使用 Dask DataFrames 对组执行任意操作的最佳方法

Question

我想使用 Dask 进行表单的操作

df.groupby(some_columns).apply(some_function)

some_function() 可以计算一些汇总统计数据、执行时间序列预测，甚至只是将组保存到 AWS S3 中的单个文件中。

Dask documentation 声明（以及其他几个 StackOverflow 答案引用） groupby-apply 不适合聚合：

Pandas 的 groupby-apply 可用于应用任意函数，包括导致每组一行的聚合。 Dask 的 groupby-apply 将对每个分区组对应用一次 func ，因此当 func 是一个缩减时，您最终会得到每个分区组对一行。要使用 Dask 应用自定义聚合，请使用 dask.dataframe.groupby.Aggregation。

尚不清楚Aggregation 是否支持对多列的操作。然而，this DataFrames tutorial 似乎完全符合我的建议，大致为some_function = lambda x: LinearRegression().fit(...)。该示例似乎按预期工作，到目前为止，我同样没有遇到任何问题，例如some_function = lambda x: x.to_csv(...).

在什么条件下我可以期望some_function 将被传递到组的所有行？如果这永远无法保证，有没有办法打破LinearRegression 的例子？最重要的是，处理这些用例的最佳方式是什么？

Answer 1

当前版本的文档和源代码似乎不同步。具体来说，在dask.groupby的源码中，有这样一条信息：

Dask groupby 支持归约，即均值、求和等，并应用。前者不会对数据进行混洗，并且可以有效地实现为树缩减。后者是通过打乱底层分区来实现的，这样一个组的所有项目都可以在同一个分区中找到。

这与文档中关于partition-group 的警告不一致。下面的 sn-p 和任务图可视化也显示了数据的洗牌，以确保分区包含同一组的所有成员：

import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame({'group': [0,0,1,1,2,0,1,2], 'npeople': [1,2,3,4,5,6,7,8]})
ddf = dd.from_pandas(df, npartitions=2)

def myfunc(df):
    return df['npeople'].sum()

results_pandas = df.groupby('group').apply(myfunc).sort_index()
results_dask = ddf.groupby('group').apply(myfunc).compute().sort_index()

print(sum(results_dask != results_pandas))
# will print 0, so there are no differences
# between dask and pandas groupby

Answer 2

这是推测性的，但也许解决partition-group 导致单个组中的行跨分区拆分的情况的一种方法是明确地重新分区数据，以确保每个组与唯一的相关联分区。

实现此目的的一种方法是创建与组标识符列相同的索引。这通常不是cheap operation，但可以通过以组标识符已排序的方式预处理数据来帮助它。