This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new eccb8faffcbd [SPARK-58993][ML] Use Java HashMap for CountVectorizer
vocabulary lookup
eccb8faffcbd is described below
commit eccb8faffcbd25cc8b4dfcac8a7b3ca215c0583b
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 26 16:25:22 2026 +0800
[SPARK-58993][ML] Use Java HashMap for CountVectorizer vocabulary lookup
### What changes were proposed in this pull request?
Replace the immutable Scala `Map[String, Int]` used by
`CountVectorizerModel` for vocabulary lookup with a pre-sized
`java.util.HashMap[String, Integer]`.
Build and broadcast the dictionary once, then dereference the broadcast
value and its size once per document instead of once per token. Register
`java.util.HashMap` with Spark Kryo so this remains compatible with
`spark.kryo.registrationRequired=true`.
### Why are the changes needed?
`CountVectorizerModel.transform` performs one vocabulary lookup per input
token. A local Java 17 benchmark with 262,144 String keys, 2 million lookups,
and a 90% hit rate measured:
- Scala immutable Map: 69.5 ns/lookup
- Spark OpenHashMap: 38.3 ns/lookup
- java.util.HashMap: 20.7 ns/lookup
For the same vocabulary, Java HashMap used about 13% less deserialized heap
than the immutable map, produced a 26% smaller JavaSerializer + LZ4 broadcast
payload, and deserialized about 47% faster. OpenHashMap used less heap and
deserialized faster than Java HashMap, but Java HashMap provides the best
tradeoff for this lookup-heavy workload.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added strict-Kryo coverage for `java.util.HashMap` and ran:
- `build/sbt "core/testOnly *KryoSerializerSuite -- -z \"basic types\""`
- `build/sbt "mllib/testOnly *CountVectorizerSuite"`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58278 from zhengruifeng/ml-countvectorizer-openhashmap-dev-4.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit ed8dbf4c8858ba542985a8ccbcac0a74282a7fc8)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../apache/spark/serializer/KryoSerializer.scala | 3 ++-
.../spark/serializer/KryoSerializerSuite.scala | 3 +++
.../apache/spark/ml/feature/CountVectorizer.scala | 22 +++++++++++++++-------
3 files changed, 20 insertions(+), 8 deletions(-)
diff --git
a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
index 65f0ffcbe18e..9b120d2f5bdb 100644
--- a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
+++ b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
@@ -555,7 +555,8 @@ private[serializer] object KryoSerializer {
classOf[SparkConf],
classOf[TaskCommitMessage],
classOf[SerializedLambda],
- classOf[BitSet]
+ classOf[BitSet],
+ classOf[java.util.HashMap[_, _]]
)
private val toRegisterSerializer = Map[Class[_], KryoClassSerializer[_]](
diff --git
a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
index 236890023364..f963d2a815e7 100644
--- a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
@@ -117,6 +117,9 @@ class KryoSerializerSuite extends SparkFunSuite with
SharedSparkContext {
check(Array.empty[Int])
check(Array(Array("1", "2"), Array("1", "2", "3", "4")))
check(Array(Array(1.toByte)))
+ val javaMap = new java.util.HashMap[String, Int]()
+ javaMap.put("one", 1)
+ check(javaMap)
}
test("pairs") {
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index 0c4a42c0fd57..a0da92fa7d7e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -17,6 +17,7 @@
package org.apache.spark.ml.feature
import java.io.{DataInputStream, DataOutputStream}
+import java.util.{HashMap => JHashMap}
import org.apache.hadoop.fs.Path
@@ -33,7 +34,7 @@ import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.SizeEstimator
-import org.apache.spark.util.collection.{OpenHashMap, Utils}
+import org.apache.spark.util.collection.OpenHashMap
/**
* Params for [[CountVectorizer]] and [[CountVectorizerModel]].
@@ -294,25 +295,32 @@ class CountVectorizerModel(
def setBinary(value: Boolean): this.type = set(binary, value)
/** Dictionary created from [[vocabulary]] and its indices, broadcast once
for [[transform()]] */
- private var broadcastDict: Option[Broadcast[Map[String, Int]]] = None
+ private var broadcastDict: Option[Broadcast[JHashMap[String, Integer]]] =
None
@Since("2.0.0")
override def transform(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
if (broadcastDict.isEmpty) {
- val dict = Utils.toMapWithIndex(vocabulary)
+ val dict = new JHashMap[String, Integer](math.ceil(vocabulary.length /
0.75).toInt)
+ var index = 0
+ while (index < vocabulary.length) {
+ dict.put(vocabulary(index), index)
+ index += 1
+ }
broadcastDict = Some(dataset.sparkSession.sparkContext.broadcast(dict))
}
val dictBr = broadcastDict.get
val localMinTF = $(minTF)
val localBinary = $(binary)
val vectorizer = udf { document: Seq[String] =>
+ val dict = dictBr.value
+ val dictSize = dict.size()
val termCounts = new OpenHashMap[Int, Int]
var tokenCount = 0L
document.foreach { term =>
- dictBr.value.get(term) match {
- case Some(index) => termCounts.changeValue(index, 1, _ + 1)
- case None => // ignore terms not in the vocabulary
+ val index = dict.get(term)
+ if (index != null) {
+ termCounts.changeValue(index.intValue(), 1, _ + 1)
}
tokenCount += 1
}
@@ -323,7 +331,7 @@ class CountVectorizerModel(
termCounts.filter(_._2 >= effectiveMinTF).map(p => (p._1,
p._2.toDouble)).toSeq
}
- Vectors.sparse(dictBr.value.size, effectiveCounts)
+ Vectors.sparse(dictSize, effectiveCounts)
}
dataset.withColumn($(outputCol), vectorizer(col($(inputCol))),
outputSchema($(outputCol)).metadata)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
