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 <[email protected]>
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 <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
