按组分组并应用：pandas与dask的比较

Question

关于 dask.dataframe 行为有些地方我还不是很理解。假设我想从 pandas 复制以下内容：

import pandas as pd
import dask.dataframe as dd
import random

s = "abcd"
lst = 10*[0]+list(range(1,6))
n = 100
df = pd.DataFrame({"col1": [random.choice(s) for i in range(n)],
                   "col2": [random.choice(lst) for i in range(n)]})
# I will need an hash in dask
df["hash"] = 2*df.col1
df = df[["hash","col1","col2"]]

def fun(data):
    if data["col2"].mean()>1:
        data["col3"]=2
    else:
        data["col3"]=1
    return(data)

df1 = df.groupby("col1").apply(fun)
df1.head()

这会返回

  hash col1  col2  col3
0   dd    d     0     1
1   aa    a     0     2
2   bb    b     0     1
3   bb    b     0     1
4   aa    a     0     2

在Dask中，我尝试了...

def fun2(data):
    if data["col2"].mean()>1:
        return 2
    else:
        return 1

ddf = df.copy()
ddf.set_index("hash",inplace=True)
ddf = dd.from_pandas(ddf, npartitions=2)

gpb = ddf.groupby("col1").apply(fun2, meta=pd.Series())

使用groupby后，结果与pandas中相同，但我很难将结果合并到新列上并保留哈希索引。

我想要以下结果：

      col1  col2  col3
hash           
aa      a     5     2
aa      a     0     2
aa      a     0     2
aa      a     0     2
aa      a     4     2

更新

在使用合并功能时，我找到了这个解决方案。

ddf1 = dd.merge(ddf, gpb.to_frame(), 
                left_on="col1",
                left_index=False, right_index=True)
ddf1 = ddf1.rename(columns={0:"col3"})

如果我需要对多列进行分组，我不确定这样做会怎样，而且还不够优雅。

- rpanai

---

您的解决方案涉及将计算的dask df与另一个df合并，这违反了文档中的两个性能建议：http://dask.pydata.org/en/latest/dataframe-performance.html#joins - kuanb

嗨@knuab，谢谢。我已经查看了那个文档，但是我找不到其他解决我的问题的方法。 - rpanai

Answer 1

使用join怎么样？

以下是您的dask代码，除了给Series命名为pd.Series(name='col3')外：

def fun2(data):
    if data["col2"].mean()>1:
        return 2
    else:
        return 1

ddf = df.copy()
ddf.set_index("hash",inplace=True)
ddf = dd.from_pandas(ddf, npartitions=2)

gpb = ddf.groupby("col1").apply(fun2, meta=pd.Series(name='col3'))

然后是连接操作。

ddf.join(gpb.to_frame(), on='col1')
print(ddf.compute().head())
      col1  col2  col3
hash                 
cc      c     0     2
cc      c     0     2
cc      c     0     2
cc      c     2     2
cc      c     0     2

- Bob Haffner