[#37994] Fix NullPointerException in Spark Runner with multiple outputs and serialization

Author: stankiewicz

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.translation;
 
+import java.util.Objects;
 import static org.apache.beam.runners.spark.translation.TranslationUtils.canAvoidRddSerialization;
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
@@ -486,6 +487,8 @@ public void evaluate(
                 TranslationUtils.getTupleTagCoders(outputs);
             all =
                 all.mapToPair(TranslationUtils.getTupleTagEncodeFunction(coderMap))
+                    .filter(Objects::nonNull) // skip nulls to save on encoding, nulls are tags that
+                    // are not read
                     .persist(level)
                     .mapToPair(TranslationUtils.getTupleTagDecodeFunction(coderMap));
           }

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java

@@ -445,8 +445,13 @@ public static Map<TupleTag<?>, Coder<WindowedValue<?>>> getTupleTagCoders(
     return tuple2 -> {
       TupleTag<?> tupleTag = tuple2._1;
       WindowedValue<?> windowedValue = tuple2._2;
-      return new Tuple2<>(
-          tupleTag, ValueAndCoderLazySerializable.of(windowedValue, coderMap.get(tupleTag)));
+      Coder<WindowedValue<?>> coder = coderMap.get(tupleTag);
+      if (coder == null) {
+        // there is no coder as this step is leaf step and is not read anywhere, so coder is pruned
+        // from coderMap
+        return null;
+      }
+      return new Tuple2<>(tupleTag, ValueAndCoderLazySerializable.of(windowedValue, coder));
     };
   }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StatefulStreamingParDoEvaluator.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.translation.streaming;
 
+import java.util.Objects;
 import static org.apache.beam.runners.spark.translation.TranslationUtils.getBatchDuration;
 import static org.apache.beam.runners.spark.translation.TranslationUtils.hasEventTimers;
 import static org.apache.beam.runners.spark.translation.TranslationUtils.hasTimers;
@@ -234,6 +235,7 @@ public void evaluate(
           TranslationUtils.getTupleTagCoders(outputs);
       all =
           all.mapToPair(TranslationUtils.getTupleTagEncodeFunction(coderMap))
+              .filter(Objects::nonNull) // skip nulls to save on encoding, nulls are tags that are not read
               .cache()
               .mapToPair(TranslationUtils.getTupleTagDecodeFunction(coderMap));
 

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.translation.streaming;
 
+import java.util.Objects;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
@@ -593,6 +594,7 @@ public void evaluate(
               TranslationUtils.getTupleTagCoders(outputs);
           all =
               all.mapToPair(TranslationUtils.getTupleTagEncodeFunction(coderMap))
+                .filter(Objects::nonNull) // skip nulls to save on encoding, nulls are tags that are not read
                   .cache()
                   .mapToPair(TranslationUtils.getTupleTagDecodeFunction(coderMap));
         }

runners/spark/src/test/java/org/apache/beam/runners/spark/translation/TransformTranslatorTest.java

@@ -40,6 +40,8 @@
 import org.apache.beam.sdk.coders.VarIntCoder;
 import org.apache.beam.sdk.transforms.Count;
 import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
@@ -48,6 +50,7 @@
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
@@ -247,4 +250,65 @@ public void testMultipleOutputParDoShouldHaveFilterWhenSideOutputIsConsumed() {
       assertTrue(parsed.stream().anyMatch(e -> e.getName().contains(tag.getId())));
     }
   }
+
+  @Test
+  public void testMultipleOutputParDoWithUnconsumedSideOutputAndSerializationStorageLevel() {
+    Pipeline p = Pipeline.create();
+    TupleTag<String> tag1 = new TupleTag<String>("tag1") {};
+    TupleTag<String> tag2 = new TupleTag<String>("tag2") {};
+    TupleTag<String> tag3 = new TupleTag<String>("tag3") {};
+
+    SparkPipelineOptions options = contextRule.createPipelineOptions();
+    // Force serialization by setting storage level to MEMORY_AND_DISK_SER
+    options.setStorageLevel("MEMORY_AND_DISK_SER");
+
+    TransformTranslator.Translator translator = new TransformTranslator.Translator();
+
+    PTransform<PBegin, PCollection<String>> createTransform = Create.of("foo", "bar");
+
+    PCollectionTuple pCollectionTuple =
+        p.apply("Create Values", createTransform)
+            .apply(
+                "Multiple Output ParDo",
+                ParDo.of(new MultiOutputDoFn(tag1, tag2, tag3))
+                    .withOutputTags(tag1, TupleTagList.of(tag2).and(tag3)));
+
+    // consume tag1 and tag2
+    pCollectionTuple.get(tag1).apply("Count1", Count.globally());
+    pCollectionTuple.get(tag2).apply("Count2", Count.globally());
+
+    p.replaceAll(SparkTransformOverrides.getDefaultOverrides(false));
+
+    EvaluationContext ctxt = new EvaluationContext(contextRule.getSparkContext(), p, options);
+    SparkRunner.initAccumulators(options, ctxt.getSparkContext());
+    SparkRunner.updateDependentTransforms(p, translator, ctxt);
+
+    // This should not throw NullPointerException
+    p.traverseTopologically(new SparkRunner.Evaluator(translator, ctxt));
+
+    // Also trigger some action on the RDD to ensure serialization happens
+    @SuppressWarnings("unchecked")
+    BoundedDataset<String> dataset =
+        (BoundedDataset<String>) ctxt.borrowDataset(pCollectionTuple.get(tag1));
+    dataset.getRDD().count();
+  }
+
+  private static class MultiOutputDoFn extends DoFn<String, String> {
+    private final TupleTag<String> tag1;
+    private final TupleTag<String> tag2;
+    private final TupleTag<String> tag3;
+
+    MultiOutputDoFn(TupleTag<String> tag1, TupleTag<String> tag2, TupleTag<String> tag3) {
+      this.tag1 = tag1;
+      this.tag2 = tag2;
+      this.tag3 = tag3;
+    }
+
+    @ProcessElement
+    public void process(@Element String input, MultiOutputReceiver outputReceiver) {
+      outputReceiver.get(tag1).output(input);
+      outputReceiver.get(tag2).output(input);
+      outputReceiver.get(tag3).output(input);
+    }
+  }
 }