如何在PySpark中使用KNNImputer？解决UDF报错问题

错误根源分析 #

你遇到的float() argument must be a string or a number, not 'list'错误，本质是两个核心问题：

1. sklearn KNNImputer的工作逻辑：它需要基于整个数据集/分区的二维特征矩阵计算邻居，逐行调用fit_transform既不符合KNN的核心逻辑，也会导致输入格式不匹配。
2. Spark UDF的参数传递错误：如果直接把多列打包成数组传入UDF，或者错误地将整列数据以列表形式传入，KNNImputer会把一维列表当成单个输入值，触发类型转换失败。

正确解决方案：用mapPartitions结合sklearn KNNImputer #

由于数据量较大，我们可以利用Spark的分区机制，在每个分区内用sklearn处理小批量数据，既保留KNN的填充逻辑，又保证处理效率。以下是具体实现步骤：

1. 准备示例数据 #

假设你的Spark数据框结构如下：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("KNNImpute").getOrCreate()
df_spark = spark.createDataFrame([
    (1.0, None, 5.0),
    (2.0, 3.0, None),
    (None, 4.0, 6.0),
    (5.0, 6.0, 7.0),
    (3.0, None, 8.0)
], ["col1", "col2", "col3"])

2. 编写分区处理函数 #

该函数会将每个分区的Spark行转换为Pandas数据框，用KNNImputer填充空值，再返回填充后的结果：

from sklearn.impute import KNNImputer
import pandas as pd
from pyspark.sql import Row

def knn_impute_partition(iterator):
    # 将分区内的行转换为Pandas DataFrame
    rows = list(iterator)
    if not rows:
        return iter([])
    # 获取特征列名
    feature_cols = rows[0].asDict().keys()
    df_part = pd.DataFrame([row.asDict() for row in rows], columns=feature_cols)
    
    # 初始化KNNImputer（可根据需求调整n_neighbors参数）
    imputer = KNNImputer(n_neighbors=2)
    # 执行填充
    filled_data = imputer.fit_transform(df_part)
    
    # 将填充后的数据转换为Spark Row对象返回
    for filled_row in filled_data:
        yield Row(**dict(zip([f"{col}_filled" for col in feature_cols], filled_row)))

3. 应用分区处理并合并结果 #

# 对原数据框的RDD应用分区处理，转换回Spark数据框
filled_rdd = df_spark.rdd.mapPartitions(knn_impute_partition)
filled_df = filled_rdd.toDF()

# 将填充后的列与原数据框关联（添加索引确保对应关系）
from pyspark.sql.functions import monotonically_increasing_id

df_with_id = df_spark.withColumn("row_id", monotonically_increasing_id())
filled_with_id = filled_df.withColumn("row_id", monotonically_increasing_id())

final_df = df_with_id.join(filled_with_id, on="row_id").drop("row_id")
final_df.show()

4. 如果只需要生成指定的pred列 #

比如你只想填充col1作为pred列，可以修改分区函数：

def knn_impute_pred(iterator):
    rows = list(iterator)
    if not rows:
        return iter([])
    df_part = pd.DataFrame([row.asDict() for row in rows], columns=["col1", "col2", "col3"])
    
    imputer = KNNImputer(n_neighbors=2)
    # 用col2、col3作为特征填充col1
    df_part[["col1", "col2", "col3"]] = imputer.fit_transform(df_part[["col1", "col2", "col3"]])
    
    for idx, row in df_part.iterrows():
        yield Row(original_col1=row["col1"], pred=row["col1"], col2=row["col2"], col3=row["col3"])

filled_pred_rdd = df_spark.rdd.mapPartitions(knn_impute_pred)
pred_df = filled_pred_rdd.toDF()
pred_df.show()

注意事项 #

• 调整n_neighbors参数：根据你的数据分布选择合适的邻居数量，避免过拟合或欠拟合。
• 分区大小：如果单个分区数据量仍然过大，可以通过repartition()调整分区数，平衡sklearn的处理压力和Spark的调度开销。
• 数据类型确保：所有特征列必须是数值类型（int/float），否则需要先做类型转换。

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