From eb43d9c203aa7fc344e8230b57b42635af5706a2 Mon Sep 17 00:00:00 2001 From: Min Yang Date: Mon, 19 Oct 2026 20:56:26 +0000 Subject: [PATCH 1/3] [SPARK-60359][SS] Remove state only when eviction iterator advances Co-authored-by: Isaac --- .../StatefulStreamlineAggregateExec.scala | 4 -- .../StreamingAggregationStateManager.scala | 2 +- .../streaming/state/EvictionIterator.scala | 20 ++++------ ...treamingAggregationStateManagerSuite.scala | 39 ++++++++++++++++++- 4 files changed, 47 insertions(+), 18 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala index a93499295789a..310f4bfd8d674 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala @@ -300,10 +300,6 @@ case class StatefulStreamlineAggregateExec( incrementalEvictionIter.foreach { evictionIter => allRemovalsTimeMs += timeTakenMs { var numRemovalsCurrRecord = 0 - // NOTE: EvictionIterator removes (and counts) a row inside hasNext, not next(), so - // a bare hasNext already deletes and bumps the metrics. To stop at exactly - // incrementalCleanupFactor removals per input we must re-check the count before - // each hasNext rather than draining the iterator. while (numRemovalsCurrRecord < incrementalCleanupFactor && evictionIter.hasNext) { // The removal happens inside of the iterator; in Update mode, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala index 63dab2c392c0f..82f868dc31333 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala @@ -64,7 +64,7 @@ sealed trait StreamingAggregationStateManager extends Serializable { /** * Return an iterator over the rows of the target state store whose event time is older than - * `evictionTimestamp`, removing each row as it is returned. See [[EvictionIterator]]. + * `evictionTimestamp`, removing each row when `next()` returns it. See [[EvictionIterator]]. */ def evictionIterator( store: StateStore, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala index e8b71adeb7466..a729ef47e53d3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala @@ -23,12 +23,11 @@ import org.apache.spark.sql.execution.streaming.operators.stateful.WatermarkSupp * An iterator over the evictable rows of a [[StateStore]], which removes each row it returns and * reports how far it has progressed. * - * Only the rows that are actually evicted are returned, so every row this iterator yields has been - * removed from the store. The removal happens in `hasNext` rather than `next` (see the note on - * `pending` below) precisely so that a caller which stops early still leaves the store consistent - * with what it observed. That lets a caller both emit the evicted rows -- streaming aggregation in - * append mode outputs a grouping key once the watermark passes it -- and count real removals rather - * than state rows scanned. + * A row is removed when `next` returns it; `hasNext` only finds the next evictable row and does + * not remove it. So a caller which stops early leaves every row it did not consume in the store, + * and `numRowsRemovedSoFar` equals the number of rows returned. That lets a caller both emit the + * evicted rows -- streaming aggregation in append mode outputs a grouping key once the watermark + * passes it -- and count real removals rather than state rows scanned. */ trait EvictionIterator extends Iterator[UnsafeRowPair] { /** Number of state rows examined so far, whether or not they were evicted. */ @@ -71,10 +70,7 @@ object EvictionIterator { override def numRowsReadDuringEvictionSoFar: Long = rowsRead override def numRowsRemovedSoFar: Long = rowsRemoved - // The row hasNext has advanced to and already removed from the store, held so next() can - // return it. Removal happens in hasNext (not next()) so that a caller which stops iterating - // after hasNext -- without the matching next() -- still leaves the store consistent with the - // rows it was told are evicted; every row this iterator surfaces has already been removed. + // The next evictable row found by hasNext, not yet removed from the store. private var pending: Option[UnsafeRowPair] = None override def hasNext: Boolean = evictionPredicate match { @@ -83,8 +79,6 @@ object EvictionIterator { val rowPair = storeIterator.next() rowsRead += 1 if (predicate.eval(rowPair.key)) { - store.remove(rowPair.key) - rowsRemoved += 1 pending = Some(rowPair) } } @@ -98,6 +92,8 @@ object EvictionIterator { } val rowPair = pending.get pending = None + store.remove(rowPair.key) + rowsRemoved += 1 rowPair } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StreamingAggregationStateManagerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StreamingAggregationStateManagerSuite.scala index 8e30f77607817..f07f6b82b44b2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StreamingAggregationStateManagerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StreamingAggregationStateManagerSuite.scala @@ -17,12 +17,14 @@ package org.apache.spark.sql.execution.streaming.state +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, SpecificInternalRow, UnsafeProjection, UnsafeRow} import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection +import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorsUtils, StatePartitionKeyExtractorFactory, StreamingAggregationStateManager} import org.apache.spark.sql.streaming.StreamTest -import org.apache.spark.sql.types.{IntegerType, StructField, StructType} +import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder, StructField, StructType} class StreamingAggregationStateManagerSuite extends StreamTest { // ============================ fields and method for test data ============================ @@ -126,6 +128,41 @@ class StreamingAggregationStateManagerSuite extends StreamTest { assert(inputRow === stateManager.get(memoryStateStore, keyRow)) } + StreamingAggregationStateManager.supportedVersions.foreach { version => + test(s"StateManager v$version - eviction iterator removes a row only on next") { + val watermarkMetadata = new MetadataBuilder() + .putLong(EventTimeWatermark.delayKey, 10) + .build() + val schema = new StructType() + .add(StructField("eventTime", LongType, nullable = false, metadata = watermarkMetadata)) + .add(StructField("value", IntegerType, nullable = false)) + val attributes = toAttributes(schema) + val stateManager = StreamingAggregationStateManager.createStateManager( + keyExpressions = Seq(attributes.head), + inputRowAttributes = attributes, + stateFormatVersion = version) + val store = new MemoryStateStore() + val inputRow = UnsafeProjection.create(schema)(InternalRow(1000L, 1)) + val key = stateManager.getKey(inputRow).copy() + stateManager.put(store, inputRow) + + val evictionIterator = stateManager.evictionIterator(store, Some(1L)) + assert(evictionIterator.numRowsReadDuringEvictionSoFar === 0) + assert(evictionIterator.numRowsRemovedSoFar === 0) + + assert(evictionIterator.hasNext) + assert(evictionIterator.hasNext) + assert(evictionIterator.numRowsReadDuringEvictionSoFar === 1) + assert(evictionIterator.numRowsRemovedSoFar === 0) + assert(store.get(key) !== null) + + assert(evictionIterator.next().value === inputRow) + assert(evictionIterator.numRowsRemovedSoFar === 1) + assert(store.get(key) === null) + assert(!evictionIterator.hasNext) + } + } + // ============================ Partition Key Extraction Tests ============================ Seq(1, 2).foreach { version => From 8361679ab89f8a1591734bec8bdbcf1503842df1 Mon Sep 17 00:00:00 2001 From: Min Yang Date: Wed, 21 Oct 2026 17:35:28 +0000 Subject: [PATCH 2/3] Remove stale comments saying EvictionIterator removes in hasNext Co-authored-by: Isaac --- .../streaming/operators/stateful/statefulOperators.scala | 4 ---- .../spark/sql/streaming/StreamingAggregationSuite.scala | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala index 799f6326f7fd3..fdc1129441785 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala @@ -1429,10 +1429,6 @@ abstract class BaseStreamingDeduplicateExec allRemovalsTimeMs += timeTakenMs { // Remove up to incrementalCleanupFactor eligible rows for every input record. var numRemovalsCurrRecord = 0 - // NOTE: the eviction iterator removes a row inside hasNext, not next(), so a bare - // hasNext already deletes it. To stop at exactly incrementalCleanupFactor removals per - // input we must re-check the count before each hasNext rather than draining the - // iterator. while (numRemovalsCurrRecord < incrementalCleanupFactor && evictionIter.hasNext) { evictionIter.next() numRemovalsCurrRecord += 1 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala index aab5cb6887ee6..a31491407efc8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala @@ -1273,7 +1273,7 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with Assertions { } // Append mode is the one mode that drives the operator's eviction path: a windowed grouping key - // is emitted only once the watermark passes it, via EvictionIterator (which removes in hasNext). + // is emitted only once the watermark passes it, via EvictionIterator. // Mirrors the stateStoreSave Append test above so the streamline operator is held to the same // watermark/eviction behaviour. testWithAllStateVersions("streamline aggregation: append mode emits windows past the watermark", From 0c1d3ef267c5ee789cbc25cf84baa803a2974f19 Mon Sep 17 00:00:00 2001 From: Min Yang Date: Wed, 21 Oct 2026 17:59:10 +0000 Subject: [PATCH 3/3] Clarify DeduplicateWithinWatermark eviction iterator removes in hasNext Co-authored-by: Isaac --- .../streaming/operators/stateful/statefulOperators.scala | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala index fdc1129441785..a0495eb6e0070 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala @@ -1694,9 +1694,10 @@ case class StreamingDeduplicateWithinWatermarkExec( // We cannot reuse `EvictionIterator` here because it reads the eviction timestamp from the // state store key, whereas this operator stores `expiresAtMicros` in the value row (the key is - // only the dedup key columns). We still match EvictionIterator's semantics by returning only - // the rows that are actually evicted, so each next() corresponds to a real removal and the - // caller's numRemovedStateRows metric counts removals rather than state rows scanned. + // only the dedup key columns). Like EvictionIterator, this returns only the rows that are + // actually evicted, so numRemovedStateRows counts removals rather than state rows scanned. + // Unlike EvictionIterator, a row is removed in hasNext (via filter); that is safe here because + // without incremental cleanup the caller always drains the iterator. store.iterator().filter { rowPair => numRowsReadDuringEviction += 1 val expiresAt = rowPair.value.getLong(0)