[SPARK-60359][SS] Remove state only when eviction iterator advances

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,

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,

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
@@ -1698,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)

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
       }
     }

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 =>

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",