[WIP][SPARK-XXXXX][PS][SS] Fix partial read bug for large state values in TransformWithStateInPySparkStateServer #52539
+145
−1
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
…ateServer Use readFully() instead of read() to ensure the entire protobuf message is read from the input stream. The read() method may only read partial results (experimentally 8KB), which can cause failures when processing large state values. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
Add test case to verify that map state can handle large string values (10 KB) without issues. This test exercises the code path that requires the readFully() fix to prevent partial reads when processing large state values. The test: - Creates a 10 KB string - Updates map state with the large string value - Verifies the retrieved value matches the original in both length and content - Tests both Pandas and Row-based implementations 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
Enhance the test to comprehensively test all three state types (value, list, map) with large values instead of just map state: - Renamed MapStateLargeStringProcessor to LargeValueStatefulProcessor to reflect broader coverage - Test now verifies ValueState, ListState, and MapState all handle 10 KB string values correctly - Simplified assertions: processors return retrieved values instead of asserting internally, test compares actual vs expected values - ListState now returns full contents (comma-separated) instead of just count for better verification 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
aee4e8c to
0ad023e
Compare
5044b4a to
fe1a90f
Compare
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?
Update the TransformWithState StateServer's proto message reading to fully read the desired message using the correct
readFullyAPI rather thanreadwhich may not fully consume all data. We should instead rely onreadFullywhich will fill up the allocated buffer with required length as per the InputStream API.Why are the changes needed?
For large state values used in the TransformWithState API,
inputStream.readis not guaranteed to readmessageLen's bytes of data as per the InputStream API. For large values,readwill return prematurely and the messageBytes will only be partially filled, yielding an incorrect and likely unparseable proto message.This is not a common scenario, as testing also indicated that the actual proto messages had to be fairly large (MB) to consistently trigger this error.
Does this PR introduce any user-facing change?
No
How was this patch tested?
The configured data size (4MB) was chosen to ensure that the fix is required and that there's a bit of buffer for machine differences:
Below is sample/testing results from using
readonly and adding a check on message length vs read bytes.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (claude-sonnet-4-5-20250929)