Groupby 并移动一个 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 文档，发现没有这样的方法（而在熊猫中）

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

谢谢

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)