Clone TransformWithState stateful processors across sibling streaming queries

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

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

@@ -30,11 +30,17 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
 import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
+import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessor;
+import org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.hadoop.fs.Path;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
+import org.apache.spark.sql.catalyst.plans.logical.TransformWithState;
+import org.apache.spark.sql.classic.Dataset$;
+import org.apache.spark.sql.streaming.StatefulProcessor;
 import org.apache.spark.sql.streaming.StreamingQuery;
 import org.apache.spark.sql.streaming.StreamingQueryException;
 import org.apache.spark.sql.streaming.StreamingQueryListener;
@@ -151,7 +157,17 @@ private StreamingQuery startQuery(
       int leafIndex,
       SparkStructuredStreamingPipelineOptions options) {
     try {
-      return dataset
+      // Leaf 0 uses the translated dataset's processor instances directly; leaves 1..n get fresh
+      // BeamStatefulProcessor copies so concurrent queries do not share StatefulProcessor.handle.
+      // TODO(https://github.com/apache/beam/issues/41198): Remove per-leaf copying once pipelines
+      // execute as a single streaming query.
+      Dataset<?> queryDataset =
+          leafIndex == 0
+              ? dataset
+              : Dataset$.MODULE$.ofRows(
+                  (org.apache.spark.sql.classic.SparkSession) dataset.sparkSession(),
+                  copyStatefulProcessors(dataset.queryExecution().logical()));
+      return queryDataset
           .writeStream()
           .format("noop")
           .outputMode("append")
@@ -166,6 +182,43 @@ private StreamingQuery startQuery(
     }
   }
 
+  /**
+   * Copies {@link TransformWithState} stateful processor instances so sibling leaf streaming
+   * queries do not concurrently mutate the same {@code StatefulProcessor} handle on the driver.
+   */
+  private static LogicalPlan copyStatefulProcessors(LogicalPlan plan) {
+    LogicalPlan mapped =
+        (LogicalPlan)
+            plan.mapChildren(ScalaInterop.fun1(StreamingEvaluationContext::copyStatefulProcessors));
+    if (mapped instanceof TransformWithState) {
+      TransformWithState tws = (TransformWithState) mapped;
+      @SuppressWarnings("unchecked")
+      StatefulProcessor<Object, Object, Object> copiedProcessor =
+          (StatefulProcessor<Object, Object, Object>)
+              (StatefulProcessor<?, ?, ?>)
+                  ((BeamStatefulProcessor<?, ?, ?>)
+                          (StatefulProcessor<?, ?, ?>) tws.statefulProcessor())
+                      .copy();
+      return tws.copy(
+          tws.keyDeserializer(),
+          tws.valueDeserializer(),
+          tws.groupingAttributes(),
+          tws.dataAttributes(),
+          copiedProcessor,
+          tws.timeMode(),
+          tws.outputMode(),
+          tws.keyEncoder(),
+          tws.outputObjAttr(),
+          tws.child(),
+          tws.hasInitialState(),
+          tws.initialStateGroupingAttrs(),
+          tws.initialStateDataAttrs(),
+          tws.initialStateDeserializer(),
+          tws.initialState());
+    }
+    return mapped;
+  }
+
   /** Blocks until every query in toAwait has terminated. Sibling queries stop on failure. */
   private void awaitTermination(List<StreamingQuery> toAwait) {
     List<StreamingQuery> active = new ArrayList<>(toAwait);

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

@@ -198,6 +198,13 @@ protected abstract void processElements(
   protected abstract void processTimers(
       K key, SparkTimerInternals timers, SparkStateInternals<K> state);
 
+  /**
+   * Returns a new uninitialized processor with the same configuration, so sibling streaming queries
+   * do not share the mutable {@link StatefulProcessorHandle} on the base class while still sharing
+   * the driver-registered {@code MetricsAccumulator}.
+   */
+  public abstract BeamStatefulProcessor<K, V, OutputT> copy();
+
   private void restoreTimers(SparkTimerInternals timerInternals) {
     ValueState<byte[]> timersState = checkStateNotNull(beamTimers);
     if (timersState.exists()) {

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

@@ -67,6 +67,11 @@ public GroupByKeyStatefulProcessor(
     this.optionsSupplier = optionsSupplier;
   }
 
+  @Override
+  public GroupByKeyStatefulProcessor<K, V> copy() {
+    return new GroupByKeyStatefulProcessor<>(inputCoder, windowingStrategy, optionsSupplier);
+  }
+
   @Override
   protected void processElements(
       K key,

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

@@ -64,6 +64,12 @@ public ParDoStatefulProcessor(
     this.metrics = metrics;
   }
 
+  @Override
+  public ParDoStatefulProcessor<K, V, OutputT> copy() {
+    return new ParDoStatefulProcessor<>(
+        doFn, runnerFactory, inputCoder, windowingStrategy, optionsSupplier, metrics);
+  }
+
   @Override
   protected void processElements(
       K key,

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

@@ -35,6 +35,11 @@
 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.metrics.Counter;
+import org.apache.beam.sdk.metrics.MetricNameFilter;
+import org.apache.beam.sdk.metrics.MetricQueryResults;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.metrics.MetricsFilter;
 import org.apache.beam.sdk.state.StateSpec;
 import org.apache.beam.sdk.state.StateSpecs;
 import org.apache.beam.sdk.state.TimeDomain;
@@ -48,6 +53,8 @@
 import org.apache.beam.sdk.transforms.windowing.FixedWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.junit.After;
@@ -73,6 +80,7 @@ public class StatefulParDoStreamingTest implements Serializable {
   private static final String TAG_CHAINED = "stateful-chained";
   private static final String TAG_NULL_KEY = "stateful-null-key";
   private static final String TAG_EXPIRY = "stateful-expiry";
+  private static final String TAG_MULTI_LEAF = "stateful-multi-leaf";
 
   @After
   public void tearDown() {
@@ -81,6 +89,7 @@ public void tearDown() {
     TestUnboundedSource.forget(TAG_CHAINED);
     TestUnboundedSource.forget(TAG_NULL_KEY);
     TestUnboundedSource.forget(TAG_EXPIRY);
+    TestUnboundedSource.forget(TAG_MULTI_LEAF);
   }
 
   private static Pipeline pipeline(
@@ -217,6 +226,42 @@ public void testExpiredKeyLeavesNoState() throws Exception {
     }
   }
 
+  @Test
+  public void testStatefulParDoSharedAcrossMultipleLeaves() throws Exception {
+    String c0 = StreamingTestUtils.newCollectorId(TAG_MULTI_LEAF + "-0");
+    String c1 = StreamingTestUtils.newCollectorId(TAG_MULTI_LEAF + "-1");
+    SparkStructuredStreamingPipelineOptions options =
+        StreamingTestUtils.streamingOptions(tempFolder);
+    options.setMaxRecordsPerBatch(1L);
+    Pipeline pipeline = Pipeline.create(options);
+    PCollection<Integer> counts =
+        pipeline
+            .apply("Read", Read.from(new TestUnboundedSource(TAG_MULTI_LEAF, 1, 5)))
+            .apply("WithKeys", WithKeys.of((Void) null))
+            .setCoder(KvCoder.of(VoidCoder.of(), StringUtf8Coder.of()))
+            .apply("StatefulParDo", ParDo.of(new NullKeyCountFn()));
+    counts.apply("Collect0", ParDo.of(new StreamingTestUtils.CollectDoFn<>(c0)));
+    counts.apply("Collect1", ParDo.of(new StreamingTestUtils.CollectDoFn<>(c1)));
+
+    PipelineResult result = StreamingTestUtils.run(pipeline);
+    assertEquals(PipelineResult.State.DONE, result.getState());
+    for (String collectorId : Arrays.asList(c0, c1)) {
+      assertTrue(
+          StreamingTestUtils.collected(collectorId).containsAll(Arrays.asList(1, 2, 3, 4, 5)));
+    }
+    MetricQueryResults metrics =
+        result
+            .metrics()
+            .queryMetrics(
+                MetricsFilter.builder()
+                    .addNameFilter(MetricNameFilter.named(NullKeyCountFn.class, "processed"))
+                    .build());
+    // TODO(https://github.com/apache/beam/issues/41198): Each leaf query currently executes the
+    // upstream stateful ParDo independently (5 elements per leaf = 10 total); update to 5L once
+    // #41198 is resolved.
+    assertEquals(10L, (long) Iterables.getOnlyElement(metrics.getCounters()).getAttempted());
+  }
+
   private static final class KeySelectorFn extends DoFn<String, KV<String, String>> {
     @ProcessElement
     public void process(@Element String element, OutputReceiver<KV<String, String>> out) {
@@ -363,12 +408,15 @@ public void process(
   }
 
   private static final class NullKeyCountFn extends DoFn<KV<Void, String>, Integer> {
+    private final Counter processed = Metrics.counter(NullKeyCountFn.class, "processed");
+
     @StateId("count")
     private final StateSpec<ValueState<Integer>> countSpec = StateSpecs.value();
 
     @ProcessElement
     public void process(
         @StateId("count") ValueState<Integer> countState, OutputReceiver<Integer> out) {
+      processed.inc();
       int count = (countState.read() == null ? 0 : countState.read()) + 1;
       countState.write(count);
       out.output(count);

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

@@ -48,6 +48,7 @@
 import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -75,6 +76,7 @@ public class WindowedGroupByKeyStreamingTest implements Serializable {
   private static final String TAG_CHAINED = "chained-timer";
   private static final String TAG_GLOBALLY = "globally-combine";
   private static final String TAG_EXPIRY = "gbk-expiry";
+  private static final String TAG_MULTI_LEAF = "gbk-multi-leaf";
 
   @After
   public void tearDown() {
@@ -84,6 +86,7 @@ public void tearDown() {
     TestUnboundedSource.forget(TAG_CHAINED);
     TestUnboundedSource.forget(TAG_GLOBALLY);
     TestUnboundedSource.forget(TAG_EXPIRY);
+    TestUnboundedSource.forget(TAG_MULTI_LEAF);
   }
 
   private SparkStructuredStreamingPipelineOptions options(String tag) throws Exception {
@@ -209,6 +212,31 @@ public void testExpiredKeyLeavesNoState() throws Exception {
     }
   }
 
+  @Test
+  public void testGroupByKeySharedAcrossMultipleLeaves() throws Exception {
+    Pipeline p = Pipeline.create(options(TAG_MULTI_LEAF));
+    PCollection<String> formatted =
+        p.apply(Read.from(new TestUnboundedSource(TAG_MULTI_LEAF, 1, 6)))
+            .apply(ParDo.of(new ToKvFn()))
+            .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+            .apply(Window.into(FixedWindows.of(Duration.standardSeconds(2))))
+            .apply(GroupByKey.create())
+            .apply(ParDo.of(new FormatGbkFn()));
+
+    String[] collectors = new String[4];
+    for (int i = 0; i < collectors.length; i++) {
+      collectors[i] = StreamingTestUtils.newCollectorId(TAG_MULTI_LEAF + "-" + i);
+      formatted.apply("Collect" + i, ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectors[i])));
+    }
+
+    assertEquals(PipelineResult.State.DONE, StreamingTestUtils.run(p).getState());
+    for (String col : collectors) {
+      Set<String> res = StreamingTestUtils.collected(col);
+      assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + ":k=2"));
+      assertTrue(res.contains(TestUnboundedSource.BASE_MILLIS + 2000 + ":k=2"));
+    }
+  }
+
   private static final class ToKvFn extends DoFn<String, KV<String, Integer>> {
     @ProcessElement
     public void process(@Element String el, OutputReceiver<KV<String, Integer>> out) {