diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json index d6a91b7e2e86..38ae1cf68222 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json +++ b/.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 } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java index b5057cd260ac..249042402698 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java +++ b/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 copiedProcessor = + (StatefulProcessor) + (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 toAwait) { List active = new ArrayList<>(toAwait); diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java index 568c41d23ec8..a1c34114fd9b 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java +++ b/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 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 copy(); + private void restoreTimers(SparkTimerInternals timerInternals) { ValueState timersState = checkStateNotNull(beamTimers); if (timersState.exists()) { diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java index cbe592af0eae..52a7d4e31c31 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/GroupByKeyStatefulProcessor.java +++ b/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 copy() { + return new GroupByKeyStatefulProcessor<>(inputCoder, windowingStrategy, optionsSupplier); + } + @Override protected void processElements( K key, diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java index f25c523eaac6..dee43d22cf37 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/ParDoStatefulProcessor.java +++ b/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 copy() { + return new ParDoStatefulProcessor<>( + doFn, runnerFactory, inputCoder, windowingStrategy, optionsSupplier, metrics); + } + @Override protected void processElements( K key, diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java index 0e30d6ed791d..b193ff03d770 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java +++ b/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 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> { @ProcessElement public void process(@Element String element, OutputReceiver> out) { @@ -363,12 +408,15 @@ public void process( } private static final class NullKeyCountFn extends DoFn, Integer> { + private final Counter processed = Metrics.counter(NullKeyCountFn.class, "processed"); + @StateId("count") private final StateSpec> countSpec = StateSpecs.value(); @ProcessElement public void process( @StateId("count") ValueState countState, OutputReceiver out) { + processed.inc(); int count = (countState.read() == null ? 0 : countState.read()) + 1; countState.write(count); out.output(count); diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java index 2b0305fa66d2..872dbeadef7e 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java +++ b/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 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 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> { @ProcessElement public void process(@Element String el, OutputReceiver> out) {