[runners-spark] Support stateful ParDo in the Structured Streaming batch runner

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json

@@ -0,0 +1,4 @@
+{
+  "comment": "Modify this file in a trivial way to cause this test suite to run",
+  "modification": 1
+}

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -7,5 +7,6 @@
   "https://github.com/apache/beam/pull/34123": "noting that PR #34123 should run this test",
   "https://github.com/apache/beam/pull/34080": "noting that PR #34080 should run this test",
   "https://github.com/apache/beam/pull/34155": "noting that PR #34155 should run this test",
-  "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface"
+  "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface",
+  "https://github.com/apache/beam/pull/39793": "noting that PR #39793 should run this test"
 }

CHANGES.md

@@ -69,6 +69,7 @@
 ## New Features / Improvements
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* (Java) Spark Structured Streaming runner: stateful ParDo with state, timers, `@RequiresTimeSortedInput` and tagged outputs is now supported in batch mode ([#39779](https://github.com/apache/beam/issues/39779)).
 
 ## Breaking Changes
 

runners/spark/spark_runner.gradle

@@ -510,15 +510,9 @@ tasks.register("validatesStructuredStreamingRunnerBatch", Test) {
     excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedPCollections'
     excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream'
     // State and Timers
-    excludeCategories 'org.apache.beam.sdk.testing.UsesStatefulParDo'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesMapState'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesMultimapState'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesSetState'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesTimersInParDo'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap'
-    excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo'
     excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration'
+    // Every UsesOrderedListState test also uses @OnWindowExpiration, which is unsupported
+    excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState'
     // Metrics
     excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics'
     excludeCategories 'org.apache.beam.sdk.testing.UsesSystemMetrics'
@@ -532,6 +526,11 @@ tasks.register("validatesStructuredStreamingRunnerBatch", Test) {
     excludeCategories 'org.apache.beam.sdk.testing.UsesTriggeredSideInputs'
   }
   filter {
+    // These build on PeriodicImpulse, so the pipeline is unbounded and rejected by this batch only
+    // runner, but they are not categorized as UsesUnboundedPCollections. Excluded by name rather
+    // than adding that category upstream, which would also stop other runners running them.
+    excludeTestsMatching 'org.apache.beam.sdk.transforms.PerKeyOrderingTest.testMultipleStatefulOrderingWithShuffle'
+    excludeTestsMatching 'org.apache.beam.sdk.transforms.PerKeyOrderingTest.testMultipleStatefulOrderingWithoutShuffle'
     // Combine with context not implemented
     excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineFnsTest.testComposedCombineWithContext'
     excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$CombineWithContextTests.testSimpleCombineWithContext'

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

@@ -25,6 +25,7 @@
 import org.apache.beam.runners.core.DoFnRunner;
 import org.apache.beam.runners.core.DoFnRunners;
 import org.apache.beam.runners.core.SideInputReader;
+import org.apache.beam.runners.core.StepContext;
 import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
 import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.CachedSideInputReader;
 import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.NoOpStepContext;
@@ -71,6 +72,20 @@ interface DoFnRunnerWithTeardown<InT, T> extends DoFnRunner<InT, T> {
   abstract DoFnRunnerWithTeardown<InT, T> create(
       PipelineOptions options, MetricsAccumulator metrics, WindowedValueMultiReceiver output);
 
+  /**
+   * Creates a runner backed by {@code stepContext} so that state and timers are available.
+   *
+   * <p>Only supported for a single, unfused {@link DoFn}: a fused runner cannot drive timers.
+   */
+  DoFnRunnerWithTeardown<InT, T> create(
+      PipelineOptions options,
+      MetricsAccumulator metrics,
+      WindowedValueMultiReceiver output,
+      StepContext stepContext) {
+    throw new UnsupportedOperationException(
+        "Stateful execution is not supported by " + getClass().getSimpleName());
+  }
+
   /**
    * Fuses the factory for the following {@link DoFnRunner} into a single factory that processes
    * both DoFns in a single step.
@@ -128,6 +143,15 @@ <T2> DoFnRunnerFactory<InT, T2> fuse(DoFnRunnerFactory<T, T2> next) {
     @Override
     DoFnRunnerWithTeardown<InT, T> create(
         PipelineOptions options, MetricsAccumulator metrics, WindowedValueMultiReceiver output) {
+      return create(options, metrics, output, new NoOpStepContext());
+    }
+
+    @Override
+    DoFnRunnerWithTeardown<InT, T> create(
+        PipelineOptions options,
+        MetricsAccumulator metrics,
+        WindowedValueMultiReceiver output,
+        StepContext stepContext) {
       DoFnRunner<InT, T> simpleRunner =
           DoFnRunners.simpleRunner(
               options,
@@ -136,7 +160,7 @@ DoFnRunnerWithTeardown<InT, T> create(
               filterMainOutput ? new FilteredOutput<>(output, mainOutput) : output,
               mainOutput,
               additionalOutputs,
-              new NoOpStepContext(),
+              stepContext,
               coder,
               outputCoders,
               windowingStrategy,

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

@@ -35,6 +35,7 @@
 import org.apache.beam.runners.core.SideInputReader;
 import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
+import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator.TranslationState;
 import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator.UnresolvedTranslation;
 import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
 import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues;
@@ -64,8 +65,11 @@
  *
  * <p>Each tag is encoded as individual column with a respective schema & encoder each.
  *
+ * <p>Stateful {@link org.apache.beam.sdk.transforms.DoFn DoFns}, those using timers, and those
+ * annotated with {@link DoFn.RequiresTimeSortedInput} are translated by {@link
+ * StatefulParDoTranslatorBatch} instead.
+ *
  * <p>TODO:
- * <li>Add support for state and timers.
  * <li>Add support for SplittableDoFn
  */
 class ParDoTranslatorBatch<InputT, OutputT>
@@ -87,18 +91,18 @@ public boolean canTranslate(ParDo.MultiOutput<InputT, OutputT> transform) {
         "Not expected to directly translate splittable DoFn, should have been overridden: %s",
         doFn);
 
-    // TODO: add support of states and timers
+    // Stateful, timer using and time sorted DoFns are routed to StatefulParDoTranslatorBatch by
+    // PipelineTranslatorBatch#getTransformTranslator. Reaching here with one means dispatch is
+    // broken, not that the feature is unsupported.
     checkState(
-        !signature.usesState() && !signature.usesTimers(),
-        "States and timers are not supported for the moment.");
+        !StatefulParDoTranslatorBatch.appliesTo(transform),
+        "Stateful / time sorted DoFn should have been translated by %s: %s",
+        StatefulParDoTranslatorBatch.class.getSimpleName(),
+        doFn);
 
     checkState(
         signature.onWindowExpiration() == null, "onWindowExpiration is not supported: %s", doFn);
 
-    checkState(
-        !signature.processElement().requiresTimeSortedInput(),
-        "@RequiresTimeSortedInput is not supported for the moment");
-
     SparkSideInputReader.validateMaterializations(transform.getSideInputs().values());
     return true;
   }
@@ -211,11 +215,11 @@ public Dataset<WindowedValue<T>> resolve(
    * <p>This can help to avoid unnecessary caching in case of multiple outputs if only {@code
    * mainTag} is consumed.
    */
-  private Map<TupleTag<?>, PCollection<?>> skipUnconsumedOutputs(
+  static Map<TupleTag<?>, PCollection<?>> skipUnconsumedOutputs(
       Map<TupleTag<?>, PCollection<?>> outputs,
       TupleTag<?> mainTag,
       TupleTagList otherTags,
-      Context cxt) {
+      TranslationState cxt) {
     switch (outputs.size()) {
       case 1:
         return outputs; // always keep main output
@@ -235,7 +239,7 @@ private Map<TupleTag<?>, PCollection<?>> skipUnconsumedOutputs(
     }
   }
 
-  private Map<String, Integer> tagsColumnIndex(Collection<TupleTag<?>> tags) {
+  static Map<String, Integer> tagsColumnIndex(Collection<TupleTag<?>> tags) {
     Map<String, Integer> index = Maps.newHashMapWithExpectedSize(tags.size());
     for (TupleTag<?> tag : tags) {
       index.put(tag.getId(), index.size());

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

@@ -82,11 +82,27 @@ public class PipelineTranslatorBatch extends PipelineTranslator {
         SplittableParDo.PrimitiveBoundedRead.class, new ReadSourceTranslatorBatch<>());
   }
 
+  /**
+   * Translators that shadow the {@link #TRANSFORM_TRANSLATORS} entry for their transform class when
+   * a predicate matches, so that a single transform class can be translated in more than one way
+   * depending on the transform instance.
+   *
+   * <p>Currently only {@link ParDo.MultiOutput} needs this, to route stateful and time sorted
+   * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link ParDoTranslatorBatch}.
+   */
+  @SuppressWarnings("rawtypes")
+  private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR =
+      new StatefulParDoTranslatorBatch<>();
+
   /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */
   @Override
   @Nullable
   protected <InT extends PInput, OutT extends POutput, TransformT extends PTransform<InT, OutT>>
       TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT transform) {
+    if (transform instanceof ParDo.MultiOutput
+        && StatefulParDoTranslatorBatch.appliesTo((ParDo.MultiOutput<?, ?>) transform)) {
+      return STATEFUL_PARDO_TRANSLATOR;
+    }
     return TRANSFORM_TRANSLATORS.get(transform.getClass());
   }
 }