(spark) branch branch-4.x updated: [SPARK-58706][ML] Use SQL aggregation in OneHotEncoder fitting

[ru...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new 6bd1f806186f [SPARK-58706][ML] Use SQL aggregation in OneHotEncoder fitting
6bd1f806186f is described below

commit 6bd1f806186fbbfe1bf3d606300c8ca4bafab0c5
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Aug 12 14:34:32 2026 +0800

    [SPARK-58706][ML] Use SQL aggregation in OneHotEncoder fitting
    
    ### What changes were proposed in this pull request?
    
    This PR replaces the RDD `treeAggregate` in `OneHotEncoder` fitting with a SQL aggregation. The SQL expressions validate category indices and compute the maximum category index for each input column.
    
    The validation expressions use `raise_error`, so an executor fails when it evaluates an invalid index. This replaces the stale RDD-specific early-stop comment. As with the old implementation, Spark cancels the query after the failure, although independently running tasks may already be processing rows.
    
    ### Why are the changes needed?
    
    The previous implementation serializes an array sized to all input columns as the aggregation zero value. SQL performs partial aggregation on executors and returns a single result row to the driver, reducing driver serialization and aggregation overhead for wide encoders.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    `build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 "mllib / Compile / compile"`
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex (GPT-5)
    
    Closes #57919 from zhengruifeng/SPARK-58706-onehot-fit-aggregation.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
    (cherry picked from commit af6fe07b8a05adc821f0b483906feeff8fa775be)
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../apache/spark/ml/feature/OneHotEncoder.scala    | 46 +++++++++-------------
 1 file changed, 18 insertions(+), 28 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
index c5120d026d47..5599aa73a4e8 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
@@ -31,8 +31,8 @@ import org.apache.spark.ml.param.shared.{HasHandleInvalid, HasInputCol, HasInput
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.{DataFrame, Dataset}
 import org.apache.spark.sql.expressions.UserDefinedFunction
-import org.apache.spark.sql.functions.{col, lit, udf}
-import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
+import org.apache.spark.sql.functions.{printf => fprintf, _}
+import org.apache.spark.sql.types.{DoubleType, IntegerType, StructField, StructType}
 import org.apache.spark.util.ArrayImplicits._
 import org.apache.spark.util.SizeEstimator
 
@@ -529,34 +529,24 @@ private[feature] object OneHotEncoderCommon {
       inputColNames: Seq[String],
       outputColNames: Seq[String],
       dropLast: Boolean): Seq[AttributeGroup] = {
-    // The RDD approach has advantage of early-stop if any values are invalid. It seems that
-    // DataFrame ops don't have equivalent functions.
     val columns = inputColNames.map { inputColName =>
-      col(inputColName).cast(DoubleType)
+      val doubleCol = col(inputColName).cast(DoubleType)
+      val intCol = doubleCol.cast(IntegerType)
+      val invalidIndexError = raise_error(fprintf(
+        lit(s"Values from column $inputColName must be indices, but got %s."), doubleCol))
+      val maxIndexError = raise_error(fprintf(
+        lit(s"OneHotEncoder only supports up to ${Int.MaxValue} indices, but got %s."),
+        doubleCol))
+      when(
+        isnull(doubleCol) || isnan(doubleCol) || doubleCol < 0.0 || doubleCol =!= intCol,
+        invalidIndexError)
+        .when(doubleCol > Int.MaxValue, maxIndexError)
+        .otherwise(intCol)
     }
-    val numOfColumns = columns.length
-
-    val numAttrsArray = dataset.select(columns: _*).rdd.map { row =>
-      (0 until numOfColumns).map(idx => row.getDouble(idx)).toArray
-    }.treeAggregate(new Array[Double](numOfColumns))(
-      (maxValues, curValues) => {
-        (0 until numOfColumns).foreach { idx =>
-          val x = curValues(idx)
-          assert(x <= Int.MaxValue,
-            s"OneHotEncoder only supports up to ${Int.MaxValue} indices, but got $x.")
-          assert(x >= 0.0 && x == x.toInt,
-            s"Values from column ${inputColNames(idx)} must be indices, but got $x.")
-          maxValues(idx) = math.max(maxValues(idx), x)
-        }
-        maxValues
-      },
-      (m0, m1) => {
-        (0 until numOfColumns).foreach { idx =>
-          m0(idx) = math.max(m0(idx), m1(idx))
-        }
-        m0
-      }
-    ).map(_.toInt + 1)
+
+    val maxValues = columns.map(c => coalesce(max(c) + 1, lit(1)))
+    val maxValuesRow = dataset.select(maxValues: _*).head()
+    val numAttrsArray = Array.tabulate(maxValues.length)(maxValuesRow.getInt)
 
     outputColNames.zip(numAttrsArray).map { case (outputColName, numAttrs) =>
       createAttrGroupForAttrNames(outputColName, numAttrs, dropLast, keepInvalid = false)


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org