Dask版本升级后降采样索引重复问题求助

背景与现象 #

在Python 3.12.8环境下，将Dask从dask[complete]==2023.12.1升级至dask[complete]==2024.12.1后出现异常：

• 将Dask DataFrame（df1）重索引匹配另一DataFrame（df2）的索引，合并后result.compute()显示正常，但result.value.compute()输出因版本而异：
  • 旧版本：索引范围为10-105，步长为5（符合预期）
  • 新版本：同一索引范围重复3次
• 新版本下绘图会因'value'列索引重复抛出IndexError: list index out of range错误

原因分析 #

原代码中使用map_partitions对每个分区执行完整的target_index重索引操作，导致每个分区都生成一份全量的目标索引数据。新版本Dask在concat时未自动合并重复索引，而旧版本的分区处理逻辑会自动去重，因此出现差异。

修复方案 #

方案1：使用Dask内置reindex方法（推荐） #

Dask的reindex方法已经内置了分区优化逻辑，会自动将目标索引分配到对应分区处理，避免重复生成数据：

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

# 生成原始数据
index1 = np.arange(0, 100, 2)
df1 = dd.from_pandas(pd.DataFrame({'value': np.random.rand(len(index1))}, index=index1), npartitions=3)

index2 = np.arange(10, 110, 5)  
df2 = dd.from_pandas(pd.DataFrame({'other_value': np.random.rand(len(index2))}, index=index2), npartitions=3)

target_index = df2.index.compute()

# 改用Dask内置reindex方法
df1_resampled = df1.reindex(target_index, method='nearest')

result = dd.concat([df1_resampled, df2], axis=1)

print(result.value.compute()) # 输出恢复为10-105步长5的索引

result.hvplot.line(
    x="index",
    y=["value", "other_value"],
    value_label="values df1 resampled, df2",
) 

方案2：手动优化map_partitions逻辑 #

如果必须保留自定义分区处理逻辑，需在每个分区仅处理自身覆盖的索引范围，避免重复生成全量数据：

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

def reindex_partition(partition, target_index, partition_bounds):
    # 筛选当前分区索引能覆盖的目标索引部分
    lower, upper = partition_bounds
    filtered_target = target_index[(target_index >= lower) & (target_index <= upper)]
    return partition.reindex(filtered_target, method='nearest')

# 生成原始数据
index1 = np.arange(0, 100, 2)
df1 = dd.from_pandas(pd.DataFrame({'value': np.random.rand(len(index1))}, index=index1), npartitions=3)

index2 = np.arange(10, 110, 5)  
df2 = dd.from_pandas(pd.DataFrame({'other_value': np.random.rand(len(index2))}, index=index2), npartitions=3)

target_index = df2.index.compute()

# 获取每个分区的索引边界
partition_bounds_list = df1.get_partition_bounds()

# 为每个分区传入对应的边界
df1_resampled = df1.map_partitions(
    reindex_partition,
    target_index,
    partition_bounds_list,
    meta=pd.DataFrame({'value': float}, index=target_index[:0])
)

result = dd.concat([df1_resampled, df2], axis=1)

print(result.value.compute())
result.hvplot.line(
    x="index",
    y=["value", "other_value"],
    value_label="values df1 resampled, df2",
) 

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