PySpark ML回归任务中如何处理唯一值超千的高基数类别特征

核心瓶颈说明 #

你当前遇到的训练效率问题、集群资源利用率不足20%的本质，并不是集群算力不够，而是方案适配性问题：

• 对树模型强行套用OneHotEncoder生成上万维稀疏列，导致RF训练时的分裂搜索空间指数级膨胀
• 超宽稀疏向量的计算无法被Spark高效调度，任务碎片化严重，单任务计算量过低，根本无法喂饱集群CPU和内存，扩容完全无法提升效率
• OHE后接PCA的方案本身逻辑自洽性不足，且计算成本更高，属于冗余操作

可直接落地的适配方案（完全兼容你的Spark 2.4.5/Databricks 6.4环境） #

1. 树模型无需对类别特征做OneHot编码 #

这是最容易被忽略的核心点：RandomForest、GBDT这类基于分裂增益计算的树模型，本身可以直接把StringIndexer输出的索引值当作数值特征参与训练，完全不需要OneHot编码。

• 之前的OHE操作是纯粹的效率浪费：每个1000+唯一值的类别经OHE后会生成上千个二值稀疏列，直接用索引的话每个类别只占1个特征维度，整体特征维度会从万级降至几十级，训练速度可提升数十倍
• 树模型直接使用类别索引的训练精度不会下降，甚至会因为减少了无意义的稀疏列分裂搜索，避免过拟合
• 你提到的StringIndexer和标准LabelEncoder的差异完全不影响使用：StringIndexer附带的类别元数据仅为ML pipeline的步骤识别提供参考，实际训练计算用的就是纯整数索引值，完全可以当LabelEncoder的输出用。

2. 控制类别基数，避免长尾膨胀 #

对必须保留类别区分度的场景，直接在StringIndexer阶段加参数控制有效唯一值数量，不要等到OHE阶段才做处理：

• 设置maxCategories=100，将出现频率排名100之后的稀有类别全部合并为“其他”类，单类别编码后的维度不会超过100
• 合并后的类别再按需选择是否做OHE，此时生成的稀疏列数仅为百级，不会出现超宽向量的问题。

3. 高基数类别适配线性模型的替代方案 #

如果后续要使用线性回归这类对特征分布敏感的模型，不需要走OHE+PCA的老路，直接用目标编码（均值编码）即可：

• 每个类别编码为对应样本的标签均值，搭配10%~20%的平滑系数（避免小样本类别过拟合），每个类别仅占1个特征维度
• 编码计算必须在交叉验证的折内完成，避免全局计算导致的标签泄露，可直接封装为Pipeline的estimator步骤嵌入现有链路
• 该方案的计算成本远低于OHE+PCA，且编码后特征的预测能力更强，不存在PCA应用于二值稀疏列的合理性问题。

现有代码的优化版本 #

直接将原有OHE链路删除，适配上述方案的核心代码如下（完全兼容你的Spark版本）：

indexers = []
for name in strings_to_index:
    # 合并长尾稀有类别，控制单类别的有效唯一值规模
    indexers.append(StringIndexer(
        inputCol=name, 
        outputCol=name+'_index', 
        handleInvalid='skip',
        maxCategories=100
    ))
    # 直接将索引列加入特征集合，无需经过OHE
    feature_list.append(name+'_index')
    feature_list.remove(name)

# 删除原有OneHotEncoderEstimator步骤
assembler = VectorAssembler(inputCols=feature_list, outputCol="features")
rf = RandomForestRegressor(
    labelCol="label", 
    featuresCol="features", 
    cacheNodeIds=True, 
    seed=42,
    subsamplingRate=0.8  # 增加子采样参数，进一步降低单棵树的计算量，几乎不损失精度
)

# 后续的paramGrid、evaluator、CrossValidator、Pipeline逻辑和你原有实现保持一致
paramGrid = ParamGridBuilder() \
      .addGrid(rf.numTrees, [10, 20]) \
      .addGrid(rf.maxDepth, [5, 10]) \
    .build()

evaluator = RegressionEvaluator(labelCol="label", predictionCol="prediction", metricName="mae")
crossval = CrossValidator(
    estimator=rf,
    estimatorParamMaps=paramGrid,
    evaluator=evaluator,   
    numFolds=2
)

pipeline = Pipeline(stages=indexers+[assembler, crossval])
cvModel_rf = pipeline.fit(data_train)

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