Spark MLlib（Java）中分类与数值特征处理方案咨询

针对你在使用Spark MLlib 2.11 Java版本时，需要给RandomForestClassifier同时传入字符串分类特征和数值特征的场景，核心思路是将不同类型的特征分别处理后合并成统一的数值特征向量——因为RandomForest只接受数值型的特征向量作为输入。下面我会拆解正确的步骤，并给出完整的可运行示例。

核心步骤说明 #

RandomForest不直接支持字符串特征，所以我们需要：

1. 拆分特征列：不要把所有特征塞进一个数组列，而是让每个特征（不管是分类还是数值）都作为DataFrame的单独列——这是后续分别处理的基础。
2. 处理标签：用StringIndexer把字符串类型的标签转成模型能识别的索引值。
3. 处理分类特征：对每个字符串分类特征，先用StringIndexer转成数值索引，再用OneHotEncoderEstimator（Spark 2.x推荐的API，替代旧的OneHotEncoder）生成独热编码向量；如果是低基数的分类特征，也可以跳过独热编码，直接用索引值（RandomForest对两种形式都能友好处理）。
4. 保留/转换数值特征：数值特征可以直接使用，确保类型为Double（Spark MLlib偏好该类型）。
5. 合并所有特征：用VectorAssembler把处理后的分类特征向量和数值特征合并成一个单一的features列。
6. 可选：特征自动识别：用VectorIndexer自动识别合并后特征向量中的分类特征（基于基数），帮助RandomForest更好地处理不同类型的特征。

完整Java示例代码 #

假设我们的数据集包含：

• 字符串标签label（比如"yes"/"no"）
• 两个字符串分类特征category1（比如"red"/"blue"/"green"）、category2（比如"small"/"medium"/"large"）
• 两个数值特征numFeature1（Double类型）、numFeature2（Double类型）

import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.PipelineStage;
import org.apache.spark.ml.classification.RandomForestClassifier;
import org.apache.spark.ml.feature.*;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import java.util.ArrayList;
import java.util.List;

public class MixedFeaturesRandomForest {
    public static void main(String[] args) {
        // 初始化SparkSession
        SparkSession spark = SparkSession.builder()
                .appName("MixedFeaturesRandomForest")
                .master("local[*]")
                .getOrCreate();

        // 1. 构建示例训练数据（实际场景可从文件读取）
        List<Row> trainingDataList = new ArrayList<>();
        trainingDataList.add(RowFactory.create("yes", "red", "small", 1.2, 3.4));
        trainingDataList.add(RowFactory.create("no", "blue", "medium", 2.3, 4.5));
        trainingDataList.add(RowFactory.create("yes", "green", "large", 3.4, 5.6));
        trainingDataList.add(RowFactory.create("no", "red", "medium", 4.5, 6.7));
        trainingDataList.add(RowFactory.create("yes", "blue", "small", 5.6, 7.8));

        StructType schema = DataTypes.createStructType(new StructField[]{
                DataTypes.createStructField("label", DataTypes.StringType, false),
                DataTypes.createStructField("category1", DataTypes.StringType, false),
                DataTypes.createStructField("category2", DataTypes.StringType, false),
                DataTypes.createStructField("numFeature1", DataTypes.DoubleType, false),
                DataTypes.createStructField("numFeature2", DataTypes.DoubleType, false)
        });

        Dataset<Row> trainingData = spark.createDataFrame(trainingDataList, schema);

        // 2. 处理标签：将字符串标签转为模型可识别的索引
        StringIndexer labelIndexer = new StringIndexer()
                .setInputCol("label")
                .setOutputCol("indexedLabel")
                .fit(trainingData);

        // 3. 处理分类特征：每个分类特征单独执行StringIndexer + OneHotEncoder
        // 处理category1
        StringIndexer cat1Indexer = new StringIndexer()
                .setInputCol("category1")
                .setOutputCol("category1Indexed")
                .fit(trainingData);

        OneHotEncoderEstimator cat1Encoder = new OneHotEncoderEstimator()
                .setInputCols(new String[]{"category1Indexed"})
                .setOutputCols(new String[]{"category1Encoded"});

        // 处理category2
        StringIndexer cat2Indexer = new StringIndexer()
                .setInputCol("category2")
                .setOutputCol("category2Indexed")
                .fit(trainingData);

        OneHotEncoderEstimator cat2Encoder = new OneHotEncoderEstimator()
                .setInputCols(new String[]{"category2Indexed"})
                .setOutputCols(new String[]{"category2Encoded"});

        // 4. 合并所有特征：将编码后的分类特征与数值特征合并为单一特征向量
        VectorAssembler assembler = new VectorAssembler()
                .setInputCols(new String[]{"category1Encoded", "category2Encoded", "numFeature1", "numFeature2"})
                .setOutputCol("rawFeatures");

        // 5. 可选：用VectorIndexer自动识别分类特征（基数<=maxCategories的视为分类特征）
        VectorIndexer featureIndexer = new VectorIndexer()
                .setInputCol("rawFeatures")
                .setOutputCol("indexedFeatures")
                .setMaxCategories(3); // 基数≤3的特征将被识别为分类特征

        // 6. 初始化RandomForestClassifier
        RandomForestClassifier rf = new RandomForestClassifier()
                .setLabelCol("indexedLabel")
                .setFeaturesCol("indexedFeatures")
                .setNumTrees(10)
                .setMaxDepth(5)
                .setSeed(42);

        // 7. 将预测结果转回原始字符串标签（方便直观查看）
        IndexToString labelConverter = new IndexToString()
                .setInputCol("prediction")
                .setOutputCol("predictedLabel")
                .setLabels(labelIndexer.labels());

        // 8. 构建Pipeline并训练模型
        Pipeline pipeline = new Pipeline()
                .setStages(new PipelineStage[]{
                        labelIndexer,
                        cat1Indexer, cat1Encoder,
                        cat2Indexer, cat2Encoder,
                        assembler,
                        featureIndexer,
                        rf,
                        labelConverter
                });

        PipelineModel model = pipeline.fit(trainingData);

        // 测试模型（示例用训练数据，实际应使用独立测试集）
        Dataset<Row> predictions = model.transform(trainingData);
        predictions.select("label", "predictedLabel", "indexedFeatures", "probability").show(false);

        spark.stop();
    }
}

针对你疑问的解答 #

1. 如何区分分类特征与数值特征？
   核心是把它们作为DataFrame的单独列：字符串类型的列用StringIndexer+OneHotEncoderEstimator处理，数值类型的列直接加入VectorAssembler即可，无需手动标记类型，通过处理逻辑自然区分。

2. 在哪里设置所有可能的类别？
   StringIndexer会自动从训练数据中学习所有类别。如果测试数据可能包含训练数据中未出现的类别，可以添加.setHandleInvalid("keep")（将未知类别映射到新索引），或者用.setCategories(...)手动指定所有可能的类别数组。

3. 关于VectorIndexer的使用
   它用于处理合并后的特征向量，自动识别低基数特征为分类特征，帮助RandomForest优化处理逻辑。如果分类特征已经做了独热编码，这一步可以省略，但保留它能让模型更精准地理解特征属性。

4. 简化方案：使用RFormula
   如果特征数量不多，可以用RFormula一步完成特征处理——它会自动将字符串分类特征转成独热编码，保留数值特征，同时处理标签。示例：

   RFormula formula = new RFormula()
           .setFormula("label ~ category1 + category2 + numFeature1 + numFeature2")
           .setLabelCol("indexedLabel")
           .setFeaturesCol("indexedFeatures");
   

   这种方式能大幅简化Pipeline步骤，但灵活性不如手动拆分处理。

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