From 1009d1f1f9f6c480f4973618b0f7f8fc3f4dca0a Mon Sep 17 00:00:00 2001 From: Akash Kumar <116457960+akashchamp@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:13:07 +0530 Subject: [PATCH] Unpersist cached Datasets after evaluation in Spark Structured Streaming runner ParDoTranslatorBatch persists an intermediate, wide-row Dataset when a multi-output ParDo needs to be divided into its per-tag Datasets, to avoid re-evaluating the parent mapPartitions transformation once per tag. It never called unpersist() on that cached Dataset, which the code even acknowledged with a "FIXME What's the strategy to unpersist Datasets / RDDs?" comment. Every multi-output ParDo therefore leaked its cached Dataset into Spark's storage memory for the lifetime of the SparkSession, causing unbounded growth, disk spilling and OOMs on large batch pipelines. Add PipelineTranslator.TranslationState#cacheDataset(Dataset, StorageLevel), a helper translators use instead of calling Dataset#persist directly when caching a dataset outside the normal putDataset path. Track every dataset cached this way and thread that collection through to EvaluationContext, which now unpersists all of them in a finally block once leaf evaluation has finished, whether it completed normally, was stopped early, or failed with an exception. ParDoTranslatorBatch now goes through cacheDataset instead of calling persist() directly, removing the FIXME. Fixes #40243 --- CHANGES.md | 1 + .../PipelineTranslatorStreaming.java | 4 +- .../StreamingEvaluationContext.java | 4 +- .../translation/EvaluationContext.java | 59 ++++++++++----- .../translation/PipelineTranslator.java | 25 ++++++- .../translation/TransformTranslator.java | 6 ++ .../batch/ParDoTranslatorBatch.java | 13 ++-- .../translation/EvaluationContextTest.java | 74 ++++++++++++++++++- .../translation/batch/ParDoTest.java | 13 +++- 9 files changed, 166 insertions(+), 33 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 24cbe75e6c71..15f39d1dc6bc 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -84,6 +84,7 @@ * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)). * (Python) Fixed `GeminiModelHandler` returning only one result per batch: each prompt is now sent as its own request (up to 32 at a time), so every input gets its own result. A custom `request_fn` must now return one response per input ([#40913](https://github.com/apache/beam/issues/40913)). +* (Java) Fixed the Spark Structured Streaming runner leaking cached intermediate `Dataset`s for multi-output `ParDo`s, which were persisted but never unpersisted in `ParDoTranslatorBatch` ([#40243](https://github.com/apache/beam/issues/40243)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java index 4d93741b6365..9c96524a14c0 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -42,6 +42,7 @@ import org.apache.beam.sdk.values.POutput; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; +import org.apache.spark.sql.Dataset; import org.apache.spark.sql.SparkSession; import org.checkerframework.checker.nullness.qual.Nullable; @@ -184,8 +185,9 @@ private static void checkNotMergingWindows( @Override protected EvaluationContext createEvaluationContext( Collection> leaves, + Collection> cachedDatasets, SparkSession session, SparkCommonPipelineOptions options) { - return new StreamingEvaluationContext(leaves, session, options); + return new StreamingEvaluationContext(leaves, cachedDatasets, session, options); } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java index b5057cd260ac..9954c30e5004 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -63,9 +63,10 @@ public class StreamingEvaluationContext extends EvaluationContext { StreamingEvaluationContext( Collection> leaves, + Collection> cachedDatasets, SparkSession session, SparkCommonPipelineOptions options) { - super(leaves, session); + super(leaves, cachedDatasets, session); this.options = options.as(SparkStructuredStreamingPipelineOptions.class); } @@ -124,6 +125,7 @@ public void evaluate() { if (idleStopListener != null) { getSparkSession().streams().removeListener(idleStopListener); } + unpersistCachedDatasets(); } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java index fb559ab8d9c1..ea078ef261ae 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java @@ -51,11 +51,16 @@ public interface NamedDataset { } private final Collection> leaves; + private final Collection> cachedDatasets; private final SparkSession session; private volatile boolean stopped = false; - protected EvaluationContext(Collection> leaves, SparkSession session) { + protected EvaluationContext( + Collection> leaves, + Collection> cachedDatasets, + SparkSession session) { this.leaves = leaves; + this.cachedDatasets = cachedDatasets; this.session = session; } @@ -64,24 +69,44 @@ protected Collection> leaves() { return leaves; } - /** Trigger evaluation of all leaf datasets. Returns early once {@link #stop()} was called. */ + /** + * Trigger evaluation of all leaf datasets. Returns early once {@link #stop()} was called. + * + *

Once all leaves have been evaluated (or evaluation was stopped), any dataset that was {@link + * PipelineTranslator.TranslationState#cacheDataset cached} during translation is unpersisted, so + * it doesn't keep occupying storage memory for the lifetime of the Spark session. + */ public void evaluate() { - for (NamedDataset ds : leaves) { - if (stopped) { - LOG.info("Evaluation stopped, skipping remaining datasets"); - return; - } - final Dataset dataset = ds.dataset(); - if (dataset == null) { - continue; - } - if (LOG.isDebugEnabled()) { - ExplainMode explainMode = ExplainMode.fromString("simple"); - String execPlan = dataset.queryExecution().explainString(explainMode); - LOG.debug("Evaluating dataset {}:\n{}", ds.name(), execPlan); + try { + for (NamedDataset ds : leaves) { + if (stopped) { + LOG.info("Evaluation stopped, skipping remaining datasets"); + return; + } + final Dataset dataset = ds.dataset(); + if (dataset == null) { + continue; + } + if (LOG.isDebugEnabled()) { + ExplainMode explainMode = ExplainMode.fromString("simple"); + String execPlan = dataset.queryExecution().explainString(explainMode); + LOG.debug("Evaluating dataset {}:\n{}", ds.name(), execPlan); + } + // force evaluation using a dummy foreach action + evaluate(ds.name(), dataset); } - // force evaluation using a dummy foreach action - evaluate(ds.name(), dataset); + } finally { + unpersistCachedDatasets(); + } + } + + /** + * Unpersists all datasets that were cached during translation. Called once leaf evaluation has + * completed (successfully, with an error, or stopped early), as none of them are needed anymore. + */ + protected final void unpersistCachedDatasets() { + for (Dataset dataset : cachedDatasets) { + dataset.unpersist(); } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java index 8470c968550b..7a0dd5b1e63c 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java @@ -129,7 +129,7 @@ public EvaluationContext translate( TranslatingVisitor translator = new TranslatingVisitor(session, options, dependencies.results); pipeline.traverseTopologically(translator); - return createEvaluationContext(translator.leaves, session, options); + return createEvaluationContext(translator.leaves, translator.cachedDatasets, session, options); } /** @@ -140,9 +140,10 @@ public EvaluationContext translate( */ protected EvaluationContext createEvaluationContext( Collection> leaves, + Collection> cachedDatasets, SparkSession session, SparkCommonPipelineOptions options) { - return new EvaluationContext(leaves, session); + return new EvaluationContext(leaves, cachedDatasets, session); } /** @@ -240,6 +241,17 @@ default void putDataset(PCollection pCollection, Dataset Broadcast> getSideInputBroadcast( PCollection pCollection, SideInputValues.Loader loader); + /** + * Persists {@code dataset} at the given {@link StorageLevel} and registers it to be unpersisted + * once the pipeline has been fully evaluated. + * + *

Translators that cache a dataset outside of {@link #putDataset} (for example to reuse it + * across multiple derived outputs within the same translation step) must use this method rather + * than calling {@link Dataset#persist(StorageLevel)} directly, so the cached data doesn't + * outlive the pipeline run it was created for. + */ + Dataset cacheDataset(Dataset dataset, StorageLevel level); + Supplier getOptionsSupplier(); PipelineOptions getOptions(); @@ -266,6 +278,7 @@ private class TranslatingVisitor extends PTransformVisitor implements Translatio private final StorageLevel storageLevel; private final Set> leaves; + private final Set> cachedDatasets; public TranslatingVisitor( SparkSession sparkSession, @@ -278,6 +291,7 @@ public TranslatingVisitor( this.storageLevel = StorageLevel.fromString(options.getStorageLevel()); this.encoders = new HashMap<>(); this.leaves = new HashSet<>(); + this.cachedDatasets = new HashSet<>(); } @Override @@ -385,6 +399,13 @@ public Broadcast> getSideInputBroadcast( return result.sideInputBroadcast; } + @Override + public Dataset cacheDataset(Dataset dataset, StorageLevel level) { + dataset.persist(level); + cachedDatasets.add(dataset); + return dataset; + } + @Override public Supplier getOptionsSupplier() { return optionsSupplier; diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TransformTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TransformTranslator.java index 24783040704e..0763c41c9ff6 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TransformTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/TransformTranslator.java @@ -45,6 +45,7 @@ import org.apache.spark.sql.Encoder; import org.apache.spark.sql.Encoders; import org.apache.spark.sql.SparkSession; +import org.apache.spark.storage.StorageLevel; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import scala.Tuple2; import scala.reflect.ClassTag; @@ -159,6 +160,11 @@ public void putDataset( state.putDataset(pCollection, dataset, cache); } + @Override + public Dataset cacheDataset(Dataset dataset, StorageLevel level) { + return state.cacheDataset(dataset, level); + } + @Override public void putUnresolved( PCollection out, UnresolvedTranslation unresolved) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java index b454fc1a34ef..a65ea4434306 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java @@ -135,15 +135,18 @@ public void translate(ParDo.MultiOutput transform, Context cxt) simple(cxt.getCurrentTransform(), input, sideInputReader, false), tagColIdx); - // FIXME What's the strategy to unpersist Datasets / RDDs? - SparkCommonPipelineOptions opts = cxt.getOptions().as(SparkCommonPipelineOptions.class); StorageLevel storageLevel = StorageLevel.fromString(opts.getStorageLevel()); - // Persist as wide rows with one column per TupleTag to support different schemas + // Persist as wide rows with one column per TupleTag to support different schemas. Caching + // avoids re-evaluating the mapPartitions above once per tag below. Persisting through + // cxt.cacheDataset (rather than calling persist directly) registers the dataset to be + // unpersisted once the pipeline has been fully evaluated, so it doesn't leak for the + // lifetime of the Spark session. Dataset>> allTagsDS = - cxt.getDataset(input).mapPartitions(doFnMapper, oneOfEncoder(encoders)); - allTagsDS.persist(storageLevel); + cxt.cacheDataset( + cxt.getDataset(input).mapPartitions(doFnMapper, oneOfEncoder(encoders)), + storageLevel); // divide into separate output datasets per tag for (TupleTag tag : outputs.keySet()) { diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContextTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContextTest.java index d1669e45c182..2a777a271918 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContextTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContextTest.java @@ -21,14 +21,23 @@ import static org.junit.Assert.assertThrows; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import java.util.Collections; +import org.apache.beam.sdk.values.WindowedValue; import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; import org.junit.Test; /** - * Unit tests for the static error-path branches of {@link EvaluationContext}. The happy-path - * branches are covered end-to-end by the structured-streaming translation tests; these tests - * specifically exercise the {@code catch} blocks that wrap and rethrow underlying Spark failures. + * Unit tests for {@link EvaluationContext}. The static-method tests specifically exercise the + * {@code catch} blocks that wrap and rethrow underlying Spark failures; the {@code evaluate()} + * tests cover the happy-path evaluation of leaf datasets end-to-end via the structured-streaming + * translation tests, and focus here on the cache cleanup added for + * https://github.com/apache/beam/issues/40243 (datasets cached with {@link + * PipelineTranslator.TranslationState#cacheDataset} must be unpersisted once evaluation of all + * leaves has finished, whether it succeeded, failed, or was stopped early). */ public class EvaluationContextTest { @@ -81,4 +90,63 @@ public void collectHandlesNullExceptionMessage() { assertThrows(RuntimeException.class, () -> EvaluationContext.collect("test-ds", ds)); assertSame(underlying, thrown.getCause()); } + + @Test + public void evaluateUnpersistsCachedDatasetsAfterLeavesAreEvaluated() { + // The per-leaf write/save mechanics (and that cached Datasets actually stop being reported + // as cached by Spark) are exercised end-to-end by ParDoTest. This test isolates the + // cache-cleanup control flow itself using a leaf with no dataset to evaluate, a case + // `evaluate()` already handles by skipping it -- `org.apache.spark.sql.DataFrameWriter` is a + // final class that Mockito cannot mock, so a leaf actually going through the write path + // cannot be stubbed here. + @SuppressWarnings("unchecked") + Dataset cached = mock(Dataset.class); + + EvaluationContext ctx = + new EvaluationContext( + Collections.singletonList(namedDataset(null)), + Collections.singletonList(cached), + mock(SparkSession.class)); + + ctx.evaluate(); + + verify(cached).unpersist(); + } + + @Test + public void evaluateUnpersistsCachedDatasetsEvenIfLeafEvaluationFails() { + // A failure evaluating one leaf must not leak datasets that were cached for other, unrelated + // parts of the pipeline: https://github.com/apache/beam/issues/40243. + @SuppressWarnings("unchecked") + Dataset> leaf = mock(Dataset.class); + when(leaf.write()).thenThrow(new RuntimeException("boom")); + + @SuppressWarnings("unchecked") + Dataset cached = mock(Dataset.class); + + EvaluationContext ctx = + new EvaluationContext( + Collections.singletonList(namedDataset(leaf)), + Collections.singletonList(cached), + mock(SparkSession.class)); + + assertThrows(RuntimeException.class, ctx::evaluate); + + verify(cached).unpersist(); + } + + private static EvaluationContext.NamedDataset namedDataset( + Dataset> dataset) { + return new EvaluationContext.NamedDataset() { + @Override + public String name() { + return "leaf"; + } + + @Override + public Dataset> dataset() { + return dataset; + } + }; + } } diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTest.java index 672a2db4fe1e..89af184af9c3 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTest.java @@ -65,15 +65,20 @@ public void testPardo() { } @Test - public void testPardoWithOutputTagsCachedRDD() { + public void testPardoWithOutputTagsUnpersistsRDDAfterRun() { + // Multiple consumed output tags force ParDoTranslatorBatch to persist an intermediate + // Dataset (regression test for https://github.com/apache/beam/issues/40243: it must not + // remain cached once the pipeline run has completed). pardoWithOutputTags("MEMORY_ONLY", true); - assertTrue("Expected cached data", SESSION.hasCachedData()); + assertTrue("Expected cached data to be unpersisted after run", !SESSION.hasCachedData()); } @Test - public void testPardoWithOutputTagsCachedDataset() { + public void testPardoWithOutputTagsUnpersistsDatasetAfterRun() { + // Same as above with a storage level that goes through the Dataset (rather than RDD) cache + // manager, see SparkSessionRule#hasCachedData. pardoWithOutputTags("MEMORY_AND_DISK", true); - assertTrue("Expected cached data", SESSION.hasCachedData()); + assertTrue("Expected cached data to be unpersisted after run", !SESSION.hasCachedData()); } @Test