You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "ericm-db (via GitHub)" <gi...@apache.org> on 2024/04/08 18:05:19 UTC
[PR] [WIP] ListStateTTL implementation [spark]
ericm-db opened a new pull request, #45932:
URL: https://github.com/apache/spark/pull/45932
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566045231
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.streaming
+
+import java.sql.Timestamp
+import java.time.Duration
+
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+case class InputEvent(
+ key: String,
+ action: String,
+ value: Int,
+ eventTime: Timestamp = null)
+
+case class OutputEvent(
+ key: String,
+ value: Int,
+ isTTLValue: Boolean,
+ ttlValue: Long)
+/**
+ * Tests that ttl works as expected for Value State for
+ * processing time and event time based ttl.
+ */
+abstract class TransformWithStateTTLTest
+ extends StreamTest {
+ import testImplicits._
+
+ def getProcessor(ttlConfig: TTLConfig): StatefulProcessor[String, InputEvent, OutputEvent]
+
+ def getStateTTLMetricName: String
+
+ test("validate state is evicted at ttl expiry") {
Review Comment:
This is true.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559866863
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,230 @@
+/*
+ * 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.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}
+import org.apache.spark.sql.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private val keySerializer = keyExprEnc.createSerializer()
+
+ private val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ 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 = {
+ 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (unsafeRowValuesIterator.hasNext) {
+ val currentRow = unsafeRowValuesIterator.next()
+ if (!stateTypesEncoder.isExpired(currentRow, batchTimestampMs)) {
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ listStatePutImpl.put(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ listStatePutImpl.appendValue(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ listStatePutImpl.appendList(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Remove this state. */
+ override def clear(): Unit = {
+ listStatePutImpl.clear()
+ }
+
+ /**
+ *
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Unit = {
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ }
+ }
+ }
+
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ new Iterator[S] {
+ override def hasNext: Boolean = {
+ unsafeRowValuesIterator.hasNext
+ }
+
+ override def next(): S = {
+ val valueUnsafeRow = unsafeRowValuesIterator.next()
+ if (valueUnsafeRow != null) {
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ } else {
+ null.asInstanceOf[S]
+ }
+ }
+ }
+ }
+
+ /**
+ * Read the ttl value associated with the grouping key.
+ */
+ private[sql] def getTTLValues(): Iterator[(S, Long)] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ new Iterator[(S, Long)] {
+ override def hasNext: Boolean = {
Review Comment:
Same here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1558571847
##########
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 values in listState which have expired after ttlDuration will not
+ * returned on get() and will be eventually removed from the state.
Review Comment:
nit: `will not be returned`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "sahnib (via GitHub)" <gi...@apache.org>.
sahnib commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1562685312
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,235 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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)
+ }
+ }
+ 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)
+ 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)
+ }
+ 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)
+ }
+ }
+
+ /**
+ *
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Unit = {
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ }
+ }
+ }
+
+ private def upsertTTLForStateKey(): Unit = {
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
Review Comment:
Can we avoid this serialization happening 2 times? We will always serialize the key to store it in the state variable column family, and then re-serialize it here to store it in ttlState. We can do this only once and that should be good enough.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,235 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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](
Review Comment:
@anishshri-db added some metrics for `TransformWithState` as part of this PR - https://github.com/apache/spark/pull/45937. I think we need to calculate expired values in this class to keep the metrics up to date.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala:
##########
@@ -237,6 +237,24 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
}
}
+ test(s"ttl States are populated for listState and ttlMode=ProcessingTime") {
Review Comment:
Can you also add listState instantiation to the test `ttl States are not populated for ttlMode=NoTTL` below, and ensure creating a ListState does not add anything to ttlStates when ttlMode=NoTTL.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala:
##########
@@ -141,14 +135,15 @@ class ValueStateImplWithTTL[S](
/**
* Read the ttl value associated with the grouping key.
*/
- private[sql] def getTTLValue(): Option[Long] = {
+ private[sql] def getTTLValue(): Option[(S, Long)] = {
val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
val retRow = store.get(encodedGroupingKey, stateName)
- if (retRow != null) {
- stateTypesEncoder.decodeTtlExpirationMs(retRow)
- } else {
- None
+ // if the returned row is not null, we want to return the value associated with the
+ // ttlExpiration
+ Option(retRow).flatMap { row =>
+ val ttlExpiration = stateTypesEncoder.decodeTtlExpirationMs(row)
+ ttlExpiration.map(expiration => (stateTypesEncoder.decodeValue(row), expiration))
Review Comment:
We can simply do
```
if (retRow != null) {
(stateTypesEncoder.decodeValue(row), stateTypesEncoder.decodeTtlExpirationMs(retRow))
} else {
None
}
```
or
```
Option(retRow).flatMap { row =>
val ttlExpiration = stateTypesEncoder.decodeTtlExpirationMs(row)
(stateTypesEncoder.decodeValue(row), expiration), ttlExpiration)
}
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala:
##########
@@ -333,7 +333,8 @@ class ValueStateSuite extends StateVariableSuiteBase {
val ttlExpirationMs = timestampMs + 60000
var ttlValue = testState.getTTLValue()
assert(ttlValue.isDefined)
- assert(ttlValue.get === ttlExpirationMs)
+ val ttlExpiration1 = ttlValue.get._2
Review Comment:
[nit] Can be renamed to `ttlExpirationMsValueInState`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,235 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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)
+ }
+ }
+ 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)
+ 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)
+ }
+ 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)
+ }
+ }
+
+ /**
+ *
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Unit = {
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ }
+ }
+ }
+
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
+ valueUnsafeRow =>
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ }
+ }
+
+ /**
+ * Read the ttl value associated with the grouping key.
+ */
+ private[sql] def getTTLValues(): Iterator[(S, Long)] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
+ valueUnsafeRow =>
+ (stateTypesEncoder.decodeValue(valueUnsafeRow),
+ stateTypesEncoder.decodeTtlExpirationMs(valueUnsafeRow).get)
+ }
+ }
+ /**
+ * Get all ttl values stored in ttl state for current implicit
+ * grouping key.
+ */
+ private[sql] def getValuesInTTLState(): Iterator[Long] = {
Review Comment:
As this function is exactly the same for `ValueStateImplWithTTL` and `ListStateImplWithTTL`, maybe we should move it to the common base class - `SingleKeyTTLStateImpl`.
Now that I think about it more, its more apt for this method to be in class `SingleKeyTTLStateImpl`.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala:
##########
@@ -27,22 +26,10 @@ import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.util.StreamManualClock
-case class InputEvent(
- key: String,
- action: String,
- value: Int,
- eventTime: Timestamp = null)
-
-case class OutputEvent(
- key: String,
- value: Int,
- isTTLValue: Boolean,
- ttlValue: Long)
-
object TTLInputProcessFunction {
def processRow(
- row: InputEvent,
- valueState: ValueStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+ row: InputEvent,
+ valueState: ValueStateImplWithTTL[Int]): Iterator[OutputEvent] = {
Review Comment:
Thanks for fixing the indent here.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeoutMode: TimeoutMode,
+ ttlMode: TTLMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 7)),
+ AddData(inputStream, InputEvent("k1", "append", 8)),
+ AddData(inputStream, InputEvent("k1", "append", 9)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // advance clock to expire the middle three elements
+ AdvanceManualClock(45 * 1000),
+ // Get all elements in the list
+ AddData(inputStream, InputEvent("k1", "get", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ // validate that the expired elements are not returned
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 7, isTTLValue = false, -1),
+ OutputEvent("k1", 8, isTTLValue = false, -1),
+ OutputEvent("k1", 9, isTTLValue = false, -1)
+ ),
+ StopStream
Review Comment:
Can we also get these elements withoutEnforcingTTL and ensure they are actually deleted from the state Store. Similarly, lets also verify the timer value for middle 3 elements is also deleted.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala:
##########
@@ -399,7 +205,7 @@ class TransformWithValueStateTTLSuite
AddData(inputStream, InputEvent(ttlKey, "get_ttl_value_from_state", -1)),
AddData(inputStream, InputEvent(noTtlKey, "get_ttl_value_from_state", -1)),
AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent(ttlKey, -1, isTTLValue = true, 61000)),
+ CheckNewAnswer(OutputEvent(ttlKey, 1, isTTLValue = true, 61000)),
Review Comment:
Was this testcase incorrect? Why do we need to change this value to 1 instead?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeoutMode: TimeoutMode,
+ ttlMode: TTLMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 7)),
+ AddData(inputStream, InputEvent("k1", "append", 8)),
+ AddData(inputStream, InputEvent("k1", "append", 9)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // advance clock to expire the middle three elements
+ AdvanceManualClock(45 * 1000),
+ // Get all elements in the list
+ AddData(inputStream, InputEvent("k1", "get", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ // validate that the expired elements are not returned
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 7, isTTLValue = false, -1),
+ OutputEvent("k1", 8, isTTLValue = false, -1),
+ OutputEvent("k1", 9, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ }
+ }
+ }
+
+ test("verify iterator works with expired values in beginning of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val inputStream = MemoryStream[InputEvent]
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(1))
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofMinutes(2))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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)
+ ),
+ StopStream
+ )
+
+ // add 3 elements with a duration of two minutes
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // 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),
+ OutputEvent("k1", 4, isTTLValue = true, 125000),
+ OutputEvent("k1", 5, isTTLValue = true, 125000),
+ OutputEvent("k1", 6, isTTLValue = true, 125000)
+ ),
+ // expire beginning values
+ AdvanceManualClock(60 * 1000),
+ AddData(inputStream, InputEvent("k1", "get", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
Review Comment:
Same as above. Can we also get these elements withoutEnforcingTTL and ensure they are actually deleted from the state Store. Similarly, lets also verify the timer value for middle 3 elements is also deleted.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566090091
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
I'm a little confused by this comment - I think using dropWhile should be okay since we are only returning one row at a time in the iterator, and `getNext` is called after every invocation of `next()`
Consider the following case (e for expired, v for valid):
[v1, v2, v3, e1, e2, e3, v4, v5, v6]
Calling `dropWhile` will drop the first elements (of which there are none in this example) and point the iterator to v1, and we set the currentRow to v1. After we return v1, we set the currentRow to v2, then v3 (as `dropWhile` doesn't drop these elements)
Now, we have returned v3, and `unsafeRowValuesIterator` is pointing at e4. We would call dropWhile on this, causing the following three elements to be dropped, and the iterator is consumed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #45932:
URL: https://github.com/apache/spark/pull/45932#issuecomment-2058102492
Thanks! Merging to master.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on PR #45932:
URL: https://github.com/apache/spark/pull/45932#issuecomment-2053489012
@HeartSaVioR PTAL, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1563063876
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,349 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+ override def getProcessor(ttlConfig: TTLConfig):
Review Comment:
nit: newline and indent seems off on the next line ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1565212846
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
Review Comment:
It doesn't look like, and we look to have an issue with this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566090091
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
I'm a little confused by this comment - I think using dropWhile should be okay since we are only returning one row at a time in the iterator.
Consider the following case (e for expired, v for valid):
[e1, e2, e3, v1, v2, v3, e4, e5, e6]
Calling `dropWhile` will drop the first three elements and point the iterator to v1, and we set the currentRow to v1. After we return v1, we set the currentRow to v2, then v3 (as `dropWhile` doesn't drop these elements)
Now, we have returned v3, and `unsafeRowValuesIterator` is pointing at e4. We would call dropWhile on this, causing the rest of the elements to be dropped.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559856458
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,230 @@
+/*
+ * 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.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}
+import org.apache.spark.sql.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private val keySerializer = keyExprEnc.createSerializer()
Review Comment:
Could we make any of these lazy vals ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on PR #45932:
URL: https://github.com/apache/spark/pull/45932#issuecomment-2046265049
Should we update the PR description ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559870514
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,365 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent]
+ with Logging {
Review Comment:
nit: indent ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566193767
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.streaming
+
+import java.sql.Timestamp
+import java.time.Duration
+
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+case class InputEvent(
+ key: String,
+ action: String,
+ value: Int,
+ eventTime: Timestamp = null)
+
+case class OutputEvent(
+ key: String,
+ value: Int,
+ isTTLValue: Boolean,
+ ttlValue: Long)
+/**
+ * Tests that ttl works as expected for Value State for
Review Comment:
Updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1562819807
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala:
##########
@@ -399,7 +205,7 @@ class TransformWithValueStateTTLSuite
AddData(inputStream, InputEvent(ttlKey, "get_ttl_value_from_state", -1)),
AddData(inputStream, InputEvent(noTtlKey, "get_ttl_value_from_state", -1)),
AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent(ttlKey, -1, isTTLValue = true, 61000)),
+ CheckNewAnswer(OutputEvent(ttlKey, 1, isTTLValue = true, 61000)),
Review Comment:
It wasn't incorrect before - I had changed the getTTLValueFromState function to return the associated value with the TTL for ListState, and had to change it for ValueState so that the APIs would match.
-1 was the placeholder value for when `isTTLValue=true`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566090091
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
I'm a little confused by this comment - I think using dropWhile should be okay since we are only returning one row at a time in the iterator, and `getNext` is called after every invocation of `next()`
Consider the following case (e for expired, v for valid):
[v1, v2, v3, e1, e2, e3, v4, v5, v6]
Calling `dropWhile` will drop the first elements (of which there are none in this example) and point the iterator to v1, and we set the currentRow to v1. After we return v1, we set the currentRow to v2, then v3 (as `dropWhile` doesn't drop these elements)
Now, we have returned v3, and `unsafeRowValuesIterator` is pointing at e4. We would call dropWhile on this, causing the following three elements to be dropped, and then we would point at v4.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566090091
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
I'm a little confused by this comment - I think using dropWhile should be okay since we are only returning one row at a time in the iterator, and `getNext` is called after every invocation of `next()`
Consider the following case (e for expired, v for valid):
[e1, e2, e3, v1, v2, v3, e4, e5, e6]
Calling `dropWhile` will drop the first three elements and point the iterator to v1, and we set the currentRow to v1. After we return v1, we set the currentRow to v2, then v3 (as `dropWhile` doesn't drop these elements)
Now, we have returned v3, and `unsafeRowValuesIterator` is pointing at e4. We would call dropWhile on this, causing the rest of the elements to be dropped.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559866658
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,230 @@
+/*
+ * 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.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}
+import org.apache.spark.sql.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private val keySerializer = keyExprEnc.createSerializer()
+
+ private val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ 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 = {
+ 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (unsafeRowValuesIterator.hasNext) {
+ val currentRow = unsafeRowValuesIterator.next()
+ if (!stateTypesEncoder.isExpired(currentRow, batchTimestampMs)) {
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ listStatePutImpl.put(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ listStatePutImpl.appendValue(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ listStatePutImpl.appendList(newState)
+ upsertTTLForStateKey()
+ }
+
+ /** Remove this state. */
+ override def clear(): Unit = {
+ listStatePutImpl.clear()
+ }
+
+ /**
+ *
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Unit = {
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ }
+ }
+ }
+
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ new Iterator[S] {
+ override def hasNext: Boolean = {
Review Comment:
Could we skip this ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559872175
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,365 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent]
+ with Logging {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeoutMode: TimeoutMode,
+ ttlMode: TTLMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ 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") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeoutMode.NoTimeouts(),
+ TTLMode.ProcessingTimeTTL(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 7)),
+ AddData(inputStream, InputEvent("k1", "append", 8)),
+ AddData(inputStream, InputEvent("k1", "append", 9)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // advance clock to expire the middle three elements
+ AdvanceManualClock(45 * 1000),
+ // Get all elements in the list
+ AddData(inputStream, InputEvent("k1", "get", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ // validate that the expired elements are not returned
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 7, isTTLValue = false, -1),
+ OutputEvent("k1", 8, isTTLValue = false, -1),
+ OutputEvent("k1", 9, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ }
+ }
+ }
+
+
Review Comment:
nit: extra newline ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #45932: [SPARK-47673][SS] Implementing TTL for ListState
URL: https://github.com/apache/spark/pull/45932
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566193244
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
Review Comment:
Yes, essentially.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala:
##########
@@ -160,4 +161,60 @@ class ListStateSuite extends StateVariableSuiteBase {
assert(listState1.get().toSeq === Seq.empty[Long])
}
}
+
+ test(s"test List state TTL") {
+ tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
+ val store = provider.getStore(0)
+ val timestampMs = 10
+ val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
+ Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]],
+ TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs))
+
+ val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
+ val testState: ListStateImplWithTTL[String] = handle.getListState[String]("testState",
+ Encoders.STRING, ttlConfig).asInstanceOf[ListStateImplWithTTL[String]]
+ ImplicitGroupingKeyTracker.setImplicitKey("test_key")
+ testState.put(Array("v1", "v2", "v3"))
+ assert(testState.get().toSeq === Seq("v1", "v2", "v3"))
+ assert(testState.getWithoutEnforcingTTL().toSeq === Seq("v1", "v2", "v3"))
+
+ val ttlExpirationMs = timestampMs + 60000
+ var ttlValues = testState.getTTLValues()
+ assert(ttlValues.nonEmpty)
+ assert(ttlValues.forall(_._2 === ttlExpirationMs))
+ var ttlStateValueIterator = testState.getValuesInTTLState()
+ assert(ttlStateValueIterator.hasNext)
+
+ // increment batchProcessingTime, or watermark and ensure expired value is not returned
+ val nextBatchHandle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
+ Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]],
+ TimeMode.ProcessingTime(), batchTimestampMs = Some(ttlExpirationMs))
+
+ val nextBatchTestState: ListStateImplWithTTL[String] =
+ nextBatchHandle.getListState[String]("testState", Encoders.STRING, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[String]]
+
+ ImplicitGroupingKeyTracker.setImplicitKey("test_key")
+
+ // ensure get does not return the expired value
+ assert(!nextBatchTestState.exists())
+ assert(nextBatchTestState.get().isEmpty)
+
+ // ttl value should still exist in state
+ ttlValues = nextBatchTestState.getTTLValues()
+ assert(ttlValues.nonEmpty)
+ assert(ttlValues.forall(_._2 === ttlExpirationMs))
+ ttlStateValueIterator = nextBatchTestState.getValuesInTTLState()
+ assert(ttlStateValueIterator.hasNext)
+ assert(ttlStateValueIterator.next() === ttlExpirationMs)
+ assert(ttlStateValueIterator.isEmpty)
+
+ // getWithoutTTL should still return the expired value
+ assert(nextBatchTestState.getWithoutEnforcingTTL().toSeq === Seq("v1", "v2", "v3"))
+
+ nextBatchTestState.clear()
+ assert(!nextBatchTestState.exists())
+ assert(nextBatchTestState.get().isEmpty)
+ }
+ }
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566191543
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
Review Comment:
Added a StreamEvent that will get the TTL values
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566192350
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
Review Comment:
Explained in a comment before the test case
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559855781
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,230 @@
+/*
+ * 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.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}
+import org.apache.spark.sql.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private val keySerializer = keyExprEnc.createSerializer()
+
+ private val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ 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 = {
+ 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ unsafeRowValuesIterator.dropWhile { row =>
Review Comment:
Should we reassign the value ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1558570672
##########
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
}
+
+/**
Review Comment:
I guess we don't need this ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "sahnib (via GitHub)" <gi...@apache.org>.
sahnib commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1557858285
##########
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
Review Comment:
scaladoc seems to point to ValueState.
`instance of ListState of type T that can be used to store state persistently`
##########
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
Review Comment:
[nit] Might be better to say -
`Any values in listState which have expired after ttlDuration will not returned on get() and will be eventually removed from the state.`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
Review Comment:
Can potentially be simplified to
```
unsafeRowValuesIterator.dropWhile { row =>
isExpired(row)
}
val next = unsafeRowValuesIterator.next()
if (next != null) {
stateTypesEncoder.decodeValue(currentRow)
} else {
finished = true
null.asInstanceOf[S]
}
```
##########
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
Review Comment:
scala doc seems to refer to ValueState.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala:
##########
@@ -81,7 +82,7 @@ abstract class SingleKeyTTLStateImpl(
stateName: String,
store: StateStore,
ttlExpirationMs: Long)
- extends TTLState {
+ extends TTLState with Logging {
Review Comment:
Is this needed?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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 {
Review Comment:
Logging does not seem to be used.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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](
Review Comment:
I wonder if there is potential to extract common logic for ListState into a base class. For ValueState, we did not find much common code to extract out, but ListState seems to have quite a lot of common code.
Overall, I would recommend consider a private function to `encodeValue`, which can add `ttlExpirationMs` for `ListStateImplWithTTL`, and reuse put/append/update logic across both ListStateImpl and ListStateImplWithTTL.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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 = {
+ 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 = {
+ val expirationMs = stateTypesEncoder.decodeTtlExpirationMs(valueRow)
+ expirationMs.exists(StateTTL.isExpired(_, batchTimestampMs))
+ }
Review Comment:
This method seems common between ValueState/ListState. Might be useful to move it to a common location.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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 = {
+ 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)
Review Comment:
To boost performance, can we potentially only remove if something is actually expired?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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.
Review Comment:
Might be useful to add specific documentation for ListState (if something is specific) rather than using same doc as parent trait.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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 = {
+ 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 = {
+ 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ new Iterator[S] {
+ override def hasNext: Boolean = {
+ unsafeRowValuesIterator.hasNext
+ }
+
+ override def next(): S = {
+ val valueUnsafeRow = unsafeRowValuesIterator.next()
+ if (valueUnsafeRow != null) {
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ } else {
+ null.asInstanceOf[S]
+ }
+ }
+ }
+ }
+
+ /**
+ * 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]] {
Review Comment:
Can this be a `Iterator[Long]`? we can just omit emitting out None values.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,267 @@
+/*
+ * 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}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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)
+
+ new NextIterator[S] {
+ private var currentRow: UnsafeRow = _
+
+ override protected def getNext(): S = {
+ if (unsafeRowValuesIterator.hasNext) {
+ currentRow = unsafeRowValuesIterator.next()
+ while (unsafeRowValuesIterator.hasNext && 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
+ }
+ } else {
+ finished = true
+ }
+ if (currentRow != null) stateTypesEncoder.decodeValue(currentRow) else null.asInstanceOf[S]
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** 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 = {
+ 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 = {
+ 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ new Iterator[S] {
+ override def hasNext: Boolean = {
+ unsafeRowValuesIterator.hasNext
+ }
+
+ override def next(): S = {
+ val valueUnsafeRow = unsafeRowValuesIterator.next()
+ if (valueUnsafeRow != null) {
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ } else {
+ null.asInstanceOf[S]
+ }
+ }
+ }
+ }
+
+ /**
+ * 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]] {
Review Comment:
Also, would it be useful to also emit actual value with expirationMs to ensure ttl is associated correctly with values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559873586
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala:
##########
@@ -356,7 +357,8 @@ class ValueStateSuite extends StateVariableSuiteBase {
// ttl value should still exist in state
ttlValue = nextBatchTestState.getTTLValue()
assert(ttlValue.isDefined)
- assert(ttlValue.get === ttlExpirationMs)
+ val ttlExpiration2 = ttlValue.get._2
+ assert(ttlExpiration2 === ttlExpirationMs)
Review Comment:
Can we add some more unit tests in ValueStateSuite/ListStateSuite/StatefulProcessorHandleSuite ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559874230
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,365 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent]
+ with Logging {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeoutMode: TimeoutMode,
+ ttlMode: TTLMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
Review Comment:
nit: newline here and add some class level comments ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1562801897
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,235 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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](
Review Comment:
yup +1 - we need to add the listStateWithTTLVars metric too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566364666
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
Review Comment:
This turned out to by my confusion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1559857448
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,230 @@
+/*
+ * 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.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}
+import org.apache.spark.sql.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private val keySerializer = keyExprEnc.createSerializer()
+
+ private val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ 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 = {
+ store.createColFamilyIfAbsent(stateName, KEY_ROW_SCHEMA, VALUE_ROW_SCHEMA_WITH_TTL,
+ NoPrefixKeyStateEncoderSpec(KEY_ROW_SCHEMA), useMultipleValuesPerKey = true)
+ }
+ /** Whether state exists or not. */
Review Comment:
nit: add new line here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566366870
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
No worries
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1563050397
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,220 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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
Review Comment:
Can we update the comment to match the args ?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TTLState.scala:
##########
@@ -75,7 +75,7 @@ trait TTLState {
*
* @return true if the state was cleared, false otherwise.
Review Comment:
Could we update this too ?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala:
##########
@@ -237,13 +237,33 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
}
}
+ test(s"ttl States are populated for listState and timeMode=ProcessingTime") {
+ tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
+ val store = provider.getStore(0)
+ val handle = new StatefulProcessorHandleImpl(store,
+ UUID.randomUUID(), keyExprEncoder, TimeMode.ProcessingTime(),
+ batchTimestampMs = Some(10))
+
+ val listStateWithTTL = handle.getListState("testState",
+ Encoders.STRING, TTLConfig(Duration.ofHours(1)))
+
+ // create another state without TTL, this should not be captured in the handle
+ handle.getListState("testState", Encoders.STRING)
+
+ assert(handle.ttlStates.size() === 1)
+ assert(handle.ttlStates.get(0) === listStateWithTTL)
+ }
+ }
+
+
Review Comment:
nit: extra newline ?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,220 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * 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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ StateStoreErrors.requireNonNullStateValue(newState, stateName)
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ store.merge(stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey),
+ stateTypesEncoder.encodeValue(newState, ttlExpirationMs), stateName)
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ newState.foreach { v =>
+ val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
+ store.merge(encodedKey, encodedValue, stateName)
+ }
+ upsertTTLForStateKey(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)
+ }
+ }
+
+ /**
+ *
Review Comment:
nit: extra empty line ?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala:
##########
@@ -237,13 +237,33 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
}
}
+ test(s"ttl States are populated for listState and timeMode=ProcessingTime") {
+ tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
+ val store = provider.getStore(0)
+ val handle = new StatefulProcessorHandleImpl(store,
+ UUID.randomUUID(), keyExprEncoder, TimeMode.ProcessingTime(),
+ batchTimestampMs = Some(10))
+
+ val listStateWithTTL = handle.getListState("testState",
+ Encoders.STRING, TTLConfig(Duration.ofHours(1)))
+
+ // create another state without TTL, this should not be captured in the handle
+ handle.getListState("testState", Encoders.STRING)
+
+ assert(handle.ttlStates.size() === 1)
+ assert(handle.ttlStates.get(0) === listStateWithTTL)
+ }
+ }
+
+
test(s"ttl States are not populated for timeMode=None") {
tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
val store = provider.getStore(0)
val handle = new StatefulProcessorHandleImpl(store,
UUID.randomUUID(), keyExprEncoder, TimeMode.None())
handle.getValueState("testState", Encoders.STRING)
+ handle.getListState("testState", Encoders.STRING)
Review Comment:
Could we use a different name for the state variable ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1563064461
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithValueStateTTLSuite.scala:
##########
@@ -171,203 +160,15 @@ case class MultipleValueStatesTTLProcessor(
}
}
-/**
- * Tests that ttl works as expected for Value State for
- * processing time and event time based ttl.
- */
-class TransformWithValueStateTTLSuite
- extends StreamTest {
- import testImplicits._
-
- test("validate state is evicted at ttl expiry") {
- withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
- classOf[RocksDBStateStoreProvider].getName) {
- withTempDir { dir =>
- val inputStream = MemoryStream[InputEvent]
- val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
- val result = inputStream.toDS()
- .groupByKey(x => x.key)
- .transformWithState(
- new ValueStateTTLProcessor(ttlConfig),
- TimeMode.ProcessingTime(),
- OutputMode.Append())
-
- val clock = new StreamManualClock
- testStream(result)(
- StartStream(
- Trigger.ProcessingTime("1 second"),
- triggerClock = clock,
- checkpointLocation = dir.getAbsolutePath),
- AddData(inputStream, InputEvent("k1", "put", 1)),
- // advance clock to trigger processing
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(),
- StopStream,
- StartStream(
- Trigger.ProcessingTime("1 second"),
- triggerClock = clock,
- checkpointLocation = dir.getAbsolutePath),
- // get this state, and make sure we get unexpired value
- AddData(inputStream, InputEvent("k1", "get", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", 1, isTTLValue = false, -1)),
- StopStream,
- StartStream(
- Trigger.ProcessingTime("1 second"),
- triggerClock = clock,
- checkpointLocation = dir.getAbsolutePath),
- // ensure ttl values were added correctly
- AddData(inputStream, InputEvent("k1", "get_ttl_value_from_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 61000)),
- AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 61000)),
- StopStream,
- StartStream(
- Trigger.ProcessingTime("1 second"),
- triggerClock = clock,
- checkpointLocation = dir.getAbsolutePath),
- // advance clock so that state expires
- AdvanceManualClock(60 * 1000),
- AddData(inputStream, InputEvent("k1", "get", -1, null)),
- AdvanceManualClock(1 * 1000),
- // validate expired value is not returned
- CheckNewAnswer(),
- // ensure this state does not exist any longer in State
- AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(),
- AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(),
- Execute { q =>
- // Filter for idle progress events and then verify the custom metrics
- // for stateful operator
- val progData = q.recentProgress.filter(prog => prog.stateOperators.size > 0)
- assert(progData.filter(prog =>
- prog.stateOperators(0).customMetrics.get("numValueStateWithTTLVars") > 0).size > 0)
- assert(progData.filter(prog =>
- prog.stateOperators(0).customMetrics
- .get("numValuesRemovedDueToTTLExpiry") > 0).size > 0)
- }
- )
- }
- }
- }
-
- test("validate state update updates the expiration timestamp") {
- withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
- classOf[RocksDBStateStoreProvider].getName) {
- val inputStream = MemoryStream[InputEvent]
- val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
- val result = inputStream.toDS()
- .groupByKey(x => x.key)
- .transformWithState(
- new ValueStateTTLProcessor(ttlConfig),
- TimeMode.ProcessingTime(),
- OutputMode.Append())
+class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
- val clock = new StreamManualClock
- testStream(result)(
- StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock),
- AddData(inputStream, InputEvent("k1", "put", 1)),
- // advance clock to trigger processing
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(),
- // get this state, and make sure we get unexpired value
- AddData(inputStream, InputEvent("k1", "get", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", 1, isTTLValue = false, -1)),
- // ensure ttl values were added correctly
- AddData(inputStream, InputEvent("k1", "get_ttl_value_from_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 61000)),
- AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 61000)),
- // advance clock and update expiration time
- AdvanceManualClock(30 * 1000),
- AddData(inputStream, InputEvent("k1", "put", 1)),
- AddData(inputStream, InputEvent("k1", "get", -1)),
- // advance clock to trigger processing
- AdvanceManualClock(1 * 1000),
- // validate value is not expired
- CheckNewAnswer(OutputEvent("k1", 1, isTTLValue = false, -1)),
- // validate ttl value is updated in the state
- AddData(inputStream, InputEvent("k1", "get_ttl_value_from_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 95000)),
- // validate ttl state has both ttl values present
- AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 61000),
- OutputEvent("k1", -1, isTTLValue = true, 95000)
- ),
- // advance clock after older expiration value
- AdvanceManualClock(30 * 1000),
- // ensure unexpired value is still present in the state
- AddData(inputStream, InputEvent("k1", "get", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", 1, isTTLValue = false, -1)),
- // validate that the older expiration value is removed from ttl state
- AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1)),
- AdvanceManualClock(1 * 1000),
- CheckNewAnswer(OutputEvent("k1", -1, isTTLValue = true, 95000))
- )
- }
+ import testImplicits._
+ override def getProcessor(ttlConfig: TTLConfig):
Review Comment:
Same here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1565148322
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+class ListStateImplWithTTL[S](
+ store: StateStore,
Review Comment:
nit: 2 more spaces (4 spaces)
##########
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 values in listState which have expired after ttlDuration will not
+ * be returned on get() and will be eventually removed from the state.
+ *
+ * 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 ListState of type T that can be used to store state persistently
+ */
+ def getListState[T](
+ stateName: String,
Review Comment:
nit: 1 more space
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
Review Comment:
Just to be 100% sure, we don't have a case where elements are stored in the direction which is out of sync with its appearance order. Do I understand correctly?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
Review Comment:
This sounds to me we support both element level TTL and TTL of the list itself. Do I understand correctly? If then looks OK.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala:
##########
@@ -245,6 +245,38 @@ class StatefulProcessorHandleImpl(
resultState
}
+ /**
+ * 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 values in listState which have expired after ttlDuration will not
+ * returned on get() and will be eventually removed from the state.
+ *
+ * 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 ListState of type T that can be used to store state persistently
+ */
+ override def getListState[T](
+ stateName: String,
Review Comment:
nit: 2 more spaces (4 spaces)
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ StateStoreErrors.requireNonNullStateValue(newState, stateName)
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ store.merge(stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey),
+ stateTypesEncoder.encodeValue(newState, ttlExpirationMs), stateName)
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ newState.foreach { v =>
+ val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
+ store.merge(encodedKey, encodedValue, stateName)
+ }
+ upsertTTLForStateKey(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)
+ }
+ }
+
+ /**
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Long = {
+ var numValuesExpired = 0L
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ } else {
+ numValuesExpired += 1
+ }
+ }
+ numValuesExpired
+ }
+
+ private def upsertTTLForStateKey(serializedGroupingKey: Array[Byte]): Unit = {
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
+ valueUnsafeRow =>
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ }
+ }
+
+ /**
+ * Read the ttl value associated with the grouping key.
+ */
+ private[sql] def getTTLValues(): Iterator[(S, Long)] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
Review Comment:
nit: same style
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala:
##########
@@ -237,13 +237,32 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
}
}
+ test(s"ttl States are populated for listState and timeMode=ProcessingTime") {
Review Comment:
nit: same
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
Review Comment:
The test is hard to understand as we change TTL config back and forth and clock is reused. Could you please add more informative comment about expected batch timestamp for the batch?
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
Wait, if we allow the change of TTL during query restart, it is possible that earlier elements in the list has later expiration time.
How this will work with `ListStateImplWithTTL.get()`? dropWhile would not work.
E.g. some elements in the middle weren't expired at batch N but going to be expired at batch N + 1, and earlier elements are somehow having later expiration time and do not expire at batch N + 1. Calling get() in batch N + 1 will produce all elements including the elements which are going to be expired in batch N + 1, because earlier elements will terminate dropWhile.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 7)),
+ AddData(inputStream, InputEvent("k1", "append", 8)),
+ AddData(inputStream, InputEvent("k1", "append", 9)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // advance clock to expire the middle three elements
+ AddData(inputStream, InputEvent("k1", "get_values_in_ttl_state", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ // validate that the expired elements are not returned
Review Comment:
This is not checking the expired elements, right? Comment confuses me.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
Review Comment:
nit: indentation is off (3 more spaces)
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ StateStoreErrors.requireNonNullStateValue(newState, stateName)
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ store.merge(stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey),
+ stateTypesEncoder.encodeValue(newState, ttlExpirationMs), stateName)
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ newState.foreach { v =>
+ val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
+ store.merge(encodedKey, encodedValue, stateName)
+ }
+ upsertTTLForStateKey(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)
+ }
+ }
+
+ /**
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Long = {
+ var numValuesExpired = 0L
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ } else {
+ numValuesExpired += 1
+ }
+ }
+ numValuesExpired
+ }
+
+ private def upsertTTLForStateKey(serializedGroupingKey: Array[Byte]): Unit = {
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
+ valueUnsafeRow =>
+ stateTypesEncoder.decodeValue(valueUnsafeRow)
+ }
+ }
+
+ /**
+ * Read the ttl value associated with the grouping key.
+ */
+ private[sql] def getTTLValues(): Iterator[(S, Long)] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
+ valueUnsafeRow =>
+ (stateTypesEncoder.decodeValue(valueUnsafeRow),
+ stateTypesEncoder.decodeTtlExpirationMs(valueUnsafeRow).get)
+ }
+ }
+ /**
Review Comment:
nit: empty line between two
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImplWithTTL.scala:
##########
@@ -0,0 +1,221 @@
+/*
+ * 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.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.streaming.{ListState, TTLConfig}
+import org.apache.spark.util.NextIterator
+
+/**
+ * Class that provides a concrete implementation for a list state state associated with state
+ * variables (with ttl expiration support) 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
+ * @param ttlConfig - TTL configuration for values stored in this state
+ * @param batchTimestampMs - current batch processing timestamp.
+ * @tparam S - data type of object that will be stored
+ */
+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] {
+
+ private lazy val keySerializer = keyExprEnc.createSerializer()
+
+ private lazy val stateTypesEncoder = StateTypesEncoder(
+ keySerializer, valEncoder, stateName, hasTtl = true)
+
+ private lazy 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 = {
+ get().nonEmpty
+ }
+
+ /**
+ * 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)
+
+ new NextIterator[S] {
+
+ override protected def getNext(): S = {
+ val iter = unsafeRowValuesIterator.dropWhile { row =>
+ stateTypesEncoder.isExpired(row, batchTimestampMs)
+ }
+
+ if (iter.hasNext) {
+ val currentRow = iter.next()
+ stateTypesEncoder.decodeValue(currentRow)
+ } else {
+ finished = true
+ null.asInstanceOf[S]
+ }
+ }
+
+ override protected def close(): Unit = {}
+ }
+ }
+
+ /** Update the value of the list. */
+ override def put(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ 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)
+ }
+ }
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entry to the list. */
+ override def appendValue(newState: S): Unit = {
+ StateStoreErrors.requireNonNullStateValue(newState, stateName)
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ store.merge(stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey),
+ stateTypesEncoder.encodeValue(newState, ttlExpirationMs), stateName)
+ upsertTTLForStateKey(serializedGroupingKey)
+ }
+
+ /** Append an entire list to the existing value. */
+ override def appendList(newState: Array[S]): Unit = {
+ validateNewState(newState)
+
+ val serializedGroupingKey = stateTypesEncoder.serializeGroupingKey()
+ val encodedKey = stateTypesEncoder.encodeSerializedGroupingKey(serializedGroupingKey)
+ newState.foreach { v =>
+ val encodedValue = stateTypesEncoder.encodeValue(v, ttlExpirationMs)
+ store.merge(encodedKey, encodedValue, stateName)
+ }
+ upsertTTLForStateKey(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)
+ }
+ }
+
+ /**
+ * Loops through all the values associated with the grouping key, and removes
+ * the expired elements from the list.
+ * @param groupingKey grouping key for which cleanup should be performed.
+ */
+ override def clearIfExpired(groupingKey: Array[Byte]): Long = {
+ var numValuesExpired = 0L
+ 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 (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) {
+ if (isFirst) {
+ store.put(encodedGroupingKey, encodedValue, stateName)
+ isFirst = false
+ } else {
+ store.merge(encodedGroupingKey, encodedValue, stateName)
+ }
+ } else {
+ numValuesExpired += 1
+ }
+ }
+ numValuesExpired
+ }
+
+ private def upsertTTLForStateKey(serializedGroupingKey: Array[Byte]): Unit = {
+ 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
+ * 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[S] = {
+ val encodedGroupingKey = stateTypesEncoder.encodeGroupingKey()
+ val unsafeRowValuesIterator = store.valuesIterator(encodedGroupingKey, stateName)
+ unsafeRowValuesIterator.map{
Review Comment:
nit: space between p and {, also valueUnsafeRow => to be after {. Saving one line and also one indentation.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ListStateSuite.scala:
##########
@@ -160,4 +161,60 @@ class ListStateSuite extends StateVariableSuiteBase {
assert(listState1.get().toSeq === Seq.empty[Long])
}
}
+
+ test(s"test List state TTL") {
+ tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider =>
+ val store = provider.getStore(0)
+ val timestampMs = 10
+ val handle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
+ Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]],
+ TimeMode.ProcessingTime(), batchTimestampMs = Some(timestampMs))
+
+ val ttlConfig = TTLConfig(ttlDuration = Duration.ofMinutes(1))
+ val testState: ListStateImplWithTTL[String] = handle.getListState[String]("testState",
+ Encoders.STRING, ttlConfig).asInstanceOf[ListStateImplWithTTL[String]]
+ ImplicitGroupingKeyTracker.setImplicitKey("test_key")
+ testState.put(Array("v1", "v2", "v3"))
+ assert(testState.get().toSeq === Seq("v1", "v2", "v3"))
+ assert(testState.getWithoutEnforcingTTL().toSeq === Seq("v1", "v2", "v3"))
+
+ val ttlExpirationMs = timestampMs + 60000
+ var ttlValues = testState.getTTLValues()
+ assert(ttlValues.nonEmpty)
+ assert(ttlValues.forall(_._2 === ttlExpirationMs))
+ var ttlStateValueIterator = testState.getValuesInTTLState()
+ assert(ttlStateValueIterator.hasNext)
+
+ // increment batchProcessingTime, or watermark and ensure expired value is not returned
+ val nextBatchHandle = new StatefulProcessorHandleImpl(store, UUID.randomUUID(),
+ Encoders.STRING.asInstanceOf[ExpressionEncoder[Any]],
+ TimeMode.ProcessingTime(), batchTimestampMs = Some(ttlExpirationMs))
+
+ val nextBatchTestState: ListStateImplWithTTL[String] =
+ nextBatchHandle.getListState[String]("testState", Encoders.STRING, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[String]]
+
+ ImplicitGroupingKeyTracker.setImplicitKey("test_key")
+
+ // ensure get does not return the expired value
+ assert(!nextBatchTestState.exists())
+ assert(nextBatchTestState.get().isEmpty)
+
+ // ttl value should still exist in state
+ ttlValues = nextBatchTestState.getTTLValues()
+ assert(ttlValues.nonEmpty)
+ assert(ttlValues.forall(_._2 === ttlExpirationMs))
+ ttlStateValueIterator = nextBatchTestState.getValuesInTTLState()
+ assert(ttlStateValueIterator.hasNext)
+ assert(ttlStateValueIterator.next() === ttlExpirationMs)
+ assert(ttlStateValueIterator.isEmpty)
+
+ // getWithoutTTL should still return the expired value
+ assert(nextBatchTestState.getWithoutEnforcingTTL().toSeq === Seq("v1", "v2", "v3"))
+
+ nextBatchTestState.clear()
+ assert(!nextBatchTestState.exists())
+ assert(nextBatchTestState.get().isEmpty)
+ }
+ }
Review Comment:
Do we want to replicate the test case for duration to be 0 or negative? See ValueStateSuite.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatefulProcessorHandleSuite.scala:
##########
@@ -219,7 +219,7 @@ class StatefulProcessorHandleSuite extends StateVariableSuiteBase {
}
}
- test(s"ttl States are populated for timeMode=ProcessingTime") {
+ test(s"ttl States are populated for valueState and timeMode=ProcessingTime") {
Review Comment:
nit: (while we are here) `s` before "" is not needed
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.streaming
+
+import java.sql.Timestamp
+import java.time.Duration
+
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+case class InputEvent(
+ key: String,
+ action: String,
+ value: Int,
+ eventTime: Timestamp = null)
+
+case class OutputEvent(
+ key: String,
+ value: Int,
+ isTTLValue: Boolean,
+ ttlValue: Long)
+/**
+ * Tests that ttl works as expected for Value State for
Review Comment:
The comment does not seem valid - there's no event time based TTL (for now), and also I guess you'll reuse this with list state as well. If you don't intend to place the comment here, please move this back to TransformWithValueStateTTLSuite.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.streaming
+
+import java.sql.Timestamp
+import java.time.Duration
+
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+case class InputEvent(
+ key: String,
+ action: String,
+ value: Int,
+ eventTime: Timestamp = null)
+
+case class OutputEvent(
+ key: String,
+ value: Int,
+ isTTLValue: Boolean,
+ ttlValue: Long)
+/**
+ * Tests that ttl works as expected for Value State for
+ * processing time and event time based ttl.
+ */
+abstract class TransformWithStateTTLTest
+ extends StreamTest {
+ import testImplicits._
+
+ def getProcessor(ttlConfig: TTLConfig): StatefulProcessor[String, InputEvent, OutputEvent]
+
+ def getStateTTLMetricName: String
+
+ test("validate state is evicted at ttl expiry") {
Review Comment:
I expect tests in here aren't changed from TransformWithValueStateTTLSuite and just extracted out. Please let me know if that's not the case, and also point out all the places for the change.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
Review Comment:
nit: 2 more spaces
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateTTLTest.scala:
##########
@@ -0,0 +1,288 @@
+/*
+ * 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.streaming
+
+import java.sql.Timestamp
+import java.time.Duration
+
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+case class InputEvent(
+ key: String,
+ action: String,
+ value: Int,
+ eventTime: Timestamp = null)
+
+case class OutputEvent(
+ key: String,
+ value: Int,
+ isTTLValue: Boolean,
+ ttlValue: Long)
+/**
Review Comment:
nit: empty line between two
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
Also, it's going back to 3 minutes, not 1 minute.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
Review Comment:
nit: 2 more spaces
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
Review Comment:
Could you please add a comment to get TTL values for 4 to 6 here? Should be much helpful for reviewer to go through.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566090091
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
I'm a little confused by this comment - I think using dropWhile should be okay since we are only returning one row at a time in the iterator, and `getNext` is called after every invocation of `next()`
Consider the following case (e for expired, v for valid):
[v1, v2, v3, e1, e2, e3, v4, v5, v6]
Calling `dropWhile` will drop the first three elements and point the iterator to v1, and we set the currentRow to v1. After we return v1, we set the currentRow to v2, then v3 (as `dropWhile` doesn't drop these elements)
Now, we have returned v3, and `unsafeRowValuesIterator` is pointing at e4. We would call dropWhile on this, causing the following three elements to be dropped, and then we would point at v4.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-47673][SS] Implementing TTL for ListState [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #45932:
URL: https://github.com/apache/spark/pull/45932#discussion_r1566363954
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithListStateTTLSuite.scala:
##########
@@ -0,0 +1,350 @@
+/*
+ * 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.streaming
+
+import java.time.Duration
+
+import org.apache.spark.sql.Encoders
+import org.apache.spark.sql.execution.streaming.{ListStateImplWithTTL, MemoryStream}
+import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.streaming.util.StreamManualClock
+
+class ListStateTTLProcessor(ttlConfig: TTLConfig)
+ extends StatefulProcessor[String, InputEvent, OutputEvent] {
+
+ @transient private var _listState: ListStateImplWithTTL[Int] = _
+
+ override def init(
+ outputMode: OutputMode,
+ timeMode: TimeMode): Unit = {
+ _listState = getHandle
+ .getListState("listState", Encoders.scalaInt, ttlConfig)
+ .asInstanceOf[ListStateImplWithTTL[Int]]
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[InputEvent],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[OutputEvent] = {
+ var results = List[OutputEvent]()
+
+ inputRows.foreach { row =>
+ val resultIter = processRow(row, _listState)
+ resultIter.foreach { r =>
+ results = r :: results
+ }
+ }
+
+ results.iterator
+ }
+
+ def processRow(
+ row: InputEvent,
+ listState: ListStateImplWithTTL[Int]): Iterator[OutputEvent] = {
+
+ var results = List[OutputEvent]()
+ val key = row.key
+ if (row.action == "get") {
+ val currState = listState.get()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_without_enforcing_ttl") {
+ val currState = listState.getWithoutEnforcingTTL()
+ currState.foreach { v =>
+ results = OutputEvent(key, v, isTTLValue = false, -1) :: results
+ }
+ } else if (row.action == "get_ttl_value_from_state") {
+ val ttlValues = listState.getTTLValues()
+ ttlValues.foreach { ttlValue =>
+ results = OutputEvent(key, ttlValue._1, isTTLValue = true, ttlValue._2) :: results
+ }
+ } else if (row.action == "put") {
+ listState.put(Array(row.value))
+ } else if (row.action == "append") {
+ listState.appendValue(row.value)
+ } else if (row.action == "get_values_in_ttl_state") {
+ val ttlValues = listState.getValuesInTTLState()
+ ttlValues.foreach { v =>
+ results = OutputEvent(key, -1, isTTLValue = true, ttlValue = v) :: results
+ }
+ }
+
+ results.iterator
+ }
+}
+
+/**
+ * Test suite for testing list state with TTL.
+ * We use the base TTL suite with a list state processor.
+ */
+class TransformWithListStateTTLSuite extends TransformWithStateTTLTest {
+
+ import testImplicits._
+
+ override def getProcessor(ttlConfig: TTLConfig):
+ StatefulProcessor[String, InputEvent, OutputEvent] = {
+ new ListStateTTLProcessor(ttlConfig)
+ }
+
+ override def getStateTTLMetricName: String = "numListStateWithTTLVars"
+
+ test("verify iterator works with expired values in middle of list") {
+ withSQLConf(SQLConf.STATE_STORE_PROVIDER_CLASS.key ->
+ classOf[RocksDBStateStoreProvider].getName,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
+ withTempDir { checkpointLocation =>
+ val ttlConfig1 = TTLConfig(ttlDuration = Duration.ofMinutes(3))
+ val inputStream = MemoryStream[InputEvent]
+ val result1 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig1),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+
+ val clock = new StreamManualClock
+ // add 3 elements with a duration of a minute
+ testStream(result1)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ 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, 181000),
+ OutputEvent("k1", 2, isTTLValue = true, 182000),
+ OutputEvent("k1", 3, isTTLValue = true, 182000)
+ ),
+ 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)
+ ),
+ StopStream
+ )
+
+ val ttlConfig2 = TTLConfig(ttlDuration = Duration.ofSeconds(15))
+ val result2 = inputStream.toDS()
+ .groupByKey(x => x.key)
+ .transformWithState(
+ getProcessor(ttlConfig2),
+ TimeMode.ProcessingTime(),
+ OutputMode.Append())
+ // add 3 elements with a duration of 15 seconds
+ testStream(result2)(
+ StartStream(Trigger.ProcessingTime("1 second"), triggerClock = clock,
+ checkpointLocation = checkpointLocation.getAbsolutePath),
+ AddData(inputStream, InputEvent("k1", "append", 4)),
+ AddData(inputStream, InputEvent("k1", "append", 5)),
+ AddData(inputStream, InputEvent("k1", "append", 6)),
+ // advance clock to trigger processing
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(),
+ // get all elements without enforcing ttl
+ AddData(inputStream, InputEvent("k1", "get_without_enforcing_ttl", -1, null)),
+ AdvanceManualClock(1 * 1000),
+ CheckNewAnswer(
+ OutputEvent("k1", 1, isTTLValue = false, -1),
+ OutputEvent("k1", 2, isTTLValue = false, -1),
+ OutputEvent("k1", 3, isTTLValue = false, -1),
+ OutputEvent("k1", 4, isTTLValue = false, -1),
+ OutputEvent("k1", 5, isTTLValue = false, -1),
+ OutputEvent("k1", 6, isTTLValue = false, -1)
+ ),
+ StopStream
+ )
+ // add 3 more elements with a duration of a minute
Review Comment:
Ah you're right, it was confusing because you're creating a new iterator based on underlying iterator. It's a logical (view) iterator which also moves the underlying iterator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org