如何修正Dask DataFrame填充NaN时取行最大值的问题？

需求：在DataFrame中使用各列的最大值填充对应列的np.nan值。经查阅文档，dask_ml.impute.SimpleImputer与sklearn.impute.SimpleImputer均不支持max填充策略，因此尝试手动通过fillna方法实现。

尝试代码如下：

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

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98, np.nan],
    'weight': [np.nan, 150, 126, 133, 164, 203]
})

df_dask = dd.from_pandas(df, npartitions=2) 
meta = [('height', 'float'),('weight', 'float')]
df_dask = df_dask.apply(lambda x: x.fillna(x.max()), axis=1, meta=meta)

df_dask.compute()

运行结果：

height  weight
0   6.21    6.21
1   5.12    150.00
2   5.85    126.00
3   5.78    133.00
4   5.98    164.00
5   203.00  203.00

设置axis=1期望按列处理，但Dask实际取了每行的最大值进行填充，结果不符合预期，请问如何修复该问题？

解决方案 #

问题核心是对Dask中apply方法的axis参数理解有误：axis=1指定的是按行处理，lambda中的x代表单行数据，因此x.max()计算的是每行的最大值，自然无法实现按列填充的需求。

要实现按列用最大值填充缺失值，正确步骤是：

1. 计算Dask DataFrame每一列的全局最大值
2. 将这些列最大值作为参数传入fillna，填充对应列的缺失值

基础实现（小数据集友好） #

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

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98, np.nan],
    'weight': [np.nan, 150, 126, 133, 164, 203]
})

df_dask = dd.from_pandas(df, npartitions=2)
# 计算每列的全局最大值
col_max = df_dask.max().compute()
# 用列最大值填充对应列的缺失值
df_dask_filled = df_dask.fillna(col_max)

print(df_dask_filled.compute())

运行结果（符合预期）：

height  weight
0    6.21   203.0
1    5.12   150.0
2    5.85   126.0
3    5.78   133.0
4    5.98   164.0
5    6.21   203.0

height列缺失值用该列最大值6.21填充，weight列缺失值用该列最大值203填充。

大数据集优化（延迟计算） #

如果数据集过大，直接compute()列最大值会将全量数据加载到内存，可通过Dask的延迟计算特性整合任务流：

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

df = pd.DataFrame({
    'height':  [6.21, 5.12, 5.85, 5.78, 5.98, np.nan],
    'weight': [np.nan, 150, 126, 133, 164, 203]
})

df_dask = dd.from_pandas(df, npartitions=2)
# 延迟计算列最大值，避免提前加载全量数据
col_max_delayed = delayed(df_dask.max)()
# 基于延迟计算的最大值填充缺失值
df_dask_filled = df_dask.fillna(col_max_delayed)

# 最后统一触发计算
print(df_dask_filled.compute())

这种方式下，列最大值计算与填充操作会被整合到Dask的任务图中，适合处理超大规模数据集。

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