diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..f22b48bf32b2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -141,6 +141,7 @@ * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). +* (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 6f2d14a69b8a..d5024d0bba90 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; @@ -194,8 +195,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 2ca54a1e9481..16747c54931b 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 @@ -62,9 +62,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); } @@ -123,6 +124,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