WIP Address a part of review comments from @tdas

* TODO list
  * replace all the usages for direct call of store.xxx whenever state manager is available
  * add iterator / remove in StreamingAggregationStateManager to remove restoreOriginRow
  * add docs

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

@@ -22,7 +22,7 @@ import org.json4s.jackson.Serialization
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.RuntimeConfig
-import org.apache.spark.sql.execution.streaming.state.FlatMapGroupsWithStateExecHelper
+import org.apache.spark.sql.execution.streaming.state.{FlatMapGroupsWithStateExecHelper, StreamingAggregationStateManager}
 import org.apache.spark.sql.internal.SQLConf.{FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION, _}
 
 /**
@@ -106,7 +106,7 @@ object OffsetSeqMetadata extends Logging {
     FLATMAPGROUPSWITHSTATE_STATE_FORMAT_VERSION.key ->
       FlatMapGroupsWithStateExecHelper.legacyVersion.toString,
     STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key ->
-      StatefulOperatorsHelper.legacyVersion.toString
+      StreamingAggregationStateManager.legacyVersion.toString
   )
 
   def apply(json: String): OffsetSeqMetadata = Serialization.read[OffsetSeqMetadata](json)

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

@@ -20,8 +20,12 @@ package org.apache.spark.sql.execution.streaming
 import scala.reflect.ClassTag
 
 import org.apache.spark.TaskContext
+import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateUnsafeProjection, GenerateUnsafeRowJoiner}
 import org.apache.spark.sql.internal.SessionState
 import org.apache.spark.sql.types.StructType
 
@@ -81,4 +85,110 @@ package object state {
         storeCoordinator)
     }
   }
+
+  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
+  }
+
+  object StreamingAggregationStateManager extends Logging {
+    val supportedVersions = Seq(1, 2)
+    val legacyVersion = 1
+
+    def createStateManager(
+        keyExpressions: Seq[Attribute],
+        inputRowAttributes: Seq[Attribute],
+        stateFormatVersion: Int): StreamingAggregationStateManager = {
+      stateFormatVersion match {
+        case 1 => new StreamingAggregationStateManagerImplV1(keyExpressions, inputRowAttributes)
+        case 2 => new StreamingAggregationStateManagerImplV2(keyExpressions, inputRowAttributes)
+        case _ => throw new IllegalArgumentException(s"Version $stateFormatVersion is invalid")
+      }
+    }
+  }
+
+  abstract class StreamingAggregationStateManagerBaseImpl(
+      protected val keyExpressions: Seq[Attribute],
+      protected val inputRowAttributes: Seq[Attribute]) extends StreamingAggregationStateManager {
+
+    @transient protected lazy val keyProjector =
+      GenerateUnsafeProjection.generate(keyExpressions, inputRowAttributes)
+
+    def getKey(row: InternalRow): UnsafeRow = keyProjector(row)
+  }
+
+  class StreamingAggregationStateManagerImplV1(
+      keyExpressions: Seq[Attribute],
+      inputRowAttributes: Seq[Attribute])
+    extends StreamingAggregationStateManagerBaseImpl(keyExpressions, inputRowAttributes) {
+
+    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)
+    }
+  }
+
+  class StreamingAggregationStateManagerImplV2(
+      keyExpressions: Seq[Attribute],
+      inputRowAttributes: Seq[Attribute])
+    extends StreamingAggregationStateManagerBaseImpl(keyExpressions, inputRowAttributes) {
+
+    private val valueExpressions: Seq[Attribute] = inputRowAttributes.diff(keyExpressions)
+    private val keyValueJoinedExpressions: Seq[Attribute] = keyExpressions ++ valueExpressions
+    private val needToProjectToRestoreValue: Boolean =
+      keyValueJoinedExpressions != inputRowAttributes
+
+    @transient private lazy val valueProjector =
+      GenerateUnsafeProjection.generate(valueExpressions, inputRowAttributes)
+
+    @transient private lazy val joiner =
+      GenerateUnsafeRowJoiner.create(StructType.fromAttributes(keyExpressions),
+        StructType.fromAttributes(valueExpressions))
+    @transient private lazy val restoreValueProjector = GenerateUnsafeProjection.generate(
+      keyValueJoinedExpressions, inputRowAttributes)
+
+    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) {
+        return savedState
+      }
+
+      val joinedRow = joiner.join(key, savedState)
+      if (needToProjectToRestoreValue) {
+        restoreValueProjector(joinedRow)
+      } else {
+        joinedRow
+      }
+    }
+
+    override def put(store: StateStore, row: UnsafeRow): Unit = {
+      val key = keyProjector(row)
+      val value = valueProjector(row)
+      store.put(key, value)
+    }
+  }
+
 }

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

@@ -30,7 +30,6 @@ import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistrib
 import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.spark.sql.execution.streaming.StatefulOperatorsHelper.StreamingAggregationStateManager
 import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.streaming.{OutputMode, StateOperatorProgress}
 import org.apache.spark.sql.types._

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

@@ -19,11 +19,10 @@ package org.apache.spark.sql.execution.streaming.state
 
 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.execution.streaming.StatefulOperatorsHelper.StreamingAggregationStateManager
 import org.apache.spark.sql.streaming.StreamTest
 import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 
-class StatefulOperatorsHelperSuite extends StreamTest {
+class StreamingAggregationStateManagerSuite extends StreamTest {
   // ============================ fields and method for test data ============================
 
   val testKeys: Seq[String] = Seq("key1", "key2")

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala

@@ -30,7 +30,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
 import org.apache.spark.sql.execution.exchange.Exchange
 import org.apache.spark.sql.execution.streaming._
-import org.apache.spark.sql.execution.streaming.state.StateStore
+import org.apache.spark.sql.execution.streaming.state.{StateStore, StreamingAggregationStateManager}
 import org.apache.spark.sql.expressions.scalalang.typed
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -65,7 +65,7 @@ class StreamingAggregationSuite extends StateStoreMetricsTest
 
   def testWithAllStateVersions(name: String, confPairs: (String, String)*)
                               (func: => Any): Unit = {
-    for (version <- StatefulOperatorsHelper.supportedVersions) {
+    for (version <- StreamingAggregationStateManager.supportedVersions) {
       test(s"$name - state format version $version") {
         executeFuncWithStateVersionSQLConf(version, confPairs, func)
       }
@@ -74,7 +74,7 @@ class StreamingAggregationSuite extends StateStoreMetricsTest
 
   def testQuietlyWithAllStateVersions(name: String, confPairs: (String, String)*)
                                      (func: => Any): Unit = {
-    for (version <- StatefulOperatorsHelper.supportedVersions) {
+    for (version <- StreamingAggregationStateManager.supportedVersions) {
       testQuietly(s"$name - state format version $version") {
         executeFuncWithStateVersionSQLConf(version, confPairs, func)
       }