Github user tdas commented on a diff in the pull request:
https://github.com/apache/spark/pull/19327#discussion_r140911245
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/SymmetricHashJoinStateManager.scala
---
@@ -87,70 +87,157 @@ class SymmetricHashJoinStateManager(
}
/**
- * Remove using a predicate on keys. See class docs for more context and
implement details.
+ * Remove using a predicate on keys.
+ *
+ * This produces an iterator over the (key, value) pairs satisfying
condition(key), where the
+ * underlying store is updated as a side-effect of producing next.
+ *
+ * This implies the iterator must be consumed fully without any other
operations on this manager
+ * or the underlying store being interleaved.
*/
- def removeByKeyCondition(condition: UnsafeRow => Boolean): Unit = {
- val allKeyToNumValues = keyToNumValues.iterator
-
- while (allKeyToNumValues.hasNext) {
- val keyToNumValue = allKeyToNumValues.next
- if (condition(keyToNumValue.key)) {
- keyToNumValues.remove(keyToNumValue.key)
- keyWithIndexToValue.removeAllValues(keyToNumValue.key,
keyToNumValue.numValue)
+ def removeByKeyCondition(condition: UnsafeRow => Boolean):
Iterator[UnsafeRowPair] = {
+ new NextIterator[UnsafeRowPair] {
+
+ private val allKeyToNumValues = keyToNumValues.iterator
+
+ private var currentKeyToNumValue: KeyAndNumValues = null
+ private var currentValues: Iterator[KeyWithIndexAndValue] = null
+
+ private def currentKey = currentKeyToNumValue.key
+
+ private val reusedPair = new UnsafeRowPair()
+
+ private def getAndRemoveValue() = {
+ val keyWithIndexAndValue = currentValues.next()
+ keyWithIndexToValue.remove(currentKey,
keyWithIndexAndValue.valueIndex)
+ reusedPair.withRows(currentKey, keyWithIndexAndValue.value)
+ }
+
+ override def getNext(): UnsafeRowPair = {
+ if (currentValues != null && currentValues.hasNext) {
+ return getAndRemoveValue()
+ } else {
+ while (allKeyToNumValues.hasNext) {
+ currentKeyToNumValue = allKeyToNumValues.next()
+ if (condition(currentKey)) {
+ currentValues = keyWithIndexToValue.getAll(
+ currentKey, currentKeyToNumValue.numValue)
+ keyToNumValues.remove(currentKey)
+
+ if (currentValues.hasNext) {
+ return getAndRemoveValue()
+ }
+ }
+ }
+ }
+
+ finished = true
+ null
}
+
+ override def close: Unit = {}
}
}
/**
- * Remove using a predicate on values. See class docs for more context
and implementation details.
+ * Remove using a predicate on values.
+ *
+ * At a high level, this produces an iterator over the (key, value)
pairs such that value
+ * satisfies the predicate, where producing an element removes the value
from the state store
+ * and producing all elements with a given key updates it accordingly.
+ *
+ * This implies the iterator must be consumed fully without any other
operations on this manager
+ * or the underlying store being interleaved.
*/
- def removeByValueCondition(condition: UnsafeRow => Boolean): Unit = {
- val allKeyToNumValues = keyToNumValues.iterator
+ def removeByValueCondition(condition: UnsafeRow => Boolean):
Iterator[UnsafeRowPair] = {
+ new NextIterator[UnsafeRowPair] {
- while (allKeyToNumValues.hasNext) {
- val keyToNumValue = allKeyToNumValues.next
- val key = keyToNumValue.key
+ // Reuse this object to avoid creation+GC overhead.
+ private val reusedPair = new UnsafeRowPair()
- var numValues: Long = keyToNumValue.numValue
- var index: Long = 0L
- var valueRemoved: Boolean = false
- var valueForIndex: UnsafeRow = null
+ private val allKeyToNumValues = keyToNumValues.iterator
- while (index < numValues) {
- if (valueForIndex == null) {
- valueForIndex = keyWithIndexToValue.get(key, index)
+ private var currentKey: UnsafeRow = null
+ private var numValues: Long = 0L
+ private var index: Long = 0L
+ private var valueRemoved: Boolean = false
+
+ // Push the data for the current key to the numValues store, and
reset the tracking variables
+ // to their empty state.
+ private def storeCurrentKey(): Unit = {
--- End diff --
nit: you are not really storing a new key, you are updating the value of an
existing key, or removing the key. how about simply, updateNumValueForKey?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
