[#36841] Spark 4 streaming: windowed GroupByKey and Combine via transformWithState

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -12,5 +12,5 @@
   "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/pull/40281": "noting that PR #40281 should run this test",
-  "https://github.com/apache/beam/issues/36841": "Spark 4 streaming null keys"
+  "https://github.com/apache/beam/issues/36841": "Spark 4 streaming GroupByKey"
 }

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

@@ -20,6 +20,7 @@
 import java.util.Collection;
 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.GroupByKeyStreamingTranslator;
 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;
@@ -33,20 +34,21 @@
 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.transforms.windowing.DefaultTrigger;
 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.beam.sdk.values.WindowingStrategy.AccumulationMode;
 import org.apache.spark.sql.SparkSession;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Pipeline translator for streaming pipelines on Spark 4. It extends the common registry to reuse
- * the stateless single output ParDo, Window.Assign, Flatten and Reshuffle translators, which are
- * safe on a streaming Dataset. Every other primitive fails at translation, the batch translators
- * for them persist or collect the Dataset, which Spark rejects for streaming plans.
+ * 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.
  */
 @Internal
 public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
@@ -61,7 +63,18 @@ public class PipelineTranslatorStreaming extends PipelineTranslatorCommon {
         protected void translate(
             ParDo.MultiOutput<KV<Object, Object>, Object> transform, Context cxt) {
           PCollection<?> input = (PCollection<?>) cxt.getInput();
-          checkNotMergingWindows(input.getWindowingStrategy());
+          checkNotMergingWindows(input.getWindowingStrategy(), "Stateful ParDo");
+          super.translate(transform, cxt);
+        }
+      };
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static final TransformTranslator GROUP_BY_KEY =
+      new GroupByKeyStreamingTranslator<Object, Object>() {
+        @Override
+        protected void translate(GroupByKey<Object, Object> transform, Context cxt) {
+          PCollection<?> input = (PCollection<?>) cxt.getInput();
+          checkSupportedGroupByKey(input);
           super.translate(transform, cxt);
         }
       };
@@ -92,11 +105,17 @@ TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT tra
     }
 
     if (transform instanceof GroupByKey) {
-      throw new UnsupportedOperationException("GroupByKey" + NOT_SUPPORTED);
+      return GROUP_BY_KEY;
+    }
+
+    if (transform instanceof Combine.Globally
+        && ((Combine.Globally<?, ?>) transform).isInsertDefault()) {
+      throw new UnsupportedOperationException(
+          "Combine.globally() with a default value (use withoutDefaults())" + NOT_SUPPORTED);
     }
 
-    if (transform instanceof Combine.PerKey) {
-      throw new UnsupportedOperationException("Combine.perKey" + NOT_SUPPORTED);
+    if (transform instanceof Combine.PerKey || transform instanceof Combine.Globally) {
+      return null;
     }
 
     if (transform instanceof ParDo.MultiOutput) {
@@ -153,10 +172,22 @@ private static void checkNoProcessingTimeTimers(DoFn<?, ?> doFn, DoFnSignature s
     }
   }
 
-  private static void checkNotMergingWindows(WindowingStrategy<?, ?> windowingStrategy) {
-    if (windowingStrategy.needsMerge()) {
+  private static void checkSupportedGroupByKey(PCollection<?> input) {
+    WindowingStrategy<?, ?> strategy = input.getWindowingStrategy();
+    checkNotMergingWindows(strategy, "GroupByKey");
+    if (!(strategy.getTrigger() instanceof DefaultTrigger)) {
       throw new UnsupportedOperationException(
-          "Stateful ParDo over merging windows" + NOT_SUPPORTED);
+          "Custom triggers (" + strategy.getTrigger() + ")" + NOT_SUPPORTED);
+    }
+    if (strategy.getMode() != AccumulationMode.DISCARDING_FIRED_PANES) {
+      throw new UnsupportedOperationException("Accumulating fired panes" + NOT_SUPPORTED);
+    }
+  }
+
+  private static void checkNotMergingWindows(
+      WindowingStrategy<?, ?> windowingStrategy, String feature) {
+    if (windowingStrategy.needsMerge()) {
+      throw new UnsupportedOperationException(feature + " over merging windows" + NOT_SUPPORTED);
     }
   }
 

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

@@ -0,0 +1,72 @@
+/*
+ * 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.coders.CoderHelpers;
+import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
+import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.GroupByKeyStatefulProcessor;
+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.GroupByKey;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+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.Encoders;
+import org.apache.spark.sql.streaming.OutputMode;
+import org.apache.spark.sql.streaming.TimeMode;
+
+/** Translates {@link GroupByKey} for Spark 4 Structured Streaming via transformWithState. */
+@Internal
+public class GroupByKeyStreamingTranslator<K, V>
+    extends TransformTranslator<
+        PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupByKey<K, V>> {
+
+  public GroupByKeyStreamingTranslator() {
+    super(0.2f);
+  }
+
+  @Override
+  protected void translate(GroupByKey<K, V> transform, Context cxt) {
+    PCollection<KV<K, V>> input = cxt.getInput();
+    WindowingStrategy<?, ?> windowing = input.getWindowingStrategy();
+    KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
+    Coder<K> keyCoder = inputCoder.getKeyCoder();
+    PCollection<KV<K, Iterable<V>>> output = cxt.getOutput();
+    Coder<KV<K, Iterable<V>>> outputCoder = output.getCoder();
+
+    GroupByKeyStatefulProcessor<K, V> processor =
+        new GroupByKeyStatefulProcessor<>(inputCoder, windowing, cxt.getOptionsSupplier());
+
+    MapFunction<WindowedValue<KV<K, V>>, byte[]> keyFn =
+        v -> CoderHelpers.toByteArray(v.getValue().getKey(), keyCoder);
+    Dataset<WindowedValue<KV<K, Iterable<V>>>> result =
+        cxt.getDataset(input)
+            .groupByKey(keyFn, Encoders.BINARY())
+            .transformWithState(
+                processor,
+                TimeMode.EventTime(),
+                OutputMode.Append(),
+                cxt.windowedEncoder(outputCoder));
+
+    cxt.putDataset(output, result);
+  }
+}

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

@@ -22,7 +22,7 @@
 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.runners.spark.structuredstreaming.translation.streaming.state.ParDoStatefulProcessor;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
@@ -68,8 +68,8 @@ protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform, Context
             cxt.getCurrentTransform(), input, SparkSideInputReader.empty(), false);
     MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession());
 
-    BeamStatefulProcessor<K, V, OutputT> processor =
-        new BeamStatefulProcessor<>(
+    ParDoStatefulProcessor<K, V, OutputT> processor =
+        new ParDoStatefulProcessor<>(
             doFn, runnerFactory, inputCoder, windowing, cxt.getOptionsSupplier(), metrics);
 
     MapFunction<WindowedValue<KV<K, V>>, byte[]> keyFn =