akashchamp opened a new pull request, #40265:
URL: https://github.com/apache/beam/pull/40265
Fixes #40243
### Problem
In the Spark Structured Streaming runner, `ParDoTranslatorBatch` explicitly
`persist()`s the intermediate, wide-row `Dataset` it builds for a
multi-output
`ParDo` (so the per-`mapPartitions` transformation backing every tagged
output
isn't re-evaluated once per tag). The code even carried a `FIXME`
acknowledging
the gap:
```java
// FIXME What's the strategy to unpersist Datasets / RDDs?
```
No corresponding `unpersist()` was ever called, so every multi-output `ParDo`
in a pipeline leaves its cached `Dataset` in Spark's storage memory for the
lifetime of the `SparkSession`. On a large batch pipeline with several such
transforms this accumulates without bound and leads to storage-memory
exhaustion, disk spilling, and OOMs.
### Fix
- Added `PipelineTranslator.TranslationState#cacheDataset(Dataset,
StorageLevel)`,
a small helper translators use instead of calling `Dataset#persist`
directly
when they need to cache a dataset that isn't going through the normal
`putDataset` path. It persists the dataset and registers it for later
cleanup.
- `PipelineTranslator.TranslatingVisitor` now tracks every dataset cached
this
way and threads that collection through to `EvaluationContext`.
- `EvaluationContext#evaluate()` now unpersists all such datasets in a
`finally` block once evaluation of the leaves has finished — whether it
completed normally, was stopped early, or an evaluation failed with an
exception — so a cached dataset never outlives the pipeline run that
created
it.
- `ParDoTranslatorBatch` now calls `cxt.cacheDataset(...)` instead of calling
`persist()` directly on the wide-row `Dataset`, removing the `FIXME`.
### Testing
- `EvaluationContextTest`: added two new tests asserting that datasets
registered via the new cache-tracking collection are `unpersist()`-ed once
`evaluate()` finishes, both on the happy path and when evaluating a leaf
throws.
- `ParDoTest`: updated the two existing tests that assert a multi-output
`ParDo` produces a cached RDD/Dataset (`testPardoWithOutputTagsCachedRDD` /
`...CachedDataset`, renamed to `...UnpersistsRDDAfterRun` /
`...UnpersistsDatasetAfterRun`) so that, after the pipeline run completes,
they assert the opposite: `SESSION.hasCachedData()` is now `false`, i.e.
the
dataset was unpersisted rather than left cached.
- Reproduced the leak first: ran the two updated `ParDoTest` regression tests
(with their new "unpersisted after run" assertions) against the unfixed
production code — both fail with
`AssertionError: Expected cached data to be unpersisted after run`,
confirming `SESSION.hasCachedData()` is still `true` after the pipeline run
completes, i.e. the `Dataset` really is left cached.
- Re-applied the fix and ran the full `EvaluationContextTest` and `ParDoTest`
suites via `./gradlew :runners:spark:3:test --tests
...EvaluationContextTest --tests ...batch.ParDoTest`: 14 tests, 14 passed,
0 failed. I did not run the `runners:spark:4` variant locally (shares the
same `src/` sources via `runners/spark/4`); CI covers it.
### Notes
This only changes the Spark Structured Streaming runner's translation-time
bookkeeping; it does not change what gets computed or written, only when the
intermediate cached `Dataset` used to fan out a multi-output `ParDo` into its
per-tag `Dataset`s is released.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
