[jvm] Support Regressor and Ranker in Pipeline with columnar input (#12058)

Author: wbo4958

jvm-packages/xgboost4j-example/src/main/java/ml/dmlc/xgboost4j/java/example/PredictLeafIndices.java

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ Copyright (c) 2014-2026 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
  */
 package ml.dmlc.xgboost4j.java.example;
 
-import java.util.Arrays;
 import java.util.HashMap;
 
 import ml.dmlc.xgboost4j.java.Booster;
@@ -53,14 +52,14 @@ public static void main(String[] args) throws XGBoostError {
 
     //predict using first 2 tree
     float[][] leafindex = booster.predictLeaf(testMat, 2);
-    for (float[] leafs : leafindex) {
-      System.out.println(Arrays.toString(leafs));
+    if (leafindex.length > 0 && leafindex[0].length > 1) {
+      System.out.println(leafindex[0][0] + ", " + leafindex[0][1]);
     }
 
     //predict all trees
     leafindex = booster.predictLeaf(testMat, 0);
-    for (float[] leafs : leafindex) {
-      System.out.println(Arrays.toString(leafs));
+    if (leafindex.length > 0 && leafindex[0].length > 1) {
+      System.out.println(leafindex[0][0] + ", " + leafindex[0][1]);
     }
   }
 }

jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/PredictLeafIndices.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014-2024 by Contributors
+ Copyright (c) 2014-2026 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -41,14 +41,14 @@ object PredictLeafIndices {
 
     // predict using first 2 tree
     val leafIndex = booster.predictLeaf(testMat, 2)
-    for (leafs <- leafIndex) {
-      println(java.util.Arrays.toString(leafs))
+    if (leafIndex.length > 0 && leafIndex(0).length > 1) {
+      println(s"${leafIndex(0)(0)}, ${leafIndex(0)(1)}")
     }
 
     // predict all trees
     val leafIndex2 = booster.predictLeaf(testMat, 0)
-    for (leafs <- leafIndex2) {
-      println(java.util.Arrays.toString(leafs))
+    if (leafIndex2.length > 0 && leafIndex2(0).length > 1) {
+      println(s"${leafIndex2(0)(0)}, ${leafIndex2(0)(1)}")
     }
   }
 }

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostEstimator.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2024-2025 by Contributors
+ Copyright (c) 2024-2026 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -29,7 +29,7 @@ import org.apache.spark.ml.functions.array_to_vector
 import org.apache.spark.ml.linalg.{SparseVector, Vector}
 import org.apache.spark.ml.param.{Param, ParamMap}
 import org.apache.spark.ml.util.{DefaultParamsWritable, MLReader, MLWritable, MLWriter}
-import org.apache.spark.ml.xgboost.{SparkUtils, XGBProbabilisticClassifierParams}
+import org.apache.spark.ml.xgboost.{SparkUtils, XGBPredictorParams, XGBProbabilisticClassifierParams}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions.{array, col, udf}
@@ -725,7 +725,8 @@ private[spark] abstract class XGBoostModelReader[M <: XGBoostModel[M]] extends M
 }
 
 // Trait for Ranker and Regressor Model
-private[spark] trait RankerRegressorBaseModel[M <: XGBoostModel[M]] extends XGBoostModel[M] {
+private[spark] trait RankerRegressorBaseModel[M <: XGBoostModel[M]] extends XGBoostModel[M]
+  with XGBPredictorParams[M] {
 
   override protected[spark] def postTransform(dataset: Dataset[_],
                                               pred: PredictedColumns): Dataset[_] = {

jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/ml/xgboost/SparkUtils.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2024 by Contributors
+ Copyright (c) 2024-2026 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 package org.apache.spark.ml.xgboost
 
 import org.apache.spark.{SparkContext, SparkException}
+import org.apache.spark.ml.PredictorParams
 import org.apache.spark.ml.classification.ProbabilisticClassifierParams
 import org.apache.spark.ml.linalg.VectorUDT
 import org.apache.spark.ml.param.Params
@@ -55,6 +56,20 @@ trait XGBProbabilisticClassifierParams[T <: Params]
   addNonXGBoostParam(rawPredictionCol, probabilityCol, thresholds)
 }
 
+trait XGBPredictorParams[T <: Params] extends PredictorParams {
+
+  /**
+   * XGBoost doesn't use validateAndTransformSchema since spark validateAndTransformSchema
+   * needs to ensure the feature is vector type
+   */
+  override protected def validateAndTransformSchema(schema: StructType,
+                                                    fitting: Boolean,
+                                                    featuresDataType: DataType): StructType = {
+    SparkUtils.appendColumn(schema, $(predictionCol), DoubleType)
+  }
+
+}
+
 /** Utils to access the spark internal functions */
 object SparkUtils {
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostEstimatorSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2024 by Contributors
+ Copyright (c) 2024-2026 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@ import java.util.Arrays
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.SparkException
+import org.apache.spark.ml.Pipeline
 import org.apache.spark.ml.linalg.{DenseVector, SparseVector, Vectors}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.xgboost.SparkUtils
@@ -741,4 +742,35 @@ class XGBoostEstimatorSuite extends AnyFunSuite with PerTest with TmpFolderPerSu
     // No exception
     transformedDf.collect()
   }
+
+  test("Pipeline with columnar input for Regressor and Ranker") {
+    val df = ss.createDataFrame(sc.parallelize(Seq(
+      (0.0, 1, 2, 1.0, 2.0, 3.0),
+      (1.0, 0, 5, 1.0, 2.0, 3.0),
+      (2.0, 2, 7, 1.0, 2.0, 3.0)
+    ))).toDF("label", "group", "weight", "c1", "c2", "c3")
+
+    // XGBoostRegressor with columnar features
+    val regressor = new XGBoostRegressor()
+      .setNumRound(1)
+      .setNumWorkers(1)
+      .setFeaturesCol(Array("c1", "c2", "c3"))
+
+    val regressorPipeline = new Pipeline().setStages(Array(regressor))
+    val regressorFit = regressorPipeline.fit(df)
+    val regressorPredictions = regressorFit.transform(df)
+    assert(regressorPredictions.count() === 3)
+
+    // XGBoostRanker with columnar features (requires group column)
+    val ranker = new XGBoostRanker()
+      .setNumRound(1)
+      .setNumWorkers(1)
+      .setGroupCol("group")
+      .setFeaturesCol(Array("c1", "c2", "c3"))
+
+    val rankerPipeline = new Pipeline().setStages(Array(ranker))
+    val rankerFit = rankerPipeline.fit(df)
+    val rankerPredictions = rankerFit.transform(df)
+    assert(rankerPredictions.count() === 3)
+  }
 }