[SPARK-55416][SS][PYTHON] Streaming Python Data Source memory leak when end-offset is not updated by vinodkc · Pull Request #54237 · apache/spark

vinodkc · 2026-02-09T22:48:33Z

What changes were proposed in this pull request?

In _SimpleStreamReaderWrapper.latestOffset(), validate that custom implementation of datasource based on SimpleDataSourceStreamReader.read() does not return a non-empty batch with end == start. If it does, raise PySparkException with error class SIMPLE_STREAM_READER_OFFSET_DID_NOT_ADVANCE before appending to the cache. Empty batches with end == start remain allowed.

Why are the changes needed?

When a user implements read(start) incorrectly and returns:

Same offset for both: end = start (e.g. both {"offset": 0}).
Non-empty iterator: e.g. 2 rows.

If a reader returns end == start with data (e.g. return (it, {"offset": start_idx})), the wrapper keeps appending to its prefetch cache on every trigger while commit(end) never trims entries (first matching index is 0). The cache grows without bound and driver (non-JVM) memory increases until OOM. Validating and raising error before appending stops this and fails fast with a clear error.

Empty batches with end == start remain allowed , it will allow the Python data source to represent that there is no data to read.

Does this PR introduce any user-facing change?

Yes. Implementations that return end == start with a non-empty iterator now get PySparkException instead of unbounded memory growth. Empty batches with end == start are unchanged.

How was this patch tested?

Added unit test in test_python_streaming_datasource.py

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

No.

HeartSaVioR

Looks good in overall, only minor and nits. Nice finding!

HeartSaVioR

+1

HeartSaVioR · 2026-02-13T07:41:14Z

Thanks! Merging to master/4.1/4.0.

HeartSaVioR · 2026-02-13T07:44:09Z

I could only merge to master.

@vinodkc There is merge conflict in 4.1/4.0. Could you please submit a PR for these branches? It'd probably good to have it in all available version line branches since it's a memory leak issue.

vinodkc · 2026-02-17T19:29:56Z

@HeartSaVioR , Thanks for the details. I've raised two PRs for branch 4.1 and 4.0.
branch 4.1 : #54321
branch 4.0 : #54317

…en end-offset is not updated ### What changes were proposed in this pull request? In `_SimpleStreamReaderWrapper.latestOffset()`, validate that custom implementation of datasource based on `SimpleDataSourceStreamReader.read()` does not return a non-empty batch with end == start. If it does, raise PySparkException with error class `SIMPLE_STREAM_READER_OFFSET_DID_NOT_ADVANCE` before appending to the cache. Empty batches with end == start remain allowed. ### Why are the changes needed? When a user implements read(start) incorrectly and returns: - Same offset for both: end = start (e.g. both {"offset": 0}). - Non-empty iterator: e.g. 2 rows. If a reader returns end == start with data (e.g. return (it, {"offset": start_idx})), the wrapper keeps appending to its prefetch cache on every trigger while commit(end) never trims entries (first matching index is 0). The cache grows without bound and driver (non-JVM) memory increases until OOM. Validating and raising error before appending stops this and fails fast with a clear error. Empty batches with end == start remain allowed , it will allow the Python data source to represent that there is no data to read. ### Does this PR introduce _any_ user-facing change? Yes. Implementations that return end == start with a non-empty iterator now get PySparkException instead of unbounded memory growth. Empty batches with end == start are unchanged. ### How was this patch tested? Added unit test in `test_python_streaming_datasource.py` ### Was this patch authored or co-authored using generative AI tooling? No. Closesapache#54237 from vinodkc/br_SPARK-55416. Authored-by: vinodkc <vinod.kc.in@gmail.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

…ak when end-offset is not updated ### What changes were proposed in this pull request? Backport #54237 to branch-4.1 In `_SimpleStreamReaderWrapper.latestOffset()`, validate that custom implementation of datasource based on `SimpleDataSourceStreamReader.read()` does not return a non-empty batch with end == start. If it does, raise PySparkException with error class `SIMPLE_STREAM_READER_OFFSET_DID_NOT_ADVANCE` before appending to the cache. Empty batches with end == start remain allowed. ### Why are the changes needed? When a user implements read(start) incorrectly and returns: - Same offset for both: end = start (e.g. both {"offset": 0}). - Non-empty iterator: e.g. 2 rows. If a reader returns end == start with data (e.g. return (it, {"offset": start_idx})), the wrapper keeps appending to its prefetch cache on every trigger while commit(end) never trims entries (first matching index is 0). The cache grows without bound and driver (non-JVM) memory increases until OOM. Validating and raising error before appending stops this and fails fast with a clear error. Empty batches with end == start remain allowed , it will allow the Python data source to represent that there is no data to read. ### Does this PR introduce _any_ user-facing change? Yes. Implementations that return end == start with a non-empty iterator now get PySparkException instead of unbounded memory growth. Empty batches with end == start are unchanged. ### How was this patch tested? Added unit test in `test_python_streaming_datasource.py` ### Was this patch authored or co-authored using generative AI tooling? No. Closes#54321 from vinodkc/br_SPARK-55416_4.1. Authored-by: vinodkc <vinod.kc.in@gmail.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

…ak when end-offset is not updated ### What changes were proposed in this pull request? Backport #54237 to branch-4.0 In `_SimpleStreamReaderWrapper.latestOffset()`, validate that custom implementation of datasource based on `SimpleDataSourceStreamReader.read()` does not return a non-empty batch with end == start. If it does, raise PySparkException with error class `SIMPLE_STREAM_READER_OFFSET_DID_NOT_ADVANCE` before appending to the cache. Empty batches with end == start remain allowed. ### Why are the changes needed? When a user implements read(start) incorrectly and returns: - Same offset for both: end = start (e.g. both {"offset": 0}). - Non-empty iterator: e.g. 2 rows. If a reader returns end == start with data (e.g. return (it, {"offset": start_idx})), the wrapper keeps appending to its prefetch cache on every trigger while commit(end) never trims entries (first matching index is 0). The cache grows without bound and driver (non-JVM) memory increases until OOM. Validating and raising error before appending stops this and fails fast with a clear error. Empty batches with end == start remain allowed , it will allow the Python data source to represent that there is no data to read. ### Does this PR introduce _any_ user-facing change? Yes. Implementations that return end == start with a non-empty iterator now get PySparkException instead of unbounded memory growth. Empty batches with end == start are unchanged. ### How was this patch tested? Added unit test in `test_python_streaming_datasource.py` ### Was this patch authored or co-authored using generative AI tooling? No. Closes#54317 from vinodkc/br_SPARK-55416_4.0. Authored-by: vinodkc <vinod.kc.in@gmail.com> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

vinodkc changed the title ~~[SPARK-55416][PYTHON][SS]Streaming Python Data Source memory leak when end-offset is not updated~~[SPARK-55416][PYTHON][SS] Streaming Python Data Source memory leak when end-offset is not updatedFeb 9, 2026

HyukjinKwon requested a review from HeartSaVioR February 9, 2026 23:10

HeartSaVioR reviewed Feb 10, 2026
View reviewed changes

vinodkc added 6 commits February 10, 2026 11:40

Validate SimpleDataSourceStreamReader end offset advances
e54702b

Cache only when start != end
78c3563

Add test for empty iter and start == end offset
3644539

Rename error class
71b4e1f

Change order of ERROR class
28cdcdc

Fix review comments
3c0b205

vinodkc force-pushed the br_SPARK-55416 branch from 81b2fd3 to 3c0b205Compare February 10, 2026 19:41

Fix test failure
a290875

vinodkc changed the title ~~[SPARK-55416][PYTHON][SS] Streaming Python Data Source memory leak when end-offset is not updated~~[SPARK-55416][SS][PYTHON] Streaming Python Data Source memory leak when end-offset is not updatedFeb 12, 2026

HeartSaVioR approved these changes Feb 13, 2026
View reviewed changes

HeartSaVioR closed this in 7542473Feb 13, 2026

This was referenced Feb 13, 2026
[SPARK-55416][SS][PYTHON][4.0] Streaming Python Data Source memory leak when end-offset is not updated #54317
Closed
[SPARK-55416][SS][PYTHON][4.1] Streaming Python Data Source memory leak when end-offset is not updated #54321
Closed

Navigation Menu

Uh oh!

[SPARK-55416][SS][PYTHON] Streaming Python Data Source memory leak when end-offset is not updated - #54237

[SPARK-55416][SS][PYTHON] Streaming Python Data Source memory leak when end-offset is not updated#54237
vinodkc wants to merge 7 commits into
apache:masterfrom
vinodkc:br_SPARK-55416

vinodkc commented Feb 9, 2026•
edited
Loading

Uh oh!

HeartSaVioR left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HeartSaVioR left a comment

Uh oh!

HeartSaVioR commented Feb 13, 2026

Uh oh!

HeartSaVioR commented Feb 13, 2026•
edited
Loading

Uh oh!

vinodkc commented Feb 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

vinodkc commented Feb 9, 2026• 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!

HeartSaVioR left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HeartSaVioR left a comment

Choose a reason for hiding this comment

Uh oh!

HeartSaVioR commented Feb 13, 2026

Uh oh!

HeartSaVioR commented Feb 13, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vinodkc commented Feb 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

vinodkc commented Feb 9, 2026•
edited
Loading

HeartSaVioR commented Feb 13, 2026•
edited
Loading