[SPARK-51351][SS] Do not materialize the output in Python worker for TWS #50110
+11
−2
Conversation
This file contains 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
anishshri-db
approved these changes
Feb 28, 2025
|
I'm going to provide the branch I was used to test this behavior. I'll update the PR description once it is ready. |
|
Also updated the PR description to explain how I figured out the issue and how I verified the fix. |
|
cc. @HyukjinKwon Would you mind taking a look? Thanks! |
HeartSaVioR
commented
Feb 28, 2025
| @@ -1223,8 +1223,17 @@ def dump_stream(self, iterator, stream): | |||
| Read through an iterator of (iterator of pandas DataFrame), serialize them to Arrow | |||
| RecordBatches, and write batches to stream. | |||
| """ | |||
| result = [(b, t) for x in iterator for y, t in x for b in y] | |||
| super().dump_stream(result, stream) | |||
|
|
|||
There was a problem hiding this comment.
nit: linter asked me to have an empty line here - probably due to inner function?
HyukjinKwon
approved these changes
Feb 28, 2025
|
https://github.com/HeartSaVioR/spark/actions/runs/13583292499/job/37988558904 One module failed due to heap space during compilation. I'm rerunning it. |
|
Thanks! Merging to master/4.0. |
HeartSaVioR
added a commit
that referenced
this pull request
Feb 28, 2025
### What changes were proposed in this pull request? This PR proposes to fix the logic of serializer in TWS PySpark version to NOT materialize the output entirely. This PR changes the logic of creating a list to create a generator instead, so that it can be lazily consumed. ### Why are the changes needed? Without this PR, all the outputs are materialized when JVM signals to Python worker that there is no further input (at task completion), which brings up two critical issues: * downstream operator can only see outputs after TWS operator processes all inputs * all the outputs are materialized into "memory" in Python worker, which could lead memory issue ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing UTs. I've confirmed manually below: * Before this PR, all the outputs are available after processing all inputs * After this PR, outputs are available during processing inputs The change I have made to verify the fix manually: HeartSaVioR@cd30db0 If we call run_test() in testcode.py in PySpark, the log messages `Spark pulls the iterators` and `The data is being retrieved from Python worker` are interleaved with this fix. Without the fix, there is a sequence of log messages, all `Spark pulls the iterators` messages come first, and then all `The data is being retrieved from Python worker` messages come later. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #50110 from HeartSaVioR/SPARK-51351. Authored-by: Jungtaek Lim <kabhwan.opensource@gmail.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com> (cherry picked from commit 496fe7a) Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>
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?
This PR proposes to fix the logic of serializer in TWS PySpark version to NOT materialize the output entirely. This PR changes the logic of creating a list to create a generator instead, so that it can be lazily consumed.
Why are the changes needed?
Without this PR, all the outputs are materialized when JVM signals to Python worker that there is no further input (at task completion), which brings up two critical issues:
Does this PR introduce any user-facing change?
No.
How was this patch tested?
Existing UTs. I've confirmed manually below:
The change I have made to verify the fix manually:
HeartSaVioR@cd30db0
If we call run_test() in testcode.py in PySpark, the log messages
Spark pulls the iteratorsandThe data is being retrieved from Python workerare interleaved with this fix. Without the fix, there is a sequence of log messages, allSpark pulls the iteratorsmessages come first, and then allThe data is being retrieved from Python workermessages come later.Was this patch authored or co-authored using generative AI tooling?
No.