zhengruifeng commented on code in PR #58278:
URL: https://github.com/apache/spark/pull/58278#discussion_r3859849365
##########
core/src/main/scala/org/apache/spark/util/collection/Utils.scala:
##########
@@ -31,6 +32,24 @@ import org.apache.spark.util.SparkCollectionUtils
*/
private[spark] object Utils extends SparkCollectionUtils {
+ /**
+ * Same function as `keys.zipWithIndex.toMap`, but uses an [[OpenHashMap]]
and allows the first
+ * index to be specified.
+ */
+ def toOpenHashMapWithIndex[K: ClassTag](
+ keys: Array[K],
+ indexOffset: Int = 0): OpenHashMap[K, Int] = {
+ val map = new OpenHashMap[K, Int](keys.length)
Review Comment:
The OpenHashMap helper has been removed. The replacement Java HashMap is
pre-sized with `ceil(vocabulary.length / 0.75)`, avoiding a construction-time
resize for the expected number of entries.
##########
mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala:
##########
@@ -310,9 +311,9 @@ class CountVectorizerModel(
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 encodedIndex = dictBr.value(term)
Review Comment:
Fixed. The UDF now reads `dictBr.value` and `dict.size()` once at the start
of each document, before the token loop. Every token lookup uses the local
dictionary reference.
##########
core/src/main/scala/org/apache/spark/util/collection/Utils.scala:
##########
@@ -31,6 +32,24 @@ import org.apache.spark.util.SparkCollectionUtils
*/
private[spark] object Utils extends SparkCollectionUtils {
+ /**
+ * Same function as `keys.zipWithIndex.toMap`, but uses an [[OpenHashMap]]
and allows the first
+ * index to be specified.
+ */
+ def toOpenHashMapWithIndex[K: ClassTag](
+ keys: Array[K],
+ indexOffset: Int = 0): OpenHashMap[K, Int] = {
Review Comment:
The helper and its index-offset API have been removed, so the ambiguous
default and sentinel contract no longer exist.
##########
mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala:
##########
@@ -294,13 +294,14 @@ 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[OpenHashMap[String, Int]]] = 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)
+ // Reserve 0 for missing terms, which avoids allocating an Option for
every lookup.
+ val dict = Utils.toOpenHashMapWithIndex(vocabulary, indexOffset = 1)
broadcastDict = Some(dataset.sparkSession.sparkContext.broadcast(dict))
Review Comment:
I tested both approaches. `java.util.HashMap` also needs registration with
strict Kryo, so the current patch registers that single class. I also
implemented the vocabulary-array plus executor-cache alternative, but it keeps
both the array references and dictionary resident and adds cache/lifecycle
complexity. I chose to broadcast one Java HashMap instead: its compressed
payload measured 2.50 MB with JavaSerializer and 1.83 MB with Kryo, and the
registration preserves `spark.kryo.registrationRequired=true`.
##########
mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala:
##########
@@ -310,9 +311,9 @@ class CountVectorizerModel(
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 encodedIndex = dictBr.value(term)
+ if (encodedIndex != 0) {
Review Comment:
Fixed by switching to `java.util.HashMap`. `get` returns a boxed Integer and
null indicates a missing term, so there is no Option allocation, zero sentinel,
offset encoding, or decode arithmetic.
##########
core/src/test/scala/org/apache/spark/util/collection/OpenHashMapSuite.scala:
##########
@@ -27,6 +27,14 @@ import org.apache.spark.util.SizeEstimator
class OpenHashMapSuite extends SparkFunSuite with Matchers {
+ test("build with index") {
Review Comment:
The OpenHashMap helper and its test have been removed. The patch now adds
strict-Kryo coverage for Java HashMap, while CountVectorizerSuite already
covers present and missing vocabulary terms; all 17 CountVectorizer tests pass.
##########
mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala:
##########
@@ -294,13 +294,14 @@ 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[OpenHashMap[String, Int]]] = None
Review Comment:
I measured this side as well for 262,144 String keys. Estimated heap was
25.6 MiB for the immutable map, 17.5 MiB for OpenHashMap, and 22.3 MiB for Java
HashMap. JavaSerializer+LZ4 payloads were 3.37, 3.39, and 2.50 MB respectively;
Kryo+LZ4 payloads were 2.73, 3.09, and 1.83 MB. I tried broadcasting the array
and lazily caching the map, but that retains both array references and the map
and adds cache lifecycle complexity. I therefore chose a single broadcast Java
HashMap, which is smaller than the old map in heap and compressed wire size.
##########
core/src/main/scala/org/apache/spark/util/collection/Utils.scala:
##########
@@ -31,6 +32,24 @@ import org.apache.spark.util.SparkCollectionUtils
*/
private[spark] object Utils extends SparkCollectionUtils {
+ /**
+ * Same function as `keys.zipWithIndex.toMap`, but uses an [[OpenHashMap]]
and allows the first
+ * index to be specified.
+ */
+ def toOpenHashMapWithIndex[K: ClassTag](
Review Comment:
The helper has been removed, so there is no new placement or reuse API in
this patch. I summarized separately that existing OpenHashMap/OpenHashSet call
sites are worth revisiting per workload; String-key insert/lookup favors JDK
collections, while primitive specialization and aggregation may still favor the
Spark collections.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
