[SPARK-53870][PYTHON][SS] Fix partial read bug for large proto messages in TransformWithStateInPySparkStateServer #52539
+138
−1
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?
Fix the TransformWithState StateServer's
parseProtoMessage
method to fully read the desired message using the correct readFully DataInputStream API rather thanread
(InputStream/FilterInputStream) which only reads all available data and may not return the full message.readFully
(DataInputStream) will continue fetching until it fills up the provided buffer.In addition to the linked API above, this StackOverflow post also illustrates the difference between the two APIs: https://stackoverflow.com/a/25900095
Why are the changes needed?
For large state values used in the TransformWithState API,
inputStream.read
is not guaranteed to readmessageLen
's bytes of data as per the InputStream API. For large values,read
will 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 somewhat large to consistently trigger this error. The test case I added uses 512KB strings in the state value updates.
Does this PR introduce any user-facing change?
No
How was this patch tested?
Added a new test case using 512KB strings:
The configured data size (512KB) triggers an incomplete read, while also completing in a reasonable time (within 30s on my laptop). I had separately tested a larger input size of 4MB which took 30min which I considered too expensive to include in the test.
Below is sample/testing results from using
read
only (i.e., no fix) and adding a check on message length vs read bytes (test code is included in this commit but reverted later for the PR). The check is no longer required after thereadFully
fix as that is handled within the provided API.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (claude-sonnet-4-5-20250929)