WIP Address a part of review comments from @tdas

* TODO list
  * add docs

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala

@@ -89,9 +89,13 @@ package object state {
   sealed trait StreamingAggregationStateManager extends Serializable {
     def getKey(row: InternalRow): UnsafeRow
     def getStateValueSchema: StructType
-    def restoreOriginRow(rowPair: UnsafeRowPair): UnsafeRow
     def get(store: StateStore, key: UnsafeRow): UnsafeRow
     def put(store: StateStore, row: UnsafeRow): Unit
+    def commit(store: StateStore): Long
+    def remove(store: StateStore, key: UnsafeRow): Unit
+    def iterator(store: StateStore): Iterator[UnsafeRowPair]
+    def keys(store: StateStore): Iterator[UnsafeRow]
+    def values(store: StateStore): Iterator[UnsafeRow]
   }
 
   object StreamingAggregationStateManager extends Logging {
@@ -118,6 +122,15 @@ package object state {
       GenerateUnsafeProjection.generate(keyExpressions, inputRowAttributes)
 
     def getKey(row: InternalRow): UnsafeRow = keyProjector(row)
+
+    override def commit(store: StateStore): Long = store.commit()
+
+    override def remove(store: StateStore, key: UnsafeRow): Unit = store.remove(key)
+
+    override def keys(store: StateStore): Iterator[UnsafeRow] = {
+      // discard and don't convert values to avoid computation
+      store.getRange(None, None).map(_.key)
+    }
   }
 
   class StreamingAggregationStateManagerImplV1(
@@ -127,17 +140,21 @@ package object state {
 
     override def getStateValueSchema: StructType = inputRowAttributes.toStructType
 
-    override def restoreOriginRow(rowPair: UnsafeRowPair): UnsafeRow = {
-      rowPair.value
-    }
-
     override def get(store: StateStore, key: UnsafeRow): UnsafeRow = {
       store.get(key)
     }
 
     override def put(store: StateStore, row: UnsafeRow): Unit = {
       store.put(getKey(row), row)
     }
+
+    override def iterator(store: StateStore): Iterator[UnsafeRowPair] = {
+      store.iterator()
+    }
+
+    override def values(store: StateStore): Iterator[UnsafeRow] = {
+      store.iterator().map(_.value)
+    }
   }
 
   class StreamingAggregationStateManagerImplV2(
@@ -161,15 +178,6 @@ package object state {
 
     override def getStateValueSchema: StructType = valueExpressions.toStructType
 
-    override def restoreOriginRow(rowPair: UnsafeRowPair): UnsafeRow = {
-      val joinedRow = joiner.join(rowPair.key, rowPair.value)
-      if (needToProjectToRestoreValue) {
-        restoreValueProjector(joinedRow)
-      } else {
-        joinedRow
-      }
-    }
-
     override def get(store: StateStore, key: UnsafeRow): UnsafeRow = {
       val savedState = store.get(key)
       if (savedState == null) {
@@ -189,6 +197,23 @@ package object state {
       val value = valueProjector(row)
       store.put(key, value)
     }
+
+    override def iterator(store: StateStore): Iterator[UnsafeRowPair] = {
+      store.iterator().map(rowPair => new UnsafeRowPair(rowPair.key, restoreOriginRow(rowPair)))
+    }
+
+    override def values(store: StateStore): Iterator[UnsafeRow] = {
+      store.iterator().map(rowPair => restoreOriginRow(rowPair))
+    }
+
+    private def restoreOriginRow(rowPair: UnsafeRowPair): UnsafeRow = {
+      val joinedRow = joiner.join(rowPair.key, rowPair.value)
+      if (needToProjectToRestoreValue) {
+        restoreValueProjector(joinedRow)
+      } else {
+        joinedRow
+      }
+    }
   }
 
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala

@@ -165,6 +165,18 @@ trait WatermarkSupport extends UnaryExecNode {
       }
     }
   }
+
+  protected def removeKeysOlderThanWatermark(storeManager: StreamingAggregationStateManager,
+                                             store: StateStore)
+    : Unit = {
+    if (watermarkPredicateForKeys.nonEmpty) {
+      storeManager.keys(store).foreach { keyRow =>
+        if (watermarkPredicateForKeys.get.eval(keyRow)) {
+          store.remove(keyRow)
+        }
+      }
+    }
+  }
 }
 
 object WatermarkSupport {
@@ -293,12 +305,12 @@ case class StateStoreSaveExec(
             }
             allRemovalsTimeMs += 0
             commitTimeMs += timeTakenMs {
-              store.commit()
+              stateManager.commit(store)
             }
             setStoreMetrics(store)
-            store.iterator().map { rowPair =>
+            stateManager.values(store).map { valueRow =>
               numOutputRows += 1
-              stateManager.restoreOriginRow(rowPair)
+              valueRow
             }
 
           // Update and output only rows being evicted from the StateStore
@@ -314,16 +326,16 @@ case class StateStoreSaveExec(
             }
 
             val removalStartTimeNs = System.nanoTime
-            val rangeIter = store.getRange(None, None)
+            val rangeIter = stateManager.iterator(store)
 
             new NextIterator[InternalRow] {
               override protected def getNext(): InternalRow = {
                 var removedValueRow: InternalRow = null
                 while(rangeIter.hasNext && removedValueRow == null) {
                   val rowPair = rangeIter.next()
                   if (watermarkPredicateForKeys.get.eval(rowPair.key)) {
-                    store.remove(rowPair.key)
-                    removedValueRow = stateManager.restoreOriginRow(rowPair)
+                    stateManager.remove(store, rowPair.key)
+                    removedValueRow = rowPair.value
                   }
                 }
                 if (removedValueRow == null) {
@@ -336,7 +348,7 @@ case class StateStoreSaveExec(
 
               override protected def close(): Unit = {
                 allRemovalsTimeMs += NANOSECONDS.toMillis(System.nanoTime - removalStartTimeNs)
-                commitTimeMs += timeTakenMs { store.commit() }
+                commitTimeMs += timeTakenMs { stateManager.commit(store) }
                 setStoreMetrics(store)
               }
             }
@@ -370,7 +382,7 @@ case class StateStoreSaveExec(
 
                 // Remove old aggregates if watermark specified
                 allRemovalsTimeMs += timeTakenMs { removeKeysOlderThanWatermark(store) }
-                commitTimeMs += timeTakenMs { store.commit() }
+                commitTimeMs += timeTakenMs { stateManager.commit(store) }
                 setStoreMetrics(store)
               }
             }

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StreamingAggregationStateManagerSuite.scala

@@ -103,6 +103,7 @@ class StreamingAggregationStateManagerSuite extends StreamTest {
     stateManager.put(memoryStateStore, inputRow)
 
     assert(memoryStateStore.iterator().size === 1)
+    assert(stateManager.iterator(memoryStateStore).size === memoryStateStore.iterator().size)
 
     val keyRow = stateManager.getKey(inputRow)
     assert(keyRow === expectedStateKey)
@@ -111,7 +112,15 @@ class StreamingAggregationStateManagerSuite extends StreamTest {
     val pair = memoryStateStore.iterator().next()
     assert(pair.key === keyRow)
     assert(pair.value === expectedStateValue)
-    assert(stateManager.restoreOriginRow(pair) === inputRow)
+
+    // iterate with state manager and see whether original rows are returned as values
+    val pairFromStateManager = stateManager.iterator(memoryStateStore).next()
+    assert(pairFromStateManager.key === keyRow)
+    assert(pairFromStateManager.value === inputRow)
+
+    // following as keys and values
+    assert(stateManager.keys(memoryStateStore).next() === keyRow)
+    assert(stateManager.values(memoryStateStore).next() === inputRow)
 
     // verify the stored value once again via get
     assert(memoryStateStore.get(keyRow) === expectedStateValue)