[Bug]: Unpersist cached Datasets after evaluation in Spark Structured Streaming runner

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

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<? extends EvaluationContext.NamedDataset<?>> leaves,
+      Collection<Dataset<?>> cachedDatasets,
       SparkSession session,
       SparkCommonPipelineOptions options) {
-    return new StreamingEvaluationContext(leaves, session, options);
+    return new StreamingEvaluationContext(leaves, cachedDatasets, session, options);
   }
 }

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<? extends NamedDataset<?>> leaves,
+      Collection<Dataset<?>> 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();
     }
   }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java

@@ -51,11 +51,16 @@ public interface NamedDataset<T> {
   }
 
   private final Collection<? extends NamedDataset<?>> leaves;
+  private final Collection<Dataset<?>> cachedDatasets;
   private final SparkSession session;
   private volatile boolean stopped = false;
 
-  protected EvaluationContext(Collection<? extends NamedDataset<?>> leaves, SparkSession session) {
+  protected EvaluationContext(
+      Collection<? extends NamedDataset<?>> leaves,
+      Collection<Dataset<?>> cachedDatasets,
+      SparkSession session) {
     this.leaves = leaves;
+    this.cachedDatasets = cachedDatasets;
     this.session = session;
   }
 
@@ -64,24 +69,44 @@ protected Collection<? extends NamedDataset<?>> 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.
+   *
+   * <p>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();
     }
   }
 

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<? extends EvaluationContext.NamedDataset<?>> leaves,
+      Collection<Dataset<?>> cachedDatasets,
       SparkSession session,
       SparkCommonPipelineOptions options) {
-    return new EvaluationContext(leaves, session);
+    return new EvaluationContext(leaves, cachedDatasets, session);
   }
 
   /**
@@ -240,6 +241,17 @@ default <T> void putDataset(PCollection<T> pCollection, Dataset<WindowedValue<T>
     <T> Broadcast<SideInputValues<T>> getSideInputBroadcast(
         PCollection<T> pCollection, SideInputValues.Loader<T> loader);
 
+    /**
+     * Persists {@code dataset} at the given {@link StorageLevel} and registers it to be unpersisted
+     * once the pipeline has been fully evaluated.
+     *
+     * <p>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.
+     */
+    <T> Dataset<T> cacheDataset(Dataset<T> dataset, StorageLevel level);
+
     Supplier<PipelineOptions> getOptionsSupplier();
 
     PipelineOptions getOptions();
@@ -266,6 +278,7 @@ private class TranslatingVisitor extends PTransformVisitor implements Translatio
     private final StorageLevel storageLevel;
 
     private final Set<TranslationResult<?, ?>> leaves;
+    private final Set<Dataset<?>> 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 <T> Broadcast<SideInputValues<T>> getSideInputBroadcast(
       return result.sideInputBroadcast;
     }
 
+    @Override
+    public <T> Dataset<T> cacheDataset(Dataset<T> dataset, StorageLevel level) {
+      dataset.persist(level);
+      cachedDatasets.add(dataset);
+      return dataset;
+    }
+
     @Override
     public Supplier<PipelineOptions> getOptionsSupplier() {
       return optionsSupplier;

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 <T> void putDataset(
       state.putDataset(pCollection, dataset, cache);
     }
 
+    @Override
+    public <T> Dataset<T> cacheDataset(Dataset<T> dataset, StorageLevel level) {
+      return state.cacheDataset(dataset, level);
+    }
+
     @Override
     public <InputT, T> void putUnresolved(
         PCollection<T> out, UnresolvedTranslation<InputT, T> unresolved) {

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java

@@ -135,15 +135,18 @@ public void translate(ParDo.MultiOutput<InputT, OutputT> 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<Tuple2<Integer, WindowedValue<Object>>> 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()) {

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<Object> 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<WindowedValue<Object>> leaf = mock(Dataset.class);
+    when(leaf.write()).thenThrow(new RuntimeException("boom"));
+
+    @SuppressWarnings("unchecked")
+    Dataset<Object> 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<Object> namedDataset(
+      Dataset<WindowedValue<Object>> dataset) {
+    return new EvaluationContext.NamedDataset<Object>() {
+      @Override
+      public String name() {
+        return "leaf";
+      }
+
+      @Override
+      public Dataset<WindowedValue<Object>> dataset() {
+        return dataset;
+      }
+    };
+  }
 }

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