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 1de68733a5b2 [SPARK-58704][ML] Optimize LSH transform closure
1de68733a5b2 is described below
commit 1de68733a5b2444331e271f8d068880e6b649ffe
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Aug 11 16:18:10 2026 +0800
[SPARK-58704][ML] Optimize LSH transform closure
### What changes were proposed in this pull request?
Create model-specific transform-function factories for LSH models. The
transform UDF no longer
captures the complete model: `MinHashLSHModel` captures its random
coefficients, and
`BucketedRandomProjectionLSHModel` captures its random matrix and bucket
length. The hash
implementation is shared with the public model method through companion
helpers.
### Why are the changes needed?
Avoiding the model capture reduces 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 with:
```
build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 'mllib / Compile /
compile'
```
Existing `MinHashLSHSuite` and `BucketedRandomProjectionLSHSuite` cover the
unchanged transform
behavior. They were not rerun locally.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57916 from zhengruifeng/SPARK-58704-lsh-transform-closure.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../ml/feature/BucketedRandomProjectionLSH.scala | 23 ++++++++++++++----
.../scala/org/apache/spark/ml/feature/LSH.scala | 8 ++++++-
.../org/apache/spark/ml/feature/MinHashLSH.scala | 28 +++++++++++++++-------
3 files changed, 46 insertions(+), 13 deletions(-)
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 ec7726a5d810..76a01fd3e346 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
@@ -95,10 +95,14 @@ class BucketedRandomProjectionLSHModel private[ml](
@Since("2.1.0")
override protected[ml] def hashFunction(elems: Vector): Array[Vector] = {
- val hashVec = new DenseVector(Array.ofDim[Double](randMatrix.numRows))
- BLAS.gemv(1.0 / $(bucketLength), randMatrix, elems, 0.0, hashVec)
- // TODO: Output vectors of dimension numHashFunctions in SPARK-18450
- hashVec.values.map(h => Vectors.dense(h.floor))
+ BucketedRandomProjectionLSHModel.hashFunction(elems, randMatrix,
$(bucketLength))
+ }
+
+ override protected[ml] def createTransformFunc: Vector => Array[Vector] = {
+ val localRandMatrix = randMatrix
+ val localBucketLength = $(bucketLength)
+ elems => BucketedRandomProjectionLSHModel.hashFunction(
+ elems, localRandMatrix, localBucketLength)
}
@Since("2.1.0")
@@ -222,6 +226,17 @@ object BucketedRandomProjectionLSH extends
DefaultParamsReadable[BucketedRandomP
@Since("2.1.0")
object BucketedRandomProjectionLSHModel extends
MLReadable[BucketedRandomProjectionLSHModel] {
+
+ private def hashFunction(
+ elems: Vector,
+ randMatrix: Matrix,
+ bucketLength: Double): Array[Vector] = {
+ val hashVec = new DenseVector(Array.ofDim[Double](randMatrix.numRows))
+ BLAS.gemv(1.0 / bucketLength, randMatrix, elems, 0.0, hashVec)
+ // TODO: Output vectors of dimension numHashFunctions in SPARK-18450
+ hashVec.values.map(h => Vectors.dense(h.floor))
+ }
+
// TODO: Save using the existing format of Array[Vector] once SPARK-12878 is
resolved.
private[ml] case class Data(randUnitVectors: Matrix)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala
index 4aa58efcbc9c..620b5b1ff8c1 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala
@@ -77,6 +77,12 @@ private[spark] abstract class LSHModel[T <: LSHModel[T]]
*/
protected[ml] def hashFunction(elems: Vector): Array[Vector]
+ /**
+ * Returns the hash function used by [[transform]]. The returned function
must not capture this
+ * model.
+ */
+ protected[ml] def createTransformFunc: Vector => Array[Vector]
+
/**
* Calculate the distance between two different keys using the distance
metric corresponding
* to the hashFunction.
@@ -97,7 +103,7 @@ private[spark] abstract class LSHModel[T <: LSHModel[T]]
override def transform(dataset: Dataset[_]): DataFrame = {
transformSchema(dataset.schema, logging = true)
- val transformUDF = udf(hashFunction(_: Vector))
+ val transformUDF = udf(createTransformFunc)
dataset.withColumn($(outputCol), transformUDF(dataset($(inputCol))))
}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala
index 46ef7dca720d..ea7b2d773e43 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala
@@ -71,14 +71,12 @@ class MinHashLSHModel private[ml](
@Since("2.1.0")
override protected[ml] def hashFunction(elems: Vector): Array[Vector] = {
- require(elems.nonZeroIterator.nonEmpty, "Must have at least 1 non zero
entry.")
- val hashValues = randCoefficients.map { case (a, b) =>
- elems.nonZeroIterator.map { case (i, _) =>
- ((1L + i) * a + b) % MinHashLSH.HASH_PRIME
- }.min.toDouble
- }
- // TODO: Output vectors of dimension numHashFunctions in SPARK-18450
- hashValues.map(Vectors.dense(_))
+ MinHashLSHModel.hashFunction(elems, randCoefficients)
+ }
+
+ override protected[ml] def createTransformFunc: Vector => Array[Vector] = {
+ val localRandCoefficients = randCoefficients
+ elems => MinHashLSHModel.hashFunction(elems, localRandCoefficients)
}
@Since("2.1.0")
@@ -220,6 +218,20 @@ object MinHashLSH extends
DefaultParamsReadable[MinHashLSH] {
@Since("2.1.0")
object MinHashLSHModel extends MLReadable[MinHashLSHModel] {
+
+ private def hashFunction(
+ elems: Vector,
+ randCoefficients: Array[(Int, Int)]): Array[Vector] = {
+ require(elems.nonZeroIterator.nonEmpty, "Must have at least 1 non zero
entry.")
+ val hashValues = randCoefficients.map { case (a, b) =>
+ elems.nonZeroIterator.map { case (i, _) =>
+ ((1L + i) * a + b) % MinHashLSH.HASH_PRIME
+ }.min.toDouble
+ }
+ // TODO: Output vectors of dimension numHashFunctions in SPARK-18450
+ hashValues.map(Vectors.dense(_))
+ }
+
private[ml] case class Data(randCoefficients: Array[Int])
private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
