zhengruifeng opened a new pull request, #58557: URL: https://github.com/apache/spark/pull/58557
### What changes were proposed in this pull request? This PR reduces the transform closure size of `NaiveBayesModel` for multinomial, complement, Bernoulli, and Gaussian models. The column-expression prediction hooks snapshot only the immutable state needed by the selected model type. Companion-object helpers perform raw prediction and probability conversion without retaining the complete model and its parameter graph. The existing single-instance prediction methods delegate to the same helpers. ### Why are the changes needed? The default probabilistic classifier hooks invoke bound model methods. Their UDF closures therefore retain the complete `NaiveBayesModel` even though scoring only needs its fitted vectors, matrices, model-specific derived state, and optional thresholds. This adds avoidable driver memory pressure for long-lived Spark Connect servers. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? The following checks passed: ``` build/sbt mllib/compile build/sbt 'mllib/testOnly org.apache.spark.ml.classification.NaiveBayesSuite' ``` `NaiveBayesSuite` ran 17 tests covering all four model types. No new tests were added because its existing model tests and `testPredictMethods` coverage exercise raw prediction, probability, and prediction output paths. A temporary local probe used a deterministic three-class model with 1,024 features for each model type. It extracted each `ScalaUDF.function` from the analyzed transform plan and serialized it with Spark's closure serializer. The former bound-model hooks were recreated in a temporary subclass. "All" serializes the functions for raw-prediction, probability, and prediction together. | Model/output | Before | After | Reduction | |---|---:|---:|---:| | Multinomial raw prediction | 30,140 B | 27,140 B | 10.0% | | Multinomial probability | 30,147 B | 27,148 B | 9.9% | | Multinomial prediction | 30,077 B | 27,184 B | 9.6% | | Multinomial all | 30,699 B | 27,380 B | 10.8% | | Complement raw prediction | 30,140 B | 27,002 B | 10.4% | | Complement probability | 30,147 B | 27,010 B | 10.4% | | Complement prediction | 30,077 B | 27,046 B | 10.1% | | Complement all | 30,699 B | 27,242 B | 11.3% | | Bernoulli raw prediction | 30,140 B | 27,145 B | 9.9% | | Bernoulli probability | 30,147 B | 27,153 B | 9.9% | | Bernoulli prediction | 30,077 B | 27,189 B | 9.6% | | Bernoulli all | 30,699 B | 27,385 B | 10.8% | | Gaussian raw prediction | 54,716 B | 51,812 B | 5.3% | | Gaussian probability | 54,723 B | 51,820 B | 5.3% | | Gaussian prediction | 54,653 B | 51,856 B | 5.1% | | Gaussian all | 55,275 B | 52,052 B | 5.8% | ### 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]
