[SPARK-50714][SQL][SS] Enable schema evolution for TransformWithState when Avro encoding is used #49277

ericm-db · 2024-12-23T19:32:53Z

What changes were proposed in this pull request?

This PR introduces stateful schema evolution for the TransformWithState operator when Avro is used.
We modified the StateStoreColumnFamilySchema and the StateSchemaV3 file to keep track of the key schema and value schema id in order to support versioning and lookups of schemas across query restarts.
The AvroStateEncoder now takes a StateSchemaProvider, which allows it to look up all of the active schemas in the StateStore for a given column family, allowing it to pass reader and writer schemas to the AvroEncoder class.
We have also added changes so that the StateDataSource can read from these rows with schema ID

Why are the changes needed?

Schema evolution is a critical feature for stateful stream processing applications that need to handle changing data schemas over time.

Does this PR introduce any user-facing change?

Yes - this change allows stateful schema evolution that was not previously possible for the TransformWithState operator.

How was this patch tested?

Unit and Integration tests in RocksDBStateStoreSuite and TransformWithStateSuite

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

.../main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala

anishshri-db · 2024-12-30T20:47:10Z

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

+      readerSchema: Schema,
+      valueProj: UnsafeProjection): UnsafeRow = {
+    if (valueBytes != null) {
+      val reader = new GenericDatumReader[Any](writerSchema, readerSchema)


lets add some comments here around the args

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

...n/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala

...re/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala

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

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

anishshri-db · 2025-01-03T18:28:27Z

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

+
+    dataType match {
+      // Basic types
+      case BooleanType => false


Are these Avro defaults too ?

anishshri-db · 2025-01-03T18:29:03Z

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

+
+      // Complex types
+      case ArrayType(elementType, _) =>
+        val defaultArray = new java.util.ArrayList[Any]()


Why not have empty collections ? i.e. empty array/map etc ?

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala

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

...re/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala

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

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

.../main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala

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

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

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

...rc/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala

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

anishshri-db · 2025-01-03T20:34:59Z

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

    assert(timeMode == TimeMode.EventTime.toString || timeMode == TimeMode.ProcessingTime.toString)
-    if (timeMode == TimeMode.EventTime.toString) {
+    val primaryIndex = if (timeMode == TimeMode.EventTime.toString) {


nit: i guess this could also be split into a separate function ?

Oh I was trying to address this comment: #49277 (comment)

...n/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala

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

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

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

...e/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala

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

anishshri-db · 2025-01-03T21:11:36Z

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

@@ -285,13 +285,13 @@ class TransformWithValueStateTTLSuite extends TransformWithStateTTLTest {
          .add("expiryTimestampMs", LongType, nullable = false)
        val schemaForValueRow: StructType = StructType(Array(StructField("__dummy__", NullType)))
        val schema0 = StateStoreColFamilySchema(
-          TimerStateUtils.getTimerStateVarName(TimeMode.ProcessingTime().toString),
-          schemaForKeyRow,
+          TimerStateUtils.getTimerStateVarNames(TimeMode.ProcessingTime().toString)._1, 0,


Lets also add atleast a couple of tests for other composite types such as list, map etc ?

anishshri-db · 2025-01-03T21:13:32Z

@ericm-db - can you also format the PR description and explain in more detail what functionality this PR adds. Thx

Also - this is a user facing change right ?

anishshri-db · 2025-01-03T21:15:39Z

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

+
+        val result2 = inputData.toDS()
+          .groupByKey(x => x)
+          .transformWithState(new RunningCountStatefulProcessorNestedLongs(),


Could we also add a test case where the newly added column is also part of the output ? That way we can ensure that the default values are being picked up correctly ? Prob also similar for dropped columns ?

github-actions bot added SQL STRUCTURED STREAMING PYTHON labels Dec 23, 2024

ericm-db added 7 commits December 27, 2024 10:06

more stuff

9063e10

not building rn

5a0ba29

it builds - probably has a ton of shit wrong with it though

eff6284

stuff

332b648

removing assertion

eb4e4df

stuff

ae82712

all sorts of stuff fails

5c7d940

ericm-db force-pushed the ssm branch from 03c6297 to 5c7d940 Compare December 27, 2024 22:43

ericm-db added 3 commits December 27, 2024 14:54

schema evolution works?

fda69c8

works

2af87b2

stuff

cd4d1ff

anishshri-db reviewed Dec 30, 2024

View reviewed changes

.../main/scala/org/apache/spark/sql/execution/streaming/StateStoreColumnFamilySchemaUtils.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

...rc/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

...n/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

...re/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala Outdated Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala Show resolved Hide resolved

anishshri-db reviewed Dec 30, 2024

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala Show resolved Hide resolved

ericm-db added 6 commits December 30, 2024 16:35

fix

4826d12

RocksDBStateStoreSuite passes

3e1d01e

some state data stuff working

e66598c

stuff

8fdce94

cleaning stuff up

9a0efc0

Tests

e95f2aa

ericm-db changed the title ~~Ssm~~ [SPARK-50714] Enable schema evolution for TransformWithState when Avro encoding is used Jan 3, 2025

more tests

fc56474

reordering test case

47f624d

HyukjinKwon changed the title ~~[SPARK-50714] Enable schema evolution for TransformWithState when Avro encoding is used~~ [SPARK-50714][SQL] Enable schema evolution for TransformWithState when Avro encoding is used Jan 3, 2025