[WIP] ListStateTTL implementation

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

@@ -72,6 +72,26 @@ private[sql] trait StatefulProcessorHandle extends Serializable {
    */
   def getListState[T](stateName: String, valEncoder: Encoder[T]): ListState[T]
 
+  /**
+   * Function to create new or return existing list state variable of given type
+   * with ttl. State values will not be returned past ttlDuration, and will be eventually removed
+   * from the state store. Any state update resets the ttl to current processing time plus
+   * ttlDuration.
+   *
+   * The user must ensure to call this function only within the `init()` method of the
+   * StatefulProcessor.
+   *
+   * @param stateName  - name of the state variable
+   * @param valEncoder - SQL encoder for state variable
+   * @param ttlConfig  - the ttl configuration (time to live duration etc.)
+   * @tparam T - type of state variable
+   * @return - instance of ValueState of type T that can be used to store state persistently
+   */
+  def getListState[T](
+     stateName: String,
+     valEncoder: Encoder[T],
+     ttlConfig: TTLConfig): ListState[T]
+
   /**
    * Creates new or returns existing map state associated with stateName.
    * The MapState persists Key-Value pairs of type [K, V].

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

@@ -0,0 +1,293 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+
+/**
+ * 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 keyEnc - 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 ListStateImplWithTTL[S](
+  store: StateStore,
+  stateName: String,
+  keyExprEnc: ExpressionEncoder[Any],
+  valEncoder: Encoder[S],
+  ttlConfig: TTLConfig,
+  batchTimestampMs: Long)
+  extends SingleKeyTTLStateImpl(stateName, store, batchTimestampMs) with ListState[S] with Logging {
+
+  private val keySerializer = keyExprEnc.createSerializer()
+
+  private val stateTypesEncoder = StateTypesEncoder(
+    keySerializer, valEncoder, stateName, hasTtl = true)
+
+  private val ttlExpirationMs =
+    StateTTL.calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs)
+
+  initialize()
+
+  private def initialize(): Unit = {
+    store.createColFamilyIfAbsent(stateName, KEY_ROW_SCHEMA, VALUE_ROW_SCHEMA_WITH_TTL,
+      NoPrefixKeyStateEncoderSpec(KEY_ROW_SCHEMA), useMultipleValuesPerKey = true)
+  }
+  /** Whether state exists or not. */
+  override def exists(): Boolean = {
+    val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+    val stateValue = store.get(encodedGroupingKey, stateName)
+    stateValue != null
+  }
+
+  /**
+   * Get the state value if it exists. If the state does not exist in state store, an
+   * empty iterator is returned.
+   */
+  override def get(): Iterator[S] = {
+    val encodedKey = stateTypesEncoder.encodeGroupingKey()
+    val unsafeRowValuesIterator = store.valuesIterator(encodedKey, stateName)
+    logError(s"### size: ${store.valuesIterator(encodedKey, stateName).size}")
+    var currentRow: UnsafeRow = null
+
+    new Iterator[S] {
+      override def hasNext: Boolean = {
+        if (currentRow == null) {
+          setNextValidRow()
+        }
+
+        currentRow != null
+      }
+
+      override def next(): S = {
+        if (currentRow == null) {
+          setNextValidRow()
+        }
+        if (currentRow == null) {
+          throw new NoSuchElementException("Iterator is at the end")
+        }
+        val result = stateTypesEncoder.decodeValue(currentRow)
+        currentRow = null
+        result
+      }
+
+      // sets currentRow to a valid state, where we are
+      // pointing to a non-expired row
+      private def setNextValidRow(): Unit = {
+        logError(s"### setNextValidRow")
+        assert(currentRow == null)
+        if (unsafeRowValuesIterator.hasNext) {
+          logError(s"### set currentRow")
+          currentRow = unsafeRowValuesIterator.next()
+        } else {
+          currentRow = null
+          return
+        }
+        while (unsafeRowValuesIterator.hasNext && isExpired(currentRow)) {
+          logError(s"### isExpired, ${isExpired(currentRow)}")
+          currentRow = unsafeRowValuesIterator.next()
+        }
+
+        // in this case, we have iterated to the end, and there are no
+        // non-expired values
+        if (currentRow != null && isExpired(currentRow)) {
+          currentRow = null
+        }
+      }
+    }
+  }
+
+  /** 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)
+  }
+
+  /** 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)
+  }
+
+  /** 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)
+  }
+
+  /** 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)
+    }
+  }
+
+  /**
+   * Clears the user state associated with this grouping key
+   * if it has expired. This function is called by Spark to perform
+   * cleanup at the end of transformWithState processing.
+   *
+   * Spark uses a secondary index to determine if the user state for
+   * this grouping key has expired. However, its possible that the user
+   * has updated the TTL and secondary index is out of date. Implementations
+   * must validate that the user State has actually expired before cleanup based
+   * on their own State data.
+   *
+   * @param groupingKey grouping key for which cleanup should be performed.
+   */
+  override def clearIfExpired(groupingKey: Array[Byte]): Unit = {
+    logError(s"### clearIfExpired")
+    val encodedGroupingKey = stateTypesEncoder.encodeSerializedGroupingKey(groupingKey)
+    val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+    // We clear the list, and use the iterator to put back all of the non-expired values
+    store.remove(encodedGroupingKey, stateName)
+    var isFirst = true
+    unsafeRowValuesIterator.foreach { encodedValue =>
+      if (!isExpired(encodedValue)) {
+        if (isFirst) {
+          store.put(encodedGroupingKey, encodedValue, stateName)
+          isFirst = false
+        } else {
+          store.merge(encodedGroupingKey, encodedValue, stateName)
+        }
+      }
+    }
+  }
+
+  private def isExpired(valueRow: UnsafeRow): Boolean = {
+    logError(s"### isExpired, ${batchTimestampMs}")
+    val expirationMs = stateTypesEncoder.decodeTtlExpirationMs(valueRow)
+    expirationMs.exists(StateTTL.isExpired(_, batchTimestampMs))
+  }
+  /*
+    * 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
+    * the underlying state store.
+    */
+
+  /**
+   * Retrieves the value from State even if its expired. This method is used
+   * in tests to read the state store value, and ensure if its cleaned up at the
+   * end of the micro-batch.
+   */
+  private[sql] def getWithoutEnforcingTTL(): Iterator[Option[S]] = {
+    val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+    val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+    new Iterator[Option[S]] {
+      override def hasNext: Boolean = {
+        unsafeRowValuesIterator.hasNext
+      }
+
+      override def next(): Option[S] = {
+        val valueUnsafeRow = unsafeRowValuesIterator.next()
+        if (valueUnsafeRow != null) {
+          val resState = stateTypesEncoder.decodeValue(valueUnsafeRow)
+          Some(resState)
+        } else {
+          None
+        }
+      }
+    }
+  }
+
+  /**
+   * Read the ttl value associated with the grouping key.
+   */
+  private[sql] def getTTLValues(): Iterator[Option[Long]] = {
+    val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+    val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+    new Iterator[Option[Long]] {
+      override def hasNext: Boolean = {
+        unsafeRowValuesIterator.hasNext
+      }
+
+      override def next(): Option[Long] = {
+        val valueUnsafeRow = unsafeRowValuesIterator.next()
+        stateTypesEncoder.decodeTtlExpirationMs(valueUnsafeRow)
+      }
+    }
+  }
+  /**
+   * Get all ttl values stored in ttl state for current implicit
+   * grouping key.
+   */
+  private[sql] def getValuesInTTLState(): Iterator[Long] = {
+    val ttlIterator = ttlIndexIterator()
+    val implicitGroupingKey = stateTypesEncoder.serializeGroupingKey()
+    var nextValue: Option[Long] = None
+
+    new Iterator[Long] {
+      override def hasNext: Boolean = {
+        while (nextValue.isEmpty && ttlIterator.hasNext) {
+          val nextTtlValue = ttlIterator.next()
+          val groupingKey = nextTtlValue.groupingKey
+          if (groupingKey sameElements implicitGroupingKey) {
+            nextValue = Some(nextTtlValue.expirationMs)
+          }
+        }
+        nextValue.isDefined
+      }
+
+      override def next(): Long = {
+        val result = nextValue.get
+        nextValue = None
+        result
+      }
+    }
+  }
+}

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

@@ -248,4 +248,35 @@ class StatefulProcessorHandleImpl(
       throw StateStoreErrors.ttlMustBePositive("update", stateName)
     }
   }
+
+  /**
+   * Function to create new or return existing list state variable of given type
+   * with ttl. State values will not be returned past ttlDuration, and will be eventually removed
+   * from the state store. Any state update resets the ttl to current processing time plus
+   * ttlDuration.
+   *
+   * The user must ensure to call this function only within the `init()` method of the
+   * StatefulProcessor.
+   *
+   * @param stateName  - name of the state variable
+   * @param valEncoder - SQL encoder for state variable
+   * @param ttlConfig  - the ttl configuration (time to live duration etc.)
+   * @tparam T - type of state variable
+   * @return - instance of ValueState of type T that can be used to store state persistently
+   */
+  override def getListState[T](
+    stateName: String,
+    valEncoder: Encoder[T],
+    ttlConfig: TTLConfig): ListState[T] = {
+
+    verifyStateVarOperations("get_list_state")
+    validateTTLConfig(ttlConfig, stateName)
+
+    assert(batchTimestampMs.isDefined)
+    val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName,
+      keyEncoder, valEncoder, ttlConfig, batchTimestampMs.get)
+    ttlStates.add(listStateWithTTL)
+
+    listStateWithTTL
+  }
 }

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

@@ -18,6 +18,7 @@ package org.apache.spark.sql.execution.streaming
 
 import java.time.Duration
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
 import org.apache.spark.sql.execution.streaming.state.{RangeKeyScanStateEncoderSpec, StateStore}
@@ -81,7 +82,7 @@ abstract class SingleKeyTTLStateImpl(
     stateName: String,
     store: StateStore,
     ttlExpirationMs: Long)
-  extends TTLState {
+  extends TTLState with Logging {
 
   import org.apache.spark.sql.execution.streaming.StateTTLSchema._
 
@@ -98,6 +99,7 @@ abstract class SingleKeyTTLStateImpl(
   def upsertTTLForStateKey(
       expirationMs: Long,
       groupingKey: Array[Byte]): Unit = {
+    logError(s"### upsertTTLForStateKey: expirationMs=$expirationMs")
     val encodedTtlKey = ttlKeyEncoder(InternalRow(expirationMs, groupingKey))
     store.put(encodedTtlKey, EMPTY_ROW, ttlColumnFamilyName)
   }
@@ -138,7 +140,7 @@ abstract class SingleKeyTTLStateImpl(
 /**
  * Helper methods for user State TTL.
  */
-object StateTTL {
+object StateTTL extends Logging {
   def calculateExpirationTimeForDuration(
       ttlDuration: Duration,
       batchTtlExpirationMs: Long): Long = {
@@ -148,6 +150,8 @@ object StateTTL {
   def isExpired(
       expirationMs: Long,
       batchTtlExpirationMs: Long): Boolean = {
+    logError(s"### isExpired: expirationMs=$expirationMs," +
+      s" batchTtlExpirationMs=$batchTtlExpirationMs")
     batchTtlExpirationMs >= expirationMs
   }
 }