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 6c865bea601e [SPARK-58280][ML][CONNECT] Estimate size of
BucketedRandomProjectionLSH and Word2Vec models
6c865bea601e is described below
commit 6c865bea601e81eeec71dc0d7c324a4e709381c2
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Jul 23 20:41:29 2026 +0800
[SPARK-58280][ML][CONNECT] Estimate size of BucketedRandomProjectionLSH and
Word2Vec models
### What changes were proposed in this pull request?
Adds explicit `estimatedSize` implementations for
`BucketedRandomProjectionLSHModel` and `Word2VecModel`.
- BRP-LSH directly accounts for its learned random matrix with
`getSizeInBytes`, including a null guard for the ML Connect constructor.
- Word2Vec accounts for its word-index map and float vector array.
Adds fitted-model regression tests in the corresponding algorithm suites,
each with a 4 KiB upper bound.
### Why are the changes needed?
The default reflection-based estimate can traverse substantially more state
than the fitted model owns. Direct accounting keeps model-size estimation
small, predictable, and focused on learned state.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- `JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt
'mllib/Test/compile'`\n\n### Was this patch authored or co-authored using
generative AI tooling?\n\nGenerated-by: Codex GPT-5
Closes #57449 from zhengruifeng/feature-composite-size-estimates-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../spark/ml/feature/BucketedRandomProjectionLSH.scala | 8 ++++++++
.../main/scala/org/apache/spark/ml/feature/Word2Vec.scala | 12 ++++++++++++
.../ml/feature/BucketedRandomProjectionLSHSuite.scala | 11 +++++++++++
.../scala/org/apache/spark/ml/feature/Word2VecSuite.scala | 14 +++++++++++++-
4 files changed, 44 insertions(+), 1 deletion(-)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala
index e4c5b91133f0..3d1765417775 100644
---
a/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala
+++
b/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala
@@ -71,6 +71,14 @@ class BucketedRandomProjectionLSHModel private[ml](
// For ml connect only
private[ml] def this() = this("", Matrices.empty)
+ private[spark] override def estimatedSize: Long = {
+ var size = estimateMatadataSize
+ if (randMatrix != null) {
+ size += randMatrix.getSizeInBytes
+ }
+ size
+ }
+
private[ml] def this(uid: String, randUnitVectors: Array[Vector]) = {
this(uid, Matrices.fromVectors(randUnitVectors.toImmutableArraySeq))
}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
index bfaf0acbde4b..2baf6260479f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
@@ -34,6 +34,7 @@ import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.util.{Utils, VersionUtils}
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
/**
* Params for [[Word2Vec]] and [[Word2VecModel]].
@@ -216,6 +217,17 @@ class Word2VecModel private[ml] (
// For ml connect only
private[ml] def this() = this("", null)
+ private[spark] override def estimatedSize: Long = {
+ var size = estimateMatadataSize
+ if (wordVectors != null) {
+ // wordIndex: Map[String, Int]
+ size += SizeEstimator.estimate(wordVectors.wordIndex)
+ // wordVectors: Array[Float]
+ size += SizeEstimator.estimate(wordVectors.wordVectors)
+ }
+ size
+ }
+
/**
* Returns a dataframe with two fields, "word" and "vector", with "word"
being a String and
* and the vector the DenseVector that it is mapped to.
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSHSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSHSuite.scala
index a7d320e8164b..eaa9238e4c5e 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSHSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSHSuite.scala
@@ -48,6 +48,17 @@ class BucketedRandomProjectionLSHSuite extends MLTest with
DefaultReadWriteTest
ParamsSuite.checkParams(model)
}
+ test("model estimated size") {
+ val model = new BucketedRandomProjectionLSH()
+ .setInputCol("keys")
+ .setOutputCol("values")
+ .setBucketLength(1.0)
+ .fit(dataset)
+ val maxSize = 1024 * 4
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("setters") {
val model = new BucketedRandomProjectionLSHModel("brp",
Array(Vectors.dense(0.0, 1.0)))
.setInputCol("testkeys")
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
index 0e38b782f507..f79bdb7dfd1f 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
@@ -35,6 +35,19 @@ class Word2VecSuite extends MLTest with DefaultReadWriteTest
{
ParamsSuite.checkParams(model)
}
+ test("model estimated size") {
+ val df = sc.parallelize(Seq(Array("a", "b", "c"), Array("a", "b",
"c"))).toDF("text")
+ val model = new Word2Vec()
+ .setVectorSize(3)
+ .setMinCount(1)
+ .setInputCol("text")
+ .setOutputCol("result")
+ .fit(df)
+ val maxSize = 1024 * 5
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("Word2Vec") {
val sentence = "a b ".repeat(100) + "a c ".repeat(10)
val numOfWords = sentence.split(" ").length
@@ -226,4 +239,3 @@ class Word2VecSuite extends MLTest with
DefaultReadWriteTest {
}
}
-
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
