如何在Dask DataFrame中获取分组的首个数据项？

我需要实现按ID分组后，获取每个分组排序后的首个（或前N个）条目，这个操作在Pandas里可以轻松完成，但Dask中直接用多列排序+groupby.head会报错——因为Dask暂不支持多列排序后的groupby.head组合，也没有实现head聚合方法。

Pandas中的实现示例 #

import pandas as pd

t = pd.DataFrame(
    [[1,2,"ij"],[1,2,"huHU"],[2,4],[2,9],[0,17],[0,2],[1,8],[1,-18]],
    columns=["particleID","distZ","someothercols"]
)
tz = ( 
        t
        .sort_values(["particleID","distZ"], axis=0)
        .groupby(["particleID"])
        .head(1)
    )
print(t)
print(tz)

Dask中的直接尝试（报错） #

直接照搬Pandas逻辑会抛出NotImplementedError：

import dask.dataframe as dd

t2 = dd.from_pandas(t, npartitions=2)
tz2 = ( 
        t2
        .sort_values(["particleID","distZ"], axis=0)
        .groupby(["particleID"])
        .head(1)
    )

print(t2.compute())

低效的临时解决方案 #

我试过两次单列排序+groupby.first能得到相同结果，但存在多余排序，且只能获取每组第一条，无法支持前N条的需求：

tz2 = ( 
        t2
        .sort_values(["distZ"], axis=0)
        .sort_values(["particleID"], axis=0)
        .groupby(["particleID"])
        .first()
    )

print(t2.compute())
print(tz2.compute())

高效解决方案 #

方案1：利用groupby.apply复用Pandas逻辑 #

通过groupby.apply将每个分组的处理交给Pandas，Dask负责拆分和合并分组，支持获取每组前N条：

import dask.dataframe as dd

def get_top_n(df, n=1):
    # 分组内按distZ排序后取前n条，和Pandas逻辑一致
    return df.sort_values(["distZ"]).head(n)

tz2 = t2.groupby("particleID").apply(get_top_n, meta=t2._meta).reset_index(drop=True)
print(tz2.compute())

注：该方案适合分组数量不是特别巨大的场景，逻辑简单易理解。

方案2：用rank过滤（大数据场景更高效） #

通过计算分组内的排名来过滤，避免全局排序，效率更高，适合超大规模数据集：

# 可选：先按particleID设置索引分区，保证同组数据在同一分区，提升计算效率
t2 = t2.set_index("particleID")

# 计算每个分组内distZ的排名（ascending=True表示取最小的distZ，和原逻辑一致）
t2["rank"] = t2.groupby("particleID")["distZ"].rank(method="first", ascending=True)

# 过滤出排名<=N的记录，这里N=1
tz2 = t2[t2["rank"] <= 1].drop("rank", axis=1).reset_index()
print(tz2.compute())

注：调整rank<=N中的N即可获取每组前N条记录，完全匹配SAS中if first的常见需求。

内容的提问来源于stack exchange，提问作者Dronakuul