dongjoon-hyun commented on code in PR #58278:
URL: https://github.com/apache/spark/pull/58278#discussion_r3855707689
##########
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:
This changes the broadcast value from `scala.collection.immutable.Map`
(registered by chill's `AllScalaRegistrar`) to `OpenHashMap`, which isn't
registered in `KryoSerializer` (only `BitSet` from this package is). Broadcast
values are serialized with the configured `spark.serializer`, so with
`spark.serializer=KryoSerializer` and `spark.kryo.registrationRequired=true`,
`transform` now fails with `IllegalArgumentException: Class is not registered:
org.apache.spark.util.collection.OpenHashMap$mcI$sp` where it previously worked.
Could we either register `OpenHashMap`/`OpenHashSet`/`Hasher` (including the
specialized variants) in `KryoSerializer`, or keep the wire format entries-only
(e.g. broadcast the `vocabulary` array and build the map lazily per executor)?
##########
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:
`dictBr.value` is called once per token, and `TorrentBroadcast.getValue()`
is `synchronized` with a `Reference` dereference on every call - a 1,000-token
document does 1,001 synchronized calls per row (including
`Vectors.sparse(dictBr.value.size, ...)` below). That per-token cost is
plausibly comparable to the `Option` allocation this PR eliminates.
Hoisting `val dict = dictBr.value` (and the size) into a local at the top of
the udf lambda, before the loop, keeps the win this PR is after.
##########
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:
This makes `OpenHashMap.apply`'s missing-key-returns-0 behavior load-bearing
across a module boundary - the exact caveat `OpenHashMap`'s class doc warns
about - with the `+1` encode in `Utils` and the `!= 0` / `- 1` decode here.
An Option-free primitive on `OpenHashMap` itself (e.g. a `getOrElse(k,
default)` doing a single specialized probe) would give the same performance
without the sentinel arithmetic, and would also serve
`StringIndexerModel.transform`, which has the same per-token
`get(...).getOrElse(...)` pattern but needs a real default index - so the `+1`
trick doesn't generalize to it.
##########
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:
Worth noting: `OpenHashMap` Java-serializes its full power-of-2 capacity
arrays (`_data`, `_values`, `BitSet`) - 1.4x-2.9x the entry count in slots -
while the immutable `Map` serialized entries only, so the broadcast payload and
per-executor footprint grow. The microbenchmark covers lookup cost but not this
side.
Broadcasting the compact `vocabulary` array and building the map lazily once
per executor would keep the wire format entries-only (and would also sidestep
the Kryo registration issue above).
##########
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:
This exercises `Utils.toOpenHashMapWithIndex` rather than `OpenHashMap`
itself, and only the `indexOffset = 1` path - the default offset, empty array,
and duplicate-key (last index wins) cases are untested, while the sibling
`toMapWithIndex` covers all of these in `SparkCollectionUtilsSuite`. Could we
mirror those cases here (or in a `Utils`-focused suite)?
##########
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:
Passing `keys.length` doesn't actually reserve room for `keys.length`
entries: `OpenHashSet` sets capacity to `nextPowerOf2(initialCapacity)` with a
0.7 grow threshold, so whenever `keys.length > 0.7 * nextPowerOf2(keys.length)`
- every power-of-2 size, including `CountVectorizer`'s default `vocabSize = 1
<< 18` - the build loop still triggers a full mid-build rehash of everything
inserted so far. Sizing by the load factor, e.g. `(keys.length / 0.7).toInt +
1`, removes the rehash and ends at the same final capacity.
Also note `OpenHashSet` caps capacity at `1 << 30` (~751M entries
effective), a limit the old immutable `Map` didn't have - theoretical for a
vocabulary, but a behavior difference vs `toMapWithIndex`.
##########
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 default `indexOffset = 0` conflicts with the missing-key sentinel the
only caller relies on: `apply` returns 0 for absent keys, so a future caller
using the default and copying the `if (map(k) != 0)` pattern would silently
treat the first key as absent. Nothing exercises the default today (the only
caller and the only test both pass 1).
Since the scaladoc doesn't mention the offset or the reserve-0 contract (and
the "Same function as `keys.zipWithIndex.toMap`" sentence is only accurate for
offset 0), could we document that contract here - or drop the default so
callers must choose? (Nit: the closing `*/` is indented one space short.)
##########
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:
Placement/reuse: the map-builder helpers (`toMap`, `toJavaMap`, plus the
inherited `toMapWithIndex`) are grouped at the bottom of this file, so this
method probably belongs next to them rather than at the top.
Also, `StringIndexerModel.transform` and `ALS.makeBlocks` already hand-roll
this exact loop (`new OpenHashMap(xs.length)` + positional `update`); migrating
them to this helper would give it (and the offset-0 default) real users and
keep one implementation to fix when e.g. the pre-sizing changes.
--
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]
