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 b63060727b35 [SPARK-58700][ML] Optimize ALSModel transform closure
b63060727b35 is described below
commit b63060727b35a1804aadba0b993ca646beac7f57
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Aug 11 16:15:19 2026 +0800
[SPARK-58700][ML] Optimize ALSModel transform closure
### What changes were proposed in this pull request?
Create ALSModel's prediction UDF during transform after snapshotting rank.
The UDF no longer captures the ALSModel instance.
The serialized prediction UDF decreases from 5,040 bytes to 2,815 bytes
(44.1%).
### Why are the changes needed?
Avoiding the model capture reduces the memory retained by the transform
closure, which is especially useful for Spark Connect server workloads.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Compiled the edited ALS source and serialized the generated prediction UDF
with Java serialization. The existing ALSSuite covers the unchanged prediction
behavior.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57913 from zhengruifeng/SPARK-58700-als-transform-closure.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/ml/recommendation/ALS.scala | 33 ++++++++++++----------
1 file changed, 18 insertions(+), 15 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index 9a54b70f7777..5ab6ffa046b5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -293,20 +293,6 @@ class ALSModel private[ml] (
@Since("3.0.0")
def setBlockSize(value: Int): this.type = set(blockSize, value)
- private val predict = udf { (featuresA: Seq[Float], featuresB: Seq[Float]) =>
- if (featuresA != null && featuresB != null) {
- var dotProduct = 0.0f
- var i = 0
- while (i < rank) {
- dotProduct += featuresA(i) * featuresB(i)
- i += 1
- }
- dotProduct
- } else {
- Float.NaN
- }
- }
-
@Since("2.0.0")
override def transform(dataset: Dataset[_]): DataFrame = {
transformSchema(dataset.schema)
@@ -326,7 +312,8 @@ class ALSModel private[ml] (
.join(itemFactors.alias(itemFactorsAlias),
col(s"${validatedInputAlias}.${$(itemCol)}") ===
col(s"${itemFactorsAlias}.id"), "left")
.select(col(s"${validatedInputAlias}.*"),
- predict(col(s"${userFactorsAlias}.features"),
col(s"${itemFactorsAlias}.features"))
+ ALSModel.getPredictUDF(rank)(
+ col(s"${userFactorsAlias}.features"),
col(s"${itemFactorsAlias}.features"))
.alias($(predictionCol)))
getColdStartStrategy match {
@@ -539,6 +526,22 @@ private[ml] case class FeatureData(id: Int, features:
Array[Float])
@Since("1.6.0")
object ALSModel extends MLReadable[ALSModel] {
+ private def getPredictUDF(rank: Int) = {
+ udf { (featuresA: Seq[Float], featuresB: Seq[Float]) =>
+ if (featuresA != null && featuresB != null) {
+ var dotProduct = 0.0f
+ var i = 0
+ while (i < rank) {
+ dotProduct += featuresA(i) * featuresB(i)
+ i += 1
+ }
+ dotProduct
+ } else {
+ Float.NaN
+ }
+ }
+ }
+
private[ml] def serializeData(data: FeatureData, dos: DataOutputStream):
Unit = {
import ReadWriteUtils._
dos.writeInt(data.id)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
