[SPARK-50017][SS] Support Avro encoding for TransformWithState operator #48401
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ericm-db
changed the title
ericm-db
force-pushed
the
avro
branch
2 times, most recently
from
ec1e07a to
1aca8f4
Compare
connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroDataSourceV2.scala
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
why do we need to move this file?
There was a problem hiding this comment.
Because it's used in AvroOptions
There was a problem hiding this comment.
Have we considered introducing a deprecated class under org.apache.spark.sql.avro that retains all the existing public methods, while moving their implementations into sql/core?
There was a problem hiding this comment.
Sure, we can do this.
connector/avro/src/main/scala/org/apache/spark/sql/avro/DeprecatedSchemaConverters.scala
Outdated
Show resolved
Hide resolved
connector/avro/src/main/scala/org/apache/spark/sql/v2/avro/AvroPartitionReaderFactory.scala
Outdated
Show resolved
Hide resolved
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ListStateImpl.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StateTypesEncoderUtils.scala
Outdated
Show resolved
Hide resolved
ericm-db
changed the title
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ValueStateImplWithTTL.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
Outdated
Show resolved
Hide resolved
|
|
||
| @deprecated("Use org.apache.spark.sql.core.avro.SchemaConverters instead", "4.0.0") | ||
| @Evolving | ||
| object DeprecatedSchemaConverters { |
There was a problem hiding this comment.
Let's keep the name SchemaConverters and don't have Deprecated in the object name
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/package.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
Outdated
Show resolved
Hide resolved
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
Outdated
Show resolved
Hide resolved
brkyvz
suggested changes
Nov 21, 2024
There was a problem hiding this comment.
The changes are SOOO much cleaner now, thank you. It can get even cleaner though:
- I feel like you can add a Serde interface for the StateEncoder code changes. That should simplify the code even further
- Any reason we just didn't extend the suites with a different SQLConf to test out the different encoding type? I feel that would remove a ton of code changes as well
| @@ -563,13 +684,233 @@ class RangeKeyScanStateEncoder( | |||
| writer.getRow() | |||
| } | |||
|
|
|||
| def encodePrefixKeyForRangeScan( | |||
There was a problem hiding this comment.
Can you add a scaladoc please?
| out.toByteArray | ||
| } | ||
|
|
||
| def decodePrefixKeyForRangeScan( |
| virtualColFamilyId: Option[Short] = None) | ||
| extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) { | ||
| virtualColFamilyId: Option[Short] = None, | ||
| avroEnc: Option[AvroEncoder] = None) |
There was a problem hiding this comment.
Instead of avroEnc, I would honestly introduce another interface:
trait Serde {
def encodeToBytes(...)
def decodeToUnsafeRow(...)
def encodePrefixKeyForRangeScan(...)
def decodePrefixKeyForRangeScan(...)
}and move the logic in there so that you don't have to keep on doing avroEnc.isDefined for these
The logic seems pretty similar except for the input data. The AvroStateSerde or whatever you want to name it would have the private lazy val remainingKeyAvroType = SchemaConverters.toAvroType(remainingKeySchema)
There was a problem hiding this comment.
Spoke offline - it doesn't look like this simplifies things an awful lot - can be a follow-up.
| virtualColFamilyId: Option[Short] = None) | ||
| extends RocksDBKeyStateEncoderBase(useColumnFamilies, virtualColFamilyId) { | ||
| virtualColFamilyId: Option[Short] = None, | ||
| avroEnc: Option[AvroEncoder] = None) |
| Some(newColFamilyId), avroEnc), RocksDBStateEncoder.getValueEncoder(valueSchema, | ||
| useMultipleValuesPerKey, avroEnc))) | ||
| } | ||
| private def getAvroSerializer(schema: StructType): AvroSerializer = { |
There was a problem hiding this comment.
nit: line before the method please
| @@ -74,10 +75,71 @@ private[sql] class RocksDBStateStoreProvider | |||
| isInternal: Boolean = false): Unit = { | |||
| verifyColFamilyCreationOrDeletion("create_col_family", colFamilyName, isInternal) | |||
| val newColFamilyId = rocksDB.createColFamilyIfAbsent(colFamilyName) | |||
| // Create cache key using store ID to avoid collisions | |||
| val avroEncCacheKey = s"${stateStoreId.operatorId}_" + | |||
There was a problem hiding this comment.
Do we have the stream runId (maybe it's available in the HadoopConf)? We should add runId, otherwise there could be collisions
Comment on lines
41
to
42
| // Avro encoder that is used by the RocksDBStateStoreProvider and RocksDBStateEncoder | ||
| // in order to serialize from UnsafeRow to a byte array of Avro encoding. |
There was a problem hiding this comment.
Can you please turn this into a proper scaladoc?
/**
* ...
*/| TestWithBothChangelogCheckpointingEnabledAndDisabled ) { colFamiliesEnabled => | ||
| val testSchema: StructType = StructType( | ||
| Seq( | ||
| StructField("ordering-1", LongType, false), |
There was a problem hiding this comment.
oh, why'd you have to change these? If these are not supported by Avro, do we have any check anywhere to disallow the usage of the Avro encoder?
There was a problem hiding this comment.
Avro code would just throw an error, saying that there are invalid characters in the field name
Comment on lines
131
to
132
| def testWithEncodingTypes(testName: String, testTags: Tag*) | ||
| (testBody: => Any): Unit = { |
There was a problem hiding this comment.
one parameter per line like below please
...e/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
Outdated
Show resolved
Hide resolved
|
oh forgot - we need to add the stream run id to the Avro encoder cache key, otherwise we may risk some unintended re-use of avro encoders. we should limit the size of that cache and add expiry to it |
brkyvz
approved these changes
Nov 22, 2024
...rc/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala
Show resolved
Hide resolved
...n/scala/org/apache/spark/sql/execution/streaming/state/StateSchemaCompatibilityChecker.scala
Outdated
Show resolved
Hide resolved
|
@brkyvz asked me to help merging this - while I could do (or just having a quick through the code change from my side before doing that), I'd like to make sure that @gengliangwang is OK with this change. @gengliangwang Do you have any further outstanding comment, or was splitting the PR your only concern? I'm going to merge this if you have no outstanding comment by tomorrow. |
|
No update. I'll proceed. Thanks! Merging to master. (DISCLAIMER: I'm merging on behalf of @brkyvz ) |
HeartSaVioR
changed the title
|
Thank you!
…On Mon, Nov 25, 2024, 8:33 PM Jungtaek Lim ***@***.***> wrote:
Closed #48401 <#48401> via 331d0bf
<331d0bf>
.
—
Reply to this email directly, view it on GitHub
<#48401 (comment)>, or
unsubscribe
<https://github.com/notifications/unsubscribe-auth/ABIAE62V4SZFXMEYAIZ72DT2CP2Z3AVCNFSM6AAAAABPU7JMVCVHI2DSMVQWIX3LMV45UABCJFZXG5LFIV3GK3TUJZXXI2LGNFRWC5DJN5XDWMJVGQZTAMBYGE4DCNQ>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
|
After this PR was merged, Maven daily test started to fail:
We can use the following method to confirm that this PR has caused similar test failures and reproduce the issue:
@ericm-db Could you help fix the issue mentioned above? |
|
@LuciferYang am looking into this, think I have to add a dependency in the sql/core pom.xml |
HeartSaVioR
pushed a commit
that referenced
this pull request
Nov 28, 2024
… in RocksDBStateStoreProvider ### What changes were proposed in this pull request? There are maven errors introduced by the guava dependency in `sql/core`, as we use the Guava cache to store the Avro encoders, outlined in this comment: #48401 (comment) Introduced a new constructor for the NonFateSharingCache and used this with the RocksDBStateStoreProvider. ### Why are the changes needed? To resolve maven build errors, so that the Avro change here: #48401 does not get reverted. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing unit tests are sufficient and maven build works on devbox ``` [INFO] Tests run: 47, Failures: 0, Errors: 0, Skipped: 0, Time elapsed: 17.64 s -- in test.org.apache.spark.sql.JavaDatasetSuite [INFO] [INFO] Results: [INFO] [INFO] Tests run: 47, Failures: 0, Errors: 0, Skipped: 0 [INFO] [INFO] [INFO] --- surefire:3.2.5:test (test) spark-sql_2.13 --- [INFO] Skipping execution of surefire because it has already been run for this configuration [INFO] [INFO] --- scalatest:2.2.0:test (test) spark-sql_2.13 --- [INFO] ScalaTest report directory: /home/eric.marnadi/spark/sql/core/target/surefire-reports WARNING: Using incubator modules: jdk.incubator.vector Discovery starting. Discovery completed in 2 seconds, 737 milliseconds. Run starting. Expected test count is: 0 DiscoverySuite: Run completed in 2 seconds, 765 milliseconds. Total number of tests run: 0 Suites: completed 1, aborted 0 Tests: succeeded 0, failed 0, canceled 0, ignored 0, pending 0 No tests were executed. [INFO] ------------------------------------------------------------------------ [INFO] BUILD SUCCESS [INFO] ------------------------------------------------------------------------ [INFO] Total time: 03:15 min [INFO] Finished at: 2024-11-28T01:10:36Z [INFO] ------------------------------------------------------------------------ ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48996 from ericm-db/chm. Authored-by: Eric Marnadi <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]>
HeartSaVioR
pushed a commit
that referenced
this pull request
Feb 21, 2025
…g to use the correct number of version bytes ### What changes were proposed in this pull request? There are currently two bugs: - The NoPrefixKeyStateEncoder adds an extra version byte to each row when UnsafeRow encoding is used: #47107 - Rows written with Avro encoding do not include a version byte: #48401 **Neither of these bugs have been released, since these bugs are only triggered with multiple column families, and transformWithState is only using it, which is going to be released for Spark 4.0.0.** This change fixes both of these bugs. ### Why are the changes needed? These changes are needed in order to conform with the expected state row encoding format. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #49996 from ericm-db/SPARK-51249. Lead-authored-by: Eric Marnadi <[email protected]> Co-authored-by: Eric Marnadi <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]>
HeartSaVioR
pushed a commit
that referenced
this pull request
Feb 21, 2025
…g to use the correct number of version bytes ### What changes were proposed in this pull request? There are currently two bugs: - The NoPrefixKeyStateEncoder adds an extra version byte to each row when UnsafeRow encoding is used: #47107 - Rows written with Avro encoding do not include a version byte: #48401 **Neither of these bugs have been released, since these bugs are only triggered with multiple column families, and transformWithState is only using it, which is going to be released for Spark 4.0.0.** This change fixes both of these bugs. ### Why are the changes needed? These changes are needed in order to conform with the expected state row encoding format. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes #49996 from ericm-db/SPARK-51249. Lead-authored-by: Eric Marnadi <[email protected]> Co-authored-by: Eric Marnadi <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]> (cherry picked from commit 42ab97a) Signed-off-by: Jungtaek Lim <[email protected]>
Pajaraja
pushed a commit
to Pajaraja/spark
that referenced
this pull request
Mar 6, 2025
…g to use the correct number of version bytes ### What changes were proposed in this pull request? There are currently two bugs: - The NoPrefixKeyStateEncoder adds an extra version byte to each row when UnsafeRow encoding is used: apache#47107 - Rows written with Avro encoding do not include a version byte: apache#48401 **Neither of these bugs have been released, since these bugs are only triggered with multiple column families, and transformWithState is only using it, which is going to be released for Spark 4.0.0.** This change fixes both of these bugs. ### Why are the changes needed? These changes are needed in order to conform with the expected state row encoding format. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Unit tests ### Was this patch authored or co-authored using generative AI tooling? No Closes apache#49996 from ericm-db/SPARK-51249. Lead-authored-by: Eric Marnadi <[email protected]> Co-authored-by: Eric Marnadi <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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