refactoring to use common utils

sql/api/src/main/scala/org/apache/spark/sql/streaming/ListState.scala

@@ -44,3 +44,23 @@ private[sql] trait ListState[S] extends Serializable {
   /** Removes this state for the given grouping key. */
   def clear(): Unit
 }
+
+/**
+ * Interface used for arbitrary stateful operations with the v2 API to capture
+ * list value state.
+ */
+private[sql] trait ListStateModify[S] extends Serializable {
+
+  /** Update the value of the list. */
+  def put(newState: Array[S]): Unit
+
+  /** Append an entry to the list */
+  def appendValue(newState: S): Unit
+
+  /** Append an entire list to the existing value */
+  def appendList(newState: Array[S]): Unit
+
+  /** Removes this state for the given grouping key. */
+  def clear(): Unit
+}
+

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

@@ -20,7 +20,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchema.{KEY_ROW_SCHEMA, VALUE_ROW_SCHEMA}
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.ListState
 
 /**
@@ -44,6 +44,9 @@ class ListStateImpl[S](
 
   private val stateTypesEncoder = StateTypesEncoder(keySerializer, valEncoder, stateName)
 
+  private val listStateModifyImpl = new ListStateModifyImpl[S](
+    store, stateName, keyExprEnc, valEncoder, stateTypesEncoder.encodeValue)
+
   store.createColFamilyIfAbsent(stateName, KEY_ROW_SCHEMA, VALUE_ROW_SCHEMA,
     NoPrefixKeyStateEncoderSpec(KEY_ROW_SCHEMA), useMultipleValuesPerKey = true)
 
@@ -75,51 +78,21 @@ class ListStateImpl[S](
 
    /** Update the value of the list. */
    override def put(newState: Array[S]): Unit = {
-     validateNewState(newState)
-
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     var isFirst = true
-
-     newState.foreach { v =>
-       val encodedValue = stateTypesEncoder.encodeValue(v)
-       if (isFirst) {
-         store.put(encodedKey, encodedValue, stateName)
-         isFirst = false
-       } else {
-          store.merge(encodedKey, encodedValue, stateName)
-       }
-     }
+     listStateModifyImpl.put(newState)
    }
 
    /** Append an entry to the list. */
    override def appendValue(newState: S): Unit = {
-     StateStoreErrors.requireNonNullStateValue(newState, stateName)
-     store.merge(stateTypesEncoder.encodeGroupingKey(),
-         stateTypesEncoder.encodeValue(newState), stateName)
+     listStateModifyImpl.appendValue(newState)
    }
 
    /** Append an entire list to the existing value. */
    override def appendList(newState: Array[S]): Unit = {
-     validateNewState(newState)
-
-     val encodedKey = stateTypesEncoder.encodeGroupingKey()
-     newState.foreach { v =>
-       val encodedValue = stateTypesEncoder.encodeValue(v)
-       store.merge(encodedKey, encodedValue, stateName)
-     }
+     listStateModifyImpl.appendList(newState)
    }
 
    /** Remove this state. */
    override def clear(): Unit = {
-     store.remove(stateTypesEncoder.encodeGroupingKey(), stateName)
-   }
-
-   private def validateNewState(newState: Array[S]): Unit = {
-     StateStoreErrors.requireNonNullStateValue(newState, stateName)
-     StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
-
-     newState.foreach { v =>
-       StateStoreErrors.requireNonNullStateValue(v, stateName)
-     }
+     listStateModifyImpl.clear()
    }
  }

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.streaming
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.execution.streaming.TransformWithStateKeyValueRowSchema.{KEY_ROW_SCHEMA, VALUE_ROW_SCHEMA_WITH_TTL}
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore, StateStoreErrors}
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateStore}
 import org.apache.spark.sql.streaming.{ListState, TTLConfig}
 import org.apache.spark.util.NextIterator
 
@@ -50,6 +50,9 @@ class ListStateImplWithTTL[S](
   private val ttlExpirationMs =
     StateTTL.calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs)
 
+  private val listStatePutImpl = new ListStateModifyImpl[S](
+    store, stateName, keyExprEnc, valEncoder, stateTypesEncoder.encodeValue(_, ttlExpirationMs))
+
   initialize()
 
   private def initialize(): Unit = {
@@ -98,58 +101,25 @@ class ListStateImplWithTTL[S](
 
   /** Update the value of the list. */
   override def put(newState: Array[S]): Unit = {
-    validateNewState(newState)
-
-    val encodedKey = stateTypesEncoder.encodeGroupingKey()
-    var isFirst = true
-
-    newState.foreach { v =>
-      val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
-      if (isFirst) {
-        store.put(encodedKey, encodedValue, stateName)
-        isFirst = false
-      } else {
-        store.merge(encodedKey, encodedValue, stateName)
-      }
-    }
-    val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
-    upsertTTLForStateKey(ttlExpirationMs, serializedGroupingKey)
+    listStatePutImpl.put(newState)
+    upsertTTLForStateKey()
   }
 
   /** Append an entry to the list. */
   override def appendValue(newState: S): Unit = {
-    StateStoreErrors.requireNonNullStateValue(newState, stateName)
-    store.merge(stateTypesEncoder.encodeGroupingKey(),
-      stateTypesEncoder.encodeValue(newState, ttlExpirationMs), stateName)
-    val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
-    upsertTTLForStateKey(ttlExpirationMs, serializedGroupingKey)
+    listStatePutImpl.appendValue(newState)
+    upsertTTLForStateKey()
   }
 
   /** Append an entire list to the existing value. */
   override def appendList(newState: Array[S]): Unit = {
-    validateNewState(newState)
-
-    val encodedKey = stateTypesEncoder.encodeGroupingKey()
-    newState.foreach { v =>
-      val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
-      store.merge(encodedKey, encodedValue, stateName)
-    }
-    val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
-    upsertTTLForStateKey(ttlExpirationMs, serializedGroupingKey)
+    listStatePutImpl.appendList(newState)
+    upsertTTLForStateKey()
   }
 
   /** Remove this state. */
   override def clear(): Unit = {
-    store.remove(stateTypesEncoder.encodeGroupingKey(), stateName)
-  }
-
-  private def validateNewState(newState: Array[S]): Unit = {
-    StateStoreErrors.requireNonNullStateValue(newState, stateName)
-    StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
-
-    newState.foreach { v =>
-      StateStoreErrors.requireNonNullStateValue(v, stateName)
-    }
+    listStatePutImpl.clear()
   }
 
   /**
@@ -176,6 +146,11 @@ class ListStateImplWithTTL[S](
     }
   }
 
+  private def upsertTTLForStateKey(): Unit = {
+    val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+    upsertTTLForStateKey(ttlExpirationMs, serializedGroupingKey)
+  }
+
   /*
     * Internal methods to probe state for testing. The below methods exist for unit tests
     * to read the state ttl values, and ensure that values are persisted correctly in

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

@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution.streaming
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.execution.streaming.state.{StateStore, StateStoreErrors}
+import org.apache.spark.sql.streaming.ListStateModify
+
+/**
+ * Provides concrete implementation for list of values associated with a state variable
+ * used in the streaming transformWithState operator.
+ *
+ * @param store - reference to the StateStore instance to be used for storing state
+ * @param stateName - name of logical state partition
+ * @param keyExprEnc - Spark SQL encoder for key
+ * @param valEncoder - Spark SQL encoder for value
+ * @tparam S - data type of object that will be stored in the list
+ */
+class ListStateModifyImpl[S](
+    store: StateStore,
+    stateName: String,
+    keyExprEnc: ExpressionEncoder[Any],
+    valEncoder: Encoder[S],
+    encodeValue: S => UnsafeRow)
+  extends ListStateModify[S] with Logging {
+
+  private val keySerializer = keyExprEnc.createSerializer()
+
+  private val stateTypesEncoder = StateTypesEncoder(keySerializer, valEncoder, stateName)
+
+  /** Update the value of the list. */
+  override def put(newState: Array[S]): Unit = {
+    validateNewState(newState)
+
+    val encodedKey = stateTypesEncoder.encodeGroupingKey()
+    var isFirst = true
+
+    newState.foreach { v =>
+      val encodedValue = encodeValue(v)
+      if (isFirst) {
+        store.put(encodedKey, encodedValue, stateName)
+        isFirst = false
+      } else {
+        store.merge(encodedKey, encodedValue, stateName)
+      }
+    }
+  }
+
+  /** Append an entry to the list. */
+  override def appendValue(newState: S): Unit = {
+    StateStoreErrors.requireNonNullStateValue(newState, stateName)
+    store.merge(stateTypesEncoder.encodeGroupingKey(),
+      encodeValue(newState), stateName)
+  }
+
+  /** Append an entire list to the existing value. */
+  override def appendList(newState: Array[S]): Unit = {
+    validateNewState(newState)
+
+    val encodedKey = stateTypesEncoder.encodeGroupingKey()
+    newState.foreach { v =>
+      val encodedValue = encodeValue(v)
+      store.merge(encodedKey, encodedValue, stateName)
+    }
+  }
+
+  /** Remove this state. */
+  override def clear(): Unit = {
+    store.remove(stateTypesEncoder.encodeGroupingKey(), stateName)
+  }
+
+  private def validateNewState(newState: Array[S]): Unit = {
+    StateStoreErrors.requireNonNullStateValue(newState, stateName)
+    StateStoreErrors.requireNonEmptyListStateValue(newState, stateName)
+
+    newState.foreach { v =>
+      StateStoreErrors.requireNonNullStateValue(v, stateName)
+    }
+  }
+}

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

@@ -101,7 +101,7 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
     new ListStateTTLProcessor(ttlConfig)
   }
 
-  test("verify iterator works with expired values in middle of list") {
+  test("verify iterator works with expired values in middle of list after restarting stream") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
@@ -208,7 +208,7 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
   }
 
 
-  test("verify iterator works with expired values in beginning of list") {
+  test("verify iterator works with expired values in beginning of list after restarting stream") {
     withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
       classOf[RocksDBStateStoreProvider].getName,
       SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
@@ -297,4 +297,69 @@ class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
       }
     }
   }
+
+  test("verify expired values are evicted from list state") {
+    withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+      classOf[RocksDBStateStoreProvider].getName,
+      SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+
+      val inputStream = MemoryStream[InputEvent]
+      val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
+      val result = inputStream.toDS()
+        .groupByKey(x => x.key)
+        .transformWithState(
+          getProcessor(ttlConfig),
+          TimeoutMode.NoTimeouts(),
+          TTLMode.ProcessingTimeTTL(),
+          OutputMode.Append())
+      val clock = new StreamManualClock
+
+      testStream(result)(
+        StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock),
+        AddData(inputStream, InputEvent("k1", "put", 1)),
+        AdvanceManualClock(1 * 1000),
+        AddData(inputStream, InputEvent("k1", "append", 2)),
+        AddData(inputStream, InputEvent("k1", "append", 3)),
+        // advance clock to trigger processing
+        AdvanceManualClock(1 * 1000),
+        CheckNewAnswer(),
+        // get ttl values
+        AddData(inputStream, InputEvent("k1", "get_ttl_value_from_state", -1, null)),
+        AdvanceManualClock(1 * 1000),
+        CheckNewAnswer(
+          OutputEvent("k1", 1, isTTLValue = true, 61000),
+          OutputEvent("k1", 2, isTTLValue = true, 62000),
+          OutputEvent("k1", 3, isTTLValue = true, 62000)
+        ),
+        AddData(inputStream, InputEvent("k1", "get", -1, null)),
+        AdvanceManualClock(1 * 1000),
+        CheckNewAnswer(
+          OutputEvent("k1", 1, isTTLValue = false, -1),
+          OutputEvent("k1", 2, isTTLValue = false, -1),
+          OutputEvent("k1", 3, isTTLValue = false, -1)
+        ),
+        AdvanceManualClock(45 * 1000),
+        AddData(inputStream, InputEvent("k1", "append", 4)),
+        AdvanceManualClock(1 * 1000),
+        AddData(inputStream, InputEvent("k1", "get_ttl_value_from_state", -1, null)),
+        AdvanceManualClock(1 * 1000),
+        CheckNewAnswer(
+          OutputEvent("k1", 1, isTTLValue = true, 61000),
+          OutputEvent("k1", 2, isTTLValue = true, 62000),
+          OutputEvent("k1", 3, isTTLValue = true, 62000),
+          OutputEvent("k1", 4, isTTLValue = true, 110000)
+        ),
+        // advance clock so data expires
+        AdvanceManualClock(30 * 1000),
+        // run a no data batch
+        CheckNewAnswer(),
+        AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1)),
+        AdvanceManualClock(1 * 1000),
+        CheckNewAnswer(
+          OutputEvent("k1", 4, isTTLValue = false, -1)
+        ),
+        StopStream
+      )
+    }
+  }
 }