[SPARK-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding #48944

ericm-db · 2024-11-23T21:36:51Z

What changes were proposed in this pull request?

Currently, we use the internal byte representation to store state for stateful streaming operators in the StateStore. This PR introduces Avro serialization and deserialization capabilities in the RocksDBStateEncoder so that we can instead use Avro encoding to store state. This is currently enabled for the TransformWithState operator via SQLConf to support all functionality supported by TWS

Why are the changes needed?

UnsafeRow is an inherently unstable format that makes no guarantees of being backwards-compatible. Therefore, if the format changes between Spark releases, this could cause StateStore corruptions. Avro is more stable, and inherently enables schema evolution.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Amended and added to unit tests

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

No

brkyvz

Overall looking good! Left some comments, we're missing some clean up that can be done

brkyvz · 2024-12-13T17:13:44Z

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

+  def encodeKey(row: UnsafeRow): Array[Byte]
+  def encodeRemainingKey(row: UnsafeRow): Array[Byte]
+  def encodePrefixKeyForRangeScan(row: UnsafeRow): Array[Byte]
+  def encodeValue(row: UnsafeRow): Array[Byte]


can you add docs for these actually? What additional information would be needed for these?

For example, when would encodeRemainingKey be used? What schema information needs to be passed in for this?

brkyvz · 2024-12-13T17:13:52Z

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

+  def decodeKey(bytes: Array[Byte]): UnsafeRow
+  def decodeRemainingKey(bytes: Array[Byte]): UnsafeRow
+  def decodePrefixKeyForRangeScan(bytes: Array[Byte]): UnsafeRow
+  def decodeValue(bytes: Array[Byte]): UnsafeRow


brkyvz · 2024-12-13T20:33:52Z

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

+        (RocksDBStateEncoder.getKeyEncoder(dataEncoder, keyStateEncoderSpec, useColumnFamilies,
+          Some(newColFamilyId)), RocksDBStateEncoder.getValueEncoder(dataEncoder, valueSchema,
+          useMultipleValuesPerKey)))


nit: can you make these one parameter per line please? maybe define them outside of the putIfAbsent first

brkyvz · 2024-12-13T20:34:08Z

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

+      (RocksDBStateEncoder.getKeyEncoder(dataEncoder, keyStateEncoderSpec,
+        useColumnFamilies, defaultColFamilyId),
+        RocksDBStateEncoder.getValueEncoder(dataEncoder, valueSchema, useMultipleValuesPerKey)))


brkyvz · 2024-12-13T20:37:09Z

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

    stateStoreEncoding match {
-      case "avro" => Some(
-        RocksDBStateStoreProvider.avroEncoderMap.get(
-          avroEncCacheKey,
-          new java.util.concurrent.Callable[AvroEncoder] {
-            override def call(): AvroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+      case "avro" =>
+        RocksDBStateStoreProvider.dataEncoderCache.get(
+          encoderCacheKey,
+          new java.util.concurrent.Callable[AvroStateEncoder] {
+            override def call(): AvroStateEncoder = {
+              val avroEncoder = createAvroEnc(keyStateEncoderSpec, valueSchema)
+              new AvroStateEncoder(keyStateEncoderSpec, valueSchema, avroEncoder)
+            }
+          }
+        )
+      case "unsaferow" =>
+        RocksDBStateStoreProvider.dataEncoderCache.get(
+          encoderCacheKey,
+          new java.util.concurrent.Callable[UnsafeRowDataEncoder] {
+            override def call(): UnsafeRowDataEncoder = {
+              new UnsafeRowDataEncoder(keyStateEncoderSpec, valueSchema)
+            }


nit: maybe change the order of the match-case with the cache.get?

assert(Set("avro", "unsaferow").contains(stateStoreEncoding)) RocksDBStateStoreProvider.dataEncoderCache.get( encoderCacheKey, new java.util.concurrent.Callable[DataEncoder] { () => if (stateStoreEncoding == "avro") { ... } else { ... } }

brkyvz · 2024-12-13T20:42:34Z

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

+      dataEncoder: RocksDBDataEncoder,
+      valueSchema: StructType,
+      useMultipleValuesPerKey: Boolean,
+      avroEnc: Option[AvroEncoder] = None): RocksDBValueStateEncoder = {


brkyvz · 2024-12-13T20:42:45Z

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

+  }
+}
+
+object RocksDBStateEncoder extends Logging {


Can you add scaladoc for all the methods below?

brkyvz · 2024-12-13T20:44:55Z

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

+    keyStateEncoderSpec match {
+      case NoPrefixKeyStateEncoderSpec(keySchema) =>
+        new NoPrefixKeyStateEncoder(dataEncoder, keySchema, useColumnFamilies, virtualColFamilyId)
+
+      case PrefixKeyScanStateEncoderSpec(keySchema, numColsPrefixKey) =>
+        new PrefixKeyScanStateEncoder(dataEncoder, keySchema, numColsPrefixKey,
+          useColumnFamilies, virtualColFamilyId)
+
+      case RangeKeyScanStateEncoderSpec(keySchema, orderingOrdinals) =>
+        new RangeKeyScanStateEncoder(dataEncoder, keySchema, orderingOrdinals,
+          useColumnFamilies, virtualColFamilyId)
+
+      case _ =>
+        throw new IllegalArgumentException(s"Unsupported key state encoder spec: " +
+          s"$keyStateEncoderSpec")
+    }
+  }


should the encoderSpec's expose a method instead:

def toEncoder(dataEncoder: DataEncoder, useColumnFamilies, virtualColFamilyId): RocksDBKeyStateEncoderBase

brkyvz · 2024-12-13T20:47:01Z

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

+  // Prefix Key schema and projection definitions used by the Avro Serializers
+  // and Deserializers
+  private val prefixKeySchema = StructType(keySchema.take(numColsPrefixKey))
+  private lazy val prefixKeyAvroType = SchemaConverters.toAvroType(prefixKeySchema)
+  private val prefixKeyProj = UnsafeProjection.create(prefixKeySchema)
+
+  // Remaining Key schema and projection definitions used by the Avro Serializers
+  // and Deserializers
+  private val remainingKeySchema = StructType(keySchema.drop(numColsPrefixKey))
+  private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
+  private val remainingKeyProj = UnsafeProjection.create(remainingKeySchema)


some of these need to be cleaned up. They're unused and no longer needed

brkyvz · 2024-12-13T20:48:26Z

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

+  private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
+    StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))


these are also no longer used. Please clean them up

brkyvz

Basically comment on scaladoc style, otherwise LGTM!

brkyvz · 2024-12-17T02:05:30Z

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

+/** Interface for encoding and decoding state store data between UnsafeRow and raw bytes.
+ *
+ * @note All encode methods expect non-null input rows. Handling of null values is left to the
+ * implementing classes.
+ */


is this in the right place?

brkyvz · 2024-12-17T02:05:50Z

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

+   */
  def encodeKey(row: UnsafeRow): Array[Byte]
+
+  /** Encodes the non-prefix portion of a key row. Used with prefix scan and


Text should start at the second line, not the first in scaladoc

brkyvz · 2024-12-17T02:06:39Z

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

+                             keyStateEncoderSpec: KeyStateEncoderSpec,
+                             valueSchema: StructType
+                           ): AvroEncoder = {


please fix the indent

brkyvz · 2024-12-17T02:06:55Z

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

+

uber nit: delete extra empty lines

brkyvz · 2024-12-17T02:17:58Z

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

  }
 }

+/** Factory object for creating state encoders used by RocksDB state store.


move text to next line

brkyvz · 2024-12-17T02:19:18Z

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

  def jsonValue: JValue
  def json: String = compact(render(jsonValue))
+
+  /** Creates a RocksDBKeyStateEncoder for this specification.


ditto on scaladoc

brkyvz · 2024-12-17T02:19:34Z

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

  if (numColsPrefixKey == 0 || numColsPrefixKey >= keySchema.length) {
    throw StateStoreErrors.incorrectNumOrderingColsForPrefixScan(numColsPrefixKey.toString)
  }
+  override def toEncoder(


nit: new line please

...apache/spark/sql/execution/datasources/v2/state/StateDataSourceTransformWithStateSuite.scala

brkyvz · 2024-12-17T04:14:49Z

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

+      case PrefixKeyScanStateEncoderSpec(_, _) =>
+        encodeUnsafeRowToAvro(row, avroEncoder.suffixKeySerializer.get, remainingKeyAvroType, out)
+      case RangeKeyScanStateEncoderSpec(_, _) =>
+        encodeUnsafeRowToAvro(row, avroEncoder.keySerializer, remainingKeyAvroType, out)


does this need to be suffixKeySerializer?

No, this is right - RangeKeyScan encoder doesn't use Avro serialization for both parts of the ke

brkyvz

LGTM, but you missed one comment: https://github.com/apache/spark/pull/48944/files#r1887807431 on scaladoc. Once done @HeartSaVioR , can you please help merge this?

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

HeartSaVioR · 2024-12-18T07:06:18Z

The comment link does not show the comment, which I assume the comment is outdated via new commit.
I'm merging on behalf of @brkyvz . Thanks! Merging to master.

…and UnsafeRow encoding ### What changes were proposed in this pull request? Currently, we use the internal byte representation to store state for stateful streaming operators in the StateStore. This PR introduces Avro serialization and deserialization capabilities in the RocksDBStateEncoder so that we can instead use Avro encoding to store state. This is currently enabled for the TransformWithState operator via SQLConf to support all functionality supported by TWS ### Why are the changes needed? UnsafeRow is an inherently unstable format that makes no guarantees of being backwards-compatible. Therefore, if the format changes between Spark releases, this could cause StateStore corruptions. Avro is more stable, and inherently enables schema evolution. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Amended and added to unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes apache#48944 from ericm-db/avro-ss. Lead-authored-by: Eric Marnadi <[email protected]> Co-authored-by: Eric Marnadi <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]>

Kimahriman · 2025-05-16T14:16:01Z

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

      storeConf: StateStoreConf,
      hadoopConf: Configuration,
      useMultipleValuesPerKey: Boolean = false): StateStore = {
+    hadoopConf.set(StreamExecution.RUN_ID_KEY, storeProviderId.queryRunId.toString)


Random late comment I noticed in the code, this is already called 4 line below

init

e939fa8

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

ericm-db added 4 commits November 23, 2024 14:24

fixed some stuff

35414a8

renaming

19f34ab

scaladoc

94acc5d

more stuff

5d5cb65

ericm-db changed the title ~~[WIP] init~~ [SPARK-50017] Support Avro encoding for TransformWithState operator Nov 23, 2024

HyukjinKwon changed the title ~~[SPARK-50017] Support Avro encoding for TransformWithState operator~~ [SPARK-50017][SQL] Support Avro encoding for TransformWithState operator Nov 24, 2024

ericm-db and others added 4 commits November 25, 2024 13:06

scalastyle

8d860ee

lazy val

d4f8b31

Merge branch 'master' into avro-ss

ed8e0a9

compiles

ca1353c

brkyvz suggested changes Dec 13, 2024

View reviewed changes

ericm-db and others added 4 commits December 13, 2024 13:26

feedback

85aa8da

Merge branch 'master' into avro-ss

bc373b0

moving avro enc creation into RocksDBStateENcoder

ff1e0f4

Update TransformWithMapStateSuite.scala

a29d14d

ericm-db requested a review from brkyvz December 13, 2024 21:54

ericm-db and others added 2 commits December 13, 2024 21:37

Merge branch 'master' into avro-ss

800f94b

fixed stream-stream join

e3dc8a5

brkyvz suggested changes Dec 17, 2024

View reviewed changes

feedback, scaladoc

8353908

ericm-db requested a review from brkyvz December 17, 2024 04:10

brkyvz reviewed Dec 17, 2024

View reviewed changes

ericm-db added 3 commits December 16, 2024 20:18

removing the avroEnc

ec9e300

removing avroenc from scaladoc

d12d877

adding scaladoc

2dcf844

ericm-db requested a review from brkyvz December 17, 2024 05:19

ericm-db changed the title ~~[SPARK-50017][SQL] Support Avro encoding for TransformWithState operator~~ [SPARK-50599][SQL] Support Avro encoding for TransformWithState operator Dec 17, 2024

ericm-db changed the title ~~[SPARK-50599][SQL] Support Avro encoding for TransformWithState operator~~ [SPARK-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding Dec 17, 2024

brkyvz approved these changes Dec 17, 2024

View reviewed changes

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

Update StateStore.scala

5028d66

HeartSaVioR closed this in 229118c Dec 18, 2024

Kimahriman reviewed May 16, 2025

View reviewed changes

		private val rangeScanAvroSchema = StateStoreColumnFamilySchemaUtils.convertForRangeScan(
		StructType(rangeScanKeyFieldsWithOrdinal.map(_._1).toArray))

[SPARK-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding #48944

[SPARK-50599][SQL] Create the DataEncoder trait that allows for Avro and UnsafeRow encoding #48944

Uh oh!

Conversation

ericm-db commented Nov 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

brkyvz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brkyvz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brkyvz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

HeartSaVioR commented Dec 18, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ericm-db commented Nov 23, 2024 •

edited

Loading