[SPARK-46928][SS] Add support for ListState in Arbitrary State API v2. by sahnib · Pull Request #44961 · apache/spark

sahnib · 2024-01-31T01:40:33Z

What changes were proposed in this pull request?

This PR adds changes for ListState implementation in State Api v2. As a list contains multiple values for a single key, we utilize RocksDB merge operator to persist multiple values.

Changes include

A new encoder/decoder to encode multiple values inside a single byte[] array (stored in RocksDB). The encoding scheme is compatible with RocksDB StringAppendOperator merge operator.
Support merge operations in ChangelogCheckpointing v2.
Extend StateStore to support merge operation, and read multiple values for a single key (via a Iterator). Note that these changes are only supported for RocksDB currently.

Why are the changes needed?

These changes are needed to support list values in the State Store. The changes are part of the work around adding new stateful streaming operator for arbitrary state mgmt that provides a bunch of new features listed in the SPIP JIRA here - https://issues.apache.org/jira/browse/SPARK-45939

Does this PR introduce any user-facing change?

Yes
This PR introduces a new state type (ListState) that users can use in their Spark streaming queries.

How was this patch tested?

Added a new test suite for ListState to ensure the state produces correct results.
Added additional testcases for input validation.
Added tests for merge operator with RocksDB.
Added tests for changelog checkpointing merge operator.
Added tests for reading merged values in RocksDBStateStore.

Was this patch authored or co-authored using generative AI tooling?

No

ericm-db

Only nits, most points were addressed on the previous PR.

anishshri-db · 2024-02-02T21:06:50Z

maybe move arguments to new lines ?

nit: the first argument does not seem to be moved

anishshri-db · 2024-02-08T22:47:12Z

Should we just remove this ? cc - @HeartSaVioR

anishshri-db

lgtm ! thanks

HeartSaVioR

First pass. Looks OK in general. Will have a second review once we merge the PR #45038 and this PR is rebased.

HeartSaVioR · 2024-02-13T06:19:45Z

    StateStoreProvider.createAndInit(
      stateStoreProviderId, keySchema, valueSchema, numColsPrefixKey,
-      useColumnFamilies = false, storeConf, hadoopConf.value)
+      useColumnFamilies = false, storeConf, hadoopConf.value, useMultipleValuesPerKey = false)


Shall we file a JIRA ticket to support transformWithState with state data source - reader? From what I understand, we are deferring on supporting this. Please correct me if I'm missing.

I think we already have this planned. @anishshri-db Do we have the JIRA created for this support?

Filed the ticket here - https://issues.apache.org/jira/browse/SPARK-47047

HeartSaVioR · 2024-02-13T06:22:23Z

nit: the first argument does not seem to be moved

HeartSaVioR · 2024-02-13T07:11:34Z

+class ListStateImpl[S](store: StateStore,
+     stateName: String,
+     keyExprEnc: ExpressionEncoder[Any])
+  extends ListState[S] with Logging {


I think you can create serializer from keyExprEnc here as field and pass serializer instead in encodeGroupingKey.

HeartSaVioR · 2024-02-13T07:13:56Z

+    val keyByteArr = toRow
+      .apply(keyOption.get).asInstanceOf[UnsafeRow].getBytes()
+
+    val keyEncoder = UnsafeProjection.create(KEY_ROW_SCHEMA)


If we are fully sure we will end up with binary key - binary value for every state type, this should go up to fields in this object. Same for valueEncoder. If we anticipate different schema, maybe keyEncoder / valueEncoder should be passed as parameter instead.

We do not anticipate a different schema. Every key, value will be encoded to a byte Array. While working on refactoring key/value Encoder to move them to object level, I discovered that UnsafeProjection is not thread safe. Hence, we cannot use it across state variables.

Refactored this logic to create a class instead of utils object, and the class instance is not shared across state variables.

HeartSaVioR · 2024-02-13T08:19:11Z

  def remove(key: UnsafeRow,
    colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit

+  def merge(key: UnsafeRow, value: UnsafeRow,


Shall we add method doc?

Sorry, missed the doc here. Added.

HeartSaVioR

Only details and optional comments.

HeartSaVioR · 2024-02-14T20:47:24Z

Thinking out loud - it'd be ideal if we can move out the context of schema for state store key/value entirely to StateTypesEncoderUtils. State impls don't need to know about anything except that StateTypesEncoderUtils will give the UnsafeRow, and all context about serde will be put to StateTypesEncoderUtils. Just a 2 cents, not mandatory to address.

If we anticipate further changes to StateTypesEncoderUtils it's also OK to leave it as it is and make such change in later phase.

This schema is currently used in TransformWithStateExec operator to create the StateStoreRDD, in the state variables to create RocksDB column family, and inside the EncoderUtils for serde. We can unify the last 2 easily, and that way state impl classes would not need to know about the schema. It would be better to make this change with MapState finalized as that class will modify the stateTypesEncoder. Let me know if that sounds good.

Yeah we can defer to the later phase - it may give better visibility upon further changes.

HeartSaVioR · 2024-02-14T20:55:24Z

I know this is beyond the scope of PR, but while we are here, I don't see any difference with Option(get()).

HeartSaVioR · 2024-02-14T21:13:37Z

This triggers overhead so we wouldn't like to do this for every value. Either we have to sample (that's why we only do for the first value) or not doing this at all. It's probably OK to do latter as get() also doesn't do this for the case where there are multiple CFs.

Removed the wrapper iterator, we don't need it anymore. Thanks for pointing this out.

HeartSaVioR · 2024-02-14T21:16:49Z

Let's make a contract about expected return value, especially about clarification of value for non-existing key. Multiple implementations should follow the contract so this needs to be clearly explained.

Sounds good. Done

HeartSaVioR · 2024-02-14T21:20:57Z

Maybe we want to explicitly document the contract on the behavior when performing merge against non-existing key, since all implementations should provide the same.

Makes sense, done.

HeartSaVioR · 2024-02-17T00:20:01Z

I talked offline but to leave record (and not forget) -

The overhead of creating UnsafeProjection is significant and we really need to reuse this, with keeping in mind that the instance returned from UnsafeProjecion.create() would be thread-unsafe. So we should either 1) use thread-local or 2) just change this object (singleton) into class and initialize per state impl.

Aggressively you can even create it in operator-partition level so that multiple state impls can share without worrying about thread-safety (benefit of 1) but even without thread-local), but maybe it won't be a significant gain.

Thanks for pointing out the overhead of UnsafeProjection. I changed the StateTypesEncoderUtils to a class and its initialized per state impl to reuse the UnsafeProjection (for a StateImpl).

HeartSaVioR

+1 pending CI.

HeartSaVioR · 2024-02-21T05:41:46Z

CI failed from Run / Build modules: pyspark-sql, pyspark-resource, pyspark-testing but none of test failures are relevant to this change.

HeartSaVioR · 2024-02-21T05:42:01Z

Thanks! Merging to master.

### What changes were proposed in this pull request? The pr aims to remove some unused error classes, include: - _LEGACY_ERROR_TEMP_3193: The code in PR(#44961) no longer uses this error class, but the related item in the file `error-classes.json` has not been deleted synchronously. - _LEGACY_ERROR_TEMP_3197: The code in PR(#44961) no longer uses this error class, but the related item in the file `error-classes.json` has not been deleted synchronously. - _LEGACY_ERROR_TEMP_3234: This should be the follow-up work of #44665 - STATE_STORE_MULTIPLE_VALUES_PER_KEY: Introduced by [SPARK-46864](#44883), but not actually used in codebase. - TWS_VALUE_SHOULD_NOT_BE_NULL: Introduced by [SPARK-46864](#44883), but not actually used in codebase. ### Why are the changes needed? Make the code cleaner. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Pass GA. - Manually check. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #45509 from panbingkun/remove_outdated_error_classes. Authored-by: panbingkun <panbingkun@baidu.com> Signed-off-by: Max Gekk <max.gekk@gmail.com>

github-actions Bot added SQL STRUCTURED STREAMING labels Jan 31, 2024

sahnib mentioned this pull request Jan 31, 2024

[SS] Add ListState implementation for State API v2. anishshri-db/spark#3

Closed

sahnib marked this pull request as ready for review January 31, 2024 01:45

sahnib force-pushed the state-api-v2-list-state branch from 69fcab9 to 4eec351 Compare January 31, 2024 17:15

ericm-db reviewed Jan 31, 2024

View reviewed changes

Comment thread ...rc/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala Outdated

ericm-db approved these changes Jan 31, 2024

View reviewed changes

Comment thread sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated

Comment thread sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala Outdated

sahnib force-pushed the state-api-v2-list-state branch from 4eec351 to 2a0602c Compare January 31, 2024 22:50