使用Dask并行过滤数据帧的块

Posted 2021-04-12

我有一个用例，需要过滤属于某个因子组的数据帧。

例如，假设数据框看起来像这样：

index    col1    col2
0        A       5
1        A       10
2        B       5
3        B       8
4        C       0
... and so on

我需要做的是在df[df['col1'] == 'A']处获取大量数据，然后应用另一个函数，根据其他条件测试行，并进一步过滤掉数据。

所以我的循环当前看起来像这样：

df_chunks = [df[df['col1'] == x] for x in col1_unique_values]
outputs = []
for df_chunk in df_chunks:
    output = do_my_custom_filter(df_chunk)
    outputs.append(output)

是否可以用Dask替换last for循环？我尝试了多处理路由，发现很难调试。

任何为此目的的设计模式都将非常有帮助！

答案

看来您的计算可能令人尴尬地是并行的。如果是这样，您可以编写一个使用并产生Pandas数据帧的函数，然后仅调用df.map_partitions(your_function)

或者，如果您希望将这些块合并在一起，那么您可能希望查看groupby-apply，如下所示：

df.groupby("col1").apply(my_custom_function)