[Spark 4] Support stateful ParDo in Structured Streaming via transformWithState

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -10,5 +10,6 @@
   "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",
   "https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test",
-  "https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially"
+  "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"
 }

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

@@ -21,17 +21,24 @@
 import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorCommon;
 import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator;
+import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.StatefulParDoStreamingTranslator;
 import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.TimerSpec;
 import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.Impulse;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
 import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
 import org.apache.beam.sdk.util.construction.SplittableParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PInput;
 import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.spark.sql.SparkSession;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -47,6 +54,18 @@ public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
   @SuppressWarnings("rawtypes")
   private static final TransformTranslator READ_UNBOUNDED = new ReadUnboundedTranslator<>();
 
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static final TransformTranslator STATEFUL_PAR_DO =
+      new StatefulParDoStreamingTranslator<Object, Object, Object>() {
+        @Override
+        protected void translate(
+            ParDo.MultiOutput<KV<Object, Object>, Object> transform, Context cxt) {
+          PCollection<?> input = (PCollection<?>) cxt.getInput();
+          checkNotMergingWindows(input.getWindowingStrategy());
+          super.translate(transform, cxt);
+        }
+      };
+
   private static final String NOT_SUPPORTED =
       " is not supported by the Spark 4 streaming runner yet, see"
           + " https://github.com/apache/beam/issues/36841";
@@ -83,10 +102,6 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
     if (transform instanceof ParDo.MultiOutput) {
       ParDo.MultiOutput<?, ?> parDo = (ParDo.MultiOutput<?, ?>) transform;
       DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn());
-      if (signature.usesState() || signature.usesTimers()) {
-        throw new UnsupportedOperationException(
-            "Stateful ParDo (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
-      }
       if (!parDo.getSideInputs().isEmpty()) {
         throw new UnsupportedOperationException(
             "ParDo with side inputs (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
@@ -98,11 +113,53 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
                 + ")"
                 + NOT_SUPPORTED);
       }
+      if (signature.processElement().requiresTimeSortedInput()) {
+        throw new UnsupportedOperationException(
+            "@RequiresTimeSortedInput (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
+      }
+      if (signature.onWindowExpiration() != null) {
+        throw new UnsupportedOperationException(
+            "@OnWindowExpiration (" + signature.fnClass().getName() + ")" + NOT_SUPPORTED);
+      }
+      checkNoProcessingTimeTimers(parDo.getFn(), signature);
+
+      if (signature.usesState() || signature.usesTimers()) {
+        return STATEFUL_PAR_DO;
+      }
     }
 
     return super.getTransformTranslator(transform);
   }
 
+  private static void checkNoProcessingTimeTimers(DoFn<?, ?> doFn, DoFnSignature signature) {
+    for (DoFnSignature.TimerDeclaration timer : signature.timerDeclarations().values()) {
+      TimerSpec spec = DoFnSignatures.getTimerSpecOrThrow(timer, doFn);
+      if (spec.getTimeDomain() != TimeDomain.EVENT_TIME) {
+        throw new UnsupportedOperationException(
+            spec.getTimeDomain() + " timer @TimerId(\"" + timer.id() + "\")" + NOT_SUPPORTED);
+      }
+    }
+    for (DoFnSignature.TimerFamilyDeclaration family :
+        signature.timerFamilyDeclarations().values()) {
+      TimerSpec spec = DoFnSignatures.getTimerFamilySpecOrThrow(family, doFn);
+      if (spec.getTimeDomain() != TimeDomain.EVENT_TIME) {
+        throw new UnsupportedOperationException(
+            spec.getTimeDomain()
+                + " timer family @TimerFamily(\""
+                + family.id()
+                + "\")"
+                + NOT_SUPPORTED);
+      }
+    }
+  }
+
+  private static void checkNotMergingWindows(WindowingStrategy<?, ?> windowingStrategy) {
+    if (windowingStrategy.needsMerge()) {
+      throw new UnsupportedOperationException(
+          "Stateful ParDo over merging windows" + NOT_SUPPORTED);
+    }
+  }
+
   @Override
   protected EvaluationContext createEvaluationContext(
       Collection<? extends EvaluationContext.NamedDataset<?>> leaves,

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

@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
+
+import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
+import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
+import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory;
+import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader;
+import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessor;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+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.WindowedValue;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.spark.api.java.function.MapFunction;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Encoder;
+import org.apache.spark.sql.streaming.OutputMode;
+import org.apache.spark.sql.streaming.TimeMode;
+
+/** Translates a stateful {@link ParDo.MultiOutput} for Spark 4 Structured Streaming. */
+@Internal
+public class StatefulParDoStreamingTranslator<K, V, OutputT>
+    extends TransformTranslator<
+        PCollection<? extends KV<K, V>>, PCollectionTuple, ParDo.MultiOutput<KV<K, V>, OutputT>> {
+
+  public StatefulParDoStreamingTranslator() {
+    super(0.2f);
+  }
+
+  @Override
+  protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform, Context cxt) {
+    @SuppressWarnings("unchecked")
+    PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) cxt.getInput();
+    WindowingStrategy<?, ?> windowing = input.getWindowingStrategy();
+    DoFn<KV<K, V>, OutputT> doFn = transform.getFn();
+
+    KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
+    Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder);
+    TupleTag<OutputT> mainOutputTag = transform.getMainOutputTag();
+    PCollection<OutputT> output = cxt.getOutput(mainOutputTag);
+    Coder<OutputT> outputCoder = output.getCoder();
+
+    DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory =
+        DoFnRunnerFactory.simple(
+            cxt.getCurrentTransform(), input, SparkSideInputReader.empty(), false);
+    MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession());
+
+    BeamStatefulProcessor<K, V, OutputT> processor =
+        new BeamStatefulProcessor<>(
+            doFn, runnerFactory, inputCoder, windowing, cxt.getOptionsSupplier(), metrics);
+
+    MapFunction<WindowedValue<KV<K, V>>, K> keyFn = v -> v.getValue().getKey();
+    Dataset<WindowedValue<OutputT>> result =
+        cxt.getDataset(input)
+            .groupByKey(keyFn, keyEnc)
+            .transformWithState(
+                processor,
+                TimeMode.EventTime(),
+                OutputMode.Append(),
+                cxt.windowedEncoder(outputCoder));
+
+    cxt.putDataset(output, result);
+  }
+}