如何在 dask.DataFrame 中获取一组的第一项？

Question

I want to get the first item of each set of different entries of a columns containing IDs. It works with pandas, but not in dask, as I cannot sort with multiple columns and the .head aggregation is not implemented. Is there another way of getting the desired result?

Here is the mimimal example for pandas, where everything works fine

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)

But in dask, see below, I get a 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())

I could get the pandas result with this code, but it seems quite inefficient, since I have a needless sort first. Also, in my real application, I need more then one row per group and head does not work with dask

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

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

Background: I want to convince everyone to with from SAS to python and pandas. However, we have some very large datasets and this is a very common application. In SAS it is quite easy with if first .

Answer 1

It's likely that the NotImplementedError is raised by .sort_values since right now dask.dataframe only implements sorting on a single column value, see docs .

Answer 2

The solution is dask.groupby.apply with a function that works on a DataFrame of each group.

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())