如何为Dask DataFrame按分区添加对应特定值的新列

核心思路是利用dask.dataframe.map_partitions同时传入两个分区数一致的Dask DataFrame，让Dask自动按分区一一配对处理，给第一个DataFrame的每行添加对应分区的特定值，全程无shuffle，性能最优。

具体实现步骤 #

1. 导入依赖库

import dask.dataframe as dd
import pandas as pd

2. 构造示例数据（可替换为你的实际数据）

# 第一个Dask DataFrame：3个分区，每个分区行数不同
df1 = dd.from_pandas(pd.DataFrame({'col1': range(10)}), npartitions=3)

# 第二个Dask DataFrame：与df1分区数相同，每个分区仅1行1列
df2 = dd.from_pandas(pd.DataFrame({'partition_tag': [100, 200, 300]}), npartitions=3)

3. 定义分区处理函数
   函数接收两个参数：df1的单个分区（Pandas DataFrame），df2的对应分区（Pandas DataFrame），提取df2分区的唯一值后，给df1分区新增列。

def attach_partition_value(df_part, val_part):
    # 提取当前分区的特定值（因为每个分区仅1个值）
    partition_val = val_part.iloc[0, 0]
    # 给当前分区的所有行添加该值
    df_part['partition_tag'] = partition_val
    return df_part

4. 执行分区合并
   调用map_partitions时同时传入两个DataFrame，并指定输出的元数据结构（meta）：

# meta参数需明确输出的列结构，这里基于df1新增一列构造
result_df = dd.map_partitions(
    attach_partition_value,
    df1,
    df2,
    meta=df1.assign(partition_tag=0)  # 0仅作为占位符，实际类型会被函数输出覆盖
)

关键说明 #

• 为什么之前的assign+map_partitions组合失败？
  如果单独对df2用map_partitions提取值，得到的是每个分区一个值的Dask Series，直接df1.assign(new_col=series)会让Dask尝试将整个Series广播到df1的所有行，而非按分区匹配对应值，导致结果不符合预期。

• 性能优势：
  这种方式是分区级别的独立操作，没有跨分区的数据移动或shuffle，完全利用Dask的并行计算能力，处理大数据集时效率极高。

注意事项 #

• 必须保证两个Dask DataFrame的分区数完全一致，否则map_partitions会抛出不匹配的错误。
• meta参数不可省略，Dask需要依赖它确定输出数据的结构和类型，避免运行时错误。
• 如果df2的列不固定，可改用val_part.squeeze()提取值，适配更通用的场景：

partition_val = val_part.squeeze()

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