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 <[email protected]>
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 <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit af6fe07b8a05adc821f0b483906feeff8fa775be)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
