在数据分析工作中，经常需要从一个包含多个类别的 DataFrame 里，为每个类别单独选出若干条记录。例如，从销售表中提取每个产品线销售额最高的三个订单，或者从用户日志中取出每位用户最近五次操作记录。这类需求本质上可以归纳为“分组后取每组前 N 行”。Pandas 的 groupby 机制结合不同函数，能给出多种实现方式，但每一种的适用条件、执行效率以及返回结果的含义都存在差异。理解这些差异，才能避免在数据处理流程中写出低效或错误的代码。

下面通过一个模拟的订单数据来演示。数据包含 category（类别）和 value（数值）两列，后续会依照不同场景逐步展开。

import pandas as pd
df = pd.DataFrame({
    'category': ['A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'C'],
    'value': [10, 30, 20, 40, 60, 50, 70, 80, 90],
    'order_id': [101, 102, 103, 104, 105, 106, 107, 108, 109]
})
print(df)

这个 DataFrame 一共有 9 行，三个类别 A、B、C，每个类别下的行数并不相同。我们的目标是为每个类别选取前 2 行，具体“前”的含义可以根据业务需求灵活调整。

一、直接使用 groupby.head(n) 获取每组原始顺序的前 N 行

最直接的方法是调用 groupby('category').head(2)。这个操作会保留每个分组内按原始行顺序排列的前 2 条记录。需要注意的是，它不会对数据做任何排序，而是原样截取。如果原始数据本身就是按照某个业务时间顺序存储的，那么这样取到的就是最早的记录；但如果数据是随机排列的，结果就未必符合“最大”“最新”等语义。

代码示例如下：

result_head = df.groupby('category').head(2)
print(result_head)

输出会保留原 DataFrame 的索引，例如 A 类的前两行是索引 0 和 1，B 类的前两行是索引 3 和 4，C 类的前两行是索引 7 和 8。这种方法代码极简，执行效率也不错，因为 Pandas 内部对 groupby.head 做了优化。但它有一个显著局限：无法自定义排序规则，只能依赖原始顺序。如果数据先经过了 shuffle 或者来自流式写入，原始顺序可能毫无业务含义，此时就需要引入排序步骤。

另外，groupby.head(n) 返回的索引是原 DataFrame 的索引，这有时会导致后续操作中出现索引重复或不对齐的问题，通常建议在需要时调用 reset_index(drop=True) 重置索引。

二、先排序再分组，获取按指定列排序后的每组前 N 行

大多数业务场景都要求“按某个值取最大/最小的前 N 条”。例如，每个类别里 value 最大的两个订单。这就需要先对整个 DataFrame 排序，再分组截取。排序的关键在于将分组键（category）和排序键（value）同时考虑进去，这样分组之后，每个组内的数据已经是按照 value 降序排列的。

下面的代码先按 category 升序、value 降序排序，然后对每个类别取前 2 行：

df_sorted = df.sort_values(['category', 'value'], ascending=[True, False])
result_sorted = df_sorted.groupby('category').head(2)
print(result_sorted)

排序后，A 类别中 value 最大的两个是 30 和 20，B 类别中是 70 和 60，C 类别中是 90 和 80。这样得到的结果符合大多数分析预期。值得注意的是，排序时把 category 放在第一位是为了保证相同类别的行聚集在一起，便于 groupby 操作。如果只按 value 排序，groupby.head 仍然能工作，但分组会变得零散，虽然最终结果相同，但内部处理可能更慢，而且代码可读性下降。

另一种常见写法是使用 groupby.apply 配合自定义函数，例如 df.groupby('category').apply(lambda x: x.nlargest(2, 'value'))。这种写法同样能得到排序后的前 N 行，但 apply 会触发 Python 级别的循环，在类别数量很多或数据量较大时，性能明显低于先全局排序再 head 的方式。因此，除非需要非常复杂的组内逻辑，否则推荐使用“排序 + head”。

三、使用 cumcount 或 rank 进行向量化过滤，提升大数据量性能

当 DataFrame 规模很大（例如数百万行、上千个类别）时，groupby.head 和全局排序的方案依然可行，但有时候我们希望更加精细地控制过程，或者需要保留排序后的完整数据而不立即截断。此时，可以为每个分组内的行生成一个序号，然后直接过滤出序号小于 N 的行。cumcount 就是为此设计的向量化方法，它按照分组内行出现的顺序从 0 开始编号。

如果不关心组内排序，只想按原始顺序取每组前 N 行，可以用：

df['row_num'] = df.groupby('category').cumcount()
result_cumcount = df[df['row_num'] < 2].drop(columns='row_num')
print(result_cumcount)

如果需要先按 value 降序排序再取每组前 N 行，则可以结合排序和 cumcount：

df_sorted = df.sort_values(['category', 'value'], ascending=[True, False])
df_sorted['row_num'] = df_sorted.groupby('category').cumcount()
result_cum_sorted = df_sorted[df_sorted['row_num'] < 2].drop(columns='row_num')
print(result_cum_sorted)

cumcount 完全在 C 层面实现，速度非常快，而且过滤操作可以通过布尔索引完成，不需要保留额外的分组对象。类似地，还可以使用 rank 方法给每个组内的行打上排名，再筛选排名靠前的行。但 rank 默认处理并列值时会导致取出的行数超过 N（因为并列排名会重复），需要设置 method='first' 来打破并列，保证每组恰好 N 行。rank 的写法如下：

df_sorted = df.sort_values(['category', 'value'], ascending=[True, False])
df_sorted['rank'] = df_sorted.groupby('category')['value'].rank(method='first', ascending=False)
result_rank = df_sorted[df_sorted['rank'] <= 2].drop(columns='rank')
print(result_rank)

与 cumcount 相比，rank 更灵活，可以基于任意列计算排名，而 cumcount 只关心行位置。不过 cumcount 通常更快，因为它不需要计算数值比较。在大数据场景下，优先考虑 cumcount 或 rank 这类向量化操作，避免使用 apply。

四、常见坑点与注意事项

分组键中存在缺失值（NaN）是一个容易被忽略的问题。Pandas 的 groupby 默认会忽略分组键为 NaN 的行，也就是说这些行不会出现在任何分组中，导致数据“丢失”。如果需要保留这些行（例如视为一个单独的类别），可以使用 dropna=False 参数。示例：

df_with_nan = df.copy()
df_with_nan.loc[0, 'category'] = None
# 默认会丢掉 None 行
print(df_with_nan.groupby('category').head(2))
# 保留 None 行作为一个分组
print(df_with_nan.groupby('category', dropna=False).head(2))

另一个常见的坑是索引对齐问题。groupby.head 返回的索引是原 DataFrame 的索引，如果后续操作依赖于连续索引，可能会出错。建议在获得结果后调用 reset_index(drop=True) 生成新的连续索引。此外，当使用多列分组时，排序和 cumcount 的配合需要确保分组键都参与排序，否则 cumcount 生成的序号可能不符合预期。

最后，注意 nth 方法与 head 的差异。nth(n) 可以精确取出每组第 n+1 行（从 0 开始计数），但如果某个组的总行数不足 n+1，会引发错误或返回空，而 head 则不会报错，只是返回实际存在的行。在实际工作中，如果只需前 N 行且不能确定每组都有足够行数，优先使用 head 或基于 cumcount 的过滤，而不是 nth。

综合来看，没有一种方法是绝对最优的。简单直接、不需要排序的场景，用 groupby.head 即可；需要按某列排序取最大或最小值，先全局排序再分组 head 是平衡的选择；数据量特别大或者需要嵌入更复杂的过滤逻辑时，cumcount 和 rank 提供了高效的向量化手段。掌握这些方法及其适用边界，能显著提升数据预处理流程的健壮性和执行效率。