如何在dask.DataFrame中获得组的第一项?

Question

我想获得包含id的列的每组不同条目的第一项。它适用于pandas，但不适用于dask，因为我不能对多列进行排序，并且没有实现.head聚合。有没有别的方法可以得到想要的结果?

这是熊猫的最小示例，其中一切正常

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中，见下面，我得到一个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())

我可以用这段代码得到pandas的结果，但它似乎效率很低，因为我首先有一个不必要的排序。此外，在我的实际应用程序中，每组需要不止一行，并且head不能与dask

一起工作。

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

背景:我想说服每个人，从SAS到蟒蛇和熊猫。然而，我们有一些非常大的数据集，这是一个非常常见的应用。在SAS中，使用if first很容易。

Answer 1

很可能NotImplementedError是由.sort_values引发的，因为现在dask.dataframe只对单个列值进行排序，请参阅文档。

Answer 2

解决方案是使用一个作用于每个组的DataFrame的函数。

import dask.dataframe as dd
t2=dd.from_pandas(t,npartitions=2)
tz2 = ( 
t2
.sort_values(["particleID"],axis=0)
.groupby(["particleID"])
.apply(lambda s: s.sort_values(["distZ"],axis=0).head(2),
meta={"particleID":"int", "distZ":"int", "someothercols":"object"})
)
print(t2.compute())
print(tz2.compute())