(spark) branch master updated: [SPARK-58634][ML] Avoid caching VectorIndexer transformation state

[ru...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new a42f5aa001c4 [SPARK-58634][ML] Avoid caching VectorIndexer transformation state
a42f5aa001c4 is described below

commit a42f5aa001c465963724c53cb89c9883c5600ad3
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Aug 10 12:10:15 2026 +0800

    [SPARK-58634][ML] Avoid caching VectorIndexer transformation state
    
    ### What changes were proposed in this pull request?
    
    This PR constructs VectorIndexerModel's per-vector transformation function for each `transform` call instead of caching it on the model. The function captures only the state needed for that transformation.
    
    It also adds a regression test that changes `handleInvalid` after a transform has been executed.
    
    ### Why are the changes needed?
    
    The cached function retained transformation-only state on the model and continued using the initial `handleInvalid` value after the parameter changed. Constructing it at transform time releases that state and honors the current parameter value.
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes. Updating `handleInvalid` on an existing VectorIndexerModel now takes effect after an earlier transformation.
    
    ### How was this patch tested?
    
    Added a regression case to `VectorIndexerSuite` and ran:
    
    ```
    build/sbt -Dsbt.supershell=false 'mllib/testOnly org.apache.spark.ml.feature.VectorIndexerSuite'
    ```
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex (GPT-5)
    
    Closes #57839 from zhengruifeng/SPARK-58634-vector-indexer-transform-state.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../main/scala/org/apache/spark/ml/feature/VectorIndexer.scala    | 8 +++-----
 .../scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala    | 7 +++++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
index d1810d8092ed..682746421994 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
@@ -365,17 +365,15 @@ class VectorIndexerModel private[ml] (
     attrs
   }
 
-  // TODO: Check more carefully about whether this whole class will be included in a closure.
-
   /** Per-vector transform function */
-  private lazy val transformFunc: Vector => Vector = {
+  private def getTransformFunc: Vector => Vector = {
     val sortedCatFeatureIndices = categoryMaps.keys.toArray.sorted
     val localVectorMap = categoryMaps
     val localNumFeatures = numFeatures
     val localHandleInvalid = getHandleInvalid
     val f: Vector => Vector = { (v: Vector) =>
       assert(v.size == localNumFeatures, "VectorIndexerModel expected vector of length" +
-        s" $numFeatures but found length ${v.size}")
+        s" $localNumFeatures but found length ${v.size}")
       v match {
         case dv: DenseVector =>
           var hasInvalid = false
@@ -448,7 +446,7 @@ class VectorIndexerModel private[ml] (
   override def transform(dataset: Dataset[_]): DataFrame = {
     transformSchema(dataset.schema, logging = true)
     val newField = prepOutputField(dataset.schema)
-    val transformUDF = udf { vector: Vector => transformFunc(vector) }
+    val transformUDF = udf(getTransformFunc)
     val newCol = transformUDF(dataset($(inputCol)))
     val ds = dataset.withColumn($(outputCol), newCol, newField.metadata)
     if (getHandleInvalid == VectorIndexer.SKIP_INVALID) {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
index 1529b4c661ce..664a1937b381 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
@@ -271,6 +271,13 @@ class VectorIndexerSuite extends MLTest with DefaultReadWriteTest with Logging {
       testTransformerByGlobalCheckFunc[FeatureData](points, model1, "indexed") { rows =>
         assert(rows.map(_(0)) == expected)
       }
+      model1.set(model1.handleInvalid, "keep")
+      testTransformerByGlobalCheckFunc[FeatureData](pointsTestInvalid, model1, "indexed") { rows =>
+        assert(rows.map(_(0)) == expected ++ Array(
+          Vectors.dense(2.0, 2.0, 0.0),
+          Vectors dense(0.0, 4.0, 2.0),
+          Vectors.dense(1.0, 3.0, 3.0)))
+      }
       val vectorIndexer2 = getIndexer.setMaxCategories(4).setHandleInvalid("keep")
       val model2 = vectorIndexer2.fit(points)
       testTransformerByGlobalCheckFunc[FeatureData](pointsTestInvalid, model2, "indexed") { rows =>


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