[#40949] Spark 4 streaming: stateless ParDo with additional outputs

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json

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

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -14,5 +14,6 @@
   "https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially",
   "https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test",
   "https://github.com/apache/beam/issues/36841": "Spark 4 streaming Impulse and PAssert",
-  "https://github.com/apache/beam/issues/40950": "Redistribute translated like Reshuffle"
+  "https://github.com/apache/beam/issues/40950": "Redistribute translated like Reshuffle",
+  "https://github.com/apache/beam/issues/40949": "ParDo with additional outputs in streaming"
 }

runners/spark/4/build.gradle

@@ -90,24 +90,12 @@ tasks.register("validatesStructuredStreamingRunnerStreaming", Test) {
     excludeTestsMatching 'org.apache.beam.sdk.testing.PAssertTest.testPAssertEqualsSingletonFailsForEmptyPCollection'
     excludeTestsMatching 'org.apache.beam.sdk.testing.PAssertTest.testPAssertEqualsSingletonFailsForNonSingletonPCollection'
 
-    // TODO(https://github.com/apache/beam/issues/40949): ParDo with additional outputs
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAllAttemptedMetrics'
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAttemptedBoundedTrieMetrics'
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAttemptedCounterMetrics'
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAttemptedDistributionMetrics'
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAttemptedGaugeMetrics'
-    excludeTestsMatching 'org.apache.beam.sdk.metrics.MetricsTest$AttemptedMetricTests.testAttemptedStringSetMetrics'
+    // TODO(https://github.com/apache/beam/issues/41199): hot key fanout computes a wrong result
     excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$BasicTests.testHotKeyCombining'
     excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$BasicTests.testHotKeyCombiningWithAccumulationMode'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultiplePCollectionsHavingMultipleConsumers'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoSchemaTest.testReadAndWriteMultiOutput'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$MultipleInputsAndOutputTests.testParDoEmptyWithTaggedOutput'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$MultipleInputsAndOutputTests.testParDoWithEmptyTaggedOutput'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$MultipleInputsAndOutputTests.testParDoWithOnlyTaggedOutput'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$MultipleInputsAndOutputTests.testParDoWithTaggedOutput'
+
+    // TODO(https://github.com/apache/beam/issues/40949): stateful ParDo with additional outputs
     excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$StateTests.testValueStateTaggedOutput'
-    excludeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest$TimestampTests.testParDoTaggedOutputWithTimestamp'
-    excludeTestsMatching 'org.apache.beam.sdk.values.PCollectionTupleTest.testComposePCollectionTuple'
 
     // TODO(https://github.com/apache/beam/issues/36841): PeriodicImpulse uses a splittable DoFn
     excludeTestsMatching 'org.apache.beam.sdk.transforms.PerKeyOrderingTest.testMultipleStatefulOrderingWithShuffle'

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

@@ -47,8 +47,8 @@
 
 /**
  * Pipeline translator for streaming pipelines on Spark 4. It extends the common registry to reuse
- * the stateless single output ParDo, Window.Assign, Flatten, Reshuffle and Combine.GroupedValues
- * translators, which are safe on a streaming Dataset. Unsupported primitives fail at translation.
+ * the stateless ParDo, Window.Assign, Flatten, Reshuffle and Combine.GroupedValues translators,
+ * which are safe on a streaming Dataset. Unsupported primitives fail at translation.
  */
 @Internal
 public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
@@ -103,6 +103,11 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
           "Combine.globally() with a default value (use withoutDefaults())" + NOT_SUPPORTED);
     }
 
+    if (transform instanceof Combine.PerKeyWithHotKeyFanout) {
+      throw new UnsupportedOperationException(
+          "Combine.perKey().withHotKeyFanout()" + NOT_SUPPORTED);
+    }
+
     if (transform instanceof Combine.PerKey || transform instanceof Combine.Globally) {
       return null;
     }
@@ -114,13 +119,6 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
         throw new UnsupportedOperationException(
             "ParDo with side inputs (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
       }
-      if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) {
-        throw new UnsupportedOperationException(
-            "ParDo with additional outputs ("
-                + signature.fnClass().getName()
-                + ")"
-                + NOT_SUPPORTED);
-      }
       if (signature.processElement().requiresTimeSortedInput()) {
         throw new UnsupportedOperationException(
             "@RequiresTimeSortedInput (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
@@ -132,6 +130,14 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
       checkNoProcessingTimeTimers(parDo.getFn(), signature);
 
       if (signature.usesState() || signature.usesTimers()) {
+        if (!parDo.getAdditionalOutputTags().getAll().isEmpty()) {
+          // BeamStatefulProcessor merges all output tags into the main output.
+          throw new UnsupportedOperationException(
+              "Stateful ParDo with additional outputs ("
+                  + signature.fnClass().getName()
+                  + ")"
+                  + NOT_SUPPORTED);
+        }
         return STATEFUL_PAR_DO;
       }
     }

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreamingTest.java

@@ -33,10 +33,12 @@
 import org.apache.beam.sdk.state.TimeDomain;
 import org.apache.beam.sdk.state.TimerSpec;
 import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Sum;
 import org.apache.beam.sdk.transforms.windowing.AfterPane;
 import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
@@ -89,6 +91,15 @@ private static void assertUnsupported(Pipeline pipeline, String expected) {
   @Test
   public void rejectsUnsupportedStatefulFeatures() throws Exception {
     Object[][] cases = {
+      {
+        kv("o")
+            .apply(
+                ParDo.of(new StatefulDoFn())
+                    .withOutputTags(
+                        new TupleTag<Integer>() {}, TupleTagList.of(new TupleTag<Integer>() {})))
+            .getPipeline(),
+        "Stateful ParDo with additional outputs"
+      },
       {kv("pt").apply(ParDo.of(new ProcessingTimeTimerDoFn())).getPipeline(), "PROCESSING_TIME"},
       {
         kv("spt").apply(ParDo.of(new SyncProcessingTimeTimerDoFn())).getPipeline(),
@@ -148,6 +159,12 @@ public void rejectsUnsupportedStatefulFeatures() throws Exception {
             .apply(Count.globally())
             .getPipeline(),
         "Combine.globally() with a default value (use withoutDefaults())"
+      },
+      {
+        kv("hk")
+            .apply(Combine.<String, Integer, Integer>perKey(Sum.ofIntegers()).withHotKeyFanout(2))
+            .getPipeline(),
+        "withHotKeyFanout()"
       }
     };
 
@@ -156,29 +173,13 @@ public void rejectsUnsupportedStatefulFeatures() throws Exception {
     }
   }
 
-  @Test
-  public void rejectsAdditionalOutputs() throws Exception {
-    TupleTag<Integer> main = new TupleTag<Integer>() {};
-    TupleTag<Integer> other = new TupleTag<Integer>() {};
-    assertUnsupported(
-        kv("o")
-            .apply(ParDo.of(new StatelessDoFn()).withOutputTags(main, TupleTagList.of(other)))
-            .getPipeline(),
-        "additional outputs");
-  }
-
   private static final class ToKvFn extends DoFn<String, KV<String, Integer>> {
     @ProcessElement
     public void process(@Element String element, OutputReceiver<KV<String, Integer>> out) {
       out.output(KV.of(element, 1));
     }
   }
 
-  private static final class StatelessDoFn extends DoFn<KV<String, Integer>, Integer> {
-    @ProcessElement
-    public void process() {}
-  }
-
   private static final class StatefulDoFn extends DoFn<KV<String, Integer>, Integer> {
     @DoFn.StateId("state")
     final StateSpec<?> spec = StateSpecs.value(VarIntCoder.of());

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

@@ -143,7 +143,10 @@ public void translate(ParDo.MultiOutput<InputT, OutputT> transform, Context cxt)
       // Persist as wide rows with one column per TupleTag to support different schemas
       Dataset<Tuple2<Integer, WindowedValue<Object>>> allTagsDS =
           cxt.getDataset(input).mapPartitions(doFnMapper, oneOfEncoder(encoders));
-      allTagsDS.persist(storageLevel);
+      // A streaming Dataset cannot be persisted, each consumed tag reruns the DoFn, see #41198.
+      if (!allTagsDS.isStreaming()) {
+        allTagsDS.persist(storageLevel);
+      }
 
       // divide into separate output datasets per tag
       for (TupleTag<?> tag : outputs.keySet()) {