在Dask中分组应用自定义函数后返回Pandas DataFrame而非Dask DataFrame的问题

问题出在你最后一行调用了.compute()——这个方法会触发Dask立即执行所有延迟计算，并把结果转换成Pandas DataFrame返回。如果想要保留Dask DataFrame对象，只需移除.compute()。

另外，Dask的groupby.apply需要指定meta参数，明确告诉Dask自定义函数返回的DataFrame结构（尤其是你新增了sum_value列），避免类型推断错误。

修改后的代码 #

import dask.dataframe as dd
import pandas as pd
import numpy as np

pack_assign_dict = {
    'clusters': ['1234', '1234', '6789', '6789', '9999', '9999'],
    'packs': ['P1', 'P2', 'P1', 'P2', 'P1', 'P2'],
    'pack_assigned': [1.95, 0.9, 1.8, 1.85, 0.25, 0.25]
}

def addition(df1):
    df1['sum_value'] = df1['pack_assigned'].sum()
    df1 = df1.sort_values(by=['pack_assigned'])
    return df1

pack_assigned_df = pd.DataFrame(pack_assign_dict)
ddf = dd.from_pandas(pack_assigned_df, npartitions=10)

# 定义meta：指定返回的DataFrame结构，包含新增的sum_value列
meta = pd.DataFrame({
    'clusters': str,
    'packs': str,
    'pack_assigned': float,
    'sum_value': float
})

# 移除.compute()，添加meta参数
ddf1 = ddf.groupby(['clusters']).apply(addition, meta=meta)

# 验证类型
print(type(ddf1))  # 输出：<class 'dask.dataframe.core.DataFrame'>

关键说明 #

• 移除.compute()：Dask默认是延迟计算，只有调用compute()才会实际执行并返回Pandas对象，保留Dask DataFrame可以继续进行后续的分布式操作。
• 指定meta参数：因为你的自定义函数新增了sum_value列，Dask无法自动推断新列的类型，手动定义meta能确保Dask正确处理返回结果，避免警告或异常。

内容的提问来源于stack exchange，提问作者Vigneshwar Thiyagarajan