zhengruifeng opened a new pull request, #58118: URL: https://github.com/apache/spark/pull/58118
### What changes were proposed in this pull request? This PR optimizes `KMeansModel.transform` by adding ML-side prediction helpers adapted from the existing `mllib.clustering.KMeansModel.predict` implementation. The transform UDF now captures local ML-side cluster centers with norms, optional center statistics, and the distance measure. It no longer calls `predict` on the enclosing `KMeansModel`, which previously kept the full ML model and its wrapped `mllib` model reachable from the task closure. The public `KMeansModel.predict(Vector)` path also uses the new ML-side helpers directly, avoiding the previous conversion from `ml.linalg.Vector` to `mllib.linalg.Vector` for each prediction. ### Why are the changes needed? `KMeansModel` is backed by the old `mllib` KMeans model. Before this PR, `transform` used a UDF which called `predict` on the ML model instance. That made the task closure retain the whole model wrapper and its `mllib` parent model, and each row prediction converted the ML vector to an old `mllib` vector. This follows SPARK-58584 by reducing unnecessary model state captured in transform closures and removing per-row vector conversion in the KMeans prediction path. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing KMeans coverage exercises the updated transform and single-instance prediction paths. ``` build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 'mllib/testOnly org.apache.spark.ml.clustering.KMeansSuite' ``` Also ran: ``` git diff --check ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Codex (GPT-5) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
