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 9b290e006e3b [SPARK-58853][ML] Optimize IDFModel transform closure
9b290e006e3b is described below
commit 9b290e006e3b3e67f4f1edef1c2a899a336d5cbc
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 19 18:39:48 2026 +0800
[SPARK-58853][ML] Optimize IDFModel transform closure
### What changes were proposed in this pull request?
This PR optimizes `IDFModel.transform` by snapshotting the underlying IDF
vector as a local
`Array[Double]` before creating the transform UDF. The UDF now calls
ML-side companion helpers:
* `IDFModel.predict`
* `IDFModel.predictDense`
* `IDFModel.predictSparse`
These helpers operate on `org.apache.spark.ml.linalg.Vector` inputs and an
array-backed IDF lookup,
so the transform UDF does not capture the full `IDFModel` wrapper and does
not dispatch through
`Vector.apply` for each term.
The legacy `mllib.feature.IDFModel` transform path now also uses
`Array[Double]` IDF values and
delegates directly to the ML-side `predictDense` / `predictSparse` helpers.
The RDD transform path
converts the broadcast IDF vector to an array once per partition.
### Why are the changes needed?
Avoiding the model capture reduces transform-closure retained memory, which
is useful for Spark
Connect server workloads and follows the closure-size optimizations tracked
under SPARK-58584.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Ran local static checks:
```
git diff --check
grep -rn -P "[^\x00-\x7F]"
mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
awk 'length>100 && $0 !~ /^[[:space:]]*(import|package) / && $0 !~
/https?:\/\// {print FILENAME":"FNR": "length" chars"}'
mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
```
Measured serialized `ScalaUDF.function` closure size with a temporary local
probe:
| Case | Before | After | Delta |
|---|---:|---:|---:|
| `idf.transform` | 3080 B | 1247 B | -1833 B, -59.5% |
No unit test was added because this is an internal refactor of existing IDF
vector transformation
logic.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58099 from zhengruifeng/spark-58853-idf-transform-closure.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../scala/org/apache/spark/ml/feature/IDF.scala | 55 ++++++++++++++++------
.../scala/org/apache/spark/mllib/feature/IDF.scala | 45 +++++-------------
2 files changed, 52 insertions(+), 48 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
index 3d677cddf8fe..6d1dae9392d0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
@@ -152,20 +152,8 @@ class IDFModel private[ml] (
override def transform(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
- val func = { vector: Vector =>
- vector match {
- case SparseVector(size, indices, values) =>
- val (newIndices, newValues) =
feature.IDFModel.transformSparse(idfModel.idf,
- indices, values)
- Vectors.sparse(size, newIndices, newValues)
- case DenseVector(values) =>
- val newValues = feature.IDFModel.transformDense(idfModel.idf, values)
- Vectors.dense(newValues)
- case other =>
- throw new UnsupportedOperationException(
- s"Only sparse and dense vectors are supported but got
${other.getClass}.")
- }
- }
+ val localIdf = idfModel.idf.toArray
+ val func = (vector: Vector) => IDFModel.predict(localIdf, vector)
val transformer = udf(func)
dataset.withColumn($(outputCol), transformer(col($(inputCol))),
@@ -213,6 +201,45 @@ class IDFModel private[ml] (
object IDFModel extends MLReadable[IDFModel] {
private[ml] case class Data(idf: Vector, docFreq: Array[Long], numDocs: Long)
+ private def predict(idf: Array[Double], v: Vector): Vector = {
+ v match {
+ case SparseVector(size, indices, values) =>
+ val (newIndices, newValues) = predictSparse(idf, indices, values)
+ Vectors.sparse(size, newIndices, newValues)
+ case DenseVector(values) =>
+ val newValues = predictDense(idf, values)
+ Vectors.dense(newValues)
+ case other =>
+ throw new UnsupportedOperationException(
+ s"Only sparse and dense vectors are supported but got
${other.getClass}.")
+ }
+ }
+
+ private[spark] def predictDense(idf: Array[Double], values: Array[Double]):
Array[Double] = {
+ val n = values.length
+ val newValues = new Array[Double](n)
+ var j = 0
+ while (j < n) {
+ newValues(j) = values(j) * idf(j)
+ j += 1
+ }
+ newValues
+ }
+
+ private[spark] def predictSparse(
+ idf: Array[Double],
+ indices: Array[Int],
+ values: Array[Double]): (Array[Int], Array[Double]) = {
+ val nnz = indices.length
+ val newValues = new Array[Double](nnz)
+ var k = 0
+ while (k < nnz) {
+ newValues(k) = values(k) * idf(indices(k))
+ k += 1
+ }
+ (indices, newValues)
+ }
+
private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
import ReadWriteUtils._
serializeVector(data.idf, dos)
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala
b/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala
index 5634f0e1d9b7..c7af3ca4a8b6 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala
@@ -21,6 +21,7 @@ import breeze.linalg.{DenseVector => BDV}
import org.apache.spark.annotation.Since
import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.ml.feature.{IDFModel => NewIDFModel}
import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector,
Vectors}
import org.apache.spark.rdd.RDD
@@ -184,7 +185,10 @@ class IDFModel private[spark](@Since("1.1.0") val idf:
Vector,
@Since("1.1.0")
def transform(dataset: RDD[Vector]): RDD[Vector] = {
val bcIdf = dataset.context.broadcast(idf)
- dataset.mapPartitions(iter => iter.map(v =>
IDFModel.transform(bcIdf.value, v)))
+ dataset.mapPartitions { iter =>
+ val localIdf = bcIdf.value.toArray
+ iter.map(v => IDFModel.transform(localIdf, v))
+ }
}
/**
@@ -194,7 +198,7 @@ class IDFModel private[spark](@Since("1.1.0") val idf:
Vector,
* @return a TF-IDF vector
*/
@Since("1.3.0")
- def transform(v: Vector): Vector = IDFModel.transform(idf, v)
+ def transform(v: Vector): Vector = IDFModel.transform(idf.toArray, v)
/**
* Transforms term frequency (TF) vectors to TF-IDF vectors (Java version).
@@ -210,50 +214,23 @@ class IDFModel private[spark](@Since("1.1.0") val idf:
Vector,
private[spark] object IDFModel {
/**
- * Transforms a term frequency (TF) vector to a TF-IDF vector with a IDF
vector
+ * Transforms a term frequency (TF) vector to a TF-IDF vector with IDF values
*
- * @param idf an IDF vector
+ * @param idf IDF values
* @param v a term frequency vector
* @return a TF-IDF vector
*/
- def transform(idf: Vector, v: Vector): Vector = {
+ private def transform(idf: Array[Double], v: Vector): Vector = {
v match {
case SparseVector(size, indices, values) =>
- val (newIndices, newValues) = transformSparse(idf, indices, values)
+ val (newIndices, newValues) = NewIDFModel.predictSparse(idf, indices,
values)
Vectors.sparse(size, newIndices, newValues)
case DenseVector(values) =>
- val newValues = transformDense(idf, values)
+ val newValues = NewIDFModel.predictDense(idf, values)
Vectors.dense(newValues)
case other =>
throw new UnsupportedOperationException(
s"Only sparse and dense vectors are supported but got
${other.getClass}.")
}
}
-
- private[spark] def transformDense(
- idf: Vector,
- values: Array[Double]): Array[Double] = {
- val n = values.length
- val newValues = new Array[Double](n)
- var j = 0
- while (j < n) {
- newValues(j) = values(j) * idf(j)
- j += 1
- }
- newValues
- }
-
- private[spark] def transformSparse(
- idf: Vector,
- indices: Array[Int],
- values: Array[Double]): (Array[Int], Array[Double]) = {
- val nnz = indices.length
- val newValues = new Array[Double](nnz)
- var k = 0
- while (k < nnz) {
- newValues(k) = values(k) * idf(indices(k))
- k += 1
- }
- (indices, newValues)
- }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
