如何在Dask中横向合并指定列生成列表列？

需要将Dask DataFrame中的b至e列合并为一个列表列，同时保留a列的原有状态。此前使用apply方法既无效又缓慢，参考Polars的concat_list实现方式，以下是Dask的高效替代方案：

现有数据 #

期望结果 #

高效实现代码 #

import dask.dataframe as dd

# 构造示例Dask DataFrame
df = dd.DataFrame.from_dict({
    "a": [0.1, 0.2, 0.3],
    "b": [1.1, 1.2, 1.3],
    "c": [2.1, 2.2, 2.3],
    "d": [3.1, 3.2, 3.3],
    "e": [4.1, 4.2, 4.3],
}, npartitions=1)

# 指定需要合并的列
cols_to_concat = ['b', 'c', 'd', 'e']

# 合并列并生成新列表列，同时删除原合并列
df_result = df.assign(value=df[cols_to_concat].agg(list, axis=1)).drop(cols_to_concat, axis=1)

# 查看结果
print(df_result.compute())

方案说明 #

• 利用Dask的agg(list, axis=1)方法，直接对指定列进行行级聚合，将每行的多列值合并为一个列表，相比逐行apply，该方法更高效，能充分利用Dask的分布式计算能力。
• 避免了手动构造列表的错误写法（如直接将列对象放入列表），确保生成的是每行对应列值组成的列表。

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