分组并应用熊猫与达斯克

Question

关于dask.dataframe的行为，我有些不太了解。假设我想从熊猫身上复制这个

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

在达斯克我试过

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"})

我不确定如果我必须在几列上进行分组，这将如何工作。加号并不完全优雅。

Answer 1

使用join怎么样？

这是您的 dask 代码，除了命名系列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