分组并移动一个dask数据帧

Question

我想使用dask 2.14对pandas数据帧进行一些操作。例如，我想在数据帧的一列上应用一个移位：

import dask.dataframe as dd
data = dd.read_csv('some_file.csv')
data.set_index('column_A')
data['column_B'] = data.groupby(['column_A'])['column_B'].shift(-1)

但我得到了AttributeError: 'SeriesGroupBy' object has no attribute 'shift'我阅读了dask文档，发现没有这样的方法(在pandas中有(

你能提出一些有效的替代方案吗？

谢谢

Answer 1

GitHub上有一个关于这方面的公开声明。从本质上讲，你必须使用apply来绕过它。我不确定这是否会对dask的性能产生影响。还有一张罚单提到了这个问题，并指出它存在于pandas中，但它已经开放了一段时间。

这应该相当于pandas操作：

import dask.dataframe as dd
import pandas as pd
import random
df = pd.DataFrame({'a': list(range(10)),
'b': random.choices(['x', 'y'], k=10)})
print("####### PANDAS ######")
print("Initial df")
print(df.head(10))
print("................")
pandas_df = df.copy()
print("Final df")
pandas_df['a'] = pandas_df.groupby(['b'])['a'].apply(lambda x: x.shift(-1))
print(pandas_df.head(10))
print()

print("####### DASK ######")
print("Initial df")
dask_df = dd.from_pandas(df, npartitions=1).reset_index()
print(dask_df.head(10))
print("................")
dask_df['a'] = dask_df.groupby(['b'])['a'].apply(lambda x: x.shift(-1))
print("Final df")
print(dask_df.head(10))

我显然无法在dask中对该方法进行基准测试，因为似乎没有其他选择。但是，我可以在pandas:中

import string
import numpy as np
import pandas as pd

df = pd.DataFrame({'a': list(range(100000)),
'b': np.random.choice(list(string.ascii_lowercase), 100000)
})
def normal_way(df):
df = df.groupby(['b'])['a'].shift(-1)
def apply_way(df):
df = df.groupby(['b'])['a'].apply(lambda x: x.shift(-1))

timeit结果为：

%timeit normal_way(df)
4.25 ms ± 98 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_way(df)
15 ms ± 446 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)