[Spark 4] Group the streaming stateful ParDo by encoded key bytes to support null keys

.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": 3
+  "modification": 4
 }

.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 watermark lag"
+  "https://github.com/apache/beam/issues/36841": "Spark 4 streaming null keys"
 }

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java

@@ -27,7 +27,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.LongSupplier;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
 import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;

runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java

@@ -29,9 +29,9 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache.CachedReader;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;

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

@@ -18,13 +18,13 @@
 package org.apache.beam.runners.spark.structuredstreaming.translation.batch;
 
 import static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY;
+import static org.apache.beam.runners.spark.coders.CoderHelpers.toByteArray;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGlobalGroupBy;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGroupByWindow;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.explodeWindowedKey;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueKey;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueValue;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.windowedKV;
-import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers.toByteArray;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.collectionEncoder;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.encoderOf;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.kvEncoder;

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

@@ -36,6 +36,7 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;

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

@@ -19,10 +19,10 @@
 
 import static org.apache.spark.sql.functions.col;
 
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset;
 import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.structuredstreaming.translation.streaming;
 
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 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;
@@ -35,7 +36,7 @@
 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.Encoders;
 import org.apache.spark.sql.streaming.OutputMode;
 import org.apache.spark.sql.streaming.TimeMode;
 
@@ -57,7 +58,7 @@ protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform, Context
     DoFn<KV<K, V>, OutputT> doFn = transform.getFn();
 
     KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();
-    Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder);
+    Coder<K> keyCoder = inputCoder.getKeyCoder();
     TupleTag<OutputT> mainOutputTag = transform.getMainOutputTag();
     PCollection<OutputT> output = cxt.getOutput(mainOutputTag);
     Coder<OutputT> outputCoder = output.getCoder();
@@ -71,10 +72,11 @@ protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform, Context
         new BeamStatefulProcessor<>(
             doFn, runnerFactory, inputCoder, windowing, cxt.getOptionsSupplier(), metrics);
 
-    MapFunction<WindowedValue<KV<K, V>>, K> keyFn = v -> v.getValue().getKey();
+    MapFunction<WindowedValue<KV<K, V>>, byte[]> keyFn =
+        v -> CoderHelpers.toByteArray(v.getValue().getKey(), keyCoder);
     Dataset<WindowedValue<OutputT>> result =
         cxt.getDataset(input)
-            .groupByKey(keyFn, keyEnc)
+            .groupByKey(keyFn, Encoders.BINARY())
             .transformWithState(
                 processor,
                 TimeMode.EventTime(),

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

@@ -41,6 +41,7 @@
 import org.apache.beam.runners.spark.structuredstreaming.translation.batch.StatefulTaskRunner;
 import org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.ListCoder;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.state.TimeDomain;
@@ -68,15 +69,15 @@
 
 /** Runs a stateful {@link DoFn} within Spark 4 Structured Streaming {@code transformWithState}. */
 public class BeamStatefulProcessor<K, V, OutputT>
-    extends StatefulProcessor<K, WindowedValue<KV<K, V>>, WindowedValue<OutputT>> {
+    extends StatefulProcessor<byte[], WindowedValue<KV<K, V>>, WindowedValue<OutputT>> {
 
   private static final String BEAM_STATE = "beamState";
   private static final String BEAM_TIMERS = "beamTimers";
   private static final String BEAM_WATERMARK = "beamWatermark";
 
   private final DoFn<KV<K, V>, OutputT> doFn;
   private final DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory;
-  private final Coder<KV<K, V>> inputCoder;
+  private final KvCoder<K, V> inputCoder;
   private final WindowingStrategy<?, ?> windowingStrategy;
   private final Supplier<PipelineOptions> optionsSupplier;
   private final MetricsAccumulator metrics;
@@ -93,7 +94,7 @@ public class BeamStatefulProcessor<K, V, OutputT>
   public BeamStatefulProcessor(
       DoFn<KV<K, V>, OutputT> doFn,
       DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory,
-      Coder<KV<K, V>> inputCoder,
+      KvCoder<K, V> inputCoder,
       WindowingStrategy<?, ?> windowingStrategy,
       Supplier<PipelineOptions> optionsSupplier,
       MetricsAccumulator metrics) {
@@ -138,7 +139,8 @@ private SparkTimerInternals timerInternals(TimerValues timerValues) {
 
   @Override
   public Iterator<WindowedValue<OutputT>> handleInputRows(
-      K key, Iterator<WindowedValue<KV<K, V>>> rows, TimerValues timerValues) {
+      byte[] keyBytes, Iterator<WindowedValue<KV<K, V>>> rows, TimerValues timerValues) {
+    K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
     List<WindowedValue<OutputT>> outputs = new ArrayList<>();
     SparkTimerInternals timerInternals = timerInternals(timerValues);
 
@@ -159,7 +161,8 @@ public Iterator<WindowedValue<OutputT>> handleInputRows(
 
   @Override
   public Iterator<WindowedValue<OutputT>> handleExpiredTimer(
-      K key, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) {
+      byte[] keyBytes, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) {
+    K key = CoderHelpers.fromByteArray(keyBytes, inputCoder.getKeyCoder());
     List<WindowedValue<OutputT>> outputs = new ArrayList<>();
     SparkTimerInternals timerInternals = timerInternals(timerValues);
     timerInternals.advanceWatermark();

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java

@@ -46,6 +46,7 @@
 import javax.annotation.Nullable;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.spark.StreamingTest;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule;
 import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamInputPartition;
@@ -60,7 +61,6 @@
 import org.apache.beam.sdk.io.UnboundedSource;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.hadoop.conf.Configuration;
@@ -209,7 +209,7 @@ public void testReadsElementsFromUnboundedSource() throws Exception {
     List<String> values = new ArrayList<>();
     for (List<Row> batch : batches(tag)) {
       for (Row row : batch) {
-        WindowedValue<String> value = CoderUtils.decodeFromByteArray(CODER, row.getAs(COL_PAYLOAD));
+        WindowedValue<String> value = CoderHelpers.fromByteArray(row.getAs(COL_PAYLOAD), CODER);
         values.add(value.getValue());
         assertEquals(
             value.getTimestamp().getMillis(), row.<Timestamp>getAs(COL_EVENT_TS).getTime());
@@ -593,11 +593,7 @@ private static List<String> values(List<List<Row>> batches) {
   }
 
   private static String decode(byte[] payload) {
-    try {
-      return CoderUtils.decodeFromByteArray(CODER, payload).getValue();
-    } catch (IOException e) {
-      throw new IllegalStateException(e);
-    }
+    return CoderHelpers.fromByteArray(payload, CODER).getValue();
   }
 
   private static List<Integer> nonEmptySizes(List<List<Row>> batches) {
@@ -748,7 +744,7 @@ private static int committedPosition(File checkpointDir, int shard) throws IOExc
         new BeamSourceCheckpoint(sourceDir(checkpointDir).getAbsolutePath(), new Configuration());
     byte[] coded = checkpoint.readMark(shard, epoch);
     assertNotNull("no mark at committed epoch " + epoch + " for shard " + shard, coded);
-    return CoderUtils.decodeFromByteArray(TestUnboundedSource.MARK_CODER, coded).next;
+    return CoderHelpers.fromByteArray(coded, TestUnboundedSource.MARK_CODER).next;
   }
 
   // ---------------------------------------------------------------------------------------------

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

@@ -33,6 +33,7 @@
 import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VoidCoder;
 import org.apache.beam.sdk.io.Read;
 import org.apache.beam.sdk.state.StateSpec;
 import org.apache.beam.sdk.state.StateSpecs;
@@ -43,6 +44,7 @@
 import org.apache.beam.sdk.state.ValueState;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
@@ -69,12 +71,14 @@ public class StatefulParDoStreamingTest implements Serializable {
   private static final String TAG_BATCHES = "stateful-batches";
   private static final String TAG_RESTART = "stateful-restart";
   private static final String TAG_CHAINED = "stateful-chained";
+  private static final String TAG_NULL_KEY = "stateful-null-key";
 
   @After
   public void tearDown() {
     TestUnboundedSource.forget(TAG_BATCHES);
     TestUnboundedSource.forget(TAG_RESTART);
     TestUnboundedSource.forget(TAG_CHAINED);
+    TestUnboundedSource.forget(TAG_NULL_KEY);
   }
 
   private static Pipeline pipeline(
@@ -167,6 +171,25 @@ public void testChainedStatefulParDoTimerOutputSurvives() throws Exception {
     assertTrue(collected.contains("timer-output"));
   }
 
+  @Test
+  public void testNullKey() throws Exception {
+    String collectorId = StreamingTestUtils.newCollectorId(TAG_NULL_KEY);
+    SparkStructuredStreamingPipelineOptions options =
+        StreamingTestUtils.streamingOptions(tempFolder);
+    options.setMaxRecordsPerBatch(1L);
+    Pipeline pipeline = Pipeline.create(options);
+    pipeline
+        .apply("Read", Read.from(new TestUnboundedSource(TAG_NULL_KEY, 1, 5)))
+        .apply("WithKeys", WithKeys.of((Void) null))
+        .setCoder(KvCoder.of(VoidCoder.of(), StringUtf8Coder.of()))
+        .apply("StatefulParDo", ParDo.of(new NullKeyCountFn()))
+        .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+    assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(pipeline).getState());
+    Set<Integer> collected = StreamingTestUtils.collected(collectorId);
+    assertTrue(collected.contains(5));
+  }
+
   private static final class KeySelectorFn extends DoFn<String, KV<String, String>> {
     @ProcessElement
     public void process(@Element String element, OutputReceiver<KV<String, String>> out) {
@@ -301,4 +324,17 @@ public void process(
       out.output(element.getValue());
     }
   }
+
+  private static final class NullKeyCountFn extends DoFn<KV<Void, String>, Integer> {
+    @StateId("count")
+    private final StateSpec<ValueState<Integer>> countSpec = StateSpecs.value();
+
+    @ProcessElement
+    public void process(
+        @StateId("count") ValueState<Integer> countState, OutputReceiver<Integer> out) {
+      int count = (countState.read() == null ? 0 : countState.read()) + 1;
+      countState.write(count);
+      out.output(count);
+    }
+  }
 }

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

@@ -18,13 +18,13 @@
 package org.apache.beam.runners.spark.structuredstreaming.translation.batch;
 
 import static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY;
+import static org.apache.beam.runners.spark.coders.CoderHelpers.toByteArray;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGlobalGroupBy;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.eligibleForGroupByWindow;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.explodeWindowedKey;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueKey;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.valueValue;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.batch.GroupByKeyHelpers.windowedKV;
-import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers.toByteArray;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.collectionEncoder;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.encoderOf;
 import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.kvEncoder;

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

@@ -34,8 +34,8 @@
 import java.util.Map;
 import java.util.function.Function;
 import javax.annotation.Nullable;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java

@@ -36,6 +36,7 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;