From 4ca982ea435affc059ddac68e0f06e1b103b8281 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 11:14:00 +0000 Subject: [PATCH 01/21] [Spark 4] Add streaming dispatch seam and lifecycle hooks Prepares the structured streaming runner for a Spark 4 streaming translator without changing batch behavior. - Add PipelineTranslatorFactory, the seam the Spark 4 module shadows to dispatch streaming pipelines. The shared base rejects streaming with a clear message instead of the previous generic checkArgument. - Open up EvaluationContext (non final, protected ctor, leaves()) and add a no-op stop() so a streaming context can override evaluation. - Add a createEvaluationContext hook to PipelineTranslator and skip the persist and lineage breaking optimizations for streaming datasets. - Plumb the EvaluationContext into SparkStructuredStreamingPipelineResult so cancel() can stop a running streaming query. - Default the state store provider to RocksDB, required by Spark 4 transformWithState and inert for batch. - Add watermarkDelayMillis, maxRecordsPerMicroBatch, maxBatchDurationMillis and streamingStopAfterIdleBatches options. - Narrow the Spark 4 test source override exclude to the legacy DStream package. The previous glob also matched structuredstreaming and would have silently dropped the structured streaming tests. --- runners/spark/4/build.gradle | 8 ++-- ...arkStructuredStreamingPipelineOptions.java | 34 +++++++++++++++ ...parkStructuredStreamingPipelineResult.java | 16 ++++++-- .../SparkStructuredStreamingRunner.java | 26 ++++++++---- .../translation/EvaluationContext.java | 18 ++++++-- .../translation/PipelineTranslator.java | 22 ++++++++-- .../PipelineTranslatorFactory.java | 41 +++++++++++++++++++ .../translation/SparkSessionFactory.java | 6 +++ 8 files changed, 150 insertions(+), 21 deletions(-) create mode 100644 runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java diff --git a/runners/spark/4/build.gradle b/runners/spark/4/build.gradle index ec1af8df38a4..a3da1d8922fd 100644 --- a/runners/spark/4/build.gradle +++ b/runners/spark/4/build.gradle @@ -53,11 +53,13 @@ tasks.validatesStructuredStreamingRunnerBatch { tasks.validatesRunner.dependsOn(validatesStructuredStreamingRunnerBatch) -// Exclude DStream-based streaming tests from the shared-base copy: the Spark 4 module -// supports only structured streaming (batch) and does not include legacy DStream support. +// Exclude legacy DStream-based streaming tests (under runners/spark/translation/streaming) +// from the shared-base copy: the Spark 4 module does not include legacy DStream support. // Streaming test utilities also depend on kafka.server.KafkaServerStartable which was // removed in Kafka 2.8.0 (the first Kafka version with a _2.13 artifact). +// Note: structured streaming tests (**/structuredstreaming/translation/streaming/**) are +// intentionally NOT excluded. tasks.named("copyTestSourceOverrides") { - exclude "**/translation/streaming/**" + exclude "**/runners/spark/translation/streaming/**" } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java index 3371a403b2c9..29cc4cb99cfb 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineOptions.java @@ -39,4 +39,38 @@ public interface SparkStructuredStreamingPipelineOptions extends SparkCommonPipe boolean getUseActiveSparkSession(); void setUseActiveSparkSession(boolean value); + + @Description( + "Watermark delay in milliseconds applied to event timestamps of streaming sources " + + "(streaming mode only).") + @Default.Long(0) + long getWatermarkDelayMillis(); + + void setWatermarkDelayMillis(long value); + + // Note: deliberately NOT named getMaxRecordsPerBatch. The legacy Spark runner's + // SparkPipelineOptions already declares Long getMaxRecordsPerBatch(); a same-name getter with a + // different return type breaks proxy generation for every registered PipelineOptions interface. + @Description( + "Maximum number of records to read per micro-batch from a streaming source " + + "(streaming mode only).") + @Default.Integer(1000) + int getMaxRecordsPerMicroBatch(); + + void setMaxRecordsPerMicroBatch(int value); + + @Description( + "Maximum duration in milliseconds of a micro-batch trigger interval (streaming mode only).") + @Default.Long(500) + long getMaxBatchDurationMillis(); + + void setMaxBatchDurationMillis(long value); + + @Description( + "Test-oriented: gracefully stop streaming queries after this many consecutive empty " + + "micro-batches. Disabled if negative (streaming mode only).") + @Default.Integer(-1) + int getStreamingStopAfterIdleBatches(); + + void setStreamingStopAfterIdleBatches(int value); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java index 9d3419e19473..b592b6fb742d 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java @@ -25,27 +25,33 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import javax.annotation.Nullable; +import java.util.function.Supplier; import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.MetricResults; import org.apache.beam.sdk.util.UserCodeException; import org.apache.spark.SparkException; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; public class SparkStructuredStreamingPipelineResult implements PipelineResult { private final Future pipelineExecution; + // Supplies the context of the translated pipeline, null until translation has completed. + private final Supplier evaluationContext; private final MetricsAccumulator metrics; - private @Nullable final Runnable onTerminalState; + private final @Nullable Runnable onTerminalState; private PipelineResult.State state; SparkStructuredStreamingPipelineResult( Future pipelineExecution, + Supplier evaluationContext, MetricsAccumulator metrics, - @Nullable final Runnable onTerminalState) { + final @Nullable Runnable onTerminalState) { this.pipelineExecution = pipelineExecution; + this.evaluationContext = evaluationContext; this.metrics = metrics; this.onTerminalState = onTerminalState; // pipelineExecution is expected to have started executing eagerly. @@ -113,6 +119,10 @@ public MetricResults metrics() { @Override public PipelineResult.State cancel() throws IOException { + EvaluationContext ctx = evaluationContext.get(); + if (ctx != null) { + ctx.stop(); + } pipelineExecution.cancel(true); offerNewState(PipelineResult.State.CANCELLED); return state; diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java index 96717f29e87f..84afa2b2a55c 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java @@ -17,12 +17,11 @@ */ package org.apache.beam.runners.spark.structuredstreaming; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; - import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.ThreadFactory; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import org.apache.beam.runners.core.metrics.MetricsPusher; import org.apache.beam.runners.core.metrics.NoOpMetricsSink; @@ -30,8 +29,8 @@ import org.apache.beam.runners.spark.structuredstreaming.metrics.SparkBeamMetricSource; import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext; import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.PipelineTranslatorFactory; import org.apache.beam.runners.spark.structuredstreaming.translation.SparkSessionFactory; -import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineRunner; import org.apache.beam.sdk.metrics.MetricsEnvironment; @@ -54,9 +53,9 @@ * href="https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html">Structured * Streaming framework). * - *

This runner is experimental, its coverage of the Beam model is still partial. Due to - * limitations of the Structured Streaming framework (e.g. lack of support for multiple stateful - * operators), streaming mode is not yet supported by this runner. + *

This runner is experimental, its coverage of the Beam model is still partial. Streaming + * mode requires the Spark 4 module (beam-runners-spark-4); the shared Spark 3 module supports + * batch pipelines only. * *

The runner translates transforms defined on a Beam pipeline to Spark `Dataset` transformations * (leveraging the high level Dataset API) and then submits these to Spark to be executed. @@ -145,17 +144,26 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { + " It is still experimental, its coverage of the Beam model is partial. ***"); PipelineTranslator.detectStreamingMode(pipeline, options); - checkArgument(!options.isStreaming(), "Streaming is not supported."); final SparkSession sparkSession = SparkSessionFactory.getOrCreateSession(options); final MetricsAccumulator metrics = MetricsAccumulator.getInstance(sparkSession); + // Set once the pipeline is translated, so the result can stop an ongoing (streaming) + // evaluation on cancel. Remains null until translation completes. + final AtomicReference ctxRef = new AtomicReference<>(); + final Future submissionFuture = - runAsync(() -> translatePipeline(sparkSession, pipeline).evaluate()); + runAsync( + () -> { + EvaluationContext ctx = translatePipeline(sparkSession, pipeline); + ctxRef.set(ctx); + ctx.evaluate(); + }); final SparkStructuredStreamingPipelineResult result = new SparkStructuredStreamingPipelineResult( submissionFuture, + ctxRef::get, metrics, sparkStopFn(sparkSession, options.getUseActiveSparkSession())); @@ -186,7 +194,7 @@ private EvaluationContext translatePipeline(SparkSession sparkSession, Pipeline PipelineTranslator.replaceTransforms(pipeline, options); - PipelineTranslator pipelineTranslator = new PipelineTranslatorBatch(); + PipelineTranslator pipelineTranslator = PipelineTranslatorFactory.create(options.isStreaming()); return pipelineTranslator.translate(pipeline, sparkSession, options); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java index b8448567eafc..0e677051fb61 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java @@ -40,10 +40,10 @@ */ @SuppressWarnings("Slf4jDoNotLogMessageOfExceptionExplicitly") @Internal -public final class EvaluationContext { +public class EvaluationContext { private static final Logger LOG = LoggerFactory.getLogger(EvaluationContext.class); - interface NamedDataset { + public interface NamedDataset { String name(); @Nullable @@ -53,11 +53,16 @@ interface NamedDataset { private final Collection> leaves; private final SparkSession session; - EvaluationContext(Collection> leaves, SparkSession session) { + protected EvaluationContext(Collection> leaves, SparkSession session) { this.leaves = leaves; this.session = session; } + /** The leaf datasets of the translated pipeline that require evaluation. */ + protected Collection> leaves() { + return leaves; + } + /** Trigger evaluation of all leaf datasets. */ public void evaluate() { for (NamedDataset ds : leaves) { @@ -113,6 +118,13 @@ public static void evaluate(String name, Dataset ds) { } } + /** + * Stops any ongoing streaming execution triggered by this context. + * + *

This is a no-op for batch pipelines. + */ + public void stop() {} + public SparkSession getSparkSession() { return session; } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java index a681dea2fde5..8470c968550b 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslator.java @@ -25,6 +25,7 @@ import java.io.IOException; import java.io.Serializable; import java.util.ArrayList; +import java.util.Collection; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -128,7 +129,20 @@ public EvaluationContext translate( TranslatingVisitor translator = new TranslatingVisitor(session, options, dependencies.results); pipeline.traverseTopologically(translator); - return new EvaluationContext(translator.leaves, session); + return createEvaluationContext(translator.leaves, session, options); + } + + /** + * Creates the {@link EvaluationContext} for the translated pipeline. + * + *

Subclasses may override this to return a specialized context, e.g. to evaluate streaming + * pipelines. + */ + protected EvaluationContext createEvaluationContext( + Collection> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + return new EvaluationContext(leaves, session); } /** @@ -311,12 +325,14 @@ public void putDataset( TranslationResult result = getResult(pCollection); result.dataset = dataset; - if (cache && result.usages() > 1) { + // Caching and lineage breaking are batch-only optimizations, streaming datasets must pass + // through untouched. + if (cache && result.usages() > 1 && !dataset.isStreaming()) { LOG.info("Dataset {} will be cached for reuse.", result.name); dataset.persist(storageLevel); // use NONE to disable } - if (result.estimatePlanComplexity() > PLAN_COMPLEXITY_THRESHOLD) { + if (!dataset.isStreaming() && result.estimatePlanComplexity() > PLAN_COMPLEXITY_THRESHOLD) { // Break linage of dataset to limit planning overhead for complex query plans. LOG.info("Breaking linage of dataset {} to limit complexity of query plan.", result.name); result.dataset = sparkSession.createDataset(dataset.rdd(), dataset.encoder()); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java new file mode 100644 index 000000000000..4609d1d380e7 --- /dev/null +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java @@ -0,0 +1,41 @@ +/* + * 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; + +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.sdk.annotations.Internal; + +/** + * Factory to create the {@link PipelineTranslator} matching the execution mode of the pipeline. + * + *

This shared base version only supports batch mode. The Spark 4 module shadows this file to + * additionally dispatch to a streaming translator. + */ +@Internal +public final class PipelineTranslatorFactory { + private PipelineTranslatorFactory() {} + + /** Creates a {@link PipelineTranslator} for the given execution mode. */ + public static PipelineTranslator create(boolean streaming) { + if (streaming) { + throw new UnsupportedOperationException( + "Streaming pipelines require the Spark 4 runner (beam-runners-spark-4)."); + } + return new PipelineTranslatorBatch(); + } +} diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 5b9e5b6fae86..822d1871b12e 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -175,6 +175,12 @@ private static SparkSession.Builder sessionBuilder( sparkConf.setIfMissing("spark.sql.shuffle.partitions", Integer.toString(partitions)); } + // Spark 4 transformWithState (used for streaming pipelines) requires the RocksDB state store. + // This is a harmless, inert configuration for batch pipelines on Spark 3. + sparkConf.setIfMissing( + "spark.sql.streaming.stateStore.providerClass", + "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider"); + return SparkSession.builder().config(sparkConf); } From 0131dbf251663d1428bd2010bf2dbcc75a5de9cc Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 12:34:54 +0000 Subject: [PATCH 02/21] [Spark 4] Add DataSourceV2 micro-batch source for Beam unbounded sources Wraps any Beam UnboundedSource as a Spark 4 DataSourceV2 micro-batch streaming source. The rows have exactly two columns, payload of type BINARY holding the element encoded with a FullWindowedValueCoder, and eventTimestamp of type TIMESTAMP holding the Beam event timestamp, so no Catalyst encoder has to be generated for Beam types. Offsets are opaque epoch counters. The driver never reads from the source and never inspects its progress, latestOffset simply returns the previous epoch plus one so Spark keeps planning micro-batches. Termination is therefore owned by the lifecycle, not by the offsets. The source is split once on the driver and the sub sources are memoized, so every micro-batch plans the same stable set of partitions. On the executor a static cache keyed by checkpoint location, source id and split id keeps the Beam UnboundedReader alive across micro-batches, mirroring the legacy MicrobatchSource. Checkpoint marks are cached in executor memory only and are never committed durably, which is documented POC scope with weak failure recovery. UnboundedSourceDataset is the translator facing entry point. It applies withWatermark exactly once, since Spark 4 forbids a second declaration further down the plan. The test confirms that the EventTimeWatermark logical node survives one and two typed maps in both the logical and the analyzed plan, and that the running query really tracks the watermark afterwards. It also reads a finite set of elements end to end and checks the decoded payloads and timestamps. --- .../io/streaming/BeamInputPartition.java | 96 ++++ .../io/streaming/BeamMicroBatchStream.java | 158 ++++++ .../io/streaming/BeamOffset.java | 79 +++ .../io/streaming/BeamPartitionReader.java | 150 ++++++ .../streaming/BeamPartitionReaderFactory.java | 34 ++ .../io/streaming/BeamReaderCache.java | 159 ++++++ .../io/streaming/BeamStreamingSource.java | 147 ++++++ .../io/streaming/BeamStreamingTable.java | 100 ++++ .../io/streaming/UnboundedSourceDataset.java | 106 ++++ .../streaming/BeamMicroBatchSourceTest.java | 477 ++++++++++++++++++ 10 files changed, 1506 insertions(+) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java new file mode 100644 index 000000000000..389ed3fec91e --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java @@ -0,0 +1,96 @@ +/* + * 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.io.streaming; + +import org.apache.spark.sql.connector.read.InputPartition; + +/** + * One split of a Beam unbounded source for one micro-batch. + * + *

Everything the executor needs travels as base64 of the Java serialized object, so the + * partition works across JVMs and is not limited to Spark local mode. + */ +public class BeamInputPartition implements InputPartition { + + private static final long serialVersionUID = 1L; + + private final String sourceB64; + private final String coderB64; + private final String pipelineOptionsB64; + private final String sourceId; + private final int splitId; + private final String checkpointLocation; + private final int maxRecordsPerMicroBatch; + private final long maxBatchDurationMillis; + + BeamInputPartition( + String sourceB64, + String coderB64, + String pipelineOptionsB64, + String sourceId, + int splitId, + String checkpointLocation, + int maxRecordsPerMicroBatch, + long maxBatchDurationMillis) { + this.sourceB64 = sourceB64; + this.coderB64 = coderB64; + this.pipelineOptionsB64 = pipelineOptionsB64; + this.sourceId = sourceId; + this.splitId = splitId; + this.checkpointLocation = checkpointLocation; + this.maxRecordsPerMicroBatch = maxRecordsPerMicroBatch; + this.maxBatchDurationMillis = maxBatchDurationMillis; + } + + String sourceB64() { + return sourceB64; + } + + String coderB64() { + return coderB64; + } + + String pipelineOptionsB64() { + return pipelineOptionsB64; + } + + String sourceId() { + return sourceId; + } + + int splitId() { + return splitId; + } + + String checkpointLocation() { + return checkpointLocation; + } + + int maxRecordsPerMicroBatch() { + return maxRecordsPerMicroBatch; + } + + long maxBatchDurationMillis() { + return maxBatchDurationMillis; + } + + @Override + public String toString() { + return "BeamInputPartition{source=" + sourceId + ", split=" + splitId + "}"; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java new file mode 100644 index 000000000000..33d9198412d6 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java @@ -0,0 +1,158 @@ +/* + * 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.io.streaming; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.connector.read.streaming.Offset; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * + *

Offsets are opaque epoch counters, see {@link BeamOffset}. {@link #latestOffset()} always + * reports a value greater than the previous one so Spark keeps scheduling micro-batches, even ones + * that turn out to be empty. Termination of a streaming pipeline is therefore never driven by the + * offsets, it is driven by the lifecycle owner (the idle batch listener of the evaluation context, + * or an explicit {@code StreamingQuery.stop()}). + * + *

The wrapped source is split exactly once, on the driver, and the resulting sub sources are + * cached so every micro-batch plans the same, stable set of partitions. Splits must be stable + * across micro-batches because the executor side reader cache is keyed by split index. + */ +public class BeamMicroBatchStream implements MicroBatchStream { + + private static final Logger LOG = LoggerFactory.getLogger(BeamMicroBatchStream.class); + + private final String sourceB64; + private final String coderB64; + private final String pipelineOptionsB64; + private final String sourceId; + private final String checkpointLocation; + private final int desiredNumSplits; + private final int maxRecordsPerMicroBatch; + private final long maxBatchDurationMillis; + + private long epoch; + private @Nullable List splitsB64; + + BeamMicroBatchStream(CaseInsensitiveStringMap options, String checkpointLocation) { + this.sourceB64 = BeamStreamingSource.required(options, BeamStreamingSource.OPT_SOURCE); + this.coderB64 = BeamStreamingSource.required(options, BeamStreamingSource.OPT_CODER); + this.pipelineOptionsB64 = + BeamStreamingSource.required(options, BeamStreamingSource.OPT_PIPELINE_OPTIONS); + this.sourceId = BeamStreamingSource.required(options, BeamStreamingSource.OPT_SOURCE_ID); + this.checkpointLocation = checkpointLocation; + this.desiredNumSplits = Math.max(1, options.getInt(BeamStreamingSource.OPT_NUM_SPLITS, 1)); + this.maxRecordsPerMicroBatch = + Math.max(1, options.getInt(BeamStreamingSource.OPT_MAX_RECORDS, 1000)); + this.maxBatchDurationMillis = + Math.max(1L, options.getLong(BeamStreamingSource.OPT_MAX_BATCH_DURATION_MILLIS, 500L)); + } + + @Override + public Offset initialOffset() { + return BeamOffset.ZERO; + } + + @Override + public synchronized Offset latestOffset() { + return new BeamOffset(++epoch); + } + + @Override + public Offset deserializeOffset(String json) { + return BeamOffset.fromJson(json); + } + + @Override + public void commit(Offset end) { + LOG.debug("Committed epoch offset {} of Beam source {}.", end, sourceId); + } + + @Override + public void stop() { + LOG.info("Stopping Beam micro-batch stream for source {}.", sourceId); + } + + @Override + public InputPartition[] planInputPartitions(Offset start, Offset end) { + List splits = splits(); + InputPartition[] partitions = new InputPartition[splits.size()]; + for (int i = 0; i < splits.size(); i++) { + partitions[i] = + new BeamInputPartition( + splits.get(i), + coderB64, + pipelineOptionsB64, + sourceId, + i, + checkpointLocation, + maxRecordsPerMicroBatch, + maxBatchDurationMillis); + } + return partitions; + } + + @Override + public PartitionReaderFactory createReaderFactory() { + return new BeamPartitionReaderFactory(); + } + + /** Splits the wrapped source once and memoizes the serialized sub sources. */ + private synchronized List splits() { + if (splitsB64 != null) { + return splitsB64; + } + UnboundedSource source = BeamStreamingSource.decode(sourceB64, "UnboundedSource"); + SerializablePipelineOptions serializableOptions = + BeamStreamingSource.decode(pipelineOptionsB64, "PipelineOptions"); + PipelineOptions options = serializableOptions.get(); + List> split; + try { + split = source.split(desiredNumSplits, options); + } catch (Exception e) { + throw new IllegalStateException( + "Failed to split UnboundedSource " + source.getClass().getCanonicalName(), e); + } + if (split.isEmpty()) { + split = Collections.singletonList(source); + } + List encoded = new ArrayList<>(split.size()); + for (UnboundedSource s : split) { + encoded.add(BeamStreamingSource.encode(s)); + } + LOG.info( + "Split Beam source {} into {} partition(s) (desired {}).", + sourceId, + encoded.size(), + desiredNumSplits); + splitsB64 = encoded; + return encoded; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java new file mode 100644 index 000000000000..f7e5e0e56027 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java @@ -0,0 +1,79 @@ +/* + * 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.io.streaming; + +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.spark.sql.connector.read.streaming.Offset; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * An opaque epoch counter used as the Spark streaming {@link Offset} of a Beam unbounded source. + * + *

The offset carries no information about the position inside the wrapped Beam source. The + * driver never reads from the source and never inspects its progress, it only needs a monotonically + * increasing value so that Spark keeps planning micro-batches. The actual read position lives in + * the executor side {@link BeamReaderCache} as a Beam {@code CheckpointMark}. + */ +public class BeamOffset extends Offset { + + /** The offset every Beam unbounded stream starts at. */ + public static final BeamOffset ZERO = new BeamOffset(0L); + + private static final Pattern EPOCH_PATTERN = Pattern.compile("-?\\d+"); + + private final long epoch; + + public BeamOffset(long epoch) { + this.epoch = epoch; + } + + /** The epoch counter value. */ + public long epoch() { + return epoch; + } + + @Override + public String json() { + return "{\"epoch\":" + epoch + "}"; + } + + /** Parses the form produced by {@link #json()}, a bare number is also accepted. */ + public static BeamOffset fromJson(String json) { + Matcher matcher = EPOCH_PATTERN.matcher(json); + if (!matcher.find()) { + throw new IllegalArgumentException("Not a valid BeamOffset: " + json); + } + return new BeamOffset(Long.parseLong(matcher.group())); + } + + @Override + public boolean equals(@Nullable Object o) { + return o instanceof BeamOffset && ((BeamOffset) o).epoch == epoch; + } + + @Override + public int hashCode() { + return Long.hashCode(epoch); + } + + @Override + public String toString() { + return json(); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java new file mode 100644 index 000000000000..ff6c319f01c8 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -0,0 +1,150 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache.CachedReader; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Reads one split of a Beam {@link UnboundedSource} for the duration of one Spark micro-batch. + * + *

The batch ends as soon as either {@code maxRecordsPerMicroBatch} elements were emitted or + * {@code maxBatchDurationMillis} of wall clock time elapsed, whichever comes first. When the source + * has no data available the reader polls with a short sleep until the deadline, so an idle source + * produces an empty micro-batch rather than blocking the query. + * + *

The underlying Beam reader is not closed at the end of the batch, it stays in {@link + * BeamReaderCache} and the next micro-batch continues from the same position. See that class for + * the failure recovery caveats. + * + * @param the element type of the wrapped source + */ +@SuppressWarnings({ + "nullness" // the current row is only read between a true next() and the following one +}) +public class BeamPartitionReader implements PartitionReader { + + private static final Logger LOG = LoggerFactory.getLogger(BeamPartitionReader.class); + + /** Sleep between two unsuccessful advance attempts while the batch deadline has not passed. */ + private static final long POLL_INTERVAL_MILLIS = 10L; + + private final String cacheKey; + private final CachedReader cached; + private final Coder> windowedValueCoder; + private final int maxRecordsPerMicroBatch; + private final long maxBatchDurationMillis; + + private long recordsRead; + private long deadlineMillis = -1L; + private @Nullable InternalRow current; + + BeamPartitionReader(BeamInputPartition partition) { + UnboundedSource source = + BeamStreamingSource.decode(partition.sourceB64(), "UnboundedSource split"); + this.windowedValueCoder = + BeamStreamingSource.decode(partition.coderB64(), "WindowedValue coder"); + SerializablePipelineOptions options = + BeamStreamingSource.decode(partition.pipelineOptionsB64(), "PipelineOptions"); + this.maxRecordsPerMicroBatch = partition.maxRecordsPerMicroBatch(); + this.maxBatchDurationMillis = partition.maxBatchDurationMillis(); + this.cacheKey = + BeamReaderCache.key( + partition.checkpointLocation(), partition.sourceId(), partition.splitId()); + this.cached = BeamReaderCache.getOrCreate(cacheKey, source, options.get()); + } + + @Override + public boolean next() throws IOException { + if (deadlineMillis < 0) { + deadlineMillis = System.currentTimeMillis() + maxBatchDurationMillis; + } + while (true) { + if (recordsRead >= maxRecordsPerMicroBatch) { + current = null; + return false; + } + long remaining = deadlineMillis - System.currentTimeMillis(); + if (remaining <= 0) { + current = null; + return false; + } + if (cached.startOrAdvance()) { + recordsRead++; + current = toRow(); + return true; + } + Uninterruptibles.sleepUninterruptibly( + Math.min(remaining, POLL_INTERVAL_MILLIS), java.util.concurrent.TimeUnit.MILLISECONDS); + } + } + + @Override + public InternalRow get() { + if (current == null) { + throw new IllegalStateException("No current row, next() did not return true."); + } + return current; + } + + /** + * Ends the micro-batch. The Beam reader deliberately stays open in {@link BeamReaderCache}, only + * its checkpoint mark is remembered and finalized. + */ + @Override + public void close() { + current = null; + try { + UnboundedSource.CheckpointMark mark = cached.reader().getCheckpointMark(); + BeamReaderCache.rememberCheckpointMark(cacheKey, mark); + mark.finalizeCheckpoint(); + } catch (Exception e) { + LOG.warn("Failed to finalize the checkpoint mark of Beam reader {}.", cacheKey, e); + } + LOG.debug("Beam reader {} emitted {} record(s) in this micro-batch.", cacheKey, recordsRead); + } + + private InternalRow toRow() { + Instant timestamp = cached.reader().getCurrentTimestamp(); + WindowedValue windowedValue = + WindowedValues.timestampedValueInGlobalWindow(cached.reader().getCurrent(), timestamp); + byte[] payload; + try { + payload = CoderUtils.encodeToByteArray(windowedValueCoder, windowedValue); + } catch (CoderException e) { + throw new IllegalStateException("Failed to encode element read from a Beam source.", e); + } + // Spark stores TimestampType as microseconds since the epoch. + return new GenericInternalRow(new Object[] {payload, timestamp.getMillis() * 1000L}); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java new file mode 100644 index 000000000000..db2573eff315 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java @@ -0,0 +1,34 @@ +/* + * 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.io.streaming; + +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; + +/** Creates a {@link BeamPartitionReader} for a {@link BeamInputPartition} on the executor. */ +public class BeamPartitionReaderFactory implements PartitionReaderFactory { + + private static final long serialVersionUID = 1L; + + @Override + public PartitionReader createReader(InputPartition partition) { + return new BeamPartitionReader<>((BeamInputPartition) partition); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java new file mode 100644 index 000000000000..9ccfc3388bd3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -0,0 +1,159 @@ +/* + * 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.io.streaming; + +import java.io.Closeable; +import java.io.IOException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.TimeUnit; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalListener; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Executor side cache of live Beam {@link UnboundedReader}s, keyed by (checkpoint location, source + * id, split id). + * + *

A Spark micro-batch creates a fresh {@link BeamPartitionReader} every batch, but a Beam + * unbounded reader is expensive to create and, more importantly, holds the read position. Keeping + * the reader alive between micro-batches lets the next batch simply continue where the previous one + * stopped. This mirrors what the legacy Spark runner does in {@code + * org.apache.beam.runners.spark.io.MicrobatchSource}. + * + *

POC scope, weak failure recovery. The {@code CheckpointMark} of every split is cached + * in executor memory only, it is never written to the Spark checkpoint location and never restored + * from it. If an executor JVM dies, or the cache entry expires, the reader is recreated from the + * last mark that happens to still be in memory, and from scratch if there is none. Data loss or + * duplication across executor failures is therefore possible. Making recovery durable means + * committing marks through the Spark offset log and is deliberately out of scope for the POC. + */ +public final class BeamReaderCache { + + private static final Logger LOG = LoggerFactory.getLogger(BeamReaderCache.class); + + /** Readers idle for longer than this are closed, releasing the underlying source connections. */ + private static final long READER_CACHE_INTERVAL_MILLIS = 10 * 60 * 1000L; + + private static final RemovalListener> CLOSE_ON_REMOVAL = + notification -> { + CachedReader reader = notification.getValue(); + String key = String.valueOf(notification.getKey()); + if (reader != null) { + LOG.info("Evicting cached Beam reader {}.", key); + try { + reader.close(); + } catch (IOException e) { + LOG.warn("Failed to close evicted Beam reader {}.", key, e); + } + } + }; + + private static final Cache> READERS = + CacheBuilder.newBuilder() + .expireAfterAccess(READER_CACHE_INTERVAL_MILLIS, TimeUnit.MILLISECONDS) + .removalListener(CLOSE_ON_REMOVAL) + .build(); + + /** Last known checkpoint mark per key, used when a reader has to be recreated. */ + private static final ConcurrentMap MARKS = new ConcurrentHashMap<>(); + + private BeamReaderCache() {} + + /** Builds the cache key of one split of one source of one streaming query. */ + public static String key(String checkpointLocation, String sourceId, int splitId) { + return checkpointLocation + '|' + sourceId + '|' + splitId; + } + + /** + * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if + * there is none. + */ + @SuppressWarnings({"unchecked", "nullness"}) // the mark type always matches the source + public static CachedReader getOrCreate( + String key, UnboundedSource source, PipelineOptions options) { + try { + return (CachedReader) + READERS.get( + key, + () -> { + CheckpointMarkT mark = (CheckpointMarkT) MARKS.get(key); + LOG.info( + "No cached Beam reader for {}, creating one at checkpoint mark {}.", key, mark); + return new CachedReader<>(source.createReader(options, mark)); + }); + } catch (Exception e) { + throw new IllegalStateException("Failed to get or create Beam unbounded reader " + key, e); + } + } + + /** Remembers the checkpoint mark of {@code key} so a recreated reader can resume from it. */ + public static void rememberCheckpointMark(String key, @Nullable CheckpointMark mark) { + if (mark != null) { + MARKS.put(key, mark); + } + } + + /** Closes and forgets every cached reader, intended for tests and for query shutdown. */ + @VisibleForTesting + public static void invalidateAll() { + READERS.invalidateAll(); + READERS.cleanUp(); + MARKS.clear(); + } + + /** A cached {@link UnboundedReader} that remembers whether it has been started already. */ + public static final class CachedReader implements Closeable { + private final UnboundedReader reader; + private boolean started; + + CachedReader(UnboundedReader reader) { + this.reader = reader; + } + + /** The wrapped Beam reader. */ + public UnboundedReader reader() { + return reader; + } + + /** + * Starts the reader on first use and advances it afterwards, returning {@code true} if an + * element is available. + */ + public synchronized boolean startOrAdvance() throws IOException { + if (!started) { + started = true; + return reader.start(); + } + return reader.advance(); + } + + @Override + public void close() throws IOException { + reader.close(); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java new file mode 100644 index 000000000000..b9a1100f11e4 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java @@ -0,0 +1,147 @@ +/* + * 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.io.streaming; + +import java.io.Serializable; +import java.util.Base64; +import java.util.Map; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableProvider; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.sources.DataSourceRegister; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.checkerframework.checker.nullness.qual.NonNull; + +/** + * Spark DataSourceV2 {@link TableProvider} exposing an arbitrary Beam {@link + * org.apache.beam.sdk.io.UnboundedSource} as a micro-batch streaming source. + * + *

The produced rows always have exactly two columns: + * + *

+ * + *

Deliberately no Catalyst encoder is generated for Beam types, everything stays opaque bytes + * until a downstream translator decodes it. + * + *

Translators should not reference this class directly, use {@link UnboundedSourceDataset#of} + * instead. + * + *

Note on the format name: this class implements {@link DataSourceRegister} and reports the + * short name {@value #SHORT_NAME}, but no {@code META-INF/services} entry is shipped, so the short + * name is not resolvable through the {@code ServiceLoader}. Use {@link #FORMAT}, the fully + * qualified class name, as the argument of {@code DataStreamReader.format(String)}. + */ +public class BeamStreamingSource implements TableProvider, DataSourceRegister { + + /** Short name of this source, see the class level note about {@code META-INF/services}. */ + public static final String SHORT_NAME = "beam-unbounded"; + + /** Format string to pass to {@code DataStreamReader.format(String)}. */ + public static final String FORMAT = + "org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamStreamingSource"; + + /** + * Base64 of the Java serialized {@link org.apache.beam.sdk.io.UnboundedSource}. + * + *

All option keys are lower case on purpose, Spark funnels DataSourceV2 options through {@link + * CaseInsensitiveStringMap}. + */ + public static final String OPT_SOURCE = "beam.source"; + + /** Base64 of the Java serialized {@code Coder>}. */ + public static final String OPT_CODER = "beam.coder"; + + /** + * Base64 of the Java serialized {@link + * org.apache.beam.runners.core.construction.SerializablePipelineOptions}. + */ + public static final String OPT_PIPELINE_OPTIONS = "beam.pipelineoptions"; + + /** Identifier making the reader cache key unique per source instance. */ + public static final String OPT_SOURCE_ID = "beam.sourceid"; + + /** Desired number of splits handed to {@code UnboundedSource.split}. */ + public static final String OPT_NUM_SPLITS = "beam.numsplits"; + + /** Maximum number of records read per split per micro-batch. */ + public static final String OPT_MAX_RECORDS = "beam.maxrecords"; + + /** Maximum wall clock duration of a single micro-batch read, in milliseconds. */ + public static final String OPT_MAX_BATCH_DURATION_MILLIS = "beam.maxbatchdurationmillis"; + + /** The fixed two column schema of this source. */ + public static final StructType SCHEMA = + new StructType() + .add(UnboundedSourceDataset.COL_PAYLOAD, DataTypes.BinaryType, false) + .add(UnboundedSourceDataset.COL_EVENT_TS, DataTypes.TimestampType, false); + + /** Required public no-arg constructor, Spark instantiates this provider reflectively. */ + public BeamStreamingSource() {} + + @Override + public String shortName() { + return SHORT_NAME; + } + + @Override + public StructType inferSchema(CaseInsensitiveStringMap options) { + return SCHEMA; + } + + @Override + public Table getTable( + StructType schema, Transform[] partitioning, Map properties) { + return new BeamStreamingTable(new CaseInsensitiveStringMap(properties)); + } + + @Override + public boolean supportsExternalMetadata() { + return false; + } + + /** Base64 encodes the Java serialized form of {@code value}. */ + static String encode(Serializable value) { + return Base64.getEncoder().encodeToString(SerializableUtils.serializeToByteArray(value)); + } + + /** Inverse of {@link #encode}, {@code description} is only used in error messages. */ + @SuppressWarnings("unchecked") + static T decode(String encoded, String description) { + return (T) + SerializableUtils.deserializeFromByteArray( + Base64.getDecoder().decode(encoded), description); + } + + /** Reads a required option, failing loudly rather than silently defaulting. */ + static String required(CaseInsensitiveStringMap options, String key) { + String value = options.get(key); + if (value == null) { + throw new IllegalArgumentException( + "Missing required option '" + key + "' for " + SHORT_NAME + " source."); + } + return value; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java new file mode 100644 index 000000000000..4ba6c24a5cef --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java @@ -0,0 +1,100 @@ +/* + * 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.io.streaming; + +import java.util.Set; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.spark.sql.connector.catalog.SupportsRead; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCapability; +import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.connector.read.ScanBuilder; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; + +/** + * The {@link Table} returned by {@link BeamStreamingSource}. It only ever declares {@link + * TableCapability#MICRO_BATCH_READ}, batch reads of an unbounded Beam source are handled elsewhere. + */ +public class BeamStreamingTable implements Table, SupportsRead { + + private final CaseInsensitiveStringMap options; + + BeamStreamingTable(CaseInsensitiveStringMap options) { + this.options = options; + } + + @Override + public String name() { + return "BeamUnboundedSource[" + + options.getOrDefault(BeamStreamingSource.OPT_SOURCE_ID, "?") + + "]"; + } + + @Override + public StructType schema() { + return BeamStreamingSource.SCHEMA; + } + + @Override + public Set capabilities() { + return ImmutableSet.of(TableCapability.MICRO_BATCH_READ); + } + + @Override + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap scanOptions) { + // Spark hands the full DataSourceV2 option map to both getTable and newScanBuilder. Prefer the + // scan options and fall back to the table properties for anything missing. + CaseInsensitiveStringMap merged = merge(options, scanOptions); + return () -> new BeamScan(merged); + } + + private static CaseInsensitiveStringMap merge( + CaseInsensitiveStringMap base, CaseInsensitiveStringMap override) { + java.util.Map map = new java.util.HashMap<>(base.asCaseSensitiveMap()); + map.putAll(override.asCaseSensitiveMap()); + return new CaseInsensitiveStringMap(map); + } + + /** The {@link Scan} of a Beam unbounded source, micro-batch only. */ + private static class BeamScan implements Scan { + private final CaseInsensitiveStringMap options; + + BeamScan(CaseInsensitiveStringMap options) { + this.options = options; + } + + @Override + public StructType readSchema() { + return BeamStreamingSource.SCHEMA; + } + + @Override + public String description() { + return "BeamUnboundedSource[" + + options.getOrDefault(BeamStreamingSource.OPT_SOURCE_ID, "?") + + "]"; + } + + @Override + public MicroBatchStream toMicroBatchStream(String checkpointLocation) { + return new BeamMicroBatchStream(options, checkpointLocation); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java new file mode 100644 index 000000000000..7e4af6687cfc --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -0,0 +1,106 @@ +/* + * 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.io.streaming; + +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +/** + * Translator facing entry point turning a Beam {@link UnboundedSource} into a streaming Spark + * {@link Dataset} of rows. + * + *

The returned dataset has exactly two columns, {@value #COL_PAYLOAD} of type {@code BINARY} + * carrying the element encoded with the supplied {@code WindowedValue} coder, and {@value + * #COL_EVENT_TS} of type {@code TIMESTAMP} carrying the event timestamp of that element. + * + *

The watermark is declared here and only here. Spark 4 rejects a second {@code + * withWatermark} declaration further down the plan, so this is the single declaration point for a + * whole Beam pipeline. Downstream translators must never call {@code withWatermark} again, they + * simply keep transforming the dataset. Both columns are still present in the returned dataset, the + * event timestamp column has to survive at least until the first stateful operator for the + * watermark to be meaningful. + */ +public final class UnboundedSourceDataset { + + /** Name of the binary column holding the encoded {@code WindowedValue}. */ + public static final String COL_PAYLOAD = "payload"; + + /** Name of the timestamp column holding the Beam event timestamp. */ + public static final String COL_EVENT_TS = "eventTimestamp"; + + /** Upper bound on the number of splits requested from a source, keeps the POC predictable. */ + private static final int MAX_DESIRED_SPLITS = 8; + + private UnboundedSourceDataset() {} + + /** + * Builds the streaming {@link Dataset} for {@code source}, with the event time watermark already + * applied. + * + * @param session the active Spark session + * @param source the Beam unbounded source to read + * @param windowedValueCoder the coder used to encode the {@value #COL_PAYLOAD} column, normally a + * {@code WindowedValues.FullWindowedValueCoder} + * @param options the pipeline options, supplying the watermark delay and the micro-batch limits + * @param the element type of the source + * @param the checkpoint mark type of the source + */ + public static Dataset of( + SparkSession session, + UnboundedSource source, + Coder> windowedValueCoder, + SparkStructuredStreamingPipelineOptions options) { + + Map readerOptions = new HashMap<>(); + readerOptions.put(BeamStreamingSource.OPT_SOURCE, BeamStreamingSource.encode(source)); + readerOptions.put( + BeamStreamingSource.OPT_CODER, BeamStreamingSource.encode(windowedValueCoder)); + readerOptions.put( + BeamStreamingSource.OPT_PIPELINE_OPTIONS, + BeamStreamingSource.encode(new SerializablePipelineOptions(options))); + readerOptions.put(BeamStreamingSource.OPT_SOURCE_ID, UUID.randomUUID().toString()); + readerOptions.put( + BeamStreamingSource.OPT_NUM_SPLITS, Integer.toString(desiredNumSplits(session))); + readerOptions.put( + BeamStreamingSource.OPT_MAX_RECORDS, + Integer.toString(options.getMaxRecordsPerMicroBatch())); + readerOptions.put( + BeamStreamingSource.OPT_MAX_BATCH_DURATION_MILLIS, + Long.toString(options.getMaxBatchDurationMillis())); + + Dataset rows = + session.readStream().format(BeamStreamingSource.FORMAT).options(readerOptions).load(); + + // Exactly one watermark declaration per pipeline, see the class javadoc. + return rows.withWatermark(COL_EVENT_TS, options.getWatermarkDelayMillis() + " milliseconds"); + } + + private static int desiredNumSplits(SparkSession session) { + int parallelism = session.sparkContext().defaultParallelism(); + return Math.max(1, Math.min(MAX_DESIRED_SPLITS, parallelism)); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java new file mode 100644 index 000000000000..203b8465a719 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -0,0 +1,477 @@ +/* + * 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.io.streaming; + +import static org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.COL_EVENT_TS; +import static org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.COL_PAYLOAD; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import java.io.Serializable; +import java.sql.Timestamp; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +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.spark.api.java.function.MapFunction; +import org.apache.spark.api.java.function.VoidFunction2; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark; +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryProgress; +import org.apache.spark.sql.streaming.Trigger; +import org.joda.time.Instant; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for the Spark 4 DataSourceV2 micro-batch source wrapping a Beam {@link UnboundedSource}. + * + *

Termination note: the epoch offsets of this source never settle, {@code + * StreamingQuery.processAllAvailable()} would therefore block forever. Every test here drives the + * query with {@code Trigger.ProcessingTime(100)} and stops it explicitly once the expected result + * arrived or the poll deadline expired. + */ +@Category(StreamingTest.class) +@RunWith(JUnit4.class) +public class BeamMicroBatchSourceTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder temp = new TemporaryFolder(); + + private static final AtomicInteger QUERY_COUNTER = new AtomicInteger(); + + /** Rows collected per query name by {@link #startCollecting}, driver side. */ + private static final Map> COLLECTED = new ConcurrentHashMap<>(); + + /** 2023-11-14T22:13:20Z, a plain modern timestamp with no rebase or DST subtleties. */ + private static final long BASE_MILLIS = 1_700_000_000_000L; + + private static final long INTERVAL_MILLIS = 1_000L; + + private static final long POLL_TIMEOUT_MILLIS = 120_000L; + + @After + public void tearDown() { + BeamReaderCache.invalidateAll(); + COLLECTED.clear(); + } + + /** + * THE TOP RISK OF THE POC: does Spark's {@code EventTimeWatermark} logical node survive a typed + * transformation producing a Beam typed dataset, so a later stateful operator can still see it? + * + *

Verdict, observed on Spark 4: it does. The node stays at the bottom of both the logical and + * the analyzed plan through one and through two typed maps, {@code DeserializeToObject / + * MapElements / SerializeFromObject} are simply stacked on top of it. + * + *

One nuance the translators must be aware of: the per attribute delay marker (rendered as + * {@code eventTimestamp#1-T1000ms} in the analyzed plan) only travels with the timestamp + * attribute itself. Once a typed map projects the timestamp column away, downstream attributes + * carry no marker. Operators that read the marker off an attribute, for example {@code + * groupBy(window(...))} or a stream to stream join, would therefore not see it. Operators that + * read the query wide watermark, which is what {@code transformWithState} in {@code + * TimeMode.EventTime} does, are unaffected, see {@link + * #testWatermarkIsTrackedAtRuntimeAfterTypedMap}. + */ + @Test(timeout = 300_000) + public void testEventTimeWatermarkSurvivesTypedMap() { + Dataset rows = rows(4, 1_000L); + assertTrue("source dataset must be streaming", rows.isStreaming()); + + assertWatermark("directly after withWatermark, logical plan", logical(rows)); + assertWatermark("directly after withWatermark, analyzed plan", analyzed(rows)); + + // A typed map producing a Beam-ish (opaque bytes) dataset, exactly what the read translator + // will do to turn rows into WindowedValue bytes. + Dataset typed = + rows.map((MapFunction) row -> row.getAs(COL_PAYLOAD), Encoders.BINARY()); + + assertWatermark("after a typed map, logical plan", logical(typed)); + assertWatermark("after a typed map, analyzed plan", analyzed(typed)); + + // And once more after a second typed stage, mimicking chained operators. + Dataset chained = + typed.map((MapFunction) bytes -> bytes, Encoders.BINARY()); + + assertWatermark("after two chained typed maps, logical plan", logical(chained)); + assertWatermark("after two chained typed maps, analyzed plan", analyzed(chained)); + } + + /** + * Runtime counterpart of the plan inspection above, the watermark must actually be tracked by the + * running query, not merely present as a plan node. + */ + @Test(timeout = 300_000) + public void testWatermarkIsTrackedAtRuntimeAfterTypedMap() throws Exception { + Dataset rows = rows(8, 0L); + Dataset typed = + rows.map((MapFunction) row -> row.getAs(COL_PAYLOAD), Encoders.BINARY()); + + String queryName = "beam_wm_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = startDiscarding(typed, queryName); + try { + String watermark = awaitWatermark(query); + assertNotNull("query never reported an event time watermark", watermark); + assertFalse( + "watermark never advanced past the epoch, it is not being tracked: " + watermark, + watermark.startsWith("1970-")); + } finally { + stopQuietly(query); + } + } + + /** Reads a finite set of elements through the DSv2 source and checks payloads and timestamps. */ + @Test(timeout = 300_000) + public void testReadsElementsFromUnboundedSource() throws Exception { + int count = 8; + Dataset rows = rows(count, 0L); + + String queryName = "beam_read_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = startCollecting(rows, queryName); + List collected; + try { + collected = awaitRows(queryName, count); + } finally { + stopQuietly(query); + } + + assertEquals("unexpected number of rows", count, collected.size()); + + Coder> coder = coder(); + List values = new ArrayList<>(); + for (Row row : collected) { + byte[] payload = row.getAs(COL_PAYLOAD); + Timestamp eventTs = row.getAs(COL_EVENT_TS); + WindowedValue windowedValue = CoderUtils.decodeFromByteArray(coder, payload); + values.add(windowedValue.getValue()); + assertEquals( + "eventTimestamp column must match the timestamp inside the encoded WindowedValue", + windowedValue.getTimestamp().getMillis(), + eventTs.getTime()); + assertEquals( + "elements are read into the global window", + Collections.singletonList(GlobalWindow.INSTANCE), + new ArrayList<>(windowedValue.getWindows())); + } + + Collections.sort(values); + List expected = new ArrayList<>(); + for (int i = 0; i < count; i++) { + expected.add(element(i)); + } + assertEquals(expected, values); + } + + /** The offset is an opaque, strictly increasing epoch counter. */ + @Test(timeout = 300_000) + public void testEpochOffsetRoundTrip() { + BeamOffset offset = new BeamOffset(42L); + assertEquals("{\"epoch\":42}", offset.json()); + assertEquals(offset, BeamOffset.fromJson(offset.json())); + assertEquals(0L, BeamOffset.ZERO.epoch()); + } + + // --------------------------------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------------------------------- + + private Dataset rows(int count, long watermarkDelayMillis) { + SparkStructuredStreamingPipelineOptions options = + org.apache.beam.sdk.options.PipelineOptionsFactory.create() + .as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(watermarkDelayMillis); + options.setMaxRecordsPerMicroBatch(1000); + options.setMaxBatchDurationMillis(200L); + return UnboundedSourceDataset.of(SESSION.getSession(), new ListSource(count), coder(), options); + } + + private static Coder> coder() { + return WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + } + + private static String element(int index) { + return "element-" + index; + } + + private static long timestampMillis(int index) { + return BASE_MILLIS + index * INTERVAL_MILLIS; + } + + /** + * Starts a query that throws its output away. + * + *

The {@code noop} sink is used on purpose. Spark test JVMs run with {@code + * spark.kryo.registrationRequired=true} (see {@code runners/spark/spark_runner.gradle}) and the + * {@code memory} sink returns a {@code MemoryWriterCommitMessage} that no Beam Kryo registrator + * knows about, so a memory sink query fails immediately. The {@code noop} sink commits {@code + * null} and its {@code DataWritingSparkTaskResult} is already registered, which is also why the + * streaming evaluation context writes to {@code noop}. + */ + private StreamingQuery startDiscarding(Dataset dataset, String queryName) throws Exception { + return dataset + .writeStream() + .format("noop") + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + /** + * Starts a query collecting every micro-batch into {@link #COLLECTED} under {@code queryName}. + */ + private StreamingQuery startCollecting(Dataset dataset, String queryName) throws Exception { + COLLECTED.put(queryName, Collections.synchronizedList(new ArrayList<>())); + return dataset + .writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + // Exactly one action per micro-batch: any second action would re-execute the + // batch and advance the Beam reader past records that were never collected. + List target = COLLECTED.get(queryName); + if (target != null) { + target.addAll(batch.collectAsList()); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + private static void stopQuietly(StreamingQuery query) { + try { + query.stop(); + } catch (Exception e) { + // Nothing useful to do while tearing a test query down. + } + } + + /** Polls the collected batches until {@code expected} rows arrived or the deadline expires. */ + private static List awaitRows(String queryName, int expected) throws Exception { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + List rows = COLLECTED.getOrDefault(queryName, Collections.emptyList()); + while (System.currentTimeMillis() < deadline) { + synchronized (rows) { + if (rows.size() >= expected) { + return new ArrayList<>(rows); + } + } + Thread.sleep(100L); + } + synchronized (rows) { + return new ArrayList<>(rows); + } + } + + /** Polls the query progress until it reports an event time watermark past the epoch. */ + private static @Nullable String awaitWatermark(StreamingQuery query) throws Exception { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + String last = null; + while (System.currentTimeMillis() < deadline) { + for (StreamingQueryProgress progress : query.recentProgress()) { + Map eventTime = progress.eventTime(); + String watermark = eventTime.get("watermark"); + if (watermark != null) { + last = watermark; + if (!watermark.startsWith("1970-")) { + return watermark; + } + } + } + Thread.sleep(100L); + } + return last; + } + + private static LogicalPlan logical(Dataset dataset) { + return ((org.apache.spark.sql.classic.Dataset) dataset).queryExecution().logical(); + } + + private static LogicalPlan analyzed(Dataset dataset) { + return ((org.apache.spark.sql.classic.Dataset) dataset).queryExecution().analyzed(); + } + + private static void assertWatermark(String what, LogicalPlan plan) { + assertTrue( + "no EventTimeWatermark node found " + what + ":\n" + plan.treeString(), + containsWatermark(plan)); + } + + private static boolean containsWatermark(LogicalPlan plan) { + if (plan instanceof EventTimeWatermark) { + return true; + } + scala.collection.Iterator children = plan.children().iterator(); + while (children.hasNext()) { + if (containsWatermark(children.next())) { + return true; + } + } + return false; + } + + // --------------------------------------------------------------------------------------------- + // a minimal in-memory UnboundedSource + // --------------------------------------------------------------------------------------------- + + /** + * A trivial single split {@link UnboundedSource} over a fixed number of synthetic elements with + * evenly spaced event timestamps. It is exhausted after the last element, further calls to {@code + * advance()} simply report that no data is available. + */ + private static class ListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + private final int count; + + ListSource(int count) { + this.count = count; + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + return Arrays.asList(this); + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark mark) { + return new ListReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return SerializableCoder.of(Mark.class); + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + /** Position of the next element to read. */ + static class Mark implements UnboundedSource.CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; + + private final int next; + + Mark(int next) { + this.next = next; + } + + @Override + public void finalizeCheckpoint() {} + } + + private static class ListReader extends UnboundedReader { + private final ListSource source; + private int next; + private int current = -1; + + ListReader(ListSource source, int next) { + this.source = source; + this.next = next; + } + + @Override + public boolean start() { + return advance(); + } + + @Override + public boolean advance() { + if (next < source.count) { + current = next++; + return true; + } + return false; + } + + @Override + public String getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return element(current); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(timestampMillis(current)); + } + + @Override + public Instant getWatermark() { + return current < 0 + ? BoundedWindow.TIMESTAMP_MIN_VALUE + : new Instant(timestampMillis(current)); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new Mark(next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() throws IOException {} + } + } +} From 3ffc438e23d0c8730b4db262f343142f26fbb08e Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 13:25:15 +0000 Subject: [PATCH 03/21] [Spark 4] Bridge Beam state and timers onto transformWithState Adds a single generic transformWithState operator that can host any keyed Beam transform, both the stateful ParDo stack and the group also by window stack that implements a windowed GroupByKey. The mode is selected by BeamStatefulProcessorConfig, everything else is shared. Keys, inputs and outputs are raw Beam coder bytes and every Spark encoder involved is BINARY or STRING, so Catalyst never has to derive a schema for a Beam type. TwsTransformFactory owns the row layout and documents it. Beam state is bridged by TwsStateInternals on top of BytesKV, a tiny string to bytes store backed by exactly one Spark MapState. One map is a requirement rather than a preference, ReduceFnRunner invents state tags at runtime so the set of state addresses is not known when init has to declare its state variables. Timers are bridged by TwsTimerInternals, which keeps the full TimerData in a second map and registers only bare wake ups with Spark, reconciling them against listTimers so an expiry can never be registered twice. Two boundary conditions needed care. Spark deletes the expiry it is firing after the callback returns, so a wake up re armed at that same millisecond is nudged one millisecond forward. Spark also expires a wake up as soon as the expiry is at or before the batch watermark, while Beam only fires an event time timer once the watermark is strictly past it, so timers sitting exactly on the watermark are withheld and re armed instead of being handed to Beam, which would decline to fire them and silently drop the on time pane. Processing time timers are rejected with a clear message, transformWithState runs in a single TimeMode and this operator uses event time. The state and timer bridges are unit tested against an in memory store, and BeamStatefulProcessorTest runs both modes through a real streaming query on a live SparkSession. --- .../streaming/TwsTransformFactory.java | 198 ++++++ .../state/BeamStatefulProcessor.java | 429 +++++++++++++ .../state/BeamStatefulProcessorConfig.java | 242 ++++++++ .../translation/streaming/state/BytesKV.java | 106 ++++ .../streaming/state/TwsStateInternals.java | 578 ++++++++++++++++++ .../streaming/state/TwsTimerInternals.java | 390 ++++++++++++ .../state/BeamStatefulProcessorTest.java | 443 ++++++++++++++ .../state/TwsStateInternalsTest.java | 359 +++++++++++ .../state/TwsTimerInternalsTest.java | 418 +++++++++++++ 9 files changed, 3163 insertions(+) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/TwsTransformFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorConfig.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BytesKV.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternals.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternalsTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternalsTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/TwsTransformFactory.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/TwsTransformFactory.java new file mode 100644 index 000000000000..232a6833259b --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/TwsTransformFactory.java @@ -0,0 +1,198 @@ +/* + * 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 java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.UncheckedIOException; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessor; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessorConfig; +import org.apache.beam.sdk.util.VarInt; +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; + +/** + * The translator facing entry point of the Spark 4 stateful streaming bridge: it groups a keyed + * dataset of raw Beam bytes by key and runs a {@link BeamStatefulProcessor} over it with Spark's + * {@code transformWithState}. + * + *

Everything on the wire is {@code byte[]} and every Spark encoder involved is {@code + * Encoders.BINARY()}. That is deliberate: it keeps Catalyst encoders, and therefore Catalyst + * schemas for Beam types, entirely out of the stateful path. + * + *

Input row encoding

+ * + *

Each element of the input {@code Dataset} is one Beam element, encoded as + * + * 

{@code
+ * varint32(keyBytes.length) || keyBytes || windowedValueBytes
+ * }
+ * + * where + * + * + * + *

Use {@link #encodeInputRow(byte[], byte[])} to build such a row. There is no length prefix on + * the payload, it simply runs to the end of the array. + * + *

Output row encoding

+ * + *

Each element of the returned {@code Dataset} is one tagged Beam output, encoded as + * + * 

{@code
+ * varint32(outputTagIndex) || windowedValueBytes
+ * }
+ * + * where {@code outputTagIndex} is the index of the emitting {@link + * org.apache.beam.sdk.values.TupleTag} in {@code config.outputTags()}, that is {@code 0} for the + * main output tag and {@code 1..n} for {@code config.additionalOutputTags()} in their configured + * order, and {@code windowedValueBytes} is the emitted {@code WindowedValue} encoded with {@code + * config.outputCoderFor(tag)}. Use {@link #outputTagIndex(byte[])} and {@link + * #outputPayload(byte[])} to take such a row apart, typically with one {@code filter} plus one + * {@code map} per output tag. + * + *

For {@link BeamStatefulProcessorConfig.Mode#GROUP_ALSO_BY_WINDOW} there is only the main + * output and its element type is {@code KV>}, so its configured output coder must be + * {@code KvCoder.of(keyCoder, IterableCoder.of(valueCoder))}. + * + *

Semantics

+ * + *

The operator always runs with {@code TimeMode.EventTime()} and {@code OutputMode.Append()}. + * The input dataset must already carry an event time watermark declared upstream with {@code + * withWatermark}, Spark forbids re-declaring it here. + */ +public final class TwsTransformFactory { + + private TwsTransformFactory() {} + + /** + * Groups {@code keyedInput} by the Beam key embedded in every row and runs the configured Beam + * transform over it inside Spark's {@code transformWithState}. + * + * @param keyedInput rows in the input encoding documented on this class + * @param config what to run and how to decode the rows + * @return rows in the output encoding documented on this class + */ + public static Dataset transform( + Dataset keyedInput, BeamStatefulProcessorConfig config) { + return keyedInput + .groupByKey((MapFunction) TwsTransformFactory::inputKey, Encoders.BINARY()) + .transformWithState( + new BeamStatefulProcessor(config), + TimeMode.EventTime(), + OutputMode.Append(), + Encoders.BINARY()); + } + + /** Builds an input row from the encoded Beam key and the encoded {@code WindowedValue}. */ + public static byte[] encodeInputRow(byte[] keyBytes, byte[] windowedValueBytes) { + ByteArrayOutputStream out = + new ByteArrayOutputStream( + VarInt.getLength(keyBytes.length) + keyBytes.length + windowedValueBytes.length); + try { + VarInt.encode(keyBytes.length, out); + out.write(keyBytes); + out.write(windowedValueBytes); + } catch (IOException e) { + throw new UncheckedIOException("Failed to encode a transformWithState input row", e); + } + return out.toByteArray(); + } + + /** Extracts the encoded Beam key, the Spark grouping key, from an input row. */ + public static byte[] inputKey(byte[] row) { + ByteArrayInputStream in = new ByteArrayInputStream(row); + int keyLength = readVarInt(in); + return readExactly(in, keyLength); + } + + /** Extracts the encoded {@code WindowedValue} payload from an input row. */ + public static byte[] inputPayload(byte[] row) { + ByteArrayInputStream in = new ByteArrayInputStream(row); + int keyLength = readVarInt(in); + skip(in, keyLength); + return readExactly(in, in.available()); + } + + /** Builds an output row from the output tag index and the encoded {@code WindowedValue}. */ + public static byte[] encodeOutputRow(int outputTagIndex, byte[] windowedValueBytes) { + ByteArrayOutputStream out = + new ByteArrayOutputStream(VarInt.getLength(outputTagIndex) + windowedValueBytes.length); + try { + VarInt.encode(outputTagIndex, out); + out.write(windowedValueBytes); + } catch (IOException e) { + throw new UncheckedIOException("Failed to encode a transformWithState output row", e); + } + return out.toByteArray(); + } + + /** Extracts the output tag index from an output row. */ + public static int outputTagIndex(byte[] row) { + return readVarInt(new ByteArrayInputStream(row)); + } + + /** Extracts the encoded {@code WindowedValue} payload from an output row. */ + public static byte[] outputPayload(byte[] row) { + ByteArrayInputStream in = new ByteArrayInputStream(row); + readVarInt(in); + return readExactly(in, in.available()); + } + + private static int readVarInt(ByteArrayInputStream in) { + try { + return VarInt.decodeInt(in); + } catch (IOException e) { + throw new UncheckedIOException("Malformed transformWithState row, bad length prefix", e); + } + } + + private static void skip(ByteArrayInputStream in, int length) { + long skipped = in.skip(length); + if (skipped != length) { + throw new IllegalArgumentException( + "Malformed transformWithState row, expected " + length + " more bytes"); + } + } + + private static byte[] readExactly(ByteArrayInputStream in, int length) { + byte[] bytes = new byte[length]; + int read = in.read(bytes, 0, length); + if (read != length && length > 0) { + throw new IllegalArgumentException( + "Malformed transformWithState row, expected " + + length + + " bytes but only " + + Math.max(read, 0) + + " were available"); + } + return bytes; + } +} 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 new file mode 100644 index 000000000000..0eca7e26a113 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java @@ -0,0 +1,429 @@ +/* + * 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.state; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.DoFnRunner; +import org.apache.beam.runners.core.DoFnRunners; +import org.apache.beam.runners.core.GroupAlsoByWindowViaWindowSetNewDoFn; +import org.apache.beam.runners.core.KeyedWorkItem; +import org.apache.beam.runners.core.KeyedWorkItems; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateInternalsFactory; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StatefulDoFnRunner; +import org.apache.beam.runners.core.StepContext; +import org.apache.beam.runners.core.SystemReduceFn; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.core.TimerInternalsFactory; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.TwsTransformFactory; +import org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; +import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.streaming.ExpiredTimerInfo; +import org.apache.spark.sql.streaming.MapState; +import org.apache.spark.sql.streaming.OutputMode; +import org.apache.spark.sql.streaming.StatefulProcessor; +import org.apache.spark.sql.streaming.TTLConfig; +import org.apache.spark.sql.streaming.TimeMode; +import org.apache.spark.sql.streaming.TimerValues; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A single Spark 4 {@code transformWithState} operator that can host any keyed Beam transform: a + * stateful {@code ParDo} or the group-also-by-window that implements a windowed {@code GroupByKey}. + * Which one is decided by {@link BeamStatefulProcessorConfig#mode()}. + * + *

Keys, inputs and outputs are all raw Beam coder bytes and every Spark encoder involved is + * {@code Encoders.BINARY()} or {@code Encoders.STRING()}, so no Catalyst schema is ever derived for + * a Beam type. The exact row layouts are documented on {@link TwsTransformFactory}. + * + *

State layout

+ * + *

Two Spark {@code MapState}s are declared, both {@code String -> byte[]}: + * + *

+ * + *

Both are declared with {@code TTLConfig.NONE()}: state lifetime is governed by Beam's own + * garbage collection timers, not by Spark's. + * + *

Bundles

+ * + *

Spark invokes {@code handleInputRows} once per key per micro-batch, and the Beam {@code + * DoFnRunner} has to be bound to that key's state, so a Beam bundle here is one key inside one + * micro-batch rather than the whole micro-batch. {@code startBundle} and {@code finishBundle} are + * therefore called around each invocation that has work to do, and skipped entirely for an + * invocation with neither elements nor due timers. {@code setup} and {@code teardown} are called + * once per Spark task, from {@link #init} and {@link #close}. + * + *

Watermarks

+ * + *

{@code TimerValues.getCurrentWatermarkInMs()} is the batch start watermark, that is the + * watermark Spark computed at the end of the previous micro-batch. An element therefore can never + * be considered late with respect to its own micro-batch, and an end-of-window timer fires in the + * micro-batch after the one whose data crossed the end of the window. That is the same one batch + * delay every micro-batch runner has. + * + *

Timers

+ * + *

Beam timers fire only in {@code handleExpiredTimer}, never in {@code handleInputRows}. A timer + * set while processing elements is picked up by Spark's own timer scan for the same micro-batch if + * it is already due, so nothing is delayed by that choice, and it removes any risk of firing a + * timer twice. Only event time timers are supported, see {@link TwsTimerInternals}. + * + *

Spark and Beam disagree on the firing boundary by one millisecond, Spark expiring a wake-up at + * {@code expiry <= watermark} and Beam requiring the watermark to be strictly past the timer. The + * gap is bridged in {@link TwsTimerInternals#removeTimersReadyToFire(long)}, which is where the + * reasoning lives; getting it wrong silently loses on-time panes rather than merely reordering + * them. + */ +@SuppressWarnings({ + "rawtypes", + "unchecked", + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class BeamStatefulProcessor extends StatefulProcessor { + + /** Name of the Spark state variable holding all Beam state. */ + public static final String BEAM_STATE_NAME = "beamState"; + + /** Name of the Spark state variable holding the encoded Beam timers. */ + public static final String BEAM_TIMER_STATE_NAME = "beamTimers"; + + private final BeamStatefulProcessorConfig config; + + private transient MapState beamState; + private transient MapState beamTimers; + private transient PipelineOptions options; + private transient @Nullable DoFn doFn; + private transient @Nullable DoFnInvoker doFnInvoker; + private transient Map, Integer> outputIndexes; + + public BeamStatefulProcessor(BeamStatefulProcessorConfig config) { + this.config = config; + } + + /** Returns the configuration this processor was built with. */ + public BeamStatefulProcessorConfig getConfig() { + return config; + } + + @Override + public void init(OutputMode outputMode, TimeMode timeMode) { + if (!TimeMode.EventTime().equals(timeMode)) { + throw new UnsupportedOperationException( + "BeamStatefulProcessor requires TimeMode.EventTime() but was initialised with " + + timeMode + + ". Beam event time timers cannot be expressed in any other Spark time mode."); + } + beamState = + getHandle() + .getMapState(BEAM_STATE_NAME, Encoders.STRING(), Encoders.BINARY(), TTLConfig.NONE()); + beamTimers = + getHandle() + .getMapState( + BEAM_TIMER_STATE_NAME, Encoders.STRING(), Encoders.BINARY(), TTLConfig.NONE()); + + options = config.optionsSupplier().get(); + + outputIndexes = new HashMap<>(); + List> tags = config.outputTags(); + for (int i = 0; i < tags.size(); i++) { + outputIndexes.put(tags.get(i), i); + } + + if (config.mode() == BeamStatefulProcessorConfig.Mode.STATEFUL_PARDO) { + DoFn cloned = SerializableUtils.clone(config.doFn()); + doFn = cloned; + doFnInvoker = DoFnInvokers.invokerFor(cloned); + DoFnInvokers.tryInvokeSetupFor(cloned, options); + } + } + + @Override + public scala.collection.Iterator handleInputRows( + byte[] key, scala.collection.Iterator rows, TimerValues timerValues) { + WindowedValues.FullWindowedValueCoder valueCoder = config.inputValueCoder(); + List> elements = new ArrayList<>(); + while (rows.hasNext()) { + byte[] payload = TwsTransformFactory.inputPayload(rows.next()); + elements.add(decode(valueCoder, payload, "input element")); + } + return process(key, elements, timerValues, null); + } + + @Override + public scala.collection.Iterator handleExpiredTimer( + byte[] key, TimerValues timerValues, ExpiredTimerInfo expiredTimerInfo) { + return process(key, Collections.emptyList(), timerValues, expiredTimerInfo.getExpiryTimeInMs()); + } + + @Override + public void close() { + if (doFnInvoker != null) { + doFnInvoker.invokeTeardown(); + doFnInvoker = null; + doFn = null; + } + } + + /** + * Runs one Beam bundle for a single key. + * + * @param encodedKey the Spark grouping key, the Beam key encoded with the key coder + * @param elements the elements of this micro-batch for that key, empty in a timer callback + * @param timerValues Spark's clock for this invocation + * @param firedExpiryMs the expiry Spark is firing, or {@code null} when processing elements + */ + private scala.collection.Iterator process( + byte[] encodedKey, + List> elements, + TimerValues timerValues, + @Nullable Long firedExpiryMs) { + + Object key = decode(config.keyCoder(), encodedKey, "key"); + Coder windowCoder = config.windowCoder(); + WindowingStrategy windowingStrategy = config.windowingStrategy(); + + TwsStateInternals stateInternals = TwsStateInternals.forKey(key, BytesKV.of(beamState)); + TwsTimerInternals timerInternals = + TwsTimerInternals.create( + BytesKV.of(beamTimers), + TwsTimerInternals.WakeupRegistry.of(getHandle()), + windowCoder, + new Instant(timerValues.getCurrentWatermarkInMs()), + new Instant(timerValues.getCurrentProcessingTimeInMs()), + firedExpiryMs); + + // Timers only ever fire from handleExpiredTimer, see the class javadoc. + List dueTimers = + firedExpiryMs == null + ? Collections.emptyList() + : timerInternals.removeTimersReadyToFire(firedExpiryMs); + + if (elements.isEmpty() && dueTimers.isEmpty()) { + // Nothing to do, but timer bookkeeping still has to be reconciled with Spark. + timerInternals.flush(); + return ScalaInterop.scalaIterator(Collections.emptyList()); + } + + StepContext stepContext = + new StepContext() { + @Override + public StateInternals stateInternals() { + return stateInternals; + } + + @Override + public TimerInternals timerInternals() { + return timerInternals; + } + }; + + List outputs = new ArrayList<>(); + WindowedValueMultiReceiver receiver = new EncodingReceiver(outputs); + + if (config.mode() == BeamStatefulProcessorConfig.Mode.GROUP_ALSO_BY_WINDOW) { + runGroupAlsoByWindow( + key, elements, dueTimers, stepContext, receiver, stateInternals, timerInternals); + } else { + runStatefulParDo(key, elements, dueTimers, stepContext, receiver); + } + + timerInternals.flush(); + return ScalaInterop.scalaIterator(outputs); + } + + private void runStatefulParDo( + Object key, + List> elements, + List dueTimers, + StepContext stepContext, + WindowedValueMultiReceiver receiver) { + + Coder> inputCoder = config.kvInputCoder(); + + DoFnRunner, Object> simpleRunner = + DoFnRunners.simpleRunner( + options, + (DoFn, Object>) doFn, + config.sideInputReader(), + receiver, + (TupleTag) config.mainOutputTag(), + config.additionalOutputTags(), + stepContext, + inputCoder, + config.outputCoders(), + config.windowingStrategy(), + config.doFnSchemaInformation(), + config.sideInputMapping()); + + DoFnRunner, Object> runner = + DoFnRunners.defaultStatefulDoFnRunner( + (DoFn, Object>) doFn, + inputCoder, + simpleRunner, + stepContext, + config.windowingStrategy(), + new StatefulDoFnRunner.TimeInternalsCleanupTimer<>( + stepContext.timerInternals(), config.windowingStrategy()), + new StatefulDoFnRunner.StateInternalsStateCleaner<>( + doFn, stepContext.stateInternals(), (Coder) config.windowCoder())); + + runner.startBundle(); + for (WindowedValue element : elements) { + runner.processElement(element.withValue(KV.of(key, element.getValue()))); + } + for (TimerInternals.TimerData timer : dueTimers) { + runner.onTimer( + timer.getTimerId(), + timer.getTimerFamilyId(), + key, + windowOf(timer.getNamespace()), + timer.getTimestamp(), + timer.getOutputTimestamp(), + timer.getDomain(), + timer.causedByDrain()); + } + runner.finishBundle(); + } + + private void runGroupAlsoByWindow( + Object key, + List> elements, + List dueTimers, + StepContext stepContext, + WindowedValueMultiReceiver receiver, + StateInternals stateInternals, + TimerInternals timerInternals) { + + StateInternalsFactory stateFactory = ignored -> stateInternals; + TimerInternalsFactory timerFactory = ignored -> timerInternals; + + DoFn, KV>> gabwDoFn = + (DoFn) + GroupAlsoByWindowViaWindowSetNewDoFn.create( + (WindowingStrategy) config.windowingStrategy(), + stateFactory, + timerFactory, + config.sideInputReader(), + (SystemReduceFn) SystemReduceFn.buffering(config.valueCoder()), + receiver, + (TupleTag) config.mainOutputTag()); + + DoFnRunner, KV>> runner = + DoFnRunners.simpleRunner( + options, + gabwDoFn, + config.sideInputReader(), + receiver, + (TupleTag) config.mainOutputTag(), + config.additionalOutputTags(), + stepContext, + null, // KeyedWorkItem has no coder here, SimpleDoFnRunner allows a null input coder + config.outputCoders(), + config.windowingStrategy(), + config.doFnSchemaInformation(), + config.sideInputMapping()); + + runner = + DoFnRunners.lateDataDroppingRunner( + runner, stepContext, (WindowingStrategy) config.windowingStrategy()); + + KeyedWorkItem workItem = KeyedWorkItems.workItem(key, dueTimers, elements); + + runner.startBundle(); + runner.processElement(WindowedValues.valueInGlobalWindow(workItem)); + runner.finishBundle(); + } + + private static BoundedWindow windowOf(StateNamespace namespace) { + if (namespace instanceof StateNamespaces.WindowNamespace) { + return ((StateNamespaces.WindowNamespace) namespace).getWindow(); + } + if (namespace instanceof StateNamespaces.WindowAndTriggerNamespace) { + return ((StateNamespaces.WindowAndTriggerNamespace) namespace).getWindow(); + } + throw new IllegalStateException( + "Cannot fire a Beam timer set in namespace " + + namespace.stringKey() + + ", it is not bound to a window."); + } + + private static T decode(Coder coder, byte[] bytes, String what) { + try { + return CoderUtils.decodeFromByteArray(coder, bytes); + } catch (Exception e) { + throw new IllegalStateException("Failed to decode a Beam " + what, e); + } + } + + /** Encodes every emitted element into the tagged output row layout and appends it. */ + private final class EncodingReceiver implements WindowedValueMultiReceiver { + private final List outputs; + + private EncodingReceiver(List outputs) { + this.outputs = outputs; + } + + @Override + public void output(TupleTag tag, WindowedValue value) { + Integer index = outputIndexes.get(tag); + if (index == null) { + throw new IllegalStateException( + "Step " + + config.stepName() + + " emitted to unknown output tag " + + tag + + ", known tags are " + + config.outputTags()); + } + WindowedValues.FullWindowedValueCoder coder = config.outputCoderFor(tag); + try { + outputs.add( + TwsTransformFactory.encodeOutputRow(index, CoderUtils.encodeToByteArray(coder, value))); + } catch (Exception e) { + throw new IllegalStateException("Failed to encode an output of tag " + tag, e); + } + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorConfig.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorConfig.java new file mode 100644 index 000000000000..d5319ab030c9 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorConfig.java @@ -0,0 +1,242 @@ +/* + * 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.state; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.function.Supplier; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Everything {@link BeamStatefulProcessor} needs in order to run a Beam transform inside Spark 4's + * {@code transformWithState}, shipped to the executors by Java serialisation. + * + *

Two modes are supported, see {@link Mode}. Both share the same wire format, the same state and + * timer bridges and the same operator, they differ only in which {@code DoFnRunner} stack is built + * on the executor. + * + *

Everything referenced from here must be {@link Serializable}. Beam coders, {@code DoFn}s, + * {@code WindowingStrategy} and {@code TupleTag} always are; the {@link #sideInputReader()} and the + * {@link #optionsSupplier()} are the two places where a caller could accidentally pass something + * that is not, so {@code SparkSideInputReader} and the broadcast options supplier of the evaluation + * context should be used. + */ +@AutoValue +@SuppressWarnings({"rawtypes", "unchecked"}) +public abstract class BeamStatefulProcessorConfig implements Serializable { + + /** Which Beam execution stack the operator hosts. */ + public enum Mode { + /** + * Runs the user's {@code DoFn} through {@code DoFnRunners.simpleRunner} wrapped in {@code + * DoFnRunners.defaultStatefulDoFnRunner}. Input elements are {@code KV}. + */ + STATEFUL_PARDO, + + /** + * Runs {@code GroupAlsoByWindowViaWindowSetNewDoFn} with {@code SystemReduceFn.buffering}, fed + * with {@code KeyedWorkItem}s assembled from the batch's elements and the fired timers. This is + * how windowed {@code GroupByKey} is implemented. {@link #doFn()} must be unset. + */ + GROUP_ALSO_BY_WINDOW + } + + /** A {@link Supplier} of {@link PipelineOptions} that survives Java serialisation. */ + public interface OptionsSupplier extends Supplier, Serializable {} + + /** The execution stack to host. */ + public abstract Mode mode(); + + /** + * The user's {@code DoFn} for {@link Mode#STATEFUL_PARDO}, {@code null} for {@link + * Mode#GROUP_ALSO_BY_WINDOW} where the operator builds the group-also-by-window {@code DoFn} + * itself. + */ + public abstract @Nullable DoFn doFn(); + + /** Coder of the Beam key {@code K}, used to decode the {@code byte[]} grouping key. */ + public abstract Coder keyCoder(); + + /** + * Coder of the element value {@code V}, that is the value side of the input {@code KV}. + * + *

It is also the element coder of the {@code SystemReduceFn.buffering} buffer in {@link + * Mode#GROUP_ALSO_BY_WINDOW}. + */ + public abstract Coder valueCoder(); + + /** Windowing strategy of the input {@code PCollection}. */ + public abstract WindowingStrategy windowingStrategy(); + + /** The main output tag, always output index {@code 0}. */ + public abstract TupleTag mainOutputTag(); + + /** Additional output tags, output indexes {@code 1..n} in this order. */ + public abstract List> additionalOutputTags(); + + /** Element coder per output tag; must contain an entry for every tag in {@link #outputTags()}. */ + public abstract Map, Coder> outputCoders(); + + /** Reader for statically broadcast side inputs, must be serializable. */ + public abstract SideInputReader sideInputReader(); + + /** Side input mapping passed to {@code DoFnRunners.simpleRunner}. */ + public abstract Map> sideInputMapping(); + + /** Schema information of the hosted {@code DoFn}. */ + public abstract DoFnSchemaInformation doFnSchemaInformation(); + + /** Supplies the pipeline options on the executor. */ + public abstract OptionsSupplier optionsSupplier(); + + /** Human readable step name, used in error messages only. */ + public abstract String stepName(); + + /** The window coder of {@link #windowingStrategy()}. */ + public final Coder windowCoder() { + return windowingStrategy().getWindowFn().windowCoder(); + } + + /** Full windowed value coder of the input element value, the input row payload coder. */ + public final WindowedValues.FullWindowedValueCoder inputValueCoder() { + return WindowedValues.getFullCoder((Coder) valueCoder(), windowCoder()); + } + + /** + * The element coder of the input as the hosted {@code DoFn} sees it in {@link + * Mode#STATEFUL_PARDO}, that is {@code KvCoder.of(keyCoder(), valueCoder())}. + */ + public final KvCoder kvInputCoder() { + return KvCoder.of((Coder) keyCoder(), (Coder) valueCoder()); + } + + /** All output tags, main first, in output index order. */ + public final List> outputTags() { + List> tags = new ArrayList<>(additionalOutputTags().size() + 1); + tags.add(mainOutputTag()); + tags.addAll(additionalOutputTags()); + return tags; + } + + /** Full windowed value coder of the elements emitted on {@code tag}. */ + public final WindowedValues.FullWindowedValueCoder outputCoderFor(TupleTag tag) { + Coder coder = outputCoders().get(tag); + if (coder == null) { + throw new IllegalArgumentException("No output coder configured for tag " + tag); + } + return WindowedValues.getFullCoder((Coder) coder, windowCoder()); + } + + /** Returns a builder with the optional properties already defaulted. */ + public static Builder builder() { + return new AutoValue_BeamStatefulProcessorConfig.Builder() + .setAdditionalOutputTags(Collections.emptyList()) + .setOutputCoders(Collections.emptyMap()) + .setSideInputReader(EmptySideInputReader.INSTANCE) + .setSideInputMapping(Collections.emptyMap()) + .setDoFnSchemaInformation(DoFnSchemaInformation.create()) + .setStepName(""); + } + + /** Builder for {@link BeamStatefulProcessorConfig}. */ + @AutoValue.Builder + public abstract static class Builder { + + public abstract Builder setMode(Mode mode); + + public abstract Builder setDoFn(@Nullable DoFn doFn); + + public abstract Builder setKeyCoder(Coder keyCoder); + + public abstract Builder setValueCoder(Coder valueCoder); + + public abstract Builder setWindowingStrategy(WindowingStrategy windowingStrategy); + + public abstract Builder setMainOutputTag(TupleTag mainOutputTag); + + public abstract Builder setAdditionalOutputTags(List> additionalOutputTags); + + public abstract Builder setOutputCoders(Map, Coder> outputCoders); + + public abstract Builder setSideInputReader(SideInputReader sideInputReader); + + public abstract Builder setSideInputMapping(Map> sideInputMapping); + + public abstract Builder setDoFnSchemaInformation(DoFnSchemaInformation doFnSchemaInformation); + + public abstract Builder setOptionsSupplier(OptionsSupplier optionsSupplier); + + public abstract Builder setStepName(String stepName); + + abstract BeamStatefulProcessorConfig autoBuild(); + + public BeamStatefulProcessorConfig build() { + BeamStatefulProcessorConfig config = autoBuild(); + if (config.mode() == Mode.STATEFUL_PARDO) { + checkArgument(config.doFn() != null, "STATEFUL_PARDO requires a DoFn"); + } else { + checkArgument( + config.doFn() == null, + "GROUP_ALSO_BY_WINDOW builds its own DoFn, no DoFn may be configured"); + } + for (TupleTag tag : config.outputTags()) { + checkArgument( + config.outputCoders().containsKey(tag), "No output coder configured for tag %s", tag); + } + return config; + } + } + + /** Serializable no side input reader, the default of {@link #sideInputReader()}. */ + private static class EmptySideInputReader implements SideInputReader, Serializable { + private static final EmptySideInputReader INSTANCE = new EmptySideInputReader(); + + @Override + public @Nullable T get(PCollectionView view, BoundedWindow window) { + throw new IllegalArgumentException( + "No side inputs were configured on this stateful operator, cannot read " + view); + } + + @Override + public boolean contains(PCollectionView view) { + return false; + } + + @Override + public boolean isEmpty() { + return true; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BytesKV.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BytesKV.java new file mode 100644 index 000000000000..5021e4714497 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BytesKV.java @@ -0,0 +1,106 @@ +/* + * 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.state; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.streaming.MapState; +import org.checkerframework.checker.nullness.qual.Nullable; +import scala.Tuple2; + +/** + * Minimal string keyed, byte array valued store, the single persistence primitive the Beam state + * and timer bridges are written against. + * + *

The production implementation, {@link #of(MapState)}, is backed by exactly one Spark + * {@code transformWithState} {@link MapState}. That is a requirement rather than a stylistic + * choice. Beam's {@code ReduceFnRunner} invents state tags at runtime, for example one buffer tag + * per active window plus the watermark hold and pane info tags that go with it, so the set of state + * addresses is not known when {@code StatefulProcessor.init} has to declare its state variables. A + * single map keyed by {@code namespace + tag} can express that, a fixed set of typed per + * {@code @StateId} state variables cannot. + * + *

Deferred optimisation: for a stateful {@code ParDo} the {@code @StateId} set is static, + * so those could be mapped onto one RocksDB column family per state id, which would let Spark push + * down range scans and drop the composite key prefix. That is a performance refinement only, it + * does not change semantics, and it is deliberately not done in this POC because the + * group-also-by-window path has to keep using the single map anyway. + * + *

Implementations are used from executor code inside one {@code handleInputRows} or {@code + * handleExpiredTimer} invocation and are therefore neither thread safe nor serializable. + */ +public interface BytesKV { + + /** Returns the value stored under {@code key}, or {@code null} if there is none. */ + byte @Nullable [] get(String key); + + /** Stores {@code value} under {@code key}, replacing any previous value. */ + void put(String key, byte[] value); + + /** Removes {@code key}, a no-op if it is absent. */ + void remove(String key); + + /** + * Returns a snapshot of all entries currently in the store. + * + *

The result is materialised eagerly, callers may safely mutate the store while iterating it. + */ + Iterable> entries(); + + /** Returns a {@link BytesKV} view over a Spark {@code transformWithState} {@link MapState}. */ + static BytesKV of(MapState mapState) { + return new MapStateBytesKV(mapState); + } + + /** A {@link BytesKV} backed by a single Spark {@link MapState}. */ + final class MapStateBytesKV implements BytesKV { + private final MapState mapState; + + private MapStateBytesKV(MapState mapState) { + this.mapState = mapState; + } + + @Override + public byte @Nullable [] get(String key) { + return mapState.containsKey(key) ? mapState.getValue(key) : null; + } + + @Override + public void put(String key, byte[] value) { + mapState.updateValue(key, value); + } + + @Override + public void remove(String key) { + mapState.removeKey(key); + } + + @Override + public Iterable> entries() { + List> all = new ArrayList<>(); + scala.collection.Iterator> it = mapState.iterator(); + while (it.hasNext()) { + Tuple2 next = it.next(); + all.add(new AbstractMap.SimpleImmutableEntry<>(next._1(), next._2())); + } + return all; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java new file mode 100644 index 000000000000..9df953b710f3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java @@ -0,0 +1,578 @@ +/* + * 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.state; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.coders.MapCoder; +import org.apache.beam.sdk.coders.SetCoder; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.CombiningState; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.MultimapState; +import org.apache.beam.sdk.state.OrderedListState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.state.ReadableStates; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.state.State; +import org.apache.beam.sdk.state.StateBinder; +import org.apache.beam.sdk.state.StateContext; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.CombineWithContext; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.CombineFnUtil; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Beam {@link StateInternals} on top of Spark 4 {@code transformWithState}, a port of the legacy + * {@code org.apache.beam.runners.spark.stateful.SparkStateInternals} with its Guava {@code + * Table} replaced by the {@link BytesKV} SPI. + * + *

Every Beam state cell is stored as one {@link BytesKV} entry under the composite key {@code + * namespace.stringKey() + "+" + tag.getId()}. Window namespaces render as {@code //} and always end in a slash, so the composite key is unambiguous. + * + *

Writes go straight through to the underlying store, there is no write buffering and therefore + * nothing to flush. Reads of aggregate cells (bag, set, map) decode the whole cell, mutate it in + * memory and write it back, exactly like the legacy Spark implementation. That is quadratic for + * very large bags and is an accepted POC limitation. + * + *

An instance is scoped to a single key and to a single {@code handleInputRows} or {@code + * handleExpiredTimer} invocation, because the underlying {@code MapState} resolves against Spark's + * implicit grouping key which is only set for the duration of that call. + * + * @param the Beam key type + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class TwsStateInternals implements StateInternals { + + /** Separator between the state namespace and the state tag id in the composite store key. */ + private static final String SEPARATOR = "+"; + + private final K key; + private final BytesKV store; + + private TwsStateInternals(K key, BytesKV store) { + this.key = key; + this.store = store; + } + + /** Creates state internals for {@code key} backed by {@code store}. */ + public static TwsStateInternals forKey(K key, BytesKV store) { + return new TwsStateInternals<>(key, store); + } + + /** Returns the composite store key used for a namespace and state tag id. */ + public static String storeKey(StateNamespace namespace, String tagId) { + return namespace.stringKey() + SEPARATOR + tagId; + } + + @Override + public K getKey() { + return key; + } + + @Override + public T state( + StateNamespace namespace, StateTag address, StateContext c) { + return address.getSpec().bind(address.getId(), new TwsStateBinder(namespace, c)); + } + + private class TwsStateBinder implements StateBinder { + private final StateNamespace namespace; + private final StateContext stateContext; + + private TwsStateBinder(StateNamespace namespace, StateContext stateContext) { + this.namespace = namespace; + this.stateContext = stateContext; + } + + @Override + public ValueState bindValue(String id, StateSpec> spec, Coder coder) { + return new TwsValueState<>(namespace, id, coder); + } + + @Override + public BagState bindBag(String id, StateSpec> spec, Coder elemCoder) { + return new TwsBagState<>(namespace, id, elemCoder); + } + + @Override + public SetState bindSet(String id, StateSpec> spec, Coder elemCoder) { + return new TwsSetState<>(namespace, id, elemCoder); + } + + @Override + public MapState bindMap( + String id, + StateSpec> spec, + Coder mapKeyCoder, + Coder mapValueCoder) { + return new TwsMapState<>(namespace, id, MapCoder.of(mapKeyCoder, mapValueCoder)); + } + + @Override + public MultimapState bindMultimap( + String id, + StateSpec> spec, + Coder keyCoder, + Coder valueCoder) { + throw new UnsupportedOperationException( + String.format("%s is not supported", MultimapState.class.getSimpleName())); + } + + @Override + public OrderedListState bindOrderedList( + String id, StateSpec> spec, Coder elemCoder) { + throw new UnsupportedOperationException( + String.format("%s is not supported", OrderedListState.class.getSimpleName())); + } + + @Override + public CombiningState bindCombining( + String id, + StateSpec> spec, + Coder accumCoder, + CombineFn combineFn) { + return new TwsCombiningState<>(namespace, id, accumCoder, combineFn); + } + + @Override + public + CombiningState bindCombiningWithContext( + String id, + StateSpec> spec, + Coder accumCoder, + CombineWithContext.CombineFnWithContext combineFn) { + return new TwsCombiningState<>( + namespace, id, accumCoder, CombineFnUtil.bindContext(combineFn, stateContext)); + } + + @Override + public WatermarkHoldState bindWatermark( + String id, StateSpec spec, TimestampCombiner timestampCombiner) { + return new TwsWatermarkHoldState(namespace, id, timestampCombiner); + } + } + + private class AbstractState { + final StateNamespace namespace; + final String id; + final Coder coder; + + private AbstractState(StateNamespace namespace, String id, Coder coder) { + this.namespace = namespace; + this.id = id; + this.coder = coder; + } + + private String cellKey() { + return storeKey(namespace, id); + } + + boolean exists() { + return store.get(cellKey()) != null; + } + + @Nullable + T readValue() { + byte[] buf = store.get(cellKey()); + if (buf == null) { + return null; + } + try { + return CoderUtils.decodeFromByteArray(coder, buf); + } catch (Exception e) { + throw new IllegalStateException("Failed to decode state cell " + cellKey(), e); + } + } + + void writeValue(T input) { + try { + store.put(cellKey(), CoderUtils.encodeToByteArray(coder, input)); + } catch (Exception e) { + throw new IllegalStateException("Failed to encode state cell " + cellKey(), e); + } + } + + public void clear() { + store.remove(cellKey()); + } + + ReadableState isEmptyState() { + return new ReadableState() { + @Override + public ReadableState readLater() { + return this; + } + + @Override + public Boolean read() { + return !exists(); + } + }; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof AbstractState)) { + return false; + } + @SuppressWarnings("unchecked") + AbstractState that = (AbstractState) o; + return namespace.equals(that.namespace) && id.equals(that.id); + } + + @Override + public int hashCode() { + int result = namespace.hashCode(); + result = 31 * result + id.hashCode(); + return result; + } + } + + private class TwsValueState extends AbstractState implements ValueState { + + private TwsValueState(StateNamespace namespace, String id, Coder coder) { + super(namespace, id, coder); + } + + @Override + public TwsValueState readLater() { + return this; + } + + @Override + public T read() { + return readValue(); + } + + @Override + public void write(T input) { + writeValue(input); + } + } + + private class TwsWatermarkHoldState extends AbstractState implements WatermarkHoldState { + + private final TimestampCombiner timestampCombiner; + + TwsWatermarkHoldState( + StateNamespace namespace, String id, TimestampCombiner timestampCombiner) { + super(namespace, id, InstantCoder.of()); + this.timestampCombiner = timestampCombiner; + } + + @Override + public TwsWatermarkHoldState readLater() { + return this; + } + + @Override + public Instant read() { + return readValue(); + } + + @Override + public void add(Instant outputTime) { + Instant combined = read(); + combined = + (combined == null) ? outputTime : getTimestampCombiner().combine(combined, outputTime); + writeValue(combined); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + + @Override + public TimestampCombiner getTimestampCombiner() { + return timestampCombiner; + } + } + + @SuppressWarnings("TypeParameterShadowing") + private class TwsCombiningState extends AbstractState + implements CombiningState { + + private final CombineFn combineFn; + + private TwsCombiningState( + StateNamespace namespace, + String id, + Coder coder, + CombineFn combineFn) { + super(namespace, id, coder); + this.combineFn = combineFn; + } + + @Override + public TwsCombiningState readLater() { + return this; + } + + @Override + public OutputT read() { + return combineFn.extractOutput(getAccum()); + } + + @Override + public void add(InputT input) { + writeValue(combineFn.addInput(getAccum(), input)); + } + + @Override + public AccumT getAccum() { + AccumT accum = readValue(); + return accum == null ? combineFn.createAccumulator() : accum; + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + + @Override + public void addAccum(AccumT accum) { + writeValue(combineFn.mergeAccumulators(Arrays.asList(getAccum(), accum))); + } + + @Override + public AccumT mergeAccumulators(Iterable accumulators) { + return combineFn.mergeAccumulators(accumulators); + } + } + + private final class TwsMapState extends AbstractState> + implements MapState { + + private TwsMapState(StateNamespace namespace, String id, Coder> coder) { + super(namespace, id, coder); + } + + @Override + public ReadableState get(MapKeyT mapKey) { + return getOrDefault(mapKey, null); + } + + @Override + public ReadableState getOrDefault(MapKeyT mapKey, @Nullable MapValueT defaultValue) { + return new ReadableState() { + @Override + public MapValueT read() { + return readAsMap().getOrDefault(mapKey, defaultValue); + } + + @Override + public ReadableState readLater() { + return this; + } + }; + } + + @Override + public void put(MapKeyT mapKey, MapValueT value) { + Map current = readAsMap(); + current.put(mapKey, value); + writeValue(current); + } + + @Override + public ReadableState computeIfAbsent( + MapKeyT mapKey, Function mappingFunction) { + Map current = readAsMap(); + MapValueT existing = current.get(mapKey); + if (existing == null) { + put(mapKey, mappingFunction.apply(mapKey)); + } + return ReadableStates.immediate(existing); + } + + private Map readAsMap() { + Map current = readValue(); + return current == null ? new HashMap<>() : current; + } + + @Override + public void remove(MapKeyT mapKey) { + Map current = readAsMap(); + current.remove(mapKey); + writeValue(current); + } + + @Override + public ReadableState> keys() { + return new ReadableState>() { + @Override + public Iterable read() { + return ImmutableList.copyOf(readAsMap().keySet()); + } + + @Override + public ReadableState> readLater() { + return this; + } + }; + } + + @Override + public ReadableState> values() { + return new ReadableState>() { + @Override + public Iterable read() { + return ImmutableList.copyOf(readAsMap().values()); + } + + @Override + public ReadableState> readLater() { + return this; + } + }; + } + + @Override + public ReadableState>> entries() { + return new ReadableState>>() { + @Override + public Iterable> read() { + return ImmutableList.copyOf(readAsMap().entrySet()); + } + + @Override + public ReadableState>> readLater() { + return this; + } + }; + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } + + private final class TwsSetState extends AbstractState> + implements SetState { + + private TwsSetState(StateNamespace namespace, String id, Coder coder) { + super(namespace, id, SetCoder.of(coder)); + } + + @Override + public ReadableState contains(InputT input) { + return ReadableStates.immediate(readAsSet().contains(input)); + } + + @Override + public ReadableState addIfAbsent(InputT input) { + Set current = readAsSet(); + boolean added = current.add(input); + writeValue(current); + return ReadableStates.immediate(added); + } + + @Override + public void remove(InputT input) { + Set current = readAsSet(); + current.remove(input); + writeValue(current); + } + + @Override + public SetState readLater() { + return this; + } + + @Override + public void add(InputT value) { + Set current = readAsSet(); + current.add(value); + writeValue(current); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + + @Override + public Iterable read() { + Set value = readValue(); + return value == null ? Collections.emptySet() : value; + } + + private Set readAsSet() { + Set value = readValue(); + return value == null ? new HashSet<>() : value; + } + } + + private final class TwsBagState extends AbstractState> implements BagState { + private TwsBagState(StateNamespace namespace, String id, Coder coder) { + super(namespace, id, ListCoder.of(coder)); + } + + @Override + public TwsBagState readLater() { + return this; + } + + @Override + public List read() { + List value = readValue(); + return value == null ? new ArrayList<>() : value; + } + + @Override + public void add(T input) { + List value = read(); + value.add(input); + writeValue(value); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternals.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternals.java new file mode 100644 index 000000000000..c78692e5670c --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternals.java @@ -0,0 +1,390 @@ +/* + * 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.state; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.spark.sql.streaming.StatefulProcessorHandle; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Beam {@link TimerInternals} on top of Spark 4 {@code transformWithState}. + * + *

Two things are stored per timer and they serve different purposes. + * + *

    + *
  • The full {@link TimerData}, encoded with {@link TimerInternals.TimerDataCoderV2}, lives in + * a {@link BytesKV} store keyed by {@link TimerData#stringKey()}. This is the source of truth + * for what fires, in which namespace, with which output timestamp. + *
  • A bare wake-up at the timer's expiry millisecond is registered with Spark through {@link + * StatefulProcessorHandle#registerTimer(long)}. Spark only knows a set of {@code long} expiry + * times per key, it carries no payload, so it can do nothing but wake us up. + *
+ * + *

Wake-up de-duplication. Many Beam timers can share one expiry millisecond, for example + * the end-of-window timer and the garbage collection timer of the same window when allowed lateness + * is zero. The Phase 0 spike found that registering the same expiry repeatedly is a real problem, + * so wake-ups are reconciled rather than registered blindly: on {@link #flush()} the set of + * expiries Spark currently holds for this key is read back with {@code listTimers()} and only the + * genuine difference is registered or deleted. Registering an expiry twice is therefore impossible + * by construction. + * + *

Same-millisecond re-arm inside a timer callback. Spark deletes the expiry it is + * currently firing after {@code handleExpiredTimer} returns and after the returned iterator + * is drained. A wake-up re-registered at exactly that expiry from inside the callback would + * therefore be silently removed again. When {@link #flush()} runs inside a timer callback, any + * wake-up that would land at or before the firing expiry is nudged to {@code firedExpiry + 1} + * instead. That is safe because the wake-up is only a wake-up: on the next callback all {@link + * TimerData} due at or before the new expiry are fired, so the timer still fires with its own + * timestamp and namespace. The only observable effect is that such a timer needs the watermark to + * reach one extra millisecond. + * + *

Processing time timers are out of scope for this POC. The operator runs in {@code + * TimeMode.EventTime()}, in which Spark's timer registry is driven by the event time watermark + * only, and a single {@code transformWithState} call cannot mix time modes. Setting a timer in the + * {@link TimeDomain#PROCESSING_TIME} or {@link TimeDomain#SYNCHRONIZED_PROCESSING_TIME} domain + * throws {@link UnsupportedOperationException}. + * + *

An instance is scoped to a single key and to a single {@code handleInputRows} or {@code + * handleExpiredTimer} invocation. Mutations are buffered in memory and only reach the store and + * Spark's timer registry when {@link #flush()} is called at the end of that invocation. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class TwsTimerInternals implements TimerInternals { + + /** + * The bare {@code long} wake-up registry Spark exposes on {@link StatefulProcessorHandle}. + * + *

Extracted as an interface so the timer bridge can be unit tested without a running Spark + * query. + */ + public interface WakeupRegistry { + + /** Registers a wake-up at {@code expiryMs} for the current key. */ + void register(long expiryMs); + + /** Deletes the wake-up at {@code expiryMs} for the current key. */ + void delete(long expiryMs); + + /** Returns the wake-ups Spark currently holds for the current key. */ + Set registered(); + + /** Returns a registry delegating to a Spark {@link StatefulProcessorHandle}. */ + static WakeupRegistry of(StatefulProcessorHandle handle) { + return new WakeupRegistry() { + @Override + public void register(long expiryMs) { + handle.registerTimer(expiryMs); + } + + @Override + public void delete(long expiryMs) { + handle.deleteTimer(expiryMs); + } + + @Override + public Set registered() { + Set timers = new HashSet<>(); + scala.collection.Iterator it = handle.listTimers(); + while (it.hasNext()) { + timers.add(((Number) it.next()).longValue()); + } + return timers; + } + }; + } + } + + private final BytesKV store; + private final WakeupRegistry registry; + private final TimerDataCoderV2 timerCoder; + private final Instant inputWatermark; + private final Instant processingTime; + private final @Nullable Long firedExpiryMs; + + /** Snapshot of the timers as they were loaded, used to compute the write-back diff. */ + private final Map loaded; + + /** Current timers, mutated by {@link #setTimer} and {@link #deleteTimer}. */ + private final Map current; + + private boolean flushed; + + private TwsTimerInternals( + BytesKV store, + WakeupRegistry registry, + Coder windowCoder, + Instant inputWatermark, + Instant processingTime, + @Nullable Long firedExpiryMs) { + this.store = store; + this.registry = registry; + this.timerCoder = TimerDataCoderV2.of(windowCoder); + this.inputWatermark = inputWatermark; + this.processingTime = processingTime; + this.firedExpiryMs = firedExpiryMs; + this.loaded = new LinkedHashMap<>(); + for (Map.Entry entry : store.entries()) { + loaded.put(entry.getKey(), decode(entry.getValue())); + } + this.current = new LinkedHashMap<>(loaded); + } + + /** + * Creates timer internals for one invocation. + * + * @param store where the encoded {@link TimerData} live + * @param registry Spark's bare wake-up registry for the current key + * @param windowCoder the window coder of the transform, needed to decode timer namespaces + * @param inputWatermark the event time watermark visible for this invocation + * @param processingTime the batch processing time + * @param firedExpiryMs the expiry Spark is currently firing, or {@code null} outside a timer + * callback + */ + public static TwsTimerInternals create( + BytesKV store, + WakeupRegistry registry, + Coder windowCoder, + Instant inputWatermark, + Instant processingTime, + @Nullable Long firedExpiryMs) { + return new TwsTimerInternals( + store, registry, windowCoder, inputWatermark, processingTime, firedExpiryMs); + } + + private TimerData decode(byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray(timerCoder, bytes); + } catch (Exception e) { + throw new IllegalStateException("Failed to decode a stored Beam timer", e); + } + } + + private byte[] encode(TimerData timer) { + try { + return CoderUtils.encodeToByteArray(timerCoder, timer); + } catch (Exception e) { + throw new IllegalStateException("Failed to encode Beam timer " + timer, e); + } + } + + private static void rejectUnsupportedDomain(TimeDomain domain) { + if (domain != TimeDomain.EVENT_TIME) { + throw new UnsupportedOperationException( + "The Spark 4 structured streaming runner only supports event time timers, but a timer " + + "in the " + + domain + + " domain was requested. Spark's transformWithState runs in a single TimeMode and " + + "this operator uses TimeMode.EventTime(); processing time timers are out of scope " + + "for the streaming POC."); + } + } + + @Override + public void setTimer( + StateNamespace namespace, + String timerId, + String timerFamilyId, + Instant target, + Instant outputTimestamp, + TimeDomain timeDomain) { + setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + } + + @Override + public void setTimer(TimerData timer) { + rejectUnsupportedDomain(timer.getDomain()); + current.put(timer.stringKey(), timer); + } + + @Override + public void deleteTimer( + StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { + current + .values() + .removeIf( + timer -> + namespace.equals(timer.getNamespace()) + && timerId.equals(timer.getTimerId()) + && timerFamilyId.equals(timer.getTimerFamilyId()) + && timeDomain.equals(timer.getDomain())); + } + + @Override + public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) { + throw new UnsupportedOperationException( + "Deleting a timer without a TimeDomain is not supported, use " + + "deleteTimer(namespace, timerId, timerFamilyId, timeDomain)."); + } + + @Override + public void deleteTimer(TimerData timer) { + current.remove(timer.stringKey()); + } + + @Override + public Instant currentProcessingTime() { + return processingTime; + } + + @Override + public @Nullable Instant currentSynchronizedProcessingTime() { + return null; + } + + /** + * Returns the event time watermark for this invocation. + * + *

This is the batch start watermark, that is the watermark Spark computed at the end of + * the previous micro-batch. Elements of the current micro-batch never advance it, so an element + * can never be late with respect to its own batch. + */ + @Override + public Instant currentInputWatermarkTime() { + return inputWatermark; + } + + @Override + public @Nullable Instant currentOutputWatermarkTime() { + return null; + } + + /** Returns the timers currently held for this key, in no particular order. */ + public Iterable getTimers() { + return Collections.unmodifiableCollection(new ArrayList<>(current.values())); + } + + /** + * Removes and returns the timers Beam considers due for the Spark wake-up at {@code + * firedExpiryMs}, in Beam's natural timer order. + * + *

The two systems disagree on the boundary and the difference is not cosmetic. + * + *

    + *
  • Spark expires a {@code transformWithState} wake-up as soon as {@code expiry <= + * batchWatermark}. + *
  • Beam fires an event time timer only once the input watermark is strictly past the + * timer's timestamp, see {@code InMemoryTimerInternals.removeNextTimer} and {@code + * SparkTimerInternals}, both of which use {@code currentTime.isAfter(timestamp)}. + *
+ * + *

Handing Beam a timer one millisecond early is silently destructive rather than merely early: + * {@code ReduceFnRunner} asks {@code AfterWatermark.pastEndOfWindow} whether to fire, that + * predicate is also strict, so the trigger declines, and because the runner is entitled to assume + * the timer will not be delivered again the pane is simply lost. The end-of-window timer of a + * fixed window sits at exactly {@code window.maxTimestamp()}, so this hits every window whose end + * coincides with a batch watermark, which in practice means most of them. + * + *

Timers are therefore only released once {@code timestamp < currentInputWatermarkTime()}. A + * timer withheld this way stays in the store, and {@link #flush()} re-registers a wake-up for it + * at {@code firedExpiryMs + 1} through {@link #wakeupFor}, so it fires on the next batch whose + * watermark has genuinely moved past it. + * + *

Firing a timer removes it, which is what makes a Spark wake-up that covers several Beam + * timers safe: the second wake-up simply finds nothing left to fire. + */ + public List removeTimersReadyToFire(long firedExpiryMs) { + // Spark guarantees firedExpiryMs <= inputWatermark, so this only ever lowers the bound, and it + // lowers it by at most one millisecond. + long bound = Math.min(firedExpiryMs, inputWatermark.getMillis() - 1); + return removeTimersAtOrBefore(new Instant(bound)); + } + + /** + * Removes and returns, in Beam's natural timer order, every event time timer whose timestamp is + * at or before {@code maxTimestamp}, without applying the watermark rule of {@link + * #removeTimersReadyToFire}. + */ + public List removeTimersAtOrBefore(Instant maxTimestamp) { + List due = new ArrayList<>(); + for (TimerData timer : current.values()) { + if (!timer.getTimestamp().isAfter(maxTimestamp)) { + due.add(timer); + } + } + Collections.sort(due); + for (TimerData timer : due) { + current.remove(timer.stringKey()); + } + return due; + } + + /** + * Persists the timer changes of this invocation and reconciles Spark's wake-ups with them. + * + *

Must be called exactly once, at the end of the {@code handleInputRows} or {@code + * handleExpiredTimer} invocation this instance belongs to. + */ + public void flush() { + if (flushed) { + throw new IllegalStateException("TwsTimerInternals.flush() called more than once"); + } + flushed = true; + + for (Map.Entry entry : current.entrySet()) { + TimerData before = loaded.get(entry.getKey()); + if (before == null || !before.equals(entry.getValue())) { + store.put(entry.getKey(), encode(entry.getValue())); + } + } + for (String key : loaded.keySet()) { + if (!current.containsKey(key)) { + store.remove(key); + } + } + + Set desired = new HashSet<>(); + for (TimerData timer : current.values()) { + desired.add(wakeupFor(timer.getTimestamp().getMillis())); + } + Set alreadyRegistered = registry.registered(); + for (Long expiry : desired) { + if (!alreadyRegistered.contains(expiry)) { + registry.register(expiry); + } + } + for (Long expiry : alreadyRegistered) { + // Spark removes the expiry it is currently firing on its own once the callback completes. + if (!desired.contains(expiry) && !expiry.equals(firedExpiryMs)) { + registry.delete(expiry); + } + } + } + + /** Maps a Beam timer timestamp onto the Spark wake-up millisecond to register for it. */ + private long wakeupFor(long timestampMs) { + if (firedExpiryMs != null && timestampMs <= firedExpiryMs) { + return firedExpiryMs + 1; + } + return timestampMs; + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java new file mode 100644 index 000000000000..5929d4b0ddf0 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java @@ -0,0 +1,443 @@ +/* + * 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.state; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.io.Serializable; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.TwsTransformFactory; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarLongCoder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.api.java.function.VoidFunction2; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.functions; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.After; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Proves that {@link BeamStatefulProcessor} really runs inside Spark 4's {@code transformWithState} + * from Java, on Scala 2.13, with a RocksDB state store and an event time watermark, and that both + * hosted execution stacks produce the Beam results they are supposed to. + * + *

The streaming source is a plain file source over JSON files with {@code maxFilesPerTrigger=1}, + * so the micro-batch boundaries and therefore the watermark progression are deterministic and no + * test depends on the wall clock. Each JSON record carries a Beam element timestamp in millis plus + * a Base64 encoded {@link TwsTransformFactory} input row. + * + *

Results are collected with {@code foreachBatch}. The {@code memory} sink cannot be used, Spark + * test JVMs run with {@code spark.kryo.registrationRequired=true} and its commit message is not a + * registered class. + */ +@Category(StreamingTest.class) +@RunWith(JUnit4.class) +public class BeamStatefulProcessorTest implements Serializable { + + /** + * Spark test JVMs set {@code spark.kryo.registrationRequired=true} (see {@code + * runners/spark/spark_runner.gradle}) to keep Beam's own serialisation honest. That cannot hold + * for a {@code transformWithState} query: Spark 4 broadcasts its own {@code + * org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata} to the executors with the + * user Kryo instance, and no Beam registrator knows that class, so the query dies before the + * first micro-batch. Beam never turns the flag on outside tests, so this only relaxes the test + * harness, but every streaming test that runs a stateful operator will have to do the same until + * the Spark 4 Kryo registrator learns about Spark's streaming state classes. + */ + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + + @Rule public transient TemporaryFolder temp = new TemporaryFolder(); + + /** Output rows collected per query name, driver side. */ + private static final Map> COLLECTED = new ConcurrentHashMap<>(); + + /** 2023-11-14T22:13:20Z, aligned on a ten second fixed window boundary. */ + private static final long BASE_MILLIS = 1_700_000_000_000L; + + private static final TupleTag MAIN_TAG = new TupleTag("main") {}; + + @After + public void tearDown() { + COLLECTED.clear(); + } + + // --------------------------------------------------------------------------------------------- + // Row codec, no Spark involved. + // --------------------------------------------------------------------------------------------- + + @Test(timeout = 60_000) + public void testInputRowCodecRoundTrip() { + byte[] key = "the-key".getBytes(UTF_8); + byte[] payload = new byte[] {1, 2, 3, 0, -7}; + + byte[] row = TwsTransformFactory.encodeInputRow(key, payload); + assertArrayEquals(key, TwsTransformFactory.inputKey(row)); + assertArrayEquals(payload, TwsTransformFactory.inputPayload(row)); + } + + @Test(timeout = 60_000) + public void testInputRowCodecHandlesEmptyKeyAndPayload() { + byte[] row = TwsTransformFactory.encodeInputRow(new byte[0], new byte[0]); + assertEquals(0, TwsTransformFactory.inputKey(row).length); + assertEquals(0, TwsTransformFactory.inputPayload(row).length); + } + + @Test(timeout = 60_000) + public void testOutputRowCodecRoundTrip() { + byte[] payload = new byte[] {9, 8, 7}; + for (int index : new int[] {0, 1, 127, 128, 100_000}) { + byte[] row = TwsTransformFactory.encodeOutputRow(index, payload); + assertEquals(index, TwsTransformFactory.outputTagIndex(row)); + assertArrayEquals(payload, TwsTransformFactory.outputPayload(row)); + } + } + + // --------------------------------------------------------------------------------------------- + // Real transformWithState queries. + // --------------------------------------------------------------------------------------------- + + /** + * A stateful {@code ParDo} in the global window: the running sum per key must survive across + * micro-batches, which is only possible if the Beam state really landed in the Spark state store + * and was read back on the next batch. + */ + @Test(timeout = 600_000) + public void testStatefulParDoRunsInTransformWithState() throws Exception { + WindowingStrategy strategy = WindowingStrategy.globalDefault(); + Coder> valueCoder = + WindowedValues.getFullCoder(VarLongCoder.of(), GlobalWindow.Coder.INSTANCE); + Coder>> outputCoder = + WindowedValues.getFullCoder( + KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()), GlobalWindow.Coder.INSTANCE); + + BeamStatefulProcessorConfig config = + BeamStatefulProcessorConfig.builder() + .setMode(BeamStatefulProcessorConfig.Mode.STATEFUL_PARDO) + .setDoFn(new RunningSumFn()) + .setKeyCoder(StringUtf8Coder.of()) + .setValueCoder(VarLongCoder.of()) + .setWindowingStrategy(strategy) + .setMainOutputTag(MAIN_TAG) + .setOutputCoders( + Collections.singletonMap( + MAIN_TAG, KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))) + .setOptionsSupplier(PipelineOptionsFactory::create) + .setStepName("running-sum") + .build(); + + List> batches = new ArrayList<>(); + batches.add( + Lists.newArrayList( + globalRecord("a", 1L, BASE_MILLIS), + globalRecord("b", 10L, BASE_MILLIS + 1), + globalRecord("a", 2L, BASE_MILLIS + 2))); + batches.add( + Lists.newArrayList( + globalRecord("a", 4L, BASE_MILLIS + 1_000), + globalRecord("b", 20L, BASE_MILLIS + 1_001))); + + List rows = runQuery("stateful-pardo", batches, config); + + List> emitted = new ArrayList<>(); + for (byte[] row : rows) { + assertEquals("only the main output tag is used", 0, TwsTransformFactory.outputTagIndex(row)); + emitted.add( + CoderUtils.decodeFromByteArray(outputCoder, TwsTransformFactory.outputPayload(row)) + .getValue()); + } + + // Per key the running sums must be exactly the prefix sums, in order. + Map> byKey = new HashMap<>(); + for (KV kv : emitted) { + byKey.computeIfAbsent(kv.getKey(), k -> new ArrayList<>()).add(kv.getValue()); + } + assertEquals("two keys expected", 2, byKey.size()); + assertEquals(Lists.newArrayList(1L, 3L, 7L), byKey.get("a")); + assertEquals(Lists.newArrayList(10L, 30L), byKey.get("b")); + assertFalse("the value coder must have been used", valueCoder.toString().isEmpty()); + } + + /** + * A windowed {@code GroupByKey}: three ten second fixed windows worth of data, driven so that the + * watermark passes the end of the first window while the query is still running. The grouped + * output can only appear if the end-of-window timer was registered with Spark, survived a + * checkpoint of the RocksDB timer state and fired through {@code handleExpiredTimer}. + */ + @Test(timeout = 600_000) + public void testGroupAlsoByWindowFiresOnTheEndOfWindowTimer() throws Exception { + WindowingStrategy strategy = + WindowingStrategy.of(FixedWindows.of(Duration.standardSeconds(10))) + .withAllowedLateness(Duration.ZERO); + Coder>>> outputCoder = + WindowedValues.getFullCoder( + KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of())), + IntervalWindow.getCoder()); + + BeamStatefulProcessorConfig config = + BeamStatefulProcessorConfig.builder() + .setMode(BeamStatefulProcessorConfig.Mode.GROUP_ALSO_BY_WINDOW) + .setKeyCoder(StringUtf8Coder.of()) + .setValueCoder(StringUtf8Coder.of()) + .setWindowingStrategy(strategy) + .setMainOutputTag(MAIN_TAG) + .setOutputCoders( + Collections.singletonMap( + MAIN_TAG, + KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of())))) + .setOptionsSupplier(PipelineOptionsFactory::create) + .setStepName("gabw") + .build(); + + IntervalWindow firstWindow = + new IntervalWindow(new Instant(BASE_MILLIS), new Instant(BASE_MILLIS + 10_000)); + + List> batches = new ArrayList<>(); + // Batch 1: the whole first window. The watermark is still at zero here. + batches.add( + Lists.newArrayList( + windowedRecord("a", "x", BASE_MILLIS, firstWindow), + windowedRecord("a", "y", BASE_MILLIS + 3_000, firstWindow), + windowedRecord("a", "z", BASE_MILLIS + 9_999, firstWindow))); + // Batch 2: a much later element, which pushes the watermark past the first window's end. + IntervalWindow lateWindow = + new IntervalWindow(new Instant(BASE_MILLIS + 20_000), new Instant(BASE_MILLIS + 30_000)); + batches.add( + Lists.newArrayList(windowedRecord("sentinel", "s", BASE_MILLIS + 20_000, lateWindow))); + // Batch 3: one more element, so there is a micro-batch that actually sees the advanced + // watermark and can therefore expire the first window's timer. + IntervalWindow lastWindow = + new IntervalWindow(new Instant(BASE_MILLIS + 30_000), new Instant(BASE_MILLIS + 40_000)); + batches.add( + Lists.newArrayList(windowedRecord("sentinel", "t", BASE_MILLIS + 30_000, lastWindow))); + + List rows = runQuery("gabw", batches, config); + + List>>> emitted = new ArrayList<>(); + for (byte[] row : rows) { + assertEquals(0, TwsTransformFactory.outputTagIndex(row)); + emitted.add( + CoderUtils.decodeFromByteArray(outputCoder, TwsTransformFactory.outputPayload(row))); + } + + // Only the first window is asserted on. Whether the sentinel's own window also fires depends on + // Spark scheduling a no-data batch after the last file, which processAllAvailable does not + // promise to wait for. + List>>> forKeyA = new ArrayList<>(); + for (WindowedValue>> candidate : emitted) { + if ("a".equals(candidate.getValue().getKey())) { + forKeyA.add(candidate); + } + } + + assertEquals("exactly one pane for the completed window, got " + emitted, 1, forKeyA.size()); + WindowedValue>> pane = forKeyA.get(0); + List grouped = Lists.newArrayList(pane.getValue().getValue()); + Collections.sort(grouped); + assertEquals(Lists.newArrayList("x", "y", "z"), grouped); + assertEquals( + "the pane must carry the window it belongs to", + Collections.singletonList(firstWindow), + Lists.newArrayList(pane.getWindows())); + assertTrue("the on-time pane must be the first one", pane.getPaneInfo().isFirst()); + assertEquals(PaneInfo.Timing.ON_TIME, pane.getPaneInfo().getTiming()); + } + + // --------------------------------------------------------------------------------------------- + // Harness. + // --------------------------------------------------------------------------------------------- + + /** A stateful DoFn keeping a running sum per key, the simplest thing that needs Beam state. */ + private static class RunningSumFn extends DoFn, KV> { + + @StateId("sum") + private final StateSpec> sumSpec = StateSpecs.value(VarLongCoder.of()); + + @ProcessElement + public void process( + @Element KV element, + @StateId("sum") ValueState sum, + OutputReceiver> out) { + Long current = sum.read(); + long updated = (current == null ? 0L : current) + element.getValue(); + sum.write(updated); + out.output(KV.of(element.getKey(), updated)); + } + } + + /** + * Runs one {@code transformWithState} query over {@code batches}, one JSON file per batch and one + * file per trigger, and returns every output row the query produced. + */ + private List runQuery( + String queryName, List> batches, BeamStatefulProcessorConfig config) + throws Exception { + + File input = temp.newFolder(queryName + "-input"); + long now = System.currentTimeMillis(); + for (int i = 0; i < batches.size(); i++) { + File file = new File(input, String.format("%03d.json", i)); + Files.write(file.toPath(), String.join("\n", batches.get(i)).getBytes(UTF_8)); + // Spark's file stream source orders files by modification time only. Three files written in + // the same millisecond tie and are then consumed in an arbitrary order, which for an event + // time test means arbitrary watermark progression. Space the timestamps out explicitly. + assertTrue( + "could not set the modification time of " + file, + file.setLastModified(now - (batches.size() - i) * 60_000L)); + } + + COLLECTED.put(queryName, Collections.synchronizedList(new ArrayList<>())); + + Dataset raw = + SESSION + .getSession() + .readStream() + .schema("ts BIGINT, payload STRING") + .option("maxFilesPerTrigger", 1) + .option("latestFirst", false) + .json(input.getAbsolutePath()); + + Dataset keyed = + raw.withColumn("eventTime", functions.expr("timestamp_millis(ts)")) + .withWatermark("eventTime", "0 seconds") + .map( + (MapFunction) + row -> Base64.getDecoder().decode(row.getAs("payload")), + Encoders.BINARY()); + + Dataset transformed = TwsTransformFactory.transform(keyed, config); + + StreamingQuery query = + transformed + .writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + List target = COLLECTED.get(queryName); + if (target != null) { + target.addAll(batch.collectAsList()); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName + "-cp").getAbsolutePath()) + .start(); + + try { + query.processAllAvailable(); + } finally { + query.stop(); + } + if (query.exception().isDefined()) { + throw new IllegalStateException( + "streaming query failed: " + query.exception().get().toString()); + } + return new ArrayList<>(COLLECTED.get(queryName)); + } + + /** One JSON record holding a global window element. */ + private static String globalRecord(String key, long value, long timestampMs) throws IOException { + WindowedValue windowedValue = + WindowedValues.of( + value, new Instant(timestampMs), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING); + return record( + timestampMs, + key, + StringUtf8Coder.of(), + windowedValue, + WindowedValues.getFullCoder(VarLongCoder.of(), GlobalWindow.Coder.INSTANCE)); + } + + /** One JSON record holding an element already assigned to {@code window}. */ + private static String windowedRecord( + String key, String value, long timestampMs, BoundedWindow window) throws IOException { + WindowedValue windowedValue = + WindowedValues.of(value, new Instant(timestampMs), window, PaneInfo.NO_FIRING); + return record( + timestampMs, + key, + StringUtf8Coder.of(), + windowedValue, + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder())); + } + + private static String record( + long timestampMs, + K key, + Coder keyCoder, + WindowedValue value, + Coder> valueCoder) + throws IOException { + byte[] row = + TwsTransformFactory.encodeInputRow( + CoderUtils.encodeToByteArray(keyCoder, key), + CoderUtils.encodeToByteArray(valueCoder, value)); + return "{\"ts\": " + + timestampMs + + ", \"payload\": \"" + + Base64.getEncoder().encodeToString(row) + + "\"}"; + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternalsTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternalsTest.java new file mode 100644 index 000000000000..a35e7bf6585c --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternalsTest.java @@ -0,0 +1,359 @@ +/* + * 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.state; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.CombiningState; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.transforms.Sum; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Unit tests for the Beam {@code StateInternals} bridge, exercised against an in-memory {@link + * BytesKV} rather than a live Spark {@code MapState}. + * + *

That is the whole point of the {@link BytesKV} seam: everything below it is Spark's problem, + * everything above it is Beam semantics and can be tested in milliseconds. + */ +@RunWith(JUnit4.class) +public class TwsStateInternalsTest { + + private static final StateNamespace NS_A = StateNamespaces.global(); + private static final StateNamespace NS_B = + StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE); + + /** A {@link BytesKV} backed by a plain {@link LinkedHashMap}, the test double for Spark state. */ + public static final class InMemoryBytesKV implements BytesKV { + private final Map map = new LinkedHashMap<>(); + + @Override + public byte @Nullable [] get(String key) { + return map.get(key); + } + + @Override + public void put(String key, byte[] value) { + map.put(key, value); + } + + @Override + public void remove(String key) { + map.remove(key); + } + + @Override + public Iterable> entries() { + return new ArrayList<>(map.entrySet()); + } + + /** Returns the raw store keys currently present, for addressing assertions. */ + public Iterable keys() { + return new ArrayList<>(map.keySet()); + } + + public int size() { + return map.size(); + } + } + + private InMemoryBytesKV store() { + return new InMemoryBytesKV(); + } + + private TwsStateInternals internals(BytesKV store) { + return TwsStateInternals.forKey("key", store); + } + + @Test(timeout = 30_000) + public void testKeyIsExposed() { + assertEquals("key", internals(store()).getKey()); + } + + @Test(timeout = 30_000) + public void testStoreKeyLayout() { + assertEquals(NS_A.stringKey() + "+" + "tag", TwsStateInternals.storeKey(NS_A, "tag")); + } + + @Test(timeout = 30_000) + public void testValueStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.value("v", StringUtf8Coder.of()); + + ValueState state = internals(store).state(NS_A, tag); + assertNull("an unwritten value state reads null", state.read()); + + state.write("hello"); + assertEquals("hello", state.read()); + + // A fresh bridge over the same store must see the same value, nothing is cached in memory. + assertEquals("hello", internals(store).state(NS_A, tag).read()); + + state.clear(); + assertNull(internals(store).state(NS_A, tag).read()); + assertEquals("clear must remove the cell, not blank it", 0, store.size()); + } + + @Test(timeout = 30_000) + public void testValueStateIsAddressedByNamespaceAndTag() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.value("v", StringUtf8Coder.of()); + + internals(store).state(NS_A, tag).write("hello"); + + assertEquals(1, store.size()); + assertEquals(TwsStateInternals.storeKey(NS_A, "v"), Lists.newArrayList(store.keys()).get(0)); + } + + @Test(timeout = 30_000) + public void testNamespacesAreIsolated() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.value("v", StringUtf8Coder.of()); + + internals(store).state(NS_A, tag).write("a"); + internals(store).state(NS_B, tag).write("b"); + + assertEquals(2, store.size()); + assertEquals("a", internals(store).state(NS_A, tag).read()); + assertEquals("b", internals(store).state(NS_B, tag).read()); + + internals(store).state(NS_A, tag).clear(); + assertNull(internals(store).state(NS_A, tag).read()); + assertEquals( + "clearing one namespace must not touch the other", + "b", + internals(store).state(NS_B, tag).read()); + } + + @Test(timeout = 30_000) + public void testTagsAreIsolatedWithinANamespace() { + InMemoryBytesKV store = store(); + internals(store).state(NS_A, StateTags.value("one", StringUtf8Coder.of())).write("1"); + internals(store).state(NS_A, StateTags.value("two", StringUtf8Coder.of())).write("2"); + + assertEquals(2, store.size()); + assertEquals( + "1", internals(store).state(NS_A, StateTags.value("one", StringUtf8Coder.of())).read()); + assertEquals( + "2", internals(store).state(NS_A, StateTags.value("two", StringUtf8Coder.of())).read()); + } + + @Test(timeout = 30_000) + public void testBagStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.bag("b", VarIntCoder.of()); + + BagState state = internals(store).state(NS_A, tag); + assertTrue("an unwritten bag is empty", state.isEmpty().read()); + assertEquals(0, Lists.newArrayList(state.read()).size()); + + state.add(1); + state.add(2); + state.add(2); + assertFalse(state.isEmpty().read()); + assertEquals( + Lists.newArrayList(1, 2, 2), Lists.newArrayList(internals(store).state(NS_A, tag).read())); + + state.clear(); + assertTrue(internals(store).state(NS_A, tag).isEmpty().read()); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testCombiningStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag> tag = + StateTags.combiningValueFromInputInternal("c", VarIntCoder.of(), Sum.ofIntegers()); + + CombiningState state = internals(store).state(NS_A, tag); + assertTrue(state.isEmpty().read()); + assertEquals(Integer.valueOf(0), state.read()); + + state.add(3); + state.add(4); + assertFalse(state.isEmpty().read()); + assertEquals( + "the accumulator must be persisted, not kept in memory", + Integer.valueOf(7), + internals(store).state(NS_A, tag).read()); + + internals(store).state(NS_A, tag).add(1); + assertEquals(Integer.valueOf(8), internals(store).state(NS_A, tag).read()); + + state.clear(); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testWatermarkHoldStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag tag = + StateTags.watermarkStateInternal("hold", TimestampCombiner.EARLIEST); + + WatermarkHoldState state = internals(store).state(NS_A, tag); + assertTrue(state.isEmpty().read()); + assertNull(state.read()); + assertEquals(TimestampCombiner.EARLIEST, state.getTimestampCombiner()); + + state.add(new Instant(50)); + state.add(new Instant(20)); + state.add(new Instant(70)); + assertEquals("EARLIEST must win", new Instant(20), internals(store).state(NS_A, tag).read()); + + state.clear(); + assertTrue(internals(store).state(NS_A, tag).isEmpty().read()); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testMapStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag> tag = + StateTags.map("m", StringUtf8Coder.of(), VarIntCoder.of()); + + MapState state = internals(store).state(NS_A, tag); + assertTrue(state.isEmpty().read()); + assertNull(state.get("a").read()); + assertEquals(Integer.valueOf(7), state.getOrDefault("a", 7).read()); + + state.put("a", 1); + state.put("b", 2); + assertEquals(Integer.valueOf(1), internals(store).state(NS_A, tag).get("a").read()); + + List keys = Lists.newArrayList(internals(store).state(NS_A, tag).keys().read()); + assertEquals(2, keys.size()); + assertTrue(keys.contains("a")); + assertTrue(keys.contains("b")); + assertEquals(2, Lists.newArrayList(internals(store).state(NS_A, tag).values().read()).size()); + assertEquals(2, Lists.newArrayList(internals(store).state(NS_A, tag).entries().read()).size()); + + assertEquals( + "computeIfAbsent must not overwrite", + Integer.valueOf(1), + internals(store).state(NS_A, tag).computeIfAbsent("a", k -> 99).read()); + assertNull(internals(store).state(NS_A, tag).computeIfAbsent("c", k -> 3).read()); + assertEquals(Integer.valueOf(3), internals(store).state(NS_A, tag).get("c").read()); + + internals(store).state(NS_A, tag).remove("a"); + assertNull(internals(store).state(NS_A, tag).get("a").read()); + + internals(store).state(NS_A, tag).clear(); + assertTrue(internals(store).state(NS_A, tag).isEmpty().read()); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testSetStateRoundTrip() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.set("s", StringUtf8Coder.of()); + + SetState state = internals(store).state(NS_A, tag); + assertTrue(state.isEmpty().read()); + assertFalse(state.contains("a").read()); + + assertTrue("addIfAbsent returns true the first time", state.addIfAbsent("a").read()); + assertFalse("addIfAbsent returns false the second time", state.addIfAbsent("a").read()); + state.add("b"); + + assertTrue(internals(store).state(NS_A, tag).contains("a").read()); + assertEquals(2, Lists.newArrayList(internals(store).state(NS_A, tag).read()).size()); + + internals(store).state(NS_A, tag).remove("a"); + assertFalse(internals(store).state(NS_A, tag).contains("a").read()); + + internals(store).state(NS_A, tag).clear(); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testDifferentKeysUseDifferentStores() { + // Spark scopes a MapState to the grouping key, so the bridge does not encode the key into the + // store key. Two keys are two stores, which this test pins down as an explicit contract. + InMemoryBytesKV storeOne = store(); + InMemoryBytesKV storeTwo = store(); + StateTag> tag = StateTags.value("v", StringUtf8Coder.of()); + + TwsStateInternals.forKey("one", storeOne).state(NS_A, tag).write("1"); + TwsStateInternals.forKey("two", storeTwo).state(NS_A, tag).write("2"); + + assertEquals("1", TwsStateInternals.forKey("one", storeOne).state(NS_A, tag).read()); + assertEquals("2", TwsStateInternals.forKey("two", storeTwo).state(NS_A, tag).read()); + assertEquals( + "the store key must not depend on the Beam key", + Lists.newArrayList(storeOne.keys()), + Lists.newArrayList(storeTwo.keys())); + } + + @Test(timeout = 30_000) + public void testWindowNamespacesAreDistinct() { + InMemoryBytesKV store = store(); + StateTag> tag = StateTags.value("v", StringUtf8Coder.of()); + StateNamespace first = + StateNamespaces.window( + IntervalWindow.getCoder(), new IntervalWindow(new Instant(0), new Instant(10))); + StateNamespace second = + StateNamespaces.window( + IntervalWindow.getCoder(), new IntervalWindow(new Instant(10), new Instant(20))); + + internals(store).state(first, tag).write("first"); + internals(store).state(second, tag).write("second"); + + assertEquals(2, store.size()); + assertEquals("first", internals(store).state(first, tag).read()); + assertEquals("second", internals(store).state(second, tag).read()); + } + + @Test(timeout = 30_000) + public void testUnsupportedStateTypesFailLoudly() { + InMemoryBytesKV store = store(); + assertThrows( + UnsupportedOperationException.class, + () -> + internals(store) + .state(NS_A, StateTags.multimap("mm", StringUtf8Coder.of(), VarIntCoder.of()))); + assertThrows( + UnsupportedOperationException.class, + () -> internals(store).state(NS_A, StateTags.orderedList("ol", VarIntCoder.of()))); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternalsTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternalsTest.java new file mode 100644 index 000000000000..5177dc902a4b --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternalsTest.java @@ -0,0 +1,418 @@ +/* + * 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.state; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.TwsStateInternalsTest.InMemoryBytesKV; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Unit tests for the Beam {@code TimerInternals} bridge, exercised against an in-memory {@link + * BytesKV} and a recording {@link TwsTimerInternals.WakeupRegistry} rather than a live Spark query. + * + *

The interesting behaviour is not "a timer can be set", it is the reconciliation between Beam's + * rich {@link TimerData} and Spark's bare set of {@code long} wake-ups: de-duplication, deletion of + * wake-ups that no timer needs any more, and the same-millisecond re-arm hazard inside a timer + * callback. + */ +@RunWith(JUnit4.class) +public class TwsTimerInternalsTest { + + private static final IntervalWindow WINDOW = + new IntervalWindow(new Instant(0), new Instant(10_000)); + private static final IntervalWindow OTHER_WINDOW = + new IntervalWindow(new Instant(10_000), new Instant(20_000)); + + private static final StateNamespace NS = + StateNamespaces.window(IntervalWindow.getCoder(), WINDOW); + private static final StateNamespace OTHER_NS = + StateNamespaces.window(IntervalWindow.getCoder(), OTHER_WINDOW); + + /** Records everything the bridge asks Spark to do with its wake-ups. */ + private static final class RecordingRegistry implements TwsTimerInternals.WakeupRegistry { + private final Set live = new LinkedHashSet<>(); + private final List registered = new ArrayList<>(); + private final List deleted = new ArrayList<>(); + + @Override + public void register(long expiryMs) { + registered.add(expiryMs); + live.add(expiryMs); + } + + @Override + public void delete(long expiryMs) { + deleted.add(expiryMs); + live.remove(expiryMs); + } + + @Override + public Set registered() { + return new HashSet<>(live); + } + } + + private TwsTimerInternals internals( + InMemoryBytesKV store, RecordingRegistry registry, Long firedExpiryMs) { + return internals(store, registry, firedExpiryMs, 0L); + } + + private TwsTimerInternals internals( + InMemoryBytesKV store, RecordingRegistry registry, Long firedExpiryMs, long watermarkMs) { + return TwsTimerInternals.create( + store, + registry, + IntervalWindow.getCoder(), + new Instant(watermarkMs), + new Instant(0), + firedExpiryMs); + } + + private static TimerData timer(String id, StateNamespace namespace, long timestampMs) { + return TimerData.of( + id, + "", + namespace, + new Instant(timestampMs), + new Instant(timestampMs), + TimeDomain.EVENT_TIME); + } + + @Test(timeout = 30_000) + public void testSetTimerPersistsAndRegistersAWakeup() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals timers = internals(store, registry, null); + timers.setTimer(timer("t", NS, 1_000)); + assertEquals("nothing may reach Spark before flush", 0, registry.registered.size()); + assertEquals(0, store.size()); + + timers.flush(); + assertEquals(Lists.newArrayList(1_000L), registry.registered); + assertEquals(1, store.size()); + + // A fresh instance over the same store sees the timer again, byte for byte. + TwsTimerInternals reloaded = internals(store, registry, null); + List loaded = Lists.newArrayList(reloaded.getTimers()); + assertEquals(1, loaded.size()); + assertEquals(timer("t", NS, 1_000), loaded.get(0)); + } + + @Test(timeout = 30_000) + public void testTimerDataSurvivesTheStoreRoundTripInFull() throws Exception { + TimerData original = + TimerData.of( + "timerId", + "familyId", + NS, + new Instant(1_234), + new Instant(1_200), + TimeDomain.EVENT_TIME); + TimerInternals.TimerDataCoderV2 coder = + TimerInternals.TimerDataCoderV2.of(IntervalWindow.getCoder()); + + TimerData decoded = + CoderUtils.decodeFromByteArray(coder, CoderUtils.encodeToByteArray(coder, original)); + + assertEquals(original, decoded); + assertEquals("familyId", decoded.getTimerFamilyId()); + assertEquals(new Instant(1_200), decoded.getOutputTimestamp()); + assertEquals(NS, decoded.getNamespace()); + } + + @Test(timeout = 30_000) + public void testWakeupsAreDeduplicated() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + // Two distinct Beam timers, in two namespaces, that expire in the same millisecond. This is the + // common case, an end-of-window timer and its garbage collection timer with zero lateness. + TwsTimerInternals timers = internals(store, registry, null); + timers.setTimer(timer("endOfWindow", NS, 9_999)); + timers.setTimer(timer("gc", OTHER_NS, 9_999)); + timers.flush(); + + assertEquals("one wake-up for two timers", Lists.newArrayList(9_999L), registry.registered); + assertEquals("both timers are persisted", 2, store.size()); + } + + @Test(timeout = 30_000) + public void testAnAlreadyRegisteredWakeupIsNotRegisteredAgain() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("t", NS, 1_000)); + first.flush(); + assertEquals(Lists.newArrayList(1_000L), registry.registered); + + // Second invocation, same key, sets the very same timer again. + TwsTimerInternals second = internals(store, registry, null); + second.setTimer(timer("t", NS, 1_000)); + second.flush(); + + assertEquals( + "no duplicate registerTimer call", Lists.newArrayList(1_000L), registry.registered); + assertEquals("and nothing was deleted either", 0, registry.deleted.size()); + } + + @Test(timeout = 30_000) + public void testDeleteTimerRemovesTheTimerAndItsWakeup() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("t", NS, 1_000)); + first.flush(); + + TwsTimerInternals second = internals(store, registry, null); + second.deleteTimer(NS, "t", "", TimeDomain.EVENT_TIME); + second.flush(); + + assertEquals(0, store.size()); + assertEquals(Lists.newArrayList(1_000L), registry.deleted); + assertTrue(registry.live.isEmpty()); + } + + @Test(timeout = 30_000) + public void testDeletingOneOfTwoTimersKeepsTheSharedWakeup() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("a", NS, 5_000)); + first.setTimer(timer("b", OTHER_NS, 5_000)); + first.flush(); + + TwsTimerInternals second = internals(store, registry, null); + second.deleteTimer(NS, "a", "", TimeDomain.EVENT_TIME); + second.flush(); + + assertEquals("the wake-up is still needed by timer b", 0, registry.deleted.size()); + assertEquals(1, store.size()); + } + + @Test(timeout = 30_000) + public void testMovingATimerMovesItsWakeup() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("t", NS, 1_000)); + first.flush(); + + // TimerData.stringKey() is namespace, domain, family and id, it does not contain the + // timestamp, so re-setting the same timer later is an in place move of the one store entry. + TwsTimerInternals second = internals(store, registry, null); + second.deleteTimer(NS, "t", "", TimeDomain.EVENT_TIME); + second.setTimer(timer("t", NS, 2_000)); + second.flush(); + + assertEquals(Lists.newArrayList(1_000L, 2_000L), registry.registered); + assertEquals(Lists.newArrayList(1_000L), registry.deleted); + assertEquals(1, store.size()); + } + + @Test(timeout = 30_000) + public void testRemoveTimersAtOrBeforeIsOrderedAndConsuming() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals timers = internals(store, registry, null); + timers.setTimer(timer("late", NS, 3_000)); + timers.setTimer(timer("early", NS, 1_000)); + timers.setTimer(timer("middle", NS, 2_000)); + + List due = timers.removeTimersAtOrBefore(new Instant(2_000)); + assertEquals(2, due.size()); + assertEquals(new Instant(1_000), due.get(0).getTimestamp()); + assertEquals(new Instant(2_000), due.get(1).getTimestamp()); + + assertEquals("fired timers are gone", 1, Lists.newArrayList(timers.getTimers()).size()); + assertEquals( + "and cannot fire twice", 0, timers.removeTimersAtOrBefore(new Instant(2_000)).size()); + + timers.flush(); + assertEquals("only the surviving timer is persisted", 1, store.size()); + assertEquals(Lists.newArrayList(3_000L), registry.registered); + } + + /** + * Spark expires a wake-up as soon as {@code expiry <= watermark}, Beam only fires an event time + * timer once the watermark is strictly past it. A timer sitting exactly on the batch watermark + * must therefore be withheld and re-armed rather than handed to Beam, which would swallow it. + */ + @Test(timeout = 30_000) + public void testTimerExactlyAtTheWatermarkIsWithheldAndReArmed() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("endOfWindow", NS, 9_999)); + first.flush(); + registry.registered.clear(); + registry.deleted.clear(); + + // Batch watermark is exactly 9999, so Spark expires the wake-up but Beam is not ready. + TwsTimerInternals early = internals(store, registry, 9_999L, 9_999L); + assertEquals( + "a timer on the watermark is not due yet", 0, early.removeTimersReadyToFire(9_999L).size()); + early.flush(); + + assertEquals("the timer survives", 1, store.size()); + assertEquals( + "and is re-armed one millisecond later", Lists.newArrayList(10_000L), registry.registered); + assertEquals("the firing expiry is left to Spark", 0, registry.deleted.size()); + + // Next batch, the watermark has genuinely moved past the timer. + TwsTimerInternals late = internals(store, registry, 10_000L, 20_000L); + List due = late.removeTimersReadyToFire(10_000L); + assertEquals(1, due.size()); + assertEquals( + "the Beam timestamp is untouched by the re-arm", + new Instant(9_999), + due.get(0).getTimestamp()); + late.flush(); + assertEquals("and it is consumed", 0, store.size()); + } + + @Test(timeout = 30_000) + public void testFiringExpiryIsNotDeletedByFlush() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("t", NS, 1_000)); + first.flush(); + registry.deleted.clear(); + + // Spark is firing 1000 and removes that wake-up itself once the callback returns. The bridge + // must not race it with a delete of its own. + TwsTimerInternals callback = internals(store, registry, 1_000L, 1_001L); + assertEquals(1, callback.removeTimersReadyToFire(1_000L).size()); + callback.flush(); + + assertEquals(0, registry.deleted.size()); + assertEquals(0, store.size()); + } + + @Test(timeout = 30_000) + public void testReArmAtTheFiringMillisecondIsNudgedForward() { + InMemoryBytesKV store = new InMemoryBytesKV(); + RecordingRegistry registry = new RecordingRegistry(); + + TwsTimerInternals first = internals(store, registry, null); + first.setTimer(timer("t", NS, 1_000)); + first.flush(); + registry.registered.clear(); + + // Inside the callback for expiry 1000 the DoFn sets a new timer at 1000 again. Spark would + // delete a wake-up registered at exactly 1000 when the callback finishes, so it has to land at + // 1001 instead. The TimerData keeps its own timestamp of 1000. + TwsTimerInternals callback = internals(store, registry, 1_000L, 1_001L); + callback.removeTimersReadyToFire(1_000L); + callback.setTimer(timer("again", NS, 1_000)); + callback.flush(); + + assertEquals(Lists.newArrayList(1_001L), registry.registered); + + TwsTimerInternals next = internals(store, registry, 1_001L, 1_002L); + List due = next.removeTimersReadyToFire(1_001L); + assertEquals(1, due.size()); + assertEquals( + "the Beam timestamp is untouched by the wake-up nudge", + new Instant(1_000), + due.get(0).getTimestamp()); + } + + @Test(timeout = 30_000) + public void testProcessingTimeTimersAreRejected() { + InMemoryBytesKV store = new InMemoryBytesKV(); + TwsTimerInternals timers = internals(store, new RecordingRegistry(), null); + + UnsupportedOperationException processing = + assertThrows( + UnsupportedOperationException.class, + () -> + timers.setTimer( + NS, "t", "", new Instant(1), new Instant(1), TimeDomain.PROCESSING_TIME)); + assertTrue(processing.getMessage().contains("event time timers")); + + assertThrows( + UnsupportedOperationException.class, + () -> + timers.setTimer( + NS, + "t", + "", + new Instant(1), + new Instant(1), + TimeDomain.SYNCHRONIZED_PROCESSING_TIME)); + } + + @Test(timeout = 30_000) + public void testDeleteTimerWithoutTimeDomainIsRejected() { + TwsTimerInternals timers = internals(new InMemoryBytesKV(), new RecordingRegistry(), null); + assertThrows(UnsupportedOperationException.class, () -> timers.deleteTimer(NS, "t", "")); + } + + @Test(timeout = 30_000) + public void testClocksAndUnsupportedWatermarks() { + TwsTimerInternals timers = + TwsTimerInternals.create( + new InMemoryBytesKV(), + new RecordingRegistry(), + IntervalWindow.getCoder(), + new Instant(7_000), + new Instant(9_000), + null); + + assertEquals(new Instant(7_000), timers.currentInputWatermarkTime()); + assertEquals(new Instant(9_000), timers.currentProcessingTime()); + assertNull(timers.currentSynchronizedProcessingTime()); + assertNull(timers.currentOutputWatermarkTime()); + } + + @Test(timeout = 30_000) + public void testFlushIsSingleUse() { + TwsTimerInternals timers = internals(new InMemoryBytesKV(), new RecordingRegistry(), null); + timers.flush(); + assertThrows(IllegalStateException.class, timers::flush); + } +} From 09be88ce2a026c3998fa1931f6e525397d089c6e Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 13:45:23 +0000 Subject: [PATCH 04/21] [Spark 4] Wire up the streaming registry and evaluation context Shadows PipelineTranslatorFactory so streaming pipelines dispatch to a new PipelineTranslatorStreaming, which extends PipelineTranslatorBatch to reuse Impulse, Window.Assign, Flatten, Reshuffle, the bounded read and stateless ParDo unchanged, while registering placeholders for the unbounded read, GroupByKey and stateful ParDo that WS-D2 will replace. Combine.PerKey is deliberately left unregistered so it auto-expands into GroupByKey plus ParDo. Adds StreamingEvaluationContext, which starts one noop-sink streaming query per leaf dataset with a checkpoint directory derived from the pipeline options, registers an idle-stop listener for test termination, and makes stop() idempotent and safe to call concurrently with evaluate() so cancel() can interrupt a running streaming pipeline from another thread. --- .../PipelineTranslatorFactory.java | 40 +++ .../PipelineTranslatorStreaming.java | 143 ++++++++ .../StreamingEvaluationContext.java | 304 ++++++++++++++++++ 3 files changed, 487 insertions(+) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java new file mode 100644 index 000000000000..ac7f6aa5d6c1 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java @@ -0,0 +1,40 @@ +/* + * 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; + +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.sdk.annotations.Internal; + +/** + * Factory to create the {@link PipelineTranslator} matching the execution mode of the pipeline. + * + *

This file shadows the shared base version of the same name found under {@code + * runners/spark/src}. The Spark 4 module compiles a merged source tree of the shared base plus + * {@code runners/spark/4/src}, with a later-wins duplicate strategy, so this copy silently replaces + * the base one for the Spark 4 module only. The base version keeps throwing for streaming, this one + * dispatches to the real streaming translator. + */ +@Internal +public final class PipelineTranslatorFactory { + private PipelineTranslatorFactory() {} + + /** Creates a {@link PipelineTranslator} for the given execution mode. */ + public static PipelineTranslator create(boolean streaming) { + return streaming ? new PipelineTranslatorStreaming() : new PipelineTranslatorBatch(); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java new file mode 100644 index 000000000000..be5107234ee3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -0,0 +1,143 @@ +/* + * 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; + +import java.util.Collection; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +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.PInput; +import org.apache.beam.sdk.values.POutput; +import org.apache.spark.sql.SparkSession; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * {@link PipelineTranslator} for executing a streaming {@link org.apache.beam.sdk.Pipeline} on + * Spark 4. + * + *

This extends {@link PipelineTranslatorBatch} purely for reuse: {@link + * PipelineTranslatorBatch#getTransformTranslator} and its private registry are the only way to + * reach the (package-private) batch translators for {@code Impulse}, {@code Window.Assign}, {@code + * Flatten}, {@code Reshuffle}, the bounded read, and stateless {@code ParDo}, all of which are + * reused completely unchanged for streaming. This class only intercepts the handful of transforms + * that need genuinely different, streaming-aware handling before falling back to {@code super}. + */ +@Internal +public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { + + // -------------------------------------------------------------------------------------------- + // Placeholders for WS-D2 + // -------------------------------------------------------------------------------------------- + // + // WS-D2 replaces each of the three constants below with a real translator instance (constructed + // in place, same as the batch registry does) and removes the corresponding TODO. No other change + // to this class should be necessary to plug those in: the interception points that select them in + // getTransformTranslator already exist and are marked below. + + // TODO(WS-D2): replace with `new ReadUnboundedTranslator<>()`. + private static final TransformTranslator READ_UNBOUNDED_PLACEHOLDER = + new UnsupportedStreamingTranslator( + "Unbounded read streaming translation not implemented yet (WS-D2)"); + + // TODO(WS-D2): replace with `new GroupByKeyStreamingTranslator<>()`. + private static final TransformTranslator GROUP_BY_KEY_PLACEHOLDER = + new UnsupportedStreamingTranslator( + "GroupByKey streaming translation not implemented yet (WS-D2)"); + + // TODO(WS-D2): replace with `new StatefulParDoStreamingTranslator<>()`. + private static final TransformTranslator STATEFUL_PAR_DO_PLACEHOLDER = + new UnsupportedStreamingTranslator( + "Stateful ParDo streaming translation not implemented yet (WS-D2)"); + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + @Nullable + protected > + TransformTranslator getTransformTranslator(TransformT transform) { + + // TODO(WS-D2): swap for `if (transform instanceof SplittableParDo.PrimitiveUnboundedRead)` + // returning `new ReadUnboundedTranslator<>()` (or a cached instance thereof). + if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { + return (TransformTranslator) READ_UNBOUNDED_PLACEHOLDER; + } + + // TODO(WS-D2): swap for `if (transform instanceof GroupByKey)` returning + // `new GroupByKeyStreamingTranslator<>()`. + if (transform instanceof GroupByKey) { + return (TransformTranslator) GROUP_BY_KEY_PLACEHOLDER; + } + + // Deliberately never registered: leaving Combine.PerKey unhandled here makes Beam auto-expand + // it into GroupByKey + ParDo, so the streaming translations above (and WS-D2's stateful ParDo + // translator) take over the expanded primitives instead of the batch + // CombinePerKeyTranslatorBatch, which has no streaming support. + if (transform instanceof Combine.PerKey) { + return null; + } + + if (transform instanceof ParDo.MultiOutput) { + DoFnSignature signature = + DoFnSignatures.signatureForDoFn(((ParDo.MultiOutput) transform).getFn()); + // TODO(WS-D2): swap this branch's return for `new StatefulParDoStreamingTranslator<>()`. + if (signature.usesState() || signature.usesTimers()) { + return (TransformTranslator) STATEFUL_PAR_DO_PLACEHOLDER; + } + // Stateless ParDo falls through to super, reusing ParDoTranslatorBatch unchanged. + } + + // Impulse, Window.Assign, Flatten, Reshuffle, the bounded read, and stateless ParDo: reused + // unchanged from the batch registry. + return super.getTransformTranslator(transform); + } + + @Override + protected EvaluationContext createEvaluationContext( + Collection> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + return new StreamingEvaluationContext(leaves, session, options); + } + + /** + * A {@link TransformTranslator} standing in for a streaming translation that WS-D2 has not + * implemented yet. Always throws {@link UnsupportedOperationException} as soon as the pipeline + * traversal tries to translate the transform it is registered for. + */ + private static final class UnsupportedStreamingTranslator + extends TransformTranslator> { + private final String message; + + UnsupportedStreamingTranslator(String message) { + super(0); + this.message = message; + } + + @Override + protected void translate(PTransform transform, Context cxt) { + throw new UnsupportedOperationException(message); + } + } +} 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 new file mode 100644 index 000000000000..a1e702dc7ed3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -0,0 +1,304 @@ +/* + * 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; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +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.sdk.annotations.Internal; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.Trigger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * The streaming counterpart of {@link EvaluationContext}: instead of forcing a one-shot batch + * evaluation of every leaf dataset, it starts one Spark Structured Streaming query per leaf and + * blocks until all of them reach a terminal state. + * + *

Sink choice

+ * + *

Every query uses the {@code noop} sink, never {@code memory}. Tests in this module run with + * {@code spark.kryo.registrationRequired=true}, and the {@code memory} sink's commit messages + * ({@code MemoryWriterCommitMessage}) are not registered with any Beam Kryo registrator, so a query + * using it dies on batch 0. This was already discovered the hard way during the Phase 0 spike, do + * not rediscover it. + * + *

Termination

+ * + *

Queries read from sources with opaque epoch offsets that never settle (see {@code + * UnboundedSourceDataset}), so {@code StreamingQueryManager#awaitAnyTermination} without outside + * help would hang forever. Two independent knobs exist to terminate a query: + * + *

    + *
  • {@link #stop()}, invoked by {@code SparkStructuredStreamingPipelineResult#cancel()}. + *
  • The idle-stop listener registered in {@link #evaluate()} when {@code + * SparkStructuredStreamingPipelineOptions#getStreamingStopAfterIdleBatches()} is {@code >= + * 0}: it counts consecutive micro-batches with zero input rows per query and gracefully stops + * that one query once the threshold is reached. This is how streaming tests in this module + * terminate on their own. + *
+ * + *

Draining on {@link #stop()}

+ * + *

{@link #stop()} does not attempt a full {@code Trigger.AvailableNow()} drain pass that + * processes every already-buffered offset before halting: doing so would require stopping the query + * and restarting it with a different trigger against the same checkpoint, which is more machinery + * than this POC's lifecycle warrants and risks checkpoint-compatibility bugs of its own. Instead, + * {@link #stop()} relies on {@link StreamingQuery#stop()}'s own graceful behaviour, which lets a + * micro-batch that is already in flight finish normally instead of interrupting it, and only then + * halts the query. Data that was not yet pulled into an in-flight micro-batch at the moment {@link + * #stop()} is called is simply left unprocessed. This is a documented limitation of the POC, not an + * oversight. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards both `queries` and `stopped` so evaluate() (which appends to `queries` as it starts + // queries) and stop() (which may run concurrently on another thread, see the class javadoc on + // thread-safety below) never race on which queries have been started or already stopped. + private final Object lock = new Object(); + private final List queries = new ArrayList<>(); + private boolean stopped = false; + + StreamingEvaluationContext( + Collection> leaves, + SparkSession session, + SparkCommonPipelineOptions options) { + super(leaves, session); + this.options = options.as(SparkStructuredStreamingPipelineOptions.class); + } + + /** + * Starts one streaming query per leaf dataset and blocks until every one of them has reached a + * terminal state, either because {@link #stop()} was called (typically via {@code cancel()}) or + * because the idle-stop listener stopped it after enough consecutive empty micro-batches. + */ + @Override + public void evaluate() { + String checkpointBaseDir = checkpointBaseDir(options); + int idleStopThreshold = options.getStreamingStopAfterIdleBatches(); + + StreamingQueryListener idleStopListener = null; + if (idleStopThreshold >= 0) { + idleStopListener = new IdleStopListener(idleStopThreshold); + getSparkSession().streams().addListener(idleStopListener); + } + + try { + int leafIndex = 0; + for (NamedDataset ds : leaves()) { + Dataset dataset = ds.dataset(); + if (dataset == null) { + continue; + } + synchronized (lock) { + if (stopped) { + // stop() already ran (e.g. an immediate cancel()); do not start further queries. + break; + } + } + if (!dataset.isStreaming()) { + // Defensive fallback: a leaf that turns out not to be streaming (e.g. a bounded side + // collection) is simply evaluated the batch way instead of starting a query for it. + EvaluationContext.evaluate(ds.name(), dataset); + continue; + } + + StreamingQuery query = startQuery(dataset, checkpointBaseDir, leafIndex++, options); + boolean alreadyStopped; + synchronized (lock) { + queries.add(query); + alreadyStopped = stopped; + } + if (alreadyStopped) { + // stop() ran in the window between the check above and this query actually starting. + stopQuery(query); + } + } + + List toAwait; + synchronized (lock) { + toAwait = new ArrayList<>(queries); + } + for (StreamingQuery query : toAwait) { + awaitTermination(query); + } + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + } + } + + /** + * Stops all queries started by {@link #evaluate()}. + * + *

Idempotent and safe to call from a thread other than the one running {@link #evaluate()}: + * {@code cancel()} calls this from the main thread while {@code evaluate()} is blocked awaiting + * termination on the pipeline execution thread. See the class javadoc for the drain limitation. + */ + @Override + public void stop() { + List toStop; + synchronized (lock) { + if (stopped) { + return; + } + stopped = true; + toStop = new ArrayList<>(queries); + } + for (StreamingQuery query : toStop) { + stopQuery(query); + } + } + + private StreamingQuery startQuery( + Dataset dataset, + String checkpointBaseDir, + int leafIndex, + SparkStructuredStreamingPipelineOptions options) { + try { + return dataset + .writeStream() + .format("noop") + .outputMode("append") + .option("checkpointLocation", checkpointBaseDir + "/" + leafIndex) + .trigger(Trigger.ProcessingTime(options.getMaxBatchDurationMillis())) + .start(); + } catch (TimeoutException e) { + throw new RuntimeException( + "Failed to start streaming query for leaf dataset index " + leafIndex, e); + } + } + + private void awaitTermination(StreamingQuery query) { + try { + query.awaitTermination(); + } catch (StreamingQueryException e) { + LOG.error("Streaming query {} terminated with an exception.", query.id(), e); + // Make sure sibling queries do not keep running once one of them has failed. + stop(); + throw new RuntimeException(e); + } + } + + /** + * Best-effort, idempotent stop of a single query, see the class javadoc for what "best-effort" + * means here. + */ + private void stopQuery(StreamingQuery query) { + try { + if (query.isActive()) { + query.stop(); + } + } catch (TimeoutException | RuntimeException e) { + LOG.warn( + "Error while stopping streaming query {}: {}", + query.id(), + String.valueOf(e.getMessage())); + } + } + + private void stopQueryById(UUID id) { + StreamingQuery match = null; + synchronized (lock) { + for (StreamingQuery query : queries) { + if (query.id().equals(id)) { + match = query; + break; + } + } + } + if (match != null) { + stopQuery(match); + } + } + + private static String checkpointBaseDir(SparkCommonPipelineOptions options) { + String dir = options.getCheckpointDir(); + if (dir == null || dir.isEmpty()) { + try { + dir = Files.createTempDirectory("beam-spark4-streaming-checkpoint").toString(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + LOG.warn("No checkpoint directory configured, falling back to temporary directory {}.", dir); + } + return dir; + } + + /** + * Counts, per query, the number of consecutive micro-batches with zero input rows, and gracefully + * stops a query once its count reaches {@code threshold}. The count for a query resets to zero as + * soon as one of its micro-batches has rows. + * + *

Stopping happens on a dedicated thread rather than inline in {@link #onQueryProgress}: + * {@link StreamingQuery#stop()} blocks until the query's execution thread has shut down, which + * should not happen on the listener bus thread that dispatches these callbacks. + */ + private final class IdleStopListener extends StreamingQueryListener { + private final int threshold; + private final Map idleCounts = new ConcurrentHashMap<>(); + + IdleStopListener(int threshold) { + this.threshold = threshold; + } + + @Override + public void onQueryStarted(QueryStartedEvent event) {} + + @Override + public void onQueryProgress(QueryProgressEvent event) { + UUID id = event.progress().id(); + if (event.progress().numInputRows() == 0) { + int count = idleCounts.computeIfAbsent(id, unused -> new AtomicInteger()).incrementAndGet(); + if (count >= threshold) { + idleCounts.remove(id); + Thread stopThread = new Thread(() -> stopQueryById(id), "beam-idle-stop-" + id); + stopThread.setDaemon(true); + stopThread.start(); + } + } else { + idleCounts.remove(id); + } + } + + @Override + public void onQueryTerminated(QueryTerminatedEvent event) { + idleCounts.remove(event.id()); + } + } +} From e2f6b0ea6c549f0150a93df335e97271135af36b Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 14:00:26 +0000 Subject: [PATCH 05/21] [Spark 4] Add streaming test scaffolding and Ignore'd translator skeletons Adds StreamingTestUtils, the ListBackedUnboundedSource plus static collector and pipeline option factories the streaming translators will be tested against, and five skeleton test classes for the stateless ParDo, windowed GroupByKey, stateful ParDo, chained stateful, and lifecycle scenarios. Every skeleton compiles and runs, reporting as skipped: each is Ignore'd with a comment naming the exact WS-D2 translator it needs once that lands. --- .../ChainedStatefulStreamingTest.java | 154 +++++++ .../streaming/StatefulParDoStreamingTest.java | 150 +++++++ .../StatelessParDoStreamingTest.java | 97 +++++ .../StreamingPipelineLifecycleTest.java | 142 ++++++ .../streaming/StreamingTestUtils.java | 404 ++++++++++++++++++ .../WindowedGroupByKeyStreamingTest.java | 205 +++++++++ 6 files changed, 1152 insertions(+) create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java new file mode 100644 index 000000000000..2bb367ac5ada --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java @@ -0,0 +1,154 @@ +/* + * 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 java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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.VarLongCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Sum; +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.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * A dedup stateful {@code ParDo} feeding a windowed {@code GroupByKey} sum: two {@code + * transformWithState} operators chained in a single query. + * + *

This is the most important test in the suite. It is the one POC scenario that actually + * exercises cross-operator watermark propagation, the key claim of the whole Phase 1-4 plan: that a + * Spark 4 micro-batch's watermark, computed once at the source, keeps meaning the same thing as it + * flows through a chain of independently-hosted stateful operators, so a downstream window can + * still correctly decide when it has seen everything it is going to see. Everything else in this + * package tests one operator at a time; this one tests that they compose. + * + *

Needs the same Kryo relaxation as {@code BeamStatefulProcessorTest}, for the same reason. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class ChainedStatefulStreamingTest implements Serializable { + + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final Duration WINDOW_SIZE = Duration.standardSeconds(10); + + /** + * Dedups by the outer {@code id} key (simulating at-least-once redelivery of the same logical + * event) and, on the first sighting of an id, passes the inner {@code KV} + * through for the downstream windowed sum. + */ + private static class DedupByIdFn extends DoFn>, KV> { + @StateId("seen") + private final StateSpec> seenSpec = StateSpecs.value(); + + @ProcessElement + public void process( + @Element KV> element, + @StateId("seen") ValueState seen, + OutputReceiver> out) { + Boolean alreadySeen = seen.read(); + if (alreadySeen == null || !alreadySeen) { + seen.write(true); + out.output(element.getValue()); + } + } + } + + @Test(timeout = 300_000) + @Ignore( + "Needs both StatefulParDoStreamingTranslator and GroupByKeyStreamingTranslator (WS-D2): this " + + "test chains a stateful ParDo into a windowed GroupByKey, so it is blocked on whichever " + + "of the two lands last. Also needs the cross-operator watermark propagation itself to " + + "actually work, which is the top risk called out in the POC plan (risk #3).") + public void dedupThenWindowedSumPropagatesWatermarkAcrossOperators() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("chained-stateful"); + StreamingTestUtils.clear(collectorId); + + List>>> elements = new ArrayList<>(); + // id "1" reported twice (redelivery), must count towards key "a" only once. + elements.add(TimestampedValue.of(KV.of("1", KV.of("a", 5L)), BASE)); + elements.add( + TimestampedValue.of(KV.of("1", KV.of("a", 5L)), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("2", KV.of("a", 3L)), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("3", KV.of("b", 10L)), BASE.plus(Duration.standardSeconds(3)))); + // Watermark rule: a much later element so the watermark passes the first window's end at both + // the dedup operator and the downstream windowed sum operator. + elements.add( + TimestampedValue.of( + KV.of("sentinel", KV.of("sentinel", 0L)), BASE.plus(Duration.standardSeconds(60)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + TestPipeline pipeline = TestPipeline.fromOptions(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, + KvCoder.of( + StringUtf8Coder.of(), + KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))))) + .apply("DedupById", ParDo.of(new DedupByIdFn())) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("SumPerKey", Sum.longsPerKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-E2, once WS-D2 lands): assert StreamingTestUtils.>getCollected( + // collectorId) contains exactly KV.of("a", 8L) (5 + 3, the duplicate "1" excluded) and + // KV.of("b", 10L) for the [0s, 10s) window. Getting exactly these values, rather than e.g. + // KV.of("a", 13L) from a missed dedup or nothing at all from a stuck watermark, is the + // observable proof that watermark and results both survived the operator chain intact. + } +} 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 new file mode 100644 index 000000000000..2ecacb81141a --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java @@ -0,0 +1,150 @@ +/* + * 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 java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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.io.Read; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.state.TimerSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.testing.TestPipeline; +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.values.KV; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * A stateful {@code ParDo} in the global window: a {@code @StateId ValueState} dedups + * repeated keys, and an event time {@code @TimerId} emits a sentinel once it expires. Hosted by the + * generic {@code transformWithState} super-operator ({@code + * BeamStatefulProcessorConfig.Mode#STATEFUL_PARDO}), so needs the same Kryo relaxation as {@code + * BeamStatefulProcessorTest}. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StatefulParDoStreamingTest implements Serializable { + + /** See {@code BeamStatefulProcessorTest} for why this relaxation is required. */ + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final String SENTINEL = "TIMER-FIRED"; + + /** + * Dedups repeated keys with a {@code ValueState} and, on the first sighting of a key, + * arms an event-time timer thirty seconds out; once that timer fires it emits {@link #SENTINEL}. + */ + private static class DedupWithExpiryFn extends DoFn, String> { + @StateId("seen") + private final StateSpec> seenSpec = StateSpecs.value(); + + @TimerId("expiry") + private final TimerSpec expirySpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + @ProcessElement + public void process( + @Element KV element, + @Timestamp Instant timestamp, + @StateId("seen") ValueState seen, + @TimerId("expiry") Timer expiryTimer, + OutputReceiver out) { + Boolean alreadySeen = seen.read(); + if (alreadySeen == null) { + expiryTimer.set(timestamp.plus(Duration.standardSeconds(30))); + } + if (alreadySeen == null || !alreadySeen) { + seen.write(true); + out.output(element.getValue()); + } + } + + @OnTimer("expiry") + public void onExpiry(OutputReceiver out) { + out.output(SENTINEL); + } + } + + @Test(timeout = 300_000) + @Ignore( + "Needs StatefulParDoStreamingTranslator (WS-D2) registered in place of the " + + "UnsupportedOperationException placeholder in " + + "PipelineTranslatorStreaming#STATEFUL_PAR_DO_PLACEHOLDER; a DoFn signature with state or " + + "timers currently has no streaming translation.") + public void dedupsRepeatedKeysAndFiresTimerSentinel() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("stateful-pardo-dedup"); + StreamingTestUtils.clear(collectorId); + + List> elements = new ArrayList<>(); + elements.add(TimestampedValue.of("a", BASE)); + // Duplicate key "a": the dedup state must suppress this one. + elements.add(TimestampedValue.of("a", BASE.plus(Duration.standardSeconds(1)))); + elements.add(TimestampedValue.of("b", BASE.plus(Duration.standardSeconds(2)))); + // Watermark rule: push well past the 30s timer deadline armed for key "a" (and "b"). + elements.add(TimestampedValue.of("c", BASE.plus(Duration.standardSeconds(90)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + TestPipeline pipeline = TestPipeline.fromOptions(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>(elements, StringUtf8Coder.of()))) + .apply("WithKeys", WithKeys.of(value -> value)) + .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) + .apply("DedupWithExpiry", ParDo.of(new DedupWithExpiryFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-D2): assert StreamingTestUtils.getCollected(collectorId) contains exactly one + // "a", one "b", one "c" (each key's first sighting) plus two SENTINEL values (one per timer + // armed for "a" and "b"; "c" arrives too late in the run for its own timer to have expired). + // Remember the one micro-batch timer latency floor documented on StreamingTestUtils. + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java new file mode 100644 index 000000000000..ceeedf70f43f --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -0,0 +1,97 @@ +/* + * 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 java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * The simplest possible streaming pipeline: an unbounded source feeding a plain, stateless {@code + * ParDo}, with every element expected to pass straight through. This is the baseline the other + * streaming tests in this package build on. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StatelessParDoStreamingTest implements Serializable { + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + + /** Doubles the input, so the assertion can tell the ParDo actually ran, not just passed data. */ + private static class DoubleFn extends DoFn { + @ProcessElement + public void process(@Element Integer element, OutputReceiver out) { + out.output(element * 2); + } + } + + @Test(timeout = 300_000) + @Ignore( + "Needs ReadUnboundedTranslator (WS-D2) registered in place of the UnsupportedOperationException " + + "placeholder in PipelineTranslatorStreaming#READ_UNBOUNDED_PLACEHOLDER; until then " + + "SplittableParDo.PrimitiveUnboundedRead has no streaming translation and pipeline.run() " + + "throws before any element is read.") + public void everyElementPassesThrough() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("stateless-pardo"); + StreamingTestUtils.clear(collectorId); + + List> elements = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + TestPipeline pipeline = TestPipeline.fromOptions(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>(elements, VarIntCoder.of()))) + .apply("Double", ParDo.of(new DoubleFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-D2): once the unbounded read translates, assert that + // StreamingTestUtils.getCollected(collectorId) contains exactly {0, 2, 4, ..., 18}, + // in any order (Spark micro-batches make no ordering guarantee across elements). + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java new file mode 100644 index 000000000000..78cb0070dd3e --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -0,0 +1,142 @@ +/* + * 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 static org.junit.Assert.assertEquals; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * {@link PipelineResult.State} transitions for a streaming pipeline: {@code RUNNING} immediately + * after {@code run()}, {@code DONE} once a naturally idle pipeline's idle-stop listener stops it, + * and {@code CANCELLED} after an explicit {@code cancel()}. + * + *

Both tests here set {@code testMode(false)} on top of {@link + * StreamingTestUtils#streamingOptions}, unlike every other test in this package: {@code + * SparkStructuredStreamingRunner#run()} calls {@code result.waitUntilFinish()} itself before + * returning when {@code testMode} is {@code true} (see its implementation), which would make {@code + * run()} block past the point these tests want to observe {@code RUNNING}. + * + *

Not tested here: that a streaming pipeline is rejected when run against Spark 3. That is + * {@code PipelineTranslatorFactory#create} in the shared base module + * (runners/spark/src/main/java/.../translation/PipelineTranslatorFactory.java) throwing {@code + * UnsupportedOperationException}, and this test module only ever compiles and runs against the + * Spark 4 classpath (this module's shadow copy of that same file dispatches to the real streaming + * translator instead). Exercising the rejection needs a Spark 3 dependency this module deliberately + * does not have; the right place for that check is a test in {@code runners/spark/src/test/...} + * against the shared base module alone. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingPipelineLifecycleTest implements Serializable { + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + + private List> tenElements() { + List> elements = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + return elements; + } + + @Test(timeout = 300_000) + @Ignore( + "Needs ReadUnboundedTranslator (WS-D2): today translation itself throws " + + "UnsupportedOperationException from PipelineTranslatorStreaming#READ_UNBOUNDED_PLACEHOLDER " + + "before a query ever starts, so waitUntilFinish() surfaces FAILED, not DONE.") + public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-done"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + // Observe RUNNING ourselves instead of letting run() block until finished, see class javadoc. + options.setTestMode(false); + TestPipeline pipeline = TestPipeline.fromOptions(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + PipelineResult.State finalState = result.waitUntilFinish(); + assertEquals(PipelineResult.State.DONE, finalState); + assertEquals(PipelineResult.State.DONE, result.getState()); + } + + @Test(timeout = 300_000) + @Ignore( + "Needs ReadUnboundedTranslator (WS-D2): cancel() only has an effect once translation has " + + "completed and StreamingEvaluationContext#stop() is reachable through the runner's " + + "ctxRef; today translation throws before that ever happens.") + public void cancelStopsTheQueryAndReportsCancelled() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-cancel"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + options.setTestMode(false); + // Disabled so the query only ever stops because of the explicit cancel() below, not because it + // happened to go idle first. + options.setStreamingStopAfterIdleBatches(-1); + TestPipeline pipeline = TestPipeline.fromOptions(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + assertEquals(PipelineResult.State.RUNNING, result.getState()); + + PipelineResult.State cancelledState = result.cancel(); + assertEquals(PipelineResult.State.CANCELLED, cancelledState); + assertEquals(PipelineResult.State.CANCELLED, result.getState()); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java new file mode 100644 index 000000000000..a1ed7f11c2ff --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -0,0 +1,404 @@ +/* + * 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 java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.junit.rules.TemporaryFolder; + +/** + * Shared test scaffolding for the Spark 4 streaming translators: an {@link UnboundedSource} over a + * fixed, in-memory list of elements, a driver-side static collector {@link DoFn}, and factories for + * the {@link SparkStructuredStreamingPipelineOptions} / {@link TestPipeline} every streaming test + * in this package needs. + * + *

Why streaming tests in this suite look the way they do

+ * + *

Two of the usual Beam testing tools do not work here, on purpose: + * + *

    + *
  • {@code PAssert} on an unbounded {@code PCollection} never fires: {@code PAssert} + * needs a final, +infinity watermark to know a window's contents are complete, and this + * runner's sources never produce one (see below). + *
  • {@code StreamingQuery#processAllAvailable()} hangs forever: {@link + * org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamMicroBatchStream} (WS-B) + * reports progress with opaque epoch offsets, not byte/row counts, so Spark can never decide + * that "all available" input has been consumed. + *
+ * + *

Instead, every test in this package follows the same recipe: + * + *

    + *
  1. Build a {@link ListBackedUnboundedSource} from a finite list of elements. The source is + * typed {@code UnboundedSource}, so the pipeline is genuinely a streaming pipeline, but it + * naturally runs out of input and Spark starts reporting empty micro-batches. + *
  2. Set {@link SparkStructuredStreamingPipelineOptions#setStreamingStopAfterIdleBatches} (via + * {@link #streamingOptions}, already set to {@code 3}) so the runner's idle-stop listener + * gracefully stops the query a few empty micro-batches after the source is exhausted, instead + * of running forever. + *
  3. Call {@code pipelineResult.waitUntilFinish()} and only then assert against the {@link + * #getCollected} static collector, never against the {@code PCollection} itself. + *
+ * + *

The watermark rule

+ * + *

A window only fires once the watermark has passed its end, and the watermark only ever + * advances on new data. Concretely, the watermark Spark computes for the source's event + * timestamp column is {@code max(eventTimestamp seen so far) - watermarkDelay}, with no idle-time + * advance: if the source stops producing elements, the watermark freezes where it was, forever. It + * does not jump to +infinity when the source is exhausted, unlike a bounded pipeline's final + * watermark. + * + *

The consequence for test authors: every window (or event-time timer) you want to assert on + * must be followed, in the input list, by at least one element timestamped after that window's end + * (or the timer's deadline). Without such a trailing element the watermark never crosses the + * threshold and the window or timer never fires, and the test will simply see nothing rather than + * failing loudly. + * + *

A second, related wrinkle carried over from WS-C's state bridge tests: the watermark visible + * inside a stateful operator (a {@code transformWithState} query) is the watermark as of the + * start of the current micro-batch, not the one just computed from the current batch's own + * rows. An end-of-window timer whose deadline the data has already crossed therefore fires one + * micro-batch later than the batch that carried the crossing data, not in that same batch. Tests + * that assert on timer firings should expect this one micro-batch latency floor and provide enough + * trailing elements (i.e. enough separate micro-batches) for it. + * + *

Local mode only

+ * + *

{@link #getCollected} works only because these tests run Spark in local mode inside the same + * JVM as the test itself: {@link CollectDoFn} appends to a plain static, synchronized, in-process + * map. It would not see anything written by executors of a real (multi-JVM) Spark cluster. + */ +public final class StreamingTestUtils { + + private StreamingTestUtils() {} + + // --------------------------------------------------------------------------------------------- + // Static, in-process collector. + // --------------------------------------------------------------------------------------------- + + /** Driver-side, per-collector-id accumulation of every element a {@link CollectDoFn} has seen. */ + private static final Map> COLLECTORS = new ConcurrentHashMap<>(); + + /** + * A {@link DoFn} that appends every element it sees to the static, in-process collector named + * {@code collectorId}, then passes the element through unchanged. Safe to use concurrently from + * multiple bundles/threads; see the class javadoc for why this only works in Spark local mode. + */ + public static final class CollectDoFn extends DoFn { + private final String collectorId; + + public CollectDoFn(String collectorId) { + this.collectorId = Preconditions.checkNotNull(collectorId); + } + + public String getCollectorId() { + return collectorId; + } + + @ProcessElement + public void processElement(@Element T element, OutputReceiver out) { + append(collectorId, element); + out.output(element); + } + } + + private static void append(String collectorId, Object value) { + COLLECTORS + .computeIfAbsent(collectorId, unused -> Collections.synchronizedList(new ArrayList<>())) + .add(value); + } + + /** Returns a snapshot of everything collected so far under {@code collectorId}. */ + @SuppressWarnings("unchecked") + public static List getCollected(String collectorId) { + List values = COLLECTORS.get(collectorId); + if (values == null) { + return Collections.emptyList(); + } + synchronized (values) { + return (List) new ArrayList<>(values); + } + } + + /** Discards everything collected so far under {@code collectorId}. */ + public static void clear(String collectorId) { + COLLECTORS.remove(collectorId); + } + + /** Convenience for a collector id that will not collide with other tests or other test runs. */ + public static String newCollectorId(String prefix) { + return prefix + "-" + UUID.randomUUID(); + } + + // --------------------------------------------------------------------------------------------- + // Pipeline options / TestPipeline factories. + // --------------------------------------------------------------------------------------------- + + /** + * Returns {@link SparkStructuredStreamingPipelineOptions} configured for a streaming test: the + * {@link SparkStructuredStreamingRunner}, test mode, streaming mode, a 3-idle-batch stop, a 200ms + * micro-batch trigger, and a checkpoint directory carved out of {@code checkpointDir}. + * + *

Callers that need a specific {@code SparkSession} (for example the relaxed-Kryo {@code + * SparkSessionRule} pattern required by any query that hosts a {@code transformWithState} + * operator, see {@code BeamStatefulProcessorTest}) should additionally call {@code + * SparkSessionRule#configure} on the returned options, which sets {@code useActiveSparkSession}. + */ + public static SparkStructuredStreamingPipelineOptions streamingOptions( + TemporaryFolder checkpointDir) throws IOException { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.as(SparkStructuredStreamingPipelineOptions.class); + options.setRunner(SparkStructuredStreamingRunner.class); + options.setTestMode(true); + options.setStreaming(true); + options.setStreamingStopAfterIdleBatches(3); + options.setMaxBatchDurationMillis(200); + options.setCheckpointDir(checkpointDir.newFolder("checkpoint").getAbsolutePath()); + return options; + } + + /** Builds a {@link TestPipeline} from {@link #streamingOptions}. */ + public static TestPipeline streamingPipeline(TemporaryFolder checkpointDir) throws IOException { + return TestPipeline.fromOptions(streamingOptions(checkpointDir)); + } + + // --------------------------------------------------------------------------------------------- + // ListBackedUnboundedSource. + // --------------------------------------------------------------------------------------------- + + /** + * An {@link UnboundedSource} over a fixed, finite {@link List} of {@link TimestampedValue}s. + * + *

Typed unbounded (so the pipeline it feeds is genuinely a streaming pipeline) but backed by a + * finite list (so it naturally goes idle once exhausted, which is how tests in this package + * terminate, see the {@link StreamingTestUtils} class javadoc). Supports explicit, possibly + * out-of-order event timestamps so tests can inject late data. + * + *

Uses {@link org.apache.beam.sdk.io.UnboundedSource.CheckpointMark#NOOP_CHECKPOINT_MARK}: + * this source never needs to resume a partially read split from a durable checkpoint, so there is + * no state worth persisting between micro-batches beyond the in-memory read position. + * + *

Elements are stored pre-encoded (via {@code coder}) as {@code byte[]} plus a {@code long} + * timestamp rather than kept as {@link TimestampedValue} objects, because this source (like any + * {@link UnboundedSource}) is shipped to executors with plain Java serialization, and {@link + * TimestampedValue} does not implement {@link java.io.Serializable}. + * + *

Splitting is round robin: split {@code i} of {@code n} gets every {@code n}-th element + * starting at offset {@code i}. A single split is returned unchanged if there are not enough + * elements to make splitting worthwhile; nothing about this source requires more than one split + * for correctness, round robin merely spreads elements across splits close to evenly. + */ + public static final class ListBackedUnboundedSource + extends UnboundedSource { + + private final List encodedElements; + private final List timestampsMillis; + private final Coder coder; + + public ListBackedUnboundedSource(List> elements, Coder coder) { + this.coder = Preconditions.checkNotNull(coder); + List encoded = new ArrayList<>(elements.size()); + List timestamps = new ArrayList<>(elements.size()); + for (TimestampedValue element : elements) { + encoded.add(encode(coder, element.getValue())); + timestamps.add(element.getTimestamp().getMillis()); + } + this.encodedElements = encoded; + this.timestampsMillis = timestamps; + } + + private ListBackedUnboundedSource( + List encodedElements, List timestampsMillis, Coder coder) { + this.encodedElements = encodedElements; + this.timestampsMillis = timestampsMillis; + this.coder = coder; + } + + @Override + public List> + split(int desiredNumSplits, PipelineOptions options) { + int numElements = encodedElements.size(); + if (numElements == 0 || desiredNumSplits <= 1) { + return Collections.singletonList(this); + } + int numSplits = Math.min(desiredNumSplits, numElements); + List> bucketedElements = new ArrayList<>(numSplits); + List> bucketedTimestamps = new ArrayList<>(numSplits); + for (int i = 0; i < numSplits; i++) { + bucketedElements.add(new ArrayList<>()); + bucketedTimestamps.add(new ArrayList<>()); + } + for (int i = 0; i < numElements; i++) { + int bucket = i % numSplits; + bucketedElements.get(bucket).add(encodedElements.get(i)); + bucketedTimestamps.get(bucket).add(timestampsMillis.get(i)); + } + List> splits = new ArrayList<>(numSplits); + for (int i = 0; i < numSplits; i++) { + splits.add( + new ListBackedUnboundedSource<>( + bucketedElements.get(i), bucketedTimestamps.get(i), coder)); + } + return splits; + } + + @Override + public UnboundedReader createReader( + PipelineOptions options, + UnboundedSource.CheckpointMark.@Nullable NoopCheckpointMark checkpointMark) { + return new ListBackedUnboundedReader<>(this); + } + + @Override + public Coder getCheckpointMarkCoder() { + return new NoopCheckpointMarkCoder(); + } + + @Override + public Coder getOutputCoder() { + return coder; + } + + private static byte[] encode(Coder coder, T value) { + try { + return CoderUtils.encodeToByteArray(coder, value); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode a ListBackedUnboundedSource element", e); + } + } + + private static T decode(Coder coder, byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray(coder, bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode a ListBackedUnboundedSource element", e); + } + } + + private static final class ListBackedUnboundedReader extends UnboundedReader { + private final ListBackedUnboundedSource source; + private int index = -1; + private Instant maxTimestampSeen = BoundedWindow.TIMESTAMP_MIN_VALUE; + + ListBackedUnboundedReader(ListBackedUnboundedSource source) { + this.source = source; + } + + @Override + public boolean start() throws IOException { + return advance(); + } + + @Override + public boolean advance() throws IOException { + int next = index + 1; + if (next >= source.encodedElements.size()) { + // Exhausted: report no more data, permanently. The source never produces more once past + // the end of the backing list. + return false; + } + index = next; + Instant timestamp = currentTimestamp(); + if (timestamp.isAfter(maxTimestampSeen)) { + maxTimestampSeen = timestamp; + } + return true; + } + + private Instant currentTimestamp() { + return new Instant(source.timestampsMillis.get(index)); + } + + @Override + public T getCurrent() throws NoSuchElementException { + if (index < 0) { + throw new NoSuchElementException(); + } + return decode(source.coder, source.encodedElements.get(index)); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (index < 0) { + throw new NoSuchElementException(); + } + return currentTimestamp(); + } + + @Override + public Instant getWatermark() { + // No idle advance, deliberately: once the list is exhausted this simply stops moving, + // rather than jumping to +infinity. See the StreamingTestUtils class javadoc. + return maxTimestampSeen; + } + + @Override + public UnboundedSource.CheckpointMark.NoopCheckpointMark getCheckpointMark() { + return UnboundedSource.CheckpointMark.NOOP_CHECKPOINT_MARK; + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() throws IOException {} + } + } + + /** + * A trivial coder for the stateless {@link UnboundedSource.CheckpointMark.NoopCheckpointMark}. + */ + private static final class NoopCheckpointMarkCoder + extends org.apache.beam.sdk.coders.AtomicCoder< + UnboundedSource.CheckpointMark.NoopCheckpointMark> { + @Override + public void encode( + UnboundedSource.CheckpointMark.NoopCheckpointMark value, java.io.OutputStream outStream) { + // Nothing to persist. + } + + @Override + public UnboundedSource.CheckpointMark.NoopCheckpointMark decode(java.io.InputStream inStream) { + return UnboundedSource.CheckpointMark.NOOP_CHECKPOINT_MARK; + } + } +} 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 new file mode 100644 index 000000000000..c2072e92dc4c --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java @@ -0,0 +1,205 @@ +/* + * 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 java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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.io.Read; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +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.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Windowed {@code GroupByKey} (via {@link Count#perKey()}, which auto-expands to {@code GroupByKey} + * + {@code Combine} since {@code Combine.PerKey} is deliberately unregistered for streaming, see + * {@code PipelineTranslatorStreaming}). This is hosted by the generic {@code transformWithState} + * super-operator in {@code BeamStatefulProcessorConfig.Mode #GROUP_ALSO_BY_WINDOW}, so every test + * here needs the same Kryo relaxation as {@code BeamStatefulProcessorTest}. + * + *

Every window this suite asserts on is followed, in the input list, by an element timestamped + * well past that window's end, per the watermark rule documented on {@link StreamingTestUtils}: the + * watermark only advances on new data and only fires a window once it has passed the window's end. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class WindowedGroupByKeyStreamingTest implements Serializable { + + /** + * See {@code BeamStatefulProcessorTest}: {@code transformWithState} broadcasts {@code + * StateSchemaMetadata} through Kryo, which is not registered anywhere, so the test JVM default of + * {@code spark.kryo.registrationRequired=true} (runners/spark/spark_runner.gradle) must be + * relaxed for any query that ends up hosting a stateful operator, windowed GroupByKey included. + */ + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final Duration WINDOW_SIZE = Duration.standardSeconds(10); + + private SparkStructuredStreamingPipelineOptions options() throws Exception { + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + return options; + } + + @Test(timeout = 300_000) + @Ignore( + "Needs GroupByKeyStreamingTranslator (WS-D2), which hosts the expanded GroupByKey via " + + "BeamStatefulProcessorConfig.Mode.GROUP_ALSO_BY_WINDOW; GroupByKey currently has no " + + "streaming translation and pipeline.run() throws before any window can fire.") + public void fixedWindowsCountPerKey() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("fixed-windows"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // All three fall in the first ten second window [0s, 10s). + elements.add(TimestampedValue.of(KV.of("a", "x"), BASE)); + elements.add(TimestampedValue.of(KV.of("a", "y"), BASE.plus(Duration.standardSeconds(1)))); + elements.add(TimestampedValue.of(KV.of("b", "z"), BASE.plus(Duration.standardSeconds(2)))); + // Watermark rule: a much later element so the watermark passes the first window's end. + elements.add( + TimestampedValue.of(KV.of("sentinel", "s"), BASE.plus(Duration.standardSeconds(60)))); + + TestPipeline pipeline = TestPipeline.fromOptions(options()); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("CountPerKey", Count.perKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains + // exactly KV.of("a", 2L) and KV.of("b", 1L) for the [0s, 10s) window. Remember the one + // micro-batch timer latency floor documented on StreamingTestUtils: the end-of-window timer + // fires one micro-batch after the sentinel's batch, not within it. + } + + @Test(timeout = 300_000) + @Ignore("Needs GroupByKeyStreamingTranslator (WS-D2), see fixedWindowsCountPerKey for why.") + public void slidingWindowsCountPerKey() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("sliding-windows"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // A five second sliding window every five seconds overlapping the ten second fixed window + // above: this element falls in two sliding windows, [-5s, 5s) and [0s, 10s). + elements.add(TimestampedValue.of(KV.of("a", "x"), BASE.plus(Duration.standardSeconds(2)))); + elements.add(TimestampedValue.of(KV.of("a", "y"), BASE.plus(Duration.standardSeconds(3)))); + // Watermark rule: push well past every window under test. + elements.add( + TimestampedValue.of(KV.of("sentinel", "s"), BASE.plus(Duration.standardSeconds(60)))); + + TestPipeline pipeline = TestPipeline.fromOptions(options()); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "SlidingWindows", + Window.into( + SlidingWindows.of(Duration.standardSeconds(10)).every(Duration.standardSeconds(5)))) + .apply("CountPerKey", Count.perKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains + // one KV.of("a", 2L) pane per overlapping sliding window the two "a" elements both fall into + // (out of scope note: this suite only ever asserts on non-merging windows; session windows are + // out of POC scope per the roadmap). + } + + @Test(timeout = 300_000) + @Ignore("Needs GroupByKeyStreamingTranslator (WS-D2), see fixedWindowsCountPerKey for why.") + public void lateDataIsDropped() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("late-data-dropped"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // On-time element in the first window [0s, 10s). + elements.add( + TimestampedValue.of(KV.of("a", "on-time"), BASE.plus(Duration.standardSeconds(1)))); + // Jump the watermark far past the first window's end (and its zero allowed lateness) before + // the late element arrives: this is the whole point of the test, the watermark is monotonic + // in the *order elements are read*, not in event time order, so a small timestamp read after a + // much larger one is unambiguously late. + elements.add( + TimestampedValue.of(KV.of("sentinel", "s"), BASE.plus(Duration.standardSeconds(60)))); + // Late: arrives after the watermark has already passed the end of the first window, and the + // default windowing strategy has zero allowed lateness, so this must be dropped, not emitted + // as a second, late pane. + elements.add(TimestampedValue.of(KV.of("a", "late"), BASE.plus(Duration.standardSeconds(2)))); + // One more push so there is a micro-batch that can observe the watermark has not moved + // backwards and the drop truly happened rather than merely not having fired yet. + elements.add( + TimestampedValue.of(KV.of("sentinel", "t"), BASE.plus(Duration.standardSeconds(90)))); + + TestPipeline pipeline = TestPipeline.fromOptions(options()); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("CountPerKey", Count.perKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains + // KV.of("a", 1L) (the on-time element only) for the first window, and never a count of 2: the + // late "a" element must be dropped, not merged into a late pane. + } +} From 3ec0d65c9c66853cee4d6026cfb8af5ca1c0dbbf Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 19:22:37 +0000 Subject: [PATCH 06/21] [Spark 4] Add the three real streaming translators Replaces the WS-D2 placeholders in PipelineTranslatorStreaming with real translators for the unbounded read, GroupByKey and stateful ParDo. ReadUnboundedTranslator wraps the Beam UnboundedSource with UnboundedSourceDataset, which already declares the single withWatermark for the whole query, and decodes the binary payload column into the Dataset> shape every other translator consumes. The event timestamp column is projected away, the EventTimeWatermark plan node survives below the projection and transformWithState reads the query wide watermark rather than a column. GroupByKeyStreamingTranslator and StatefulParDoStreamingTranslator both encode their input into the byte[] row layout of TwsTransformFactory, hand it to the generic transformWithState operator in GROUP_ALSO_BY_WINDOW and STATEFUL_PARDO mode respectively, and decode the tagged output rows back out. The stateful ParDo splits the output by tag index and skips additional outputs that nothing consumes, so an unused tag does not start a second streaming query. StreamingTranslationHelpers holds the shared guards and row conversions. Pipelines outside the POC scope are rejected at translation time with a message naming the feature, specifically merging windows, custom triggers, accumulating panes, processing time timers, non deterministic key coders, non KV input to a stateful ParDo, side inputs on a stateful ParDo, @OnWindowExpiration and @RequiresTimeSortedInput. --- .../PipelineTranslatorStreaming.java | 64 +---- .../GroupByKeyStreamingTranslator.java | 120 +++++++++ .../streaming/ReadUnboundedTranslator.java | 111 ++++++++ .../StatefulParDoStreamingTranslator.java | 186 ++++++++++++++ .../StreamingTranslationHelpers.java | 237 ++++++++++++++++++ 5 files changed, 663 insertions(+), 55 deletions(-) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java index be5107234ee3..7901fb9f4a31 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -20,6 +20,9 @@ import java.util.Collection; import org.apache.beam.runners.spark.SparkCommonPipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; +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; import org.apache.beam.sdk.transforms.Combine; import org.apache.beam.sdk.transforms.GroupByKey; @@ -47,30 +50,6 @@ @Internal public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { - // -------------------------------------------------------------------------------------------- - // Placeholders for WS-D2 - // -------------------------------------------------------------------------------------------- - // - // WS-D2 replaces each of the three constants below with a real translator instance (constructed - // in place, same as the batch registry does) and removes the corresponding TODO. No other change - // to this class should be necessary to plug those in: the interception points that select them in - // getTransformTranslator already exist and are marked below. - - // TODO(WS-D2): replace with `new ReadUnboundedTranslator<>()`. - private static final TransformTranslator READ_UNBOUNDED_PLACEHOLDER = - new UnsupportedStreamingTranslator( - "Unbounded read streaming translation not implemented yet (WS-D2)"); - - // TODO(WS-D2): replace with `new GroupByKeyStreamingTranslator<>()`. - private static final TransformTranslator GROUP_BY_KEY_PLACEHOLDER = - new UnsupportedStreamingTranslator( - "GroupByKey streaming translation not implemented yet (WS-D2)"); - - // TODO(WS-D2): replace with `new StatefulParDoStreamingTranslator<>()`. - private static final TransformTranslator STATEFUL_PAR_DO_PLACEHOLDER = - new UnsupportedStreamingTranslator( - "Stateful ParDo streaming translation not implemented yet (WS-D2)"); - /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ @Override @SuppressWarnings({"rawtypes", "unchecked"}) @@ -78,22 +57,18 @@ public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { protected > TransformTranslator getTransformTranslator(TransformT transform) { - // TODO(WS-D2): swap for `if (transform instanceof SplittableParDo.PrimitiveUnboundedRead)` - // returning `new ReadUnboundedTranslator<>()` (or a cached instance thereof). if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { - return (TransformTranslator) READ_UNBOUNDED_PLACEHOLDER; + return (TransformTranslator) new ReadUnboundedTranslator<>(); } - // TODO(WS-D2): swap for `if (transform instanceof GroupByKey)` returning - // `new GroupByKeyStreamingTranslator<>()`. if (transform instanceof GroupByKey) { - return (TransformTranslator) GROUP_BY_KEY_PLACEHOLDER; + return (TransformTranslator) new GroupByKeyStreamingTranslator<>(); } // Deliberately never registered: leaving Combine.PerKey unhandled here makes Beam auto-expand - // it into GroupByKey + ParDo, so the streaming translations above (and WS-D2's stateful ParDo - // translator) take over the expanded primitives instead of the batch - // CombinePerKeyTranslatorBatch, which has no streaming support. + // it into GroupByKey + ParDo, so the streaming translations above take over the expanded + // primitives instead of the batch CombinePerKeyTranslatorBatch, which has no streaming + // support. if (transform instanceof Combine.PerKey) { return null; } @@ -101,9 +76,8 @@ TransformTranslator getTransformTranslator(TransformT tra if (transform instanceof ParDo.MultiOutput) { DoFnSignature signature = DoFnSignatures.signatureForDoFn(((ParDo.MultiOutput) transform).getFn()); - // TODO(WS-D2): swap this branch's return for `new StatefulParDoStreamingTranslator<>()`. if (signature.usesState() || signature.usesTimers()) { - return (TransformTranslator) STATEFUL_PAR_DO_PLACEHOLDER; + return (TransformTranslator) new StatefulParDoStreamingTranslator<>(); } // Stateless ParDo falls through to super, reusing ParDoTranslatorBatch unchanged. } @@ -120,24 +94,4 @@ protected EvaluationContext createEvaluationContext( SparkCommonPipelineOptions options) { return new StreamingEvaluationContext(leaves, session, options); } - - /** - * A {@link TransformTranslator} standing in for a streaming translation that WS-D2 has not - * implemented yet. Always throws {@link UnsupportedOperationException} as soon as the pipeline - * traversal tries to translate the transform it is registered for. - */ - private static final class UnsupportedStreamingTranslator - extends TransformTranslator> { - private final String message; - - UnsupportedStreamingTranslator(String message) { - super(0); - this.message = message; - } - - @Override - protected void translate(PTransform transform, Context cxt) { - throw new UnsupportedOperationException(message); - } - } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java new file mode 100644 index 000000000000..d166c663321d --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java @@ -0,0 +1,120 @@ +/* + * 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 java.util.Collections; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessorConfig; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; + +/** + * Streaming translator for {@link GroupByKey}, implemented as a group-also-by-window hosted by the + * generic {@code transformWithState} operator. + * + *

Unlike the batch translator, which may ignore triggering and simply collect every value of a + * key, a streaming {@code GroupByKey} has to respect the watermark: a window's single on-time pane + * is emitted only once the watermark has passed the end of that window, and values arriving after + * that are dropped as late. Both are the responsibility of the Beam {@code ReduceFnRunner} that + * {@code BeamStatefulProcessorConfig.Mode#GROUP_ALSO_BY_WINDOW} sets up inside the operator, so all + * this translator does is put the data into and take it back out of the operator's byte[] row + * layout, see {@link TwsTransformFactory}. + * + *

{@code Combine.PerKey} is deliberately not registered for streaming, so combines reach this + * translator already expanded into {@code GroupByKey} plus a plain {@code ParDo}. + */ +public class GroupByKeyStreamingTranslator + extends TransformTranslator< + PCollection>, PCollection>>, GroupByKey> { + + /** Output tag of the group-also-by-window {@code DoFn}, index 0, its only output. */ + private static final String MAIN_OUTPUT_TAG_ID = "gbk-main-output"; + + public GroupByKeyStreamingTranslator() { + super(0.2f); + } + + @Override + @SuppressWarnings("unchecked") + protected void translate(GroupByKey transform, Context cxt) { + PCollection> input = cxt.getInput(); + String stepName = cxt.getCurrentTransform().getFullName(); + + WindowingStrategy windowing = input.getWindowingStrategy(); + StreamingTranslationHelpers.checkSupportedWindowing(windowing, stepName); + + if (!(input.getCoder() instanceof KvCoder)) { + throw StreamingTranslationHelpers.unsupported( + stepName, "the non KV input coder " + input.getCoder()); + } + KvCoder inputCoder = (KvCoder) input.getCoder(); + Coder keyCoder = inputCoder.getKeyCoder(); + Coder valueCoder = inputCoder.getValueCoder(); + StreamingTranslationHelpers.checkDeterministicKeyCoder(keyCoder, stepName); + + Coder windowCoder = windowing.getWindowFn().windowCoder(); + KvCoder> outputCoder = KvCoder.of(keyCoder, IterableCoder.of(valueCoder)); + TupleTag>> mainOutputTag = new TupleTag<>(MAIN_OUTPUT_TAG_ID); + + Dataset keyedRows = + cxt.getDataset(input) + .map( + new StreamingTranslationHelpers.EncodeKeyedRow<>( + keyCoder, WindowedValues.getFullCoder(valueCoder, windowCoder)), + Encoders.BINARY()); + + BeamStatefulProcessorConfig config = + BeamStatefulProcessorConfig.builder() + .setMode(BeamStatefulProcessorConfig.Mode.GROUP_ALSO_BY_WINDOW) + .setKeyCoder(keyCoder) + .setValueCoder(valueCoder) + .setWindowingStrategy(windowing) + .setMainOutputTag(mainOutputTag) + .setOutputCoders( + Collections., Coder>singletonMap(mainOutputTag, outputCoder)) + .setOptionsSupplier( + StreamingTranslationHelpers.optionsSupplier(cxt.getOptionsSupplier())) + .setStepName(stepName) + .build(); + + // GROUP_ALSO_BY_WINDOW has a single output tag, so every row carries index 0 and no filtering + // by tag is needed on the way out. + Dataset outputRows = TwsTransformFactory.transform(keyedRows, config); + + Encoder>>> encoder = cxt.windowedEncoder(outputCoder); + Dataset>>> result = + outputRows.map( + new StreamingTranslationHelpers.DecodeTaggedOutput<>( + WindowedValues.getFullCoder(outputCoder, windowCoder)), + encoder); + + cxt.putDataset(cxt.getOutput(), result); + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java new file mode 100644 index 000000000000..36cf305e3376 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java @@ -0,0 +1,111 @@ +/* + * 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 static org.apache.spark.sql.functions.col; + +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; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +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.Row; + +/** + * Translator for {@link SplittableParDo.PrimitiveUnboundedRead}, the streaming counterpart of the + * batch {@code ReadSourceTranslatorBatch}. + * + *

The heavy lifting is done by {@link UnboundedSourceDataset}, which wraps the Beam {@link + * UnboundedSource} in a DataSourceV2 micro-batch stream and returns a two column {@code + * Dataset}: the element encoded as a {@code WindowedValue}, plus its event timestamp. All this + * translator adds is the typed decode back into the {@code Dataset>} shape every + * other translator consumes. + * + *

Two things about the result are worth spelling out. + * + *

    + *
  • The watermark is already declared by {@link UnboundedSourceDataset} and must never + * be re-declared, Spark rejects a second {@code withWatermark} in the same plan. The {@code + * EventTimeWatermark} plan node sits below the projection and the typed map applied here and + * survives both, and a {@code transformWithState} operator downstream reads the query wide + * watermark rather than a column, so dropping the timestamp column here is safe. + *
  • Elements arrive in the global window, timestamped with the reader's record + * timestamp. A windowed pipeline therefore still needs its {@code Window.Assign}, which is + * translated by the reused batch translator. + *
+ */ +public class ReadUnboundedTranslator + extends TransformTranslator, SplittableParDo.PrimitiveUnboundedRead> { + + public ReadUnboundedTranslator() { + super(0.05f); + } + + @Override + protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Context cxt) { + PCollection output = cxt.getOutput(); + UnboundedSource source = transform.getSource(); + Coder elementCoder = output.getCoder(); + + // Matches what the partition readers emit: a value in the global window, timestamped with the + // record's own event timestamp. + WindowedValues.FullWindowedValueCoder payloadCoder = + WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + SparkStructuredStreamingPipelineOptions options = + cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); + + Dataset rows = + UnboundedSourceDataset.of(cxt.getSparkSession(), source, payloadCoder, options); + + Encoder> encoder = + cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + Dataset> dataset = + rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) + .as(Encoders.BINARY()) + .map(new DecodePayload<>(payloadCoder), encoder); + + cxt.putDataset(output, dataset); + } + + /** Decodes the binary payload column back into a Beam {@code WindowedValue}. */ + private static final class DecodePayload implements MapFunction> { + private final Coder> coder; + + DecodePayload(Coder> coder) { + this.coder = coder; + } + + @Override + public WindowedValue call(byte[] payload) { + return CoderHelpers.fromByteArray(payload, coder); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java new file mode 100644 index 000000000000..01f2dd63c557 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java @@ -0,0 +1,186 @@ +/* + * 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 static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessorConfig; +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.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.ParDoTranslation; +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.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; + +/** + * Streaming translator for a stateful {@link ParDo.MultiOutput}, that is a {@code ParDo} whose + * {@code DoFn} declares {@code @StateId} state or {@code @TimerId} timers. Stateless {@code ParDo}s + * keep using the batch translator unchanged, the streaming registry only routes stateful ones here. + * + *

The user's {@code DoFn} is hosted verbatim by the generic {@code transformWithState} operator + * in {@code BeamStatefulProcessorConfig.Mode#STATEFUL_PARDO}, which runs it through {@code + * DoFnRunners.simpleRunner} wrapped in {@code DoFnRunners.defaultStatefulDoFnRunner}, so Beam's own + * state, timer and window garbage collection semantics apply. This translator's whole job is the + * byte[] row plumbing documented on {@link TwsTransformFactory}: encode {@code WindowedValue>} into keyed input rows, and split the tagged output rows back into one {@code + * Dataset>} per {@link TupleTag}. + * + *

Note on multiple outputs: each additional tag adds one {@code filter} plus one {@code map} on + * top of the same operator, which Spark plans as a separate branch. Only outputs that are actually + * consumed downstream, plus the main output, get a dataset at all. + */ +public class StatefulParDoStreamingTranslator + extends TransformTranslator< + PCollection>, PCollectionTuple, ParDo.MultiOutput, OutputT>> { + + public StatefulParDoStreamingTranslator() { + super(0.2f); + } + + @Override + protected boolean canTranslate(ParDo.MultiOutput, OutputT> transform) { + DoFn, OutputT> doFn = transform.getFn(); + String stepName = doFn.getClass().getName(); + DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn); + + checkState( + !signature.processElement().isSplittable(), + "Not expected to directly translate splittable DoFn, should have been overridden: %s", + doFn); + + StreamingTranslationHelpers.checkNoProcessingTimeTimers(doFn, signature, stepName); + + if (signature.onWindowExpiration() != null) { + throw StreamingTranslationHelpers.unsupported(stepName, "@OnWindowExpiration"); + } + if (signature.processElement().requiresTimeSortedInput()) { + throw StreamingTranslationHelpers.unsupported(stepName, "@RequiresTimeSortedInput"); + } + if (!transform.getSideInputs().isEmpty()) { + throw StreamingTranslationHelpers.unsupported( + stepName, + "side inputs on a stateful ParDo. Broadcasting a side input requires collecting its " + + "PCollection, which is not possible while the pipeline is streaming"); + } + return true; + } + + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + protected void translate(ParDo.MultiOutput, OutputT> transform, Context cxt) { + PCollection> input = (PCollection>) cxt.getInput(); + String stepName = cxt.getCurrentTransform().getFullName(); + + WindowingStrategy windowing = input.getWindowingStrategy(); + StreamingTranslationHelpers.checkSupportedWindowing(windowing, stepName); + + // Beam guarantees a stateful DoFn is applied to a keyed PCollection, but say so clearly rather + // than failing with a ClassCastException deep in the operator. + if (!(input.getCoder() instanceof KvCoder)) { + throw StreamingTranslationHelpers.unsupported( + stepName, + "state or timers on the non KV input coder " + + input.getCoder() + + ". A stateful ParDo must be applied to a PCollection of KVs"); + } + KvCoder inputCoder = (KvCoder) input.getCoder(); + Coder keyCoder = inputCoder.getKeyCoder(); + Coder valueCoder = inputCoder.getValueCoder(); + StreamingTranslationHelpers.checkDeterministicKeyCoder(keyCoder, stepName); + + Coder windowCoder = windowing.getWindowFn().windowCoder(); + + TupleTag mainOutputTag = transform.getMainOutputTag(); + List> additionalOutputTags = + new ArrayList<>(transform.getAdditionalOutputTags().getAll()); + + // One coder per tag the DoFn may emit to, taken from the PCollection behind that tag. + Map, Coder> outputCoders = new LinkedHashMap<>(); + outputCoders.put(mainOutputTag, cxt.getOutput(mainOutputTag).getCoder()); + for (TupleTag tag : additionalOutputTags) { + outputCoders.put(tag, cxt.getOutput((TupleTag) tag).getCoder()); + } + + Dataset keyedRows = + cxt.getDataset(input) + .map( + new StreamingTranslationHelpers.EncodeKeyedRow<>( + keyCoder, WindowedValues.getFullCoder(valueCoder, windowCoder)), + Encoders.BINARY()); + + BeamStatefulProcessorConfig config = + BeamStatefulProcessorConfig.builder() + .setMode(BeamStatefulProcessorConfig.Mode.STATEFUL_PARDO) + .setDoFn(transform.getFn()) + .setKeyCoder(keyCoder) + .setValueCoder(valueCoder) + .setWindowingStrategy(windowing) + .setMainOutputTag(mainOutputTag) + .setAdditionalOutputTags(additionalOutputTags) + .setOutputCoders(outputCoders) + .setDoFnSchemaInformation( + ParDoTranslation.getSchemaInformation(cxt.getCurrentTransform())) + .setOptionsSupplier( + StreamingTranslationHelpers.optionsSupplier(cxt.getOptionsSupplier())) + .setStepName(stepName) + .build(); + + Dataset outputRows = TwsTransformFactory.transform(keyedRows, config); + + List> allTags = config.outputTags(); + boolean singleTag = allTags.size() == 1; + for (int tagIndex = 0; tagIndex < allTags.size(); tagIndex++) { + TupleTag tag = (TupleTag) allTags.get(tagIndex); + PCollection outputPCollection = cxt.getOutput(tag); + if (tagIndex > 0 && cxt.isLeaf(outputPCollection)) { + // An additional output nobody consumes: emitting it would start a whole extra streaming + // query re-running this operator for rows that are then thrown away. + continue; + } + Coder outputCoder = outputPCollection.getCoder(); + Encoder> encoder = cxt.windowedEncoder(outputCoder); + Dataset taggedRows = + singleTag + ? outputRows + : outputRows.filter(new StreamingTranslationHelpers.TagIndexFilter(tagIndex)); + cxt.putDataset( + outputPCollection, + taggedRows.map( + new StreamingTranslationHelpers.DecodeTaggedOutput<>( + WindowedValues.getFullCoder(outputCoder, windowCoder)), + encoder)); + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java new file mode 100644 index 000000000000..6b5e4878613f --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java @@ -0,0 +1,237 @@ +/* + * 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 java.util.function.Supplier; +import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state.BeamStatefulProcessorConfig; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.AfterWatermark; +import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.transforms.windowing.Trigger; +import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; +import org.apache.spark.api.java.function.FilterFunction; +import org.apache.spark.api.java.function.MapFunction; + +/** + * Shared plumbing of the three streaming translators: the guards that reject Beam features this POC + * deliberately does not implement, and the small serializable functions that convert between the + * runner's {@code Dataset>} representation and the raw {@code byte[]} row layouts + * of {@link TwsTransformFactory}. + * + *

Guards throw {@link UnsupportedOperationException} naming the offending feature. They run at + * translation time, before any Spark query is started, so an unsupported pipeline fails immediately + * and loudly rather than producing quietly wrong results. + */ +final class StreamingTranslationHelpers { + + private StreamingTranslationHelpers() {} + + // ----------------------------------------------------------------------------------------- + // Guards + // ----------------------------------------------------------------------------------------- + + /** + * Rejects windowing strategies whose semantics the {@code transformWithState} bridge cannot + * reproduce: merging (session) windows, custom triggers and accumulating panes. + */ + static void checkSupportedWindowing(WindowingStrategy strategy, String stepName) { + WindowFn windowFn = strategy.getWindowFn(); + if (!windowFn.isNonMerging()) { + throw unsupported( + stepName, + "merging windows (" + + windowFn.getClass().getSimpleName() + + "). Session windows and any other merging WindowFn are out of scope for the " + + "Spark 4 streaming runner"); + } + Trigger trigger = strategy.getTrigger(); + if (!isDefaultTrigger(trigger)) { + throw unsupported( + stepName, + "the custom trigger " + + trigger + + ". Only the default trigger (one on-time pane per window when the watermark passes " + + "its end) is implemented"); + } + if (strategy.getMode() != AccumulationMode.DISCARDING_FIRED_PANES) { + throw unsupported( + stepName, + "accumulating panes (" + strategy.getMode() + "). Only discarding panes are implemented"); + } + } + + /** + * The default trigger, as Beam models it, is either {@link DefaultTrigger} itself or the + * equivalent {@code AfterWatermark.pastEndOfWindow()} without early or late firings, which is + * what {@code WindowingStrategy} may normalise it to. + */ + private static boolean isDefaultTrigger(Trigger trigger) { + if (trigger instanceof DefaultTrigger) { + return true; + } + if (trigger instanceof AfterWatermark.FromEndOfWindow) { + // FromEndOfWindow without early/late firings is exactly the default trigger; with firings + // configured Beam represents it as one of the AfterWatermarkEarlyAndLate subclasses instead. + return true; + } + return false; + } + + /** + * The Beam key is the Spark grouping key and is compared as raw bytes, so two equal keys must + * always encode identically. + */ + static void checkDeterministicKeyCoder(Coder keyCoder, String stepName) { + try { + keyCoder.verifyDeterministic(); + } catch (Coder.NonDeterministicException e) { + throw new UnsupportedOperationException( + "Cannot translate " + + stepName + + " for streaming: the key coder " + + keyCoder + + " is not deterministic. Keys are grouped by their encoded bytes, so a " + + "non-deterministic key coder would silently split a single Beam key across " + + "several Spark state entries.", + e); + } + } + + /** + * Only event time timers reach the {@code transformWithState} operator, which runs in {@code + * TimeMode.EventTime()}; a processing time timer would never fire. + */ + static void checkNoProcessingTimeTimers( + DoFn doFn, DoFnSignature signature, String stepName) { + for (DoFnSignature.TimerDeclaration timer : signature.timerDeclarations().values()) { + TimerSpec spec = DoFnSignatures.getTimerSpecOrThrow(timer, doFn); + if (spec.getTimeDomain() != TimeDomain.EVENT_TIME) { + throw unsupported( + stepName, + "the " + + spec.getTimeDomain() + + " timer @TimerId(\"" + + timer.id() + + "\"). Only event time timers are implemented"); + } + } + for (DoFnSignature.TimerFamilyDeclaration family : + signature.timerFamilyDeclarations().values()) { + TimerSpec spec = DoFnSignatures.getTimerFamilySpecOrThrow(family, doFn); + if (spec.getTimeDomain() != TimeDomain.EVENT_TIME) { + throw unsupported( + stepName, + "the " + + spec.getTimeDomain() + + " timer family @TimerFamily(\"" + + family.id() + + "\"). Only event time timers are implemented"); + } + } + } + + static UnsupportedOperationException unsupported(String stepName, String feature) { + return new UnsupportedOperationException( + "Cannot translate " + stepName + " for streaming, it uses " + feature + "."); + } + + // ----------------------------------------------------------------------------------------- + // Row conversions + // ----------------------------------------------------------------------------------------- + + /** Adapts the translation context's options supplier to the shape the operator config wants. */ + static BeamStatefulProcessorConfig.OptionsSupplier optionsSupplier( + Supplier supplier) { + return new DelegatingOptionsSupplier(supplier); + } + + private static final class DelegatingOptionsSupplier + implements BeamStatefulProcessorConfig.OptionsSupplier { + private final Supplier delegate; + + DelegatingOptionsSupplier(Supplier delegate) { + this.delegate = delegate; + } + + @Override + public PipelineOptions get() { + return delegate.get(); + } + } + + /** + * Turns {@code WindowedValue>} into a {@link TwsTransformFactory} input row: the encoded + * key followed by the {@code WindowedValue} of the value side only. + */ + static final class EncodeKeyedRow implements MapFunction>, byte[]> { + private final Coder keyCoder; + private final Coder> payloadCoder; + + EncodeKeyedRow(Coder keyCoder, Coder> payloadCoder) { + this.keyCoder = keyCoder; + this.payloadCoder = payloadCoder; + } + + @Override + public byte[] call(WindowedValue> element) { + KV kv = element.getValue(); + return TwsTransformFactory.encodeInputRow( + CoderHelpers.toByteArray(kv.getKey(), keyCoder), + CoderHelpers.toByteArray(element.withValue(kv.getValue()), payloadCoder)); + } + } + + /** Decodes the {@code WindowedValue} payload of a {@link TwsTransformFactory} output row. */ + static final class DecodeTaggedOutput implements MapFunction> { + private final Coder> coder; + + DecodeTaggedOutput(Coder> coder) { + this.coder = coder; + } + + @Override + public WindowedValue call(byte[] row) { + return CoderHelpers.fromByteArray(TwsTransformFactory.outputPayload(row), coder); + } + } + + /** Keeps only the {@link TwsTransformFactory} output rows carrying one specific tag index. */ + static final class TagIndexFilter implements FilterFunction { + private final int tagIndex; + + TagIndexFilter(int tagIndex) { + this.tagIndex = tagIndex; + } + + @Override + public boolean call(byte[] row) { + return TwsTransformFactory.outputTagIndex(row) == tagIndex; + } + } +} From 1e157500c0d5e58a276214fd8ae975b7452fe9ed Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 19:43:02 +0000 Subject: [PATCH 07/21] [Spark 4] Enable the streaming end to end tests Turns the five Ignore'd streaming test skeletons into real, asserting tests against the WS-D2 translators. Every skeleton declared its pipeline as a local TestPipeline variable, which TestPipeline itself rejects at run() because it is not a Rule field. These tests cannot use PAssert anyway, a streaming pipeline here never gets a final watermark so panes never finalize, and each test needs its own options, so TestPipeline buys nothing. They now build a plain Pipeline.create(options) and assert after waitUntilFinish against the static collector in StreamingTestUtils. The unused TestPipeline factory is dropped from StreamingTestUtils and the reasoning is recorded in its javadoc. StatelessParDoStreamingTest and StreamingPipelineLifecycleTest gain the same relaxed Kryo SparkSessionRule as the rest of the suite, so that the whole package shares one session and a cancelled or idle stopped query cannot take that session down with it. Asserted results, all derived from each test's own input list. Stateless ParDo, exactly 0, 2, 4 up to 18. Fixed 10s window count per key, a=2 and b=1. Sliding 10s every 5s, two a=2 panes, one per shared window. Late data dropped, a=1 and sentinel=1, never a=2. Stateful ParDo, a, b, c once each plus two timer sentinels. Chained dedup into windowed sum, a=8 and b=10. The chained case is the one that matters, a=13 would mean the dedup state was lost and an empty result would mean the watermark got stuck between the two operators, so exactly a=8 and b=10 is the observable proof of cross operator watermark propagation. The late data test is the only one that depends on how the source round robins elements across splits, so it pins maxRecordsPerMicroBatch to 1, spells out the resulting per batch schedule in a comment, and asserts the split count assumption up front rather than leaving it silent. The lifecycle tests wait for a query to actually be active before cancelling, since translation is asynchronous and an early cancel would find a null evaluation context, and they check the query really stopped rather than only that the state was relabelled. --- .../ChainedStatefulStreamingTest.java | 31 +++++--- .../streaming/StatefulParDoStreamingTest.java | 28 ++++--- .../StatelessParDoStreamingTest.java | 42 +++++++--- .../StreamingPipelineLifecycleTest.java | 79 ++++++++++++++++--- .../streaming/StreamingTestUtils.java | 21 +++-- .../WindowedGroupByKeyStreamingTest.java | 79 +++++++++++++------ 6 files changed, 196 insertions(+), 84 deletions(-) diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java index 2bb367ac5ada..da1ea8df2982 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java @@ -17,12 +17,16 @@ */ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; +import static org.junit.Assert.assertEquals; + import java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.spark.StreamingTest; import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; @@ -31,7 +35,6 @@ import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.ValueState; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Sum; @@ -42,7 +45,6 @@ import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.ClassRule; -import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -99,11 +101,6 @@ public void process( } @Test(timeout = 300_000) - @Ignore( - "Needs both StatefulParDoStreamingTranslator and GroupByKeyStreamingTranslator (WS-D2): this " - + "test chains a stateful ParDo into a windowed GroupByKey, so it is blocked on whichever " - + "of the two lands last. Also needs the cross-operator watermark propagation itself to " - + "actually work, which is the top risk called out in the POC plan (risk #3).") public void dedupThenWindowedSumPropagatesWatermarkAcrossOperators() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("chained-stateful"); StreamingTestUtils.clear(collectorId); @@ -126,7 +123,7 @@ public void dedupThenWindowedSumPropagatesWatermarkAcrossOperators() throws Exce SparkStructuredStreamingPipelineOptions options = StreamingTestUtils.streamingOptions(checkpointDir); SESSION.configure(options); - TestPipeline pipeline = TestPipeline.fromOptions(options); + Pipeline pipeline = Pipeline.create(options); pipeline .apply( @@ -145,10 +142,18 @@ public void dedupThenWindowedSumPropagatesWatermarkAcrossOperators() throws Exce PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-E2, once WS-D2 lands): assert StreamingTestUtils.>getCollected( - // collectorId) contains exactly KV.of("a", 8L) (5 + 3, the duplicate "1" excluded) and - // KV.of("b", 10L) for the [0s, 10s) window. Getting exactly these values, rather than e.g. - // KV.of("a", 13L) from a missed dedup or nothing at all from a stuck watermark, is the - // observable proof that watermark and results both survived the operator chain intact. + // a=8 is 5 + 3, the redelivered id "1" excluded by the upstream dedup operator; b=10 is the + // single "3" element. Both are the [0s, 10s) window firing in the downstream operator, which + // only happens if the watermark computed once at the source still reaches the second stateful + // operator intact. a=13 would mean the dedup state was lost, and an empty result would mean the + // watermark got stuck between the two operators; getting exactly [a=8, b=10] is the observable + // proof that neither happened. The sentinel's own [60s, 70s) window never fires, nothing + // arrives after it to push the watermark past 70s. + List collected = new ArrayList<>(); + for (KV kv : StreamingTestUtils.>getCollected(collectorId)) { + collected.add(kv.getKey() + "=" + kv.getValue()); + } + Collections.sort(collected); + assertEquals("pipeline state=" + result.getState(), "[a=8, b=10]", collected.toString()); } } 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 2ecacb81141a..b0eff04d2979 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 @@ -17,12 +17,16 @@ */ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; +import static org.junit.Assert.assertEquals; + import java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.spark.StreamingTest; import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; @@ -34,7 +38,6 @@ import org.apache.beam.sdk.state.TimerSpec; import org.apache.beam.sdk.state.TimerSpecs; import org.apache.beam.sdk.state.ValueState; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.WithKeys; @@ -43,7 +46,6 @@ import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.ClassRule; -import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -107,11 +109,6 @@ public void onExpiry(OutputReceiver out) { } @Test(timeout = 300_000) - @Ignore( - "Needs StatefulParDoStreamingTranslator (WS-D2) registered in place of the " - + "UnsupportedOperationException placeholder in " - + "PipelineTranslatorStreaming#STATEFUL_PAR_DO_PLACEHOLDER; a DoFn signature with state or " - + "timers currently has no streaming translation.") public void dedupsRepeatedKeysAndFiresTimerSentinel() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("stateful-pardo-dedup"); StreamingTestUtils.clear(collectorId); @@ -127,7 +124,7 @@ public void dedupsRepeatedKeysAndFiresTimerSentinel() throws Exception { SparkStructuredStreamingPipelineOptions options = StreamingTestUtils.streamingOptions(checkpointDir); SESSION.configure(options); - TestPipeline pipeline = TestPipeline.fromOptions(options); + Pipeline pipeline = Pipeline.create(options); pipeline .apply( @@ -142,9 +139,16 @@ public void dedupsRepeatedKeysAndFiresTimerSentinel() throws Exception { PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-D2): assert StreamingTestUtils.getCollected(collectorId) contains exactly one - // "a", one "b", one "c" (each key's first sighting) plus two SENTINEL values (one per timer - // armed for "a" and "b"; "c" arrives too late in the run for its own timer to have expired). - // Remember the one micro-batch timer latency floor documented on StreamingTestUtils. + // One "a" (the second sighting is suppressed by the dedup state), one "b", one "c", plus two + // SENTINELs: the timers armed at 0s+30s for "a" and at 2s+30s for "b" both expire under the + // final watermark of 90s, while "c"'s own timer at 90s+30s = 120s never does. The sentinels + // arrive one micro-batch after the batch carrying "c", per the timer latency floor documented + // on StreamingTestUtils. + List collected = new ArrayList<>(StreamingTestUtils.getCollected(collectorId)); + Collections.sort(collected); + assertEquals( + "pipeline state=" + result.getState(), + "[" + SENTINEL + ", " + SENTINEL + ", a, b, c]", + collected.toString()); } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java index ceeedf70f43f..e352014bbd3f 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -17,21 +17,26 @@ */ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; +import static org.junit.Assert.assertEquals; + import java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.io.Read; -import org.apache.beam.sdk.testing.TestPipeline; 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.TimestampedValue; import org.joda.time.Duration; import org.joda.time.Instant; -import org.junit.Ignore; +import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -48,6 +53,16 @@ @Category(StreamingTest.class) public class StatelessParDoStreamingTest implements Serializable { + /** + * This pipeline hosts no {@code transformWithState} operator of its own, but it still shares the + * session with the rest of the streaming suite and is configured identically to it, so that the + * baseline test differs from the stateful ones in the pipeline under test and nothing else. See + * {@code BeamStatefulProcessorTest} for why the relaxation is needed at all. + */ + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); private static final Instant BASE = new Instant(0); @@ -61,11 +76,6 @@ public void process(@Element Integer element, OutputReceiver out) { } @Test(timeout = 300_000) - @Ignore( - "Needs ReadUnboundedTranslator (WS-D2) registered in place of the UnsupportedOperationException " - + "placeholder in PipelineTranslatorStreaming#READ_UNBOUNDED_PLACEHOLDER; until then " - + "SplittableParDo.PrimitiveUnboundedRead has no streaming translation and pipeline.run() " - + "throws before any element is read.") public void everyElementPassesThrough() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("stateless-pardo"); StreamingTestUtils.clear(collectorId); @@ -77,7 +87,8 @@ public void everyElementPassesThrough() throws Exception { SparkStructuredStreamingPipelineOptions options = StreamingTestUtils.streamingOptions(checkpointDir); - TestPipeline pipeline = TestPipeline.fromOptions(options); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); pipeline .apply( @@ -90,8 +101,17 @@ public void everyElementPassesThrough() throws Exception { PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-D2): once the unbounded read translates, assert that - // StreamingTestUtils.getCollected(collectorId) contains exactly {0, 2, 4, ..., 18}, - // in any order (Spark micro-batches make no ordering guarantee across elements). + // Nothing here is windowed or stateful, so every element is emitted as soon as its micro-batch + // is processed and no watermark has to cross anything. Micro-batch boundaries make the order + // arbitrary, hence the sort. + List collected = + new ArrayList<>(StreamingTestUtils.getCollected(collectorId)); + Collections.sort(collected); + + List expected = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + expected.add(i * 2); + } + assertEquals("pipeline state=" + result.getState(), expected, collected); } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java index 78cb0070dd3e..ca67c06e9f3b 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -18,21 +18,25 @@ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.io.Read; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TimestampedValue; import org.joda.time.Duration; import org.joda.time.Instant; -import org.junit.Ignore; +import org.junit.ClassRule; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -49,7 +53,8 @@ * StreamingTestUtils#streamingOptions}, unlike every other test in this package: {@code * SparkStructuredStreamingRunner#run()} calls {@code result.waitUntilFinish()} itself before * returning when {@code testMode} is {@code true} (see its implementation), which would make {@code - * run()} block past the point these tests want to observe {@code RUNNING}. + * run()} block past the point these tests want to observe {@code RUNNING}. Note that {@code + * SparkSessionRule#configure} sets {@code testMode(true)}, so the override has to come after it. * *

Not tested here: that a streaming pipeline is rejected when run against Spark 3. That is * {@code PipelineTranslatorFactory#create} in the shared base module @@ -64,10 +69,25 @@ @Category(StreamingTest.class) public class StreamingPipelineLifecycleTest implements Serializable { + /** + * These pipelines host no {@code transformWithState} operator, but they do need {@code + * useActiveSparkSession} so that a cancelled or idle-stopped query does not take the shared + * session down with it: {@code SparkStructuredStreamingRunner#sparkStopFn} only stops the session + * on a terminal state when the session was not provided from outside. Configuring the + * session the same relaxed way as the rest of the suite keeps a single session shared across the + * whole streaming test run. + */ + @ClassRule + public static final SparkSessionRule SESSION = + new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); private static final Instant BASE = new Instant(0); + /** How long to wait for a query to actually start before giving up on it. */ + private static final long QUERY_START_TIMEOUT_MILLIS = 60_000L; + private List> tenElements() { List> elements = new ArrayList<>(); for (int i = 0; i < 10; i++) { @@ -76,20 +96,33 @@ private List> tenElements() { return elements; } + /** + * Blocks until at least one streaming query is active on the shared session. Translation happens + * asynchronously on the runner's submission thread, so {@code run()} returns before any query + * exists, and {@code cancel()} before that point would find a {@code null} evaluation context and + * silently have nothing to stop. + */ + private static void awaitQueryStarted() throws InterruptedException { + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length == 0) { + assertTrue( + "no streaming query started within " + QUERY_START_TIMEOUT_MILLIS + "ms", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + @Test(timeout = 300_000) - @Ignore( - "Needs ReadUnboundedTranslator (WS-D2): today translation itself throws " - + "UnsupportedOperationException from PipelineTranslatorStreaming#READ_UNBOUNDED_PLACEHOLDER " - + "before a query ever starts, so waitUntilFinish() surfaces FAILED, not DONE.") public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("lifecycle-done"); StreamingTestUtils.clear(collectorId); SparkStructuredStreamingPipelineOptions options = StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); // Observe RUNNING ourselves instead of letting run() block until finished, see class javadoc. options.setTestMode(false); - TestPipeline pipeline = TestPipeline.fromOptions(options); + Pipeline pipeline = Pipeline.create(options); pipeline .apply( @@ -105,24 +138,32 @@ public void idlePipelineGoesFromRunningToDoneOnceIdle() throws Exception { PipelineResult.State finalState = result.waitUntilFinish(); assertEquals(PipelineResult.State.DONE, finalState); assertEquals(PipelineResult.State.DONE, result.getState()); + + // DONE has to mean the idle-stop listener stopped a query that had actually drained its input, + // not that the query fell over early, so check the data came through too. + List collected = + new ArrayList<>(StreamingTestUtils.getCollected(collectorId)); + Collections.sort(collected); + List expected = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + expected.add(i); + } + assertEquals(expected, collected); } @Test(timeout = 300_000) - @Ignore( - "Needs ReadUnboundedTranslator (WS-D2): cancel() only has an effect once translation has " - + "completed and StreamingEvaluationContext#stop() is reachable through the runner's " - + "ctxRef; today translation throws before that ever happens.") public void cancelStopsTheQueryAndReportsCancelled() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("lifecycle-cancel"); StreamingTestUtils.clear(collectorId); SparkStructuredStreamingPipelineOptions options = StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); options.setTestMode(false); // Disabled so the query only ever stops because of the explicit cancel() below, not because it // happened to go idle first. options.setStreamingStopAfterIdleBatches(-1); - TestPipeline pipeline = TestPipeline.fromOptions(options); + Pipeline pipeline = Pipeline.create(options); pipeline .apply( @@ -135,8 +176,20 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception { PipelineResult result = pipeline.run(); assertEquals(PipelineResult.State.RUNNING, result.getState()); + awaitQueryStarted(); + PipelineResult.State cancelledState = result.cancel(); assertEquals(PipelineResult.State.CANCELLED, cancelledState); assertEquals(PipelineResult.State.CANCELLED, result.getState()); + + // cancel() has to have actually stopped the query, not just relabelled the result: with + // idle-stop disabled this query would otherwise run until the JUnit timeout. + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "the streaming query was still active " + QUERY_START_TIMEOUT_MILLIS + "ms after cancel", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java index a1ed7f11c2ff..1bf14084f5ec 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -32,7 +32,6 @@ import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.CoderUtils; @@ -44,13 +43,12 @@ /** * Shared test scaffolding for the Spark 4 streaming translators: an {@link UnboundedSource} over a - * fixed, in-memory list of elements, a driver-side static collector {@link DoFn}, and factories for - * the {@link SparkStructuredStreamingPipelineOptions} / {@link TestPipeline} every streaming test - * in this package needs. + * fixed, in-memory list of elements, a driver-side static collector {@link DoFn}, and a factory for + * the {@link SparkStructuredStreamingPipelineOptions} every streaming test in this package needs. * *

Why streaming tests in this suite look the way they do

* - *

Two of the usual Beam testing tools do not work here, on purpose: + *

Three of the usual Beam testing tools do not work here, on purpose: * *

    *
  • {@code PAssert} on an unbounded {@code PCollection} never fires: {@code PAssert} @@ -60,6 +58,12 @@ * org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamMicroBatchStream} (WS-B) * reports progress with opaque epoch offsets, not byte/row counts, so Spark can never decide * that "all available" input has been consumed. + *
  • {@code TestPipeline} is not used at all: it insists on being declared as a + * {@code @Rule} field and otherwise fails {@code run()} with "Is your TestPipeline + * declaration missing a @Rule annotation?". Its two benefits, {@code PAssert} bookkeeping and + * enforcing that the pipeline was actually run, are worthless here because these tests cannot + * use {@code PAssert} anyway and each needs its own per-test options. Every test in this + * package therefore builds a plain {@code Pipeline.create(options)} instead. *
* *

Instead, every test in this package follows the same recipe: @@ -168,7 +172,7 @@ public static String newCollectorId(String prefix) { } // --------------------------------------------------------------------------------------------- - // Pipeline options / TestPipeline factories. + // Pipeline options factory. // --------------------------------------------------------------------------------------------- /** @@ -194,11 +198,6 @@ public static SparkStructuredStreamingPipelineOptions streamingOptions( return options; } - /** Builds a {@link TestPipeline} from {@link #streamingOptions}. */ - public static TestPipeline streamingPipeline(TemporaryFolder checkpointDir) throws IOException { - return TestPipeline.fromOptions(streamingOptions(checkpointDir)); - } - // --------------------------------------------------------------------------------------------- // ListBackedUnboundedSource. // --------------------------------------------------------------------------------------------- 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 c2072e92dc4c..c3fb37003da3 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 @@ -17,17 +17,20 @@ */ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; +import static org.junit.Assert.assertEquals; + import java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import org.apache.beam.runners.spark.StreamingTest; import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; 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.io.Read; -import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.windowing.FixedWindows; @@ -38,7 +41,6 @@ import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.ClassRule; -import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -56,6 +58,9 @@ *

Every window this suite asserts on is followed, in the input list, by an element timestamped * well past that window's end, per the watermark rule documented on {@link StreamingTestUtils}: the * watermark only advances on new data and only fires a window once it has passed the window's end. + * The mirror image of that rule is that the trailing "sentinel" elements' own windows are never + * asserted on, because nothing arrives after them to push the watermark past their ends, so they + * simply never fire. */ @RunWith(JUnit4.class) @Category(StreamingTest.class) @@ -83,11 +88,17 @@ private SparkStructuredStreamingPipelineOptions options() throws Exception { return options; } + /** Renders the collected panes as a sorted {@code key=count} list, for a readable assertion. */ + private static String collectedCounts(String collectorId) { + List rendered = new ArrayList<>(); + for (KV kv : StreamingTestUtils.>getCollected(collectorId)) { + rendered.add(kv.getKey() + "=" + kv.getValue()); + } + Collections.sort(rendered); + return rendered.toString(); + } + @Test(timeout = 300_000) - @Ignore( - "Needs GroupByKeyStreamingTranslator (WS-D2), which hosts the expanded GroupByKey via " - + "BeamStatefulProcessorConfig.Mode.GROUP_ALSO_BY_WINDOW; GroupByKey currently has no " - + "streaming translation and pipeline.run() throws before any window can fire.") public void fixedWindowsCountPerKey() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("fixed-windows"); StreamingTestUtils.clear(collectorId); @@ -101,7 +112,7 @@ public void fixedWindowsCountPerKey() throws Exception { elements.add( TimestampedValue.of(KV.of("sentinel", "s"), BASE.plus(Duration.standardSeconds(60)))); - TestPipeline pipeline = TestPipeline.fromOptions(options()); + Pipeline pipeline = Pipeline.create(options()); pipeline .apply( "ReadUnbounded", @@ -115,28 +126,26 @@ public void fixedWindowsCountPerKey() throws Exception { PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains - // exactly KV.of("a", 2L) and KV.of("b", 1L) for the [0s, 10s) window. Remember the one - // micro-batch timer latency floor documented on StreamingTestUtils: the end-of-window timer - // fires one micro-batch after the sentinel's batch, not within it. + // Only [0s, 10s) ever fires: the sentinel's own window [60s, 70s) has nothing after it to push + // the watermark past 70s. + assertEquals("pipeline state=" + result.getState(), "[a=2, b=1]", collectedCounts(collectorId)); } @Test(timeout = 300_000) - @Ignore("Needs GroupByKeyStreamingTranslator (WS-D2), see fixedWindowsCountPerKey for why.") public void slidingWindowsCountPerKey() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("sliding-windows"); StreamingTestUtils.clear(collectorId); List>> elements = new ArrayList<>(); - // A five second sliding window every five seconds overlapping the ten second fixed window - // above: this element falls in two sliding windows, [-5s, 5s) and [0s, 10s). + // A ten second sliding window every five seconds: both these elements fall in exactly two + // sliding windows, [-5s, 5s) and [0s, 10s). elements.add(TimestampedValue.of(KV.of("a", "x"), BASE.plus(Duration.standardSeconds(2)))); elements.add(TimestampedValue.of(KV.of("a", "y"), BASE.plus(Duration.standardSeconds(3)))); // Watermark rule: push well past every window under test. elements.add( TimestampedValue.of(KV.of("sentinel", "s"), BASE.plus(Duration.standardSeconds(60)))); - TestPipeline pipeline = TestPipeline.fromOptions(options()); + Pipeline pipeline = Pipeline.create(options()); pipeline .apply( "ReadUnbounded", @@ -153,14 +162,15 @@ public void slidingWindowsCountPerKey() throws Exception { PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains - // one KV.of("a", 2L) pane per overlapping sliding window the two "a" elements both fall into - // (out of scope note: this suite only ever asserts on non-merging windows; session windows are - // out of POC scope per the roadmap). + // One a=2 pane per sliding window the two "a" elements share, so exactly two of them. The + // sentinel's own windows [55s, 65s) and [60s, 70s) both end after the final watermark of 60s + // and therefore never fire. Out of scope note: this suite only ever asserts on non-merging + // windows, session windows are out of POC scope per the roadmap and are rejected outright by + // GroupByKeyStreamingTranslator#canTranslate. + assertEquals("pipeline state=" + result.getState(), "[a=2, a=2]", collectedCounts(collectorId)); } @Test(timeout = 300_000) - @Ignore("Needs GroupByKeyStreamingTranslator (WS-D2), see fixedWindowsCountPerKey for why.") public void lateDataIsDropped() throws Exception { String collectorId = StreamingTestUtils.newCollectorId("late-data-dropped"); StreamingTestUtils.clear(collectorId); @@ -184,7 +194,26 @@ public void lateDataIsDropped() throws Exception { elements.add( TimestampedValue.of(KV.of("sentinel", "t"), BASE.plus(Duration.standardSeconds(90)))); - TestPipeline pipeline = TestPipeline.fromOptions(options()); + SparkStructuredStreamingPipelineOptions options = options(); + // Pin one record per split per micro-batch. Without this the whole four element list lands in + // a single micro-batch, whose start watermark is still -infinity, and the "late" element is + // then perfectly on time. See the comment below on how the splitting interacts with this. + options.setMaxRecordsPerMicroBatch(1); + + // This test, alone in the suite, depends on how ListBackedUnboundedSource round robins its + // elements across splits, so make that dependency loud rather than silent. The session is + // local[2], so UnboundedSourceDataset asks for two splits and gets + // split 0: [a@1s, a@2s] split 1: [sentinel@60s, sentinel@90s] + // With one record per split per micro-batch that gives batch 1 = {a@1s, sentinel@60s} (start + // watermark -infinity, both on time, end watermark 60s) and batch 2 = {a@2s, sentinel@90s} + // (start watermark 60s, so a@2s in window [0s, 10s) is late and dropped, while the same + // batch's start watermark fires that window with the single on-time element in it). + assertEquals( + "this test assumes a two split source, see the comment above", + 2, + SESSION.getSession().sparkContext().defaultParallelism()); + + Pipeline pipeline = Pipeline.create(options); pipeline .apply( "ReadUnbounded", @@ -198,8 +227,10 @@ public void lateDataIsDropped() throws Exception { PipelineResult result = pipeline.run(); result.waitUntilFinish(); - // TODO(WS-D2): assert StreamingTestUtils.>getCollected(collectorId) contains - // KV.of("a", 1L) (the on-time element only) for the first window, and never a count of 2: the - // late "a" element must be dropped, not merged into a late pane. + // a=1, never a=2: the late "a" was dropped rather than merged into a late pane. sentinel=1 is + // the sentinel's [60s, 70s) window, which the trailing sentinel@90s pushes the watermark past; + // its [90s, 100s) window has nothing after it and never fires. + assertEquals( + "pipeline state=" + result.getState(), "[a=1, sentinel=1]", collectedCounts(collectorId)); } } From 5783125f028b6d05a9d12e540c1dcea6925b3403 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 20:02:40 +0000 Subject: [PATCH 08/21] [Spark 4] Register Spark's streaming internals with Kryo Spark 4 broadcasts a StateSchemaMetadata to the executors for every transformWithState query, using the user Kryo instance. Nothing registered that class, so any Beam streaming pipeline with state or timers died on its first micro-batch as soon as spark.kryo.registrationRequired was on, which is the default for this module's tests and a perfectly reasonable production setting. The same applies to MemoryWriterCommitMessage, the commit message of Spark's memory sink, which is nested inside the already registered DataWritingSparkTaskResult. Both are registered by name, because the shared runner base also compiles against Spark 3 where neither class exists, and both are registered with a JavaSerializer rather than Kryo's default field serializer. Both are Scala case classes holding further Scala and Spark types, immutable.Map, StructType, avro Schema and Row, none of which are registered either, so going through Java serialization covers the whole object graph at once instead of forcing this list to track Spark's internal field layout. Neither object is on a hot path. Every streaming test now runs with the module default of spark.kryo.registrationRequired=true, the per test relaxations are gone. SparkKryoRegistratorStreamingTest names both classes at compile time against the Spark 4 classpath, so a future rename becomes a compile error rather than a silently dropped registration. --- .../StreamingEvaluationContext.java | 12 +- .../streaming/BeamMicroBatchSourceTest.java | 10 +- .../SparkKryoRegistratorStreamingTest.java | 135 ++++++++++++++++++ .../ChainedStatefulStreamingTest.java | 4 +- .../streaming/StatefulParDoStreamingTest.java | 9 +- .../StatelessParDoStreamingTest.java | 13 +- .../StreamingPipelineLifecycleTest.java | 10 +- .../WindowedGroupByKeyStreamingTest.java | 12 +- .../state/BeamStatefulProcessorTest.java | 19 ++- .../translation/SparkSessionFactory.java | 52 ++++++- 10 files changed, 225 insertions(+), 51 deletions(-) create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java 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 a1e702dc7ed3..c45ec0de30ab 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 @@ -47,11 +47,13 @@ * *

Sink choice

* - *

Every query uses the {@code noop} sink, never {@code memory}. Tests in this module run with - * {@code spark.kryo.registrationRequired=true}, and the {@code memory} sink's commit messages - * ({@code MemoryWriterCommitMessage}) are not registered with any Beam Kryo registrator, so a query - * using it dies on batch 0. This was already discovered the hard way during the Phase 0 spike, do - * not rediscover it. + *

Every query uses the {@code noop} sink, never {@code memory}. A Beam pipeline emits through + * its own sinks inside the leaf {@code DoFn}s, so a leaf dataset's rows have already served their + * purpose by the time the Spark sink sees them; the {@code memory} sink would accumulate every one + * of them in driver memory for nobody to read. Note that {@code MemoryWriterCommitMessage} is + * registered by {@code SparkSessionFactory.SparkKryoRegistrator} anyway, so switching a query to + * the {@code memory} sink for debugging no longer trips {@code + * spark.kryo.registrationRequired=true}. * *

Termination

* diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java index 203b8465a719..38b27db1f4b5 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -247,12 +247,10 @@ private static long timestampMillis(int index) { /** * Starts a query that throws its output away. * - *

The {@code noop} sink is used on purpose. Spark test JVMs run with {@code - * spark.kryo.registrationRequired=true} (see {@code runners/spark/spark_runner.gradle}) and the - * {@code memory} sink returns a {@code MemoryWriterCommitMessage} that no Beam Kryo registrator - * knows about, so a memory sink query fails immediately. The {@code noop} sink commits {@code - * null} and its {@code DataWritingSparkTaskResult} is already registered, which is also why the - * streaming evaluation context writes to {@code noop}. + *

The {@code noop} sink is used on purpose: these tests observe the source through a {@code + * foreachBatch} or through the query's own progress, never through the sink, so there is no + * reason to buffer rows anywhere. That matches what the streaming evaluation context does for + * real pipelines. */ private StreamingQuery startDiscarding(Dataset dataset, String queryName) throws Exception { return dataset diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java new file mode 100644 index 000000000000..3506daf75d85 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkKryoRegistratorStreamingTest.java @@ -0,0 +1,135 @@ +/* + * 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; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.seqOf; +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.tuple; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.io.Input; +import com.esotericsoftware.kryo.io.Output; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.spark.SparkConf; +import org.apache.spark.serializer.KryoSerializer; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.execution.streaming.sources.MemoryWriterCommitMessage; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadataKey; +import org.apache.spark.sql.execution.streaming.state.StateSchemaMetadataValue; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import scala.collection.immutable.Map$; + +/** + * Guards the Spark 4 streaming entries of {@link SparkSessionFactory.SparkKryoRegistrator}: the + * classes Spark itself pushes through the user Kryo instance during a Structured Streaming query, + * which the runner never names anywhere else. + * + *

These are registered by name in the shared runner base, because that base also compiles + * against Spark 3 where neither class exists. A rename or a package move on a future Spark version + * would therefore not break the build, it would silently drop the registration and only surface as + * a streaming query dying on its first micro-batch. This test names the classes at compile time + * against the Spark 4 classpath, so that failure mode becomes a compile error instead. + * + * @see SparkSessionFactory.SparkKryoRegistrator + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class SparkKryoRegistratorStreamingTest { + + /** A Kryo configured exactly the way the runner configures it, strict registration included. */ + private static Kryo strictKryo() { + SparkConf conf = + new SparkConf(false) + .set("spark.serializer", KryoSerializer.class.getName()) + .set("spark.kryo.registrationRequired", "true") + .set( + "spark.kryo.registrator", SparkSessionFactory.SparkKryoRegistrator.class.getName()); + return new KryoSerializer(conf).newKryo(); + } + + private static Object roundTrip(Kryo kryo, Object value) { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (Output output = new Output(bytes)) { + kryo.writeClassAndObject(output, value); + } + try (Input input = new Input(new ByteArrayInputStream(bytes.toByteArray()))) { + return kryo.readClassAndObject(input); + } + } + + /** + * Spark 4 broadcasts a {@link StateSchemaMetadata} to the executors for every {@code + * transformWithState} query, so this is the registration that decides whether a Beam streaming + * pipeline with state or timers runs at all under {@code spark.kryo.registrationRequired=true}. + * + *

The instance below is deliberately not empty. It carries the nested {@code StructType} and + * {@code org.apache.avro.Schema} that make the difference between a registration that only + * survives a trivial payload and one that survives a real one. + */ + @Test + public void stateSchemaMetadataRoundTripsWithRegistrationRequired() { + StructType sqlSchema = + new StructType().add("key", DataTypes.StringType).add("value", DataTypes.BinaryType, false); + StateSchemaMetadataKey key = new StateSchemaMetadataKey("default", (short) 1, true); + StateSchemaMetadataValue value = + new StateSchemaMetadataValue( + sqlSchema, org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING)); + StateSchemaMetadata metadata = + new StateSchemaMetadata(Map$.MODULE$.from(seqOf(tuple(key, value)))); + + Kryo kryo = strictKryo(); + assertNotNull( + "StateSchemaMetadata must be registered, see SparkKryoRegistrator", + kryo.getRegistration(StateSchemaMetadata.class)); + + StateSchemaMetadata back = (StateSchemaMetadata) roundTrip(kryo, metadata); + assertEquals(1, back.activeSchemas().size()); + assertEquals(value, back.activeSchemas().apply(key)); + } + + /** + * The commit message of Spark's {@code memory} sink, nested inside the already registered {@code + * DataWritingSparkTaskResult}. The runner writes to {@code noop}, but the {@code memory} sink is + * the obvious thing to reach for when inspecting a query, and it used to fail on batch 0. + */ + @Test + public void memoryWriterCommitMessageRoundTripsWithRegistrationRequired() { + Row row = RowFactory.create("a", 1); + MemoryWriterCommitMessage message = new MemoryWriterCommitMessage(3, seqOf(row)); + + Kryo kryo = strictKryo(); + assertNotNull( + "MemoryWriterCommitMessage must be registered, see SparkKryoRegistrator", + kryo.getRegistration(MemoryWriterCommitMessage.class)); + + MemoryWriterCommitMessage back = (MemoryWriterCommitMessage) roundTrip(kryo, message); + assertEquals(3, back.partition()); + assertEquals(1, back.data().size()); + assertEquals(row, back.data().apply(0)); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java index da1ea8df2982..4b428b8739eb 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java @@ -69,9 +69,7 @@ @Category(StreamingTest.class) public class ChainedStatefulStreamingTest implements Serializable { - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); 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 b0eff04d2979..d33f7e3d4e4a 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 @@ -64,10 +64,11 @@ @Category(StreamingTest.class) public class StatefulParDoStreamingTest implements Serializable { - /** See {@code BeamStatefulProcessorTest} for why this relaxation is required. */ - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + /** + * Runs with the module default of {@code spark.kryo.registrationRequired=true}, see {@code + * BeamStatefulProcessorTest}. + */ + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java index e352014bbd3f..a69409681eda 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -32,7 +32,6 @@ import org.apache.beam.sdk.io.Read; 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.TimestampedValue; import org.joda.time.Duration; import org.joda.time.Instant; @@ -54,14 +53,12 @@ public class StatelessParDoStreamingTest implements Serializable { /** - * This pipeline hosts no {@code transformWithState} operator of its own, but it still shares the - * session with the rest of the streaming suite and is configured identically to it, so that the - * baseline test differs from the stateful ones in the pipeline under test and nothing else. See - * {@code BeamStatefulProcessorTest} for why the relaxation is needed at all. + * This pipeline hosts no {@code transformWithState} operator of its own, but it is configured + * identically to the stateful tests, including the module default of {@code + * spark.kryo.registrationRequired=true}, so that the baseline test differs from them in the + * pipeline under test and nothing else. */ - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java index ca67c06e9f3b..f58a809d591b 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -32,7 +32,6 @@ import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.transforms.ParDo; -import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.TimestampedValue; import org.joda.time.Duration; import org.joda.time.Instant; @@ -74,12 +73,11 @@ public class StreamingPipelineLifecycleTest implements Serializable { * useActiveSparkSession} so that a cancelled or idle-stopped query does not take the shared * session down with it: {@code SparkStructuredStreamingRunner#sparkStopFn} only stops the session * on a terminal state when the session was not provided from outside. Configuring the - * session the same relaxed way as the rest of the suite keeps a single session shared across the - * whole streaming test run. + * session the same way as the rest of the suite, including the module default of {@code + * spark.kryo.registrationRequired=true}, keeps a single session shared across the whole streaming + * test run. */ - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); 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 c3fb37003da3..b292def89799 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 @@ -67,14 +67,12 @@ public class WindowedGroupByKeyStreamingTest implements Serializable { /** - * See {@code BeamStatefulProcessorTest}: {@code transformWithState} broadcasts {@code - * StateSchemaMetadata} through Kryo, which is not registered anywhere, so the test JVM default of - * {@code spark.kryo.registrationRequired=true} (runners/spark/spark_runner.gradle) must be - * relaxed for any query that ends up hosting a stateful operator, windowed GroupByKey included. + * Runs with the module default of {@code spark.kryo.registrationRequired=true}, see {@code + * BeamStatefulProcessorTest} for why a {@code transformWithState} query needs {@code + * SparkSessionFactory.SparkKryoRegistrator} to know about {@code StateSchemaMetadata} for that to + * hold. */ - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java index 5929d4b0ddf0..bdfee6cf2113 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java @@ -96,18 +96,15 @@ public class BeamStatefulProcessorTest implements Serializable { /** - * Spark test JVMs set {@code spark.kryo.registrationRequired=true} (see {@code - * runners/spark/spark_runner.gradle}) to keep Beam's own serialisation honest. That cannot hold - * for a {@code transformWithState} query: Spark 4 broadcasts its own {@code - * org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata} to the executors with the - * user Kryo instance, and no Beam registrator knows that class, so the query dies before the - * first micro-batch. Beam never turns the flag on outside tests, so this only relaxes the test - * harness, but every streaming test that runs a stateful operator will have to do the same until - * the Spark 4 Kryo registrator learns about Spark's streaming state classes. + * Deliberately runs with the module default of {@code spark.kryo.registrationRequired=true} (see + * {@code runners/spark/spark_runner.gradle}). Spark 4 broadcasts its own {@code + * org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata} to the executors through + * the user Kryo instance for every {@code transformWithState} query, so a stateful query only + * survives its first micro-batch because {@code SparkSessionFactory.SparkKryoRegistrator} + * registers that class. Keeping the strict flag on here is what stops that registration from + * silently rotting. */ - @ClassRule - public static final SparkSessionRule SESSION = - new SparkSessionRule(KV.of("spark.kryo.registrationRequired", "false")); + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); @Rule public transient TemporaryFolder temp = new TemporaryFolder(); diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 822d1871b12e..81e421026499 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -22,6 +22,7 @@ import static org.apache.commons.lang3.math.NumberUtils.toInt; import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.Serializer; import com.esotericsoftware.kryo.serializers.JavaSerializer; import java.util.ArrayList; import java.util.Collection; @@ -235,6 +236,11 @@ public void registerClasses(Kryo kryo) { tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroCoder"); tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroGenericCoder"); + // Spark internals only present when running streaming pipelines on Spark 4. These are + // registered by name because the shared runner base also compiles against Spark 3, where + // none of these classes exist. See registerSparkStreamingInternals for the details. + registerSparkStreamingInternals(kryo); + // standard coders of org.apache.beam.sdk.coders kryo.register(BigDecimalCoder.class); kryo.register(BigEndianIntegerCoder.class); @@ -290,9 +296,53 @@ public void registerClasses(Kryo kryo) { kryo.register(TupleTagList.class); } + /** + * Registers the Spark internals that a Structured Streaming query serializes behind the + * runner's back, so streaming pipelines also work with {@code + * spark.kryo.registrationRequired=true}. + * + *

    + *
  • {@code StateSchemaMetadata} is broadcast by Spark 4 for every {@code + * transformWithState} query, so every streaming pipeline using Beam state or timers hits + * it, and hits it on the very first micro-batch. + *
  • {@code MemoryWriterCommitMessage} is the commit message of Spark's {@code memory} sink. + * The runner itself writes to the {@code noop} sink, but the {@code memory} sink is what + * one reaches for when inspecting a query's output, and it is nested inside the already + * registered {@link DataWritingSparkTaskResult}. + *
+ * + *

Both are Scala case classes holding further Scala and Spark types ({@code + * immutable.Map}, {@code StructType}, {@code org.apache.avro.Schema}, {@code Row}), none of + * which are registered either. Registering them with a {@link JavaSerializer} rather than + * Kryo's default field serializer covers that whole object graph in one go, since both classes + * are {@link java.io.Serializable}. That keeps this list from having to track Spark's internal + * field layout across versions. Neither object is on a hot path, one is broadcast once per + * query and the other is one message per task commit, so the cost of Java serialization here + * does not matter. + */ + private void registerSparkStreamingInternals(Kryo kryo) { + tryToRegister( + kryo, + "org.apache.spark.sql.execution.streaming.state.StateSchemaMetadata", + new JavaSerializer()); + tryToRegister( + kryo, + "org.apache.spark.sql.execution.streaming.sources.MemoryWriterCommitMessage", + new JavaSerializer()); + } + private void tryToRegister(Kryo kryo, String className) { + tryToRegister(kryo, className, null); + } + + private void tryToRegister(Kryo kryo, String className, @Nullable Serializer serializer) { try { - kryo.register(Class.forName(className)); + Class cls = Class.forName(className); + if (serializer == null) { + kryo.register(cls); + } else { + kryo.register(cls, serializer); + } } catch (ClassNotFoundException e) { LOG.info("Class {}} was not found on classpath", className); } From c8de94f0eadfb9dd34c06ef8e83f8ad4d3ae4594 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 20:09:47 +0000 Subject: [PATCH 09/21] [Spark 4] Reformat the shared base to satisfy spotless Two javadoc paragraphs in the shared base drifted from google-java-format, one from the streaming note added to SparkStructuredStreamingRunner and one from the new Kryo registrations. The shared base sources under runners/spark/src are checked by the :runners:spark project rather than by :runners:spark:3 or :runners:spark:4, whose own spotless targets only see their per version override directories, so it is easy to miss. --- .../SparkStructuredStreamingRunner.java | 4 ++-- .../translation/SparkSessionFactory.java | 16 ++++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java index 84afa2b2a55c..f78026847fad 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java @@ -54,8 +54,8 @@ * Streaming framework). * *

This runner is experimental, its coverage of the Beam model is still partial. Streaming - * mode requires the Spark 4 module (beam-runners-spark-4); the shared Spark 3 module supports - * batch pipelines only. + * mode requires the Spark 4 module (beam-runners-spark-4); the shared Spark 3 module supports batch + * pipelines only. * *

The runner translates transforms defined on a Beam pipeline to Spark `Dataset` transformations * (leveraging the high level Dataset API) and then submits these to Spark to be executed. diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 81e421026499..ffa0bb563ec1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -311,14 +311,14 @@ public void registerClasses(Kryo kryo) { * registered {@link DataWritingSparkTaskResult}. * * - *

Both are Scala case classes holding further Scala and Spark types ({@code - * immutable.Map}, {@code StructType}, {@code org.apache.avro.Schema}, {@code Row}), none of - * which are registered either. Registering them with a {@link JavaSerializer} rather than - * Kryo's default field serializer covers that whole object graph in one go, since both classes - * are {@link java.io.Serializable}. That keeps this list from having to track Spark's internal - * field layout across versions. Neither object is on a hot path, one is broadcast once per - * query and the other is one message per task commit, so the cost of Java serialization here - * does not matter. + *

Both are Scala case classes holding further Scala and Spark types ({@code immutable.Map}, + * {@code StructType}, {@code org.apache.avro.Schema}, {@code Row}), none of which are + * registered either. Registering them with a {@link JavaSerializer} rather than Kryo's default + * field serializer covers that whole object graph in one go, since both classes are {@link + * java.io.Serializable}. That keeps this list from having to track Spark's internal field + * layout across versions. Neither object is on a hot path, one is broadcast once per query and + * the other is one message per task commit, so the cost of Java serialization here does not + * matter. */ private void registerSparkStreamingInternals(Kryo kryo) { tryToRegister( From d9639f68c474a62480a8e30133d76ddb1a8095d8 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 20:23:23 +0000 Subject: [PATCH 10/21] [Spark 4] Assert the chained pipeline against Spark's query progress ChainedStatefulStreamingTest asserts what the chained pipeline computes. It cannot distinguish a single query holding two transformWithState operators from two queries that happen to add up to the same numbers, and it cannot show whether the watermark moved at all or whether the whole input simply landed in one micro-batch where everything is trivially on time. ChainedStatefulStreamingEvidenceTest asserts the run itself, against Spark's own StreamingQueryProgress. It pins one record per split per micro-batch so the watermark has to climb in steps, records every progress event through a listener, and then asserts that all of them carry one query id, that one micro-batch reports exactly two transformWithStateExec state operators, and that the watermark takes at least three strictly increasing values. The second test adds a tap between the two operators, so the late record is observed leaving the dedup operator and absent from the windowed sum. That is the difference between a record excluded for lateness and a record that never arrived. Both tests print the raw per batch progress they recorded, since that output is the evidence the phase gate report quotes. Also drops four javadoc references to a Kryo relaxation that no longer exists anywhere in the suite. --- .../ChainedStatefulStreamingEvidenceTest.java | 394 ++++++++++++++++++ .../ChainedStatefulStreamingTest.java | 4 +- .../streaming/StatefulParDoStreamingTest.java | 3 +- .../streaming/StreamingTestUtils.java | 5 +- .../WindowedGroupByKeyStreamingTest.java | 3 +- 5 files changed, 401 insertions(+), 8 deletions(-) create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java new file mode 100644 index 000000000000..c1183f4acb14 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java @@ -0,0 +1,394 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; +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.VarLongCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +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.Sum; +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.TimestampedValue; +import org.apache.spark.sql.streaming.StateOperatorProgress; +import org.apache.spark.sql.streaming.StreamingQueryListener; +import org.apache.spark.sql.streaming.StreamingQueryProgress; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * The chained stateful pipeline, asserted against Spark's own {@link StreamingQueryProgress} rather + * than only against the pipeline's output. + * + *

{@link ChainedStatefulStreamingTest} asserts what the chained pipeline computes. This + * one asserts how Spark ran it, which is the part of the POC claim a correct output alone + * does not evidence: + * + *

    + *
  1. Two distinct {@code transformWithState} operators live inside one single Spark streaming + * query, rather than the pipeline being cut into two queries that would each carry their own + * independent watermark. + *
  2. The event time watermark of that one query genuinely advances over successive + * micro-batches, instead of the whole input landing in one batch where every element is + * trivially on time. + *
  3. A record whose window the watermark has already passed is excluded from the result rather + * than quietly folded into it, and it is excluded at the downstream windowed operator, after + * having passed cleanly through the upstream stateful one. + *
+ * + *

Both tests print the raw per micro-batch progress they recorded. That printout is the evidence + * the phase gate report quotes, so keep it printing. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class ChainedStatefulStreamingEvidenceTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final org.joda.time.Instant BASE = new org.joda.time.Instant(0); + private static final Duration WINDOW_SIZE = Duration.standardSeconds(10); + + /** Spark's short name for the physical operator a {@code transformWithState} compiles down to. */ + private static final String TWS_OPERATOR_NAME = "transformWithStateExec"; + + /** Records every {@link StreamingQueryProgress} a run produced, in order. */ + private static final class ProgressRecorder extends StreamingQueryListener { + private final List progresses = + Collections.synchronizedList(new ArrayList<>()); + + @Override + public void onQueryStarted(QueryStartedEvent event) {} + + @Override + public void onQueryProgress(QueryProgressEvent event) { + progresses.add(event.progress()); + } + + @Override + public void onQueryTerminated(QueryTerminatedEvent event) {} + + List snapshot() { + synchronized (progresses) { + return new ArrayList<>(progresses); + } + } + } + + /** + * Dedups by the outer {@code id} key and passes the inner {@code KV} on. + * Deliberately identical to the one in {@link ChainedStatefulStreamingTest}, so both tests + * describe the same pipeline. + */ + private static class DedupByIdFn extends DoFn>, KV> { + @StateId("seen") + private final StateSpec> seenSpec = StateSpecs.value(); + + @ProcessElement + public void process( + @Element KV> element, + @StateId("seen") ValueState seen, + OutputReceiver> out) { + Boolean alreadySeen = seen.read(); + if (alreadySeen == null || !alreadySeen) { + seen.write(true); + out.output(element.getValue()); + } + } + } + + private static String render(String collectorId) { + List rendered = new ArrayList<>(); + for (KV kv : StreamingTestUtils.>getCollected(collectorId)) { + rendered.add(kv.getKey() + "=" + kv.getValue()); + } + Collections.sort(rendered); + return rendered.toString(); + } + + /** Prints one line per micro-batch: batch id, input rows, watermark, and each state operator. */ + private static void printProgress(String label, List progresses) { + StringBuilder out = new StringBuilder(); + out.append(System.lineSeparator()).append("===== ").append(label).append(" ====="); + for (StreamingQueryProgress progress : progresses) { + out.append(System.lineSeparator()) + .append("queryId=") + .append(progress.id()) + .append(" batchId=") + .append(progress.batchId()) + .append(" numInputRows=") + .append(progress.numInputRows()) + .append(" eventTime=") + .append(progress.eventTime()); + for (StateOperatorProgress operator : progress.stateOperators()) { + out.append(System.lineSeparator()) + .append(" stateOperator name=") + .append(operator.operatorName()) + .append(" numRowsTotal=") + .append(operator.numRowsTotal()) + .append(" numRowsUpdated=") + .append(operator.numRowsUpdated()) + .append(" numRowsRemoved=") + .append(operator.numRowsRemoved()) + .append(" numRowsDroppedByWatermark=") + .append(operator.numRowsDroppedByWatermark()) + .append(" numStateStoreInstances=") + .append(operator.numStateStoreInstances()); + } + } + out.append(System.lineSeparator()).append("===== end ").append(label).append(" ====="); + // Deliberately System.out: this printout is an artefact the phase gate report quotes, and it + // has to survive whatever log configuration the module happens to run with. + System.out.println(out); + } + + /** The event time watermark Spark used for a micro-batch, or null if it had none yet. */ + private static @Nullable Instant watermarkOf(StreamingQueryProgress progress) { + String watermark = progress.eventTime().get("watermark"); + return watermark == null ? null : Instant.parse(watermark); + } + + /** The distinct watermark values a run went through, in the order they first appeared. */ + private static List distinctWatermarks(List progresses) { + List watermarks = new ArrayList<>(); + for (StreamingQueryProgress progress : progresses) { + Instant watermark = watermarkOf(progress); + if (watermark != null + && (watermarks.isEmpty() || !watermark.equals(watermarks.get(watermarks.size() - 1)))) { + watermarks.add(watermark); + } + } + return watermarks; + } + + /** Asserts every recorded progress belongs to one and the same streaming query. */ + private static void assertSingleQuery(List progresses) { + Set ids = new LinkedHashSet<>(); + for (StreamingQueryProgress progress : progresses) { + ids.add(progress.id()); + } + assertEquals( + "expected the whole pipeline to run as one streaming query, saw " + ids, 1, ids.size()); + } + + private SparkStructuredStreamingPipelineOptions oneRecordPerSplitPerBatchOptions() + throws Exception { + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + // One record per split per micro-batch, so the watermark climbs in visible steps instead of + // reaching its final value inside a single batch. Both tests here are about what happens + // between micro-batches, which a single batch run cannot show at all. + options.setMaxRecordsPerMicroBatch(1); + return options; + } + + private static Read.Unbounded>> readOf( + List>>> elements) { + return Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, + KvCoder.of(StringUtf8Coder.of(), KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of())))); + } + + /** + * The on-time chained pipeline, asserted against Spark's own view of the run: one query, two + * {@code transformWithState} operators inside it, and a watermark that moves. + */ + @Test(timeout = 300_000) + public void twoStatefulOperatorsShareOneQueryAndItsAdvancingWatermark() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("evidence-chained"); + StreamingTestUtils.clear(collectorId); + + // Same data as ChainedStatefulStreamingTest: id "1" is redelivered and must count only once. + List>>> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(KV.of("1", KV.of("a", 5L)), BASE)); + elements.add( + TimestampedValue.of(KV.of("1", KV.of("a", 5L)), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("2", KV.of("a", 3L)), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("3", KV.of("b", 10L)), BASE.plus(Duration.standardSeconds(3)))); + elements.add( + TimestampedValue.of( + KV.of("sentinel", KV.of("sentinel", 0L)), BASE.plus(Duration.standardSeconds(60)))); + + Pipeline pipeline = Pipeline.create(oneRecordPerSplitPerBatchOptions()); + pipeline + .apply("ReadUnbounded", readOf(elements)) + .apply("DedupById", ParDo.of(new DedupByIdFn())) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("SumPerKey", Sum.longsPerKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + ProgressRecorder recorder = new ProgressRecorder(); + SESSION.getSession().streams().addListener(recorder); + PipelineResult result; + try { + result = pipeline.run(); + result.waitUntilFinish(); + } finally { + SESSION.getSession().streams().removeListener(recorder); + } + + List progresses = recorder.snapshot(); + printProgress("chained stateful, on time", progresses); + List watermarks = distinctWatermarks(progresses); + System.out.println("distinct watermarks in order: " + watermarks); + + assertEquals("pipeline state=" + result.getState(), "[a=8, b=10]", render(collectorId)); + + // (1) One query, and inside it two transformWithState operators reported together in the same + // micro-batch. A progress record is scoped to exactly one query, so two entries in one record + // cannot be two separate queries being conflated. + assertSingleQuery(progresses); + StreamingQueryProgress twoOperators = null; + for (StreamingQueryProgress progress : progresses) { + if (progress.stateOperators().length == 2) { + twoOperators = progress; + break; + } + } + assertNotNull( + "no micro-batch reported two state operators, see the printed progress above", + twoOperators); + for (StateOperatorProgress operator : twoOperators.stateOperators()) { + assertEquals(TWS_OPERATOR_NAME, operator.operatorName()); + } + + // (2) The watermark of that one query advances over micro-batches, and never moves backwards. + assertTrue( + "expected the watermark to take at least three distinct values over the run, saw " + + watermarks, + watermarks.size() >= 3); + for (int i = 1; i < watermarks.size(); i++) { + assertTrue( + "watermark moved backwards: " + watermarks, + watermarks.get(i).isAfter(watermarks.get(i - 1))); + } + } + + /** + * A record whose event time falls in a window the watermark has already passed is excluded from + * that window's result. The tap between the two operators is what makes this a statement about + * lateness rather than about the record having gone missing somewhere upstream: the very same + * record is observed leaving the dedup operator and absent from the windowed sum. + * + *

Deterministic only because of the split and batch arithmetic spelled out below. This test + * and {@code WindowedGroupByKeyStreamingTest#lateDataIsDropped} are the only two in the suite + * that depend on it. + */ + @Test(timeout = 300_000) + public void lateRecordIsExcludedByTheDownstreamWindowNotLostUpstream() throws Exception { + String tapId = StreamingTestUtils.newCollectorId("evidence-late-tap"); + String collectorId = StreamingTestUtils.newCollectorId("evidence-late"); + StreamingTestUtils.clear(tapId); + StreamingTestUtils.clear(collectorId); + + // ListBackedUnboundedSource round robins, so split 0 gets indices 0 and 2, split 1 gets 1 and + // 3. With one record per split per micro-batch that gives: + // batch 1 = {a@0s, z@60s} start watermark -infinity, both on time, end watermark 60s + // batch 2 = {a@2s, z@90s} start watermark 60s, so a@2s in window [0s, 10s) is already late + // and is dropped, while the same batch's start watermark fires that + // window with the single on-time element it holds + // batch 3 = {} start watermark 90s, fires window [60s, 70s) + List>>> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(KV.of("1", KV.of("a", 5L)), BASE)); + elements.add( + TimestampedValue.of(KV.of("s1", KV.of("z", 7L)), BASE.plus(Duration.standardSeconds(60)))); + elements.add( + TimestampedValue.of(KV.of("2", KV.of("a", 3L)), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("s2", KV.of("z", 11L)), BASE.plus(Duration.standardSeconds(90)))); + + SparkStructuredStreamingPipelineOptions options = oneRecordPerSplitPerBatchOptions(); + assertEquals( + "this test assumes a two split source, see the comment above", + 2, + SESSION.getSession().sparkContext().defaultParallelism()); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("ReadUnbounded", readOf(elements)) + .apply("DedupById", ParDo.of(new DedupByIdFn())) + // Tap between the two stateful operators: whatever this sees did leave operator one and + // did reach operator two's input. + .apply("TapAfterDedup", ParDo.of(new StreamingTestUtils.CollectDoFn<>(tapId))) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("SumPerKey", Sum.longsPerKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + ProgressRecorder recorder = new ProgressRecorder(); + SESSION.getSession().streams().addListener(recorder); + PipelineResult result; + try { + result = pipeline.run(); + result.waitUntilFinish(); + } finally { + SESSION.getSession().streams().removeListener(recorder); + } + + printProgress("chained stateful, late record", recorder.snapshot()); + System.out.println("after dedup: " + render(tapId)); + System.out.println("windowed sums: " + render(collectorId)); + + // The late record was emitted by the dedup operator, so it did reach the windowed operator. + assertEquals( + "the late record never made it past the dedup operator, so this test would prove nothing" + + " about lateness", + "[a=3, a=5, z=11, z=7]", + render(tapId)); + + // And it is still not in the result. a=5, never a=8: the a=3 that arrived after the watermark + // had passed the end of [0s, 10s) was excluded rather than folded into the sum. z=7 is window + // [60s, 70s); the z=11 in [90s, 100s) has nothing after it to push the watermark past its end, + // so that window never fires. + assertEquals("pipeline state=" + result.getState(), "[a=5, z=7]", render(collectorId)); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java index 4b428b8739eb..dd3ec0a9f782 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java @@ -63,7 +63,9 @@ * still correctly decide when it has seen everything it is going to see. Everything else in this * package tests one operator at a time; this one tests that they compose. * - *

Needs the same Kryo relaxation as {@code BeamStatefulProcessorTest}, for the same reason. + *

See {@code ChainedStatefulStreamingEvidenceTest} for the same pipeline asserted against + * Spark's own {@code StreamingQueryProgress}, which is what evidences that the two operators really + * do share one query and one advancing watermark rather than merely producing the right numbers. */ @RunWith(JUnit4.class) @Category(StreamingTest.class) 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 d33f7e3d4e4a..1bd925b4c78d 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 @@ -57,8 +57,7 @@ * A stateful {@code ParDo} in the global window: a {@code @StateId ValueState} dedups * repeated keys, and an event time {@code @TimerId} emits a sentinel once it expires. Hosted by the * generic {@code transformWithState} super-operator ({@code - * BeamStatefulProcessorConfig.Mode#STATEFUL_PARDO}), so needs the same Kryo relaxation as {@code - * BeamStatefulProcessorTest}. + * BeamStatefulProcessorConfig.Mode#STATEFUL_PARDO}). */ @RunWith(JUnit4.class) @Category(StreamingTest.class) diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java index 1bf14084f5ec..2da0458dbc07 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -180,9 +180,8 @@ public static String newCollectorId(String prefix) { * {@link SparkStructuredStreamingRunner}, test mode, streaming mode, a 3-idle-batch stop, a 200ms * micro-batch trigger, and a checkpoint directory carved out of {@code checkpointDir}. * - *

Callers that need a specific {@code SparkSession} (for example the relaxed-Kryo {@code - * SparkSessionRule} pattern required by any query that hosts a {@code transformWithState} - * operator, see {@code BeamStatefulProcessorTest}) should additionally call {@code + *

Callers that need the test to run against a specific {@code SparkSession}, for example the + * one held by a {@code SparkSessionRule}, should additionally call {@code * SparkSessionRule#configure} on the returned options, which sets {@code useActiveSparkSession}. */ public static SparkStructuredStreamingPipelineOptions streamingOptions( 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 b292def89799..32d0fd4544e6 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 @@ -52,8 +52,7 @@ * Windowed {@code GroupByKey} (via {@link Count#perKey()}, which auto-expands to {@code GroupByKey} * + {@code Combine} since {@code Combine.PerKey} is deliberately unregistered for streaming, see * {@code PipelineTranslatorStreaming}). This is hosted by the generic {@code transformWithState} - * super-operator in {@code BeamStatefulProcessorConfig.Mode #GROUP_ALSO_BY_WINDOW}, so every test - * here needs the same Kryo relaxation as {@code BeamStatefulProcessorTest}. + * super-operator in {@code BeamStatefulProcessorConfig.Mode #GROUP_ALSO_BY_WINDOW}. * *

Every window this suite asserts on is followed, in the input list, by an element timestamped * well past that window's end, per the watermark rule documented on {@link StreamingTestUtils}: the From fed52a15550836a4a4214119f7d9177cc4949e91 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Tue, 4 Aug 2026 18:08:03 +0000 Subject: [PATCH 11/21] [Spark 4] Satisfy ErrorProne on the streaming sources Two checks fire under -Werror in CI but not in a local build. Passing -Pjava17Home makes the compile fork to Java 17, and the Beam build adds -XepDisableAllChecks whenever it forks, so ErrorProne is silently off locally. Reproduce with JAVA_HOME set to a Java 17 JDK instead. TypeParameterUnusedInFormals on BeamStreamingSource.decode is suppressed. The decoded types include Coder>, which no Class token can express, so a checked variant is not available. The unused windowingStrategy local in BeamStatefulProcessor.process is removed, every later use reads config.windowingStrategy() directly. --- .../io/streaming/BeamStreamingSource.java | 10 ++++++++-- .../streaming/state/BeamStatefulProcessor.java | 1 - 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java index b9a1100f11e4..d057d39bac55 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java @@ -127,8 +127,14 @@ static String encode(Serializable value) { return Base64.getEncoder().encodeToString(SerializableUtils.serializeToByteArray(value)); } - /** Inverse of {@link #encode}, {@code description} is only used in error messages. */ - @SuppressWarnings("unchecked") + /** + * Inverse of {@link #encode}, {@code description} is only used in error messages. + * + *

The return type is inferred from the call site. Deserialization cannot check it, and the + * decoded types include generic ones such as {@code Coder>} that no {@code + * Class} token can express, so a checked variant is not available here. + */ + @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"}) static T decode(String encoded, String description) { return (T) SerializableUtils.deserializeFromByteArray( 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 0eca7e26a113..d61f4582a263 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 @@ -222,7 +222,6 @@ private scala.collection.Iterator process( Object key = decode(config.keyCoder(), encodedKey, "key"); Coder windowCoder = config.windowCoder(); - WindowingStrategy windowingStrategy = config.windowingStrategy(); TwsStateInternals stateInternals = TwsStateInternals.forKey(key, BytesKV.of(beamState)); TwsTimerInternals timerInternals = From f036681c567b2d6b6dea458d1ee48c4e40828cb6 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 19 Aug 2026 16:23:16 +0000 Subject: [PATCH 12/21] [Spark 4] Register the streaming internals at the end of the Kryo registrator The conditional, by-name registrations only resolve on a Spark 4 classpath. Registering them last keeps the auto assigned Kryo ids of every unconditional registration identical on Spark 3 and Spark 4. --- .../translation/SparkSessionFactory.java | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index ffa0bb563ec1..9b6265ab6cf8 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -236,11 +236,6 @@ public void registerClasses(Kryo kryo) { tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroCoder"); tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroGenericCoder"); - // Spark internals only present when running streaming pipelines on Spark 4. These are - // registered by name because the shared runner base also compiles against Spark 3, where - // none of these classes exist. See registerSparkStreamingInternals for the details. - registerSparkStreamingInternals(kryo); - // standard coders of org.apache.beam.sdk.coders kryo.register(BigDecimalCoder.class); kryo.register(BigEndianIntegerCoder.class); @@ -294,6 +289,14 @@ public void registerClasses(Kryo kryo) { kryo.register(CoGbkResultSchema.class); kryo.register(TupleTag.class); kryo.register(TupleTagList.class); + + // Spark internals only present when running streaming pipelines on Spark 4. These are + // registered by name because the shared runner base also compiles against Spark 3, where + // none of these classes exist. This call must stay last: Kryo auto assigns registration ids + // sequentially, so registering these conditional, by-name classes after everything else + // keeps the auto assigned ids of all the registrations above identical on Spark 3 and + // Spark 4 classpaths. See registerSparkStreamingInternals for the details. + registerSparkStreamingInternals(kryo); } /** From 497a531411ba08db99fa0ffe6583fe6e8e5bdb37 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 19 Aug 2026 16:23:17 +0000 Subject: [PATCH 13/21] [Spark 4] Clean up the fallback temporary checkpoint directory When no checkpointDir is configured the streaming evaluation context creates a temporary directory and never removed it. It is now deleted recursively, best effort, once evaluate() finishes. A user configured checkpointDir is never touched. --- .../StreamingEvaluationContext.java | 48 ++++++++++++++++- .../StreamingPipelineLifecycleTest.java | 53 +++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) 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 c45ec0de30ab..34bdd8131f13 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 @@ -20,14 +20,17 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.nio.file.Files; +import java.nio.file.Path; import java.util.ArrayList; import java.util.Collection; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; import org.apache.beam.runners.spark.SparkCommonPipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; import org.apache.beam.sdk.annotations.Internal; @@ -37,6 +40,7 @@ import org.apache.spark.sql.streaming.StreamingQueryException; import org.apache.spark.sql.streaming.StreamingQueryListener; import org.apache.spark.sql.streaming.Trigger; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -95,6 +99,13 @@ public class StreamingEvaluationContext extends EvaluationContext { private final List queries = new ArrayList<>(); private boolean stopped = false; + // Set by checkpointBaseDir() when no checkpointDir was configured, so a temporary directory was + // created as a fallback. Written once in evaluate() before any query starts and read back in the + // finally block of that same method on the same thread, so unlike `queries` and `stopped` it does + // not need `lock`. Never set when the user configured a checkpointDir, which evaluate() must + // never delete. + private @Nullable Path tempCheckpointDir; + StreamingEvaluationContext( Collection> leaves, SparkSession session, @@ -162,6 +173,9 @@ public void evaluate() { if (idleStopListener != null) { getSparkSession().streams().removeListener(idleStopListener); } + if (tempCheckpointDir != null) { + deleteTempCheckpointDir(tempCheckpointDir); + } } } @@ -249,19 +263,49 @@ private void stopQueryById(UUID id) { } } - private static String checkpointBaseDir(SparkCommonPipelineOptions options) { + private String checkpointBaseDir(SparkCommonPipelineOptions options) { String dir = options.getCheckpointDir(); if (dir == null || dir.isEmpty()) { try { - dir = Files.createTempDirectory("beam-spark4-streaming-checkpoint").toString(); + tempCheckpointDir = Files.createTempDirectory("beam-spark4-streaming-checkpoint"); } catch (IOException e) { throw new UncheckedIOException(e); } + dir = tempCheckpointDir.toString(); LOG.warn("No checkpoint directory configured, falling back to temporary directory {}.", dir); } return dir; } + /** + * Best-effort, recursive delete of the fallback temporary checkpoint directory created by {@link + * #checkpointBaseDir}. RocksDB and Spark write nested state and log files under it, so plain + * {@link Files#delete} is not enough. Failures are logged rather than thrown: cleanup is a + * courtesy, not something that should fail a pipeline that otherwise ran to completion. + */ + private static void deleteTempCheckpointDir(Path dir) { + try (Stream paths = Files.walk(dir)) { + paths + .sorted(Comparator.reverseOrder()) + .forEach( + path -> { + try { + Files.delete(path); + } catch (IOException e) { + LOG.warn( + "Failed to delete temporary checkpoint path {}: {}", + path, + String.valueOf(e.getMessage())); + } + }); + } catch (IOException e) { + LOG.warn( + "Failed to clean up temporary checkpoint directory {}: {}", + dir, + String.valueOf(e.getMessage())); + } + } + /** * Counts, per query, the number of consecutive micro-batches with zero input rows, and gracefully * stops a query once its count reaches {@code threshold}. The count for a query resets to zero as diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java index f58a809d591b..d10c0730a232 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -20,10 +20,14 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import java.io.File; import java.io.Serializable; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import org.apache.beam.runners.spark.StreamingTest; import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; @@ -190,4 +194,53 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception { Thread.sleep(50L); } } + + /** + * When no {@code checkpointDir} is configured, {@code StreamingEvaluationContext} falls back to a + * {@code beam-spark4-streaming-checkpoint*} temporary directory under {@code java.io.tmpdir}. + * That fallback directory must not survive the pipeline: it is cleaned up in {@code evaluate()}'s + * {@code finally} block once every query has reached a terminal state. + */ + @Test(timeout = 300_000) + public void tempCheckpointDirIsCleanedUpWhenNoneConfigured() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-tempdir"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + options.setTestMode(false); + // Unset the checkpointDir that streamingOptions() configured, to exercise the fallback path. + options.setCheckpointDir(null); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + Set tempDirsBefore = tempCheckpointDirNames(); + + PipelineResult result = pipeline.run(); + PipelineResult.State finalState = result.waitUntilFinish(); + assertEquals(PipelineResult.State.DONE, finalState); + + Set survivors = tempCheckpointDirNames(); + survivors.removeAll(tempDirsBefore); + assertTrue("leftover temporary checkpoint directories: " + survivors, survivors.isEmpty()); + } + + /** + * Names of {@code beam-spark4-streaming-checkpoint*} entries currently under {@code + * java.io.tmpdir}. + */ + private static Set tempCheckpointDirNames() { + File tmpDir = new File(System.getProperty("java.io.tmpdir")); + String[] names = + tmpDir.list((dir, name) -> name.startsWith("beam-spark4-streaming-checkpoint")); + return names == null ? new HashSet<>() : new HashSet<>(Arrays.asList(names)); + } } From 647cbe18f28b599feed15ca1e6970764cfcbcb18 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 08:54:26 +0000 Subject: [PATCH 14/21] [Spark 4] Name the throwing validation in the stateful streaming translator canTranslate is a predicate in every batch translator, but the streaming stateful translator deliberately throws instead of returning false, a silent decline would route the DoFn to the batch translator and lose the streaming semantics. The throwing checks now live in rejectUnsupported and the override documents the contract. No error message changed. --- .../StatefulParDoStreamingTranslator.java | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java index 01f2dd63c557..ef6df4777c47 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java @@ -69,8 +69,24 @@ public StatefulParDoStreamingTranslator() { super(0.2f); } + /** + * Unlike the batch translators, this override never returns {@code false}: {@link + * #rejectUnsupported} throws with the offending feature named instead. Silently declining here + * would make {@code PipelineTranslator#getSupportedTranslator} fall back to the batch {@code + * ParDo} translator, which would run the stateful {@code DoFn} without the streaming state and + * timer semantics and produce quietly wrong results. + */ @Override protected boolean canTranslate(ParDo.MultiOutput, OutputT> transform) { + rejectUnsupported(transform); + return true; + } + + /** + * Throws when the stateful {@code ParDo} uses a feature this translator does not implement, + * naming that feature; returns normally otherwise. + */ + private void rejectUnsupported(ParDo.MultiOutput, OutputT> transform) { DoFn, OutputT> doFn = transform.getFn(); String stepName = doFn.getClass().getName(); DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn); @@ -94,7 +110,6 @@ protected boolean canTranslate(ParDo.MultiOutput, OutputT> transform) { "side inputs on a stateful ParDo. Broadcasting a side input requires collecting its " + "PCollection, which is not possible while the pipeline is streaming"); } - return true; } @Override From aa30afb89addb105838f950f84b82a9140b3a8ab Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 08:54:27 +0000 Subject: [PATCH 15/21] [Spark 4] Surface a failing leaf query while its siblings still run awaitTermination previously awaited the leaf queries in list order, so a failure in a later query stayed hidden until every earlier query terminated. The queries are now polled round robin with a short timeout, the first failure stops the siblings and is rethrown. Adds a lifecycle test with a healthy leaf that never idles and a poisoned leaf, which hangs under the old sequential await. --- .../StreamingEvaluationContext.java | 49 +++++++++--- .../StreamingPipelineLifecycleTest.java | 79 ++++++++++++++++++- 2 files changed, 115 insertions(+), 13 deletions(-) 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 34bdd8131f13..51f9bfaa6310 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 @@ -24,6 +24,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Comparator; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.UUID; @@ -90,6 +91,11 @@ public class StreamingEvaluationContext extends EvaluationContext { private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + // How long one awaitTermination poll blocks on a single query before moving on to the next one, + // see awaitTermination(List). Short enough to surface a failure promptly, long enough to not + // busy-spin while every query is healthy. + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + private final SparkStructuredStreamingPipelineOptions options; // Guards both `queries` and `stopped` so evaluate() (which appends to `queries` as it starts @@ -166,9 +172,7 @@ public void evaluate() { synchronized (lock) { toAwait = new ArrayList<>(queries); } - for (StreamingQuery query : toAwait) { - awaitTermination(query); - } + awaitTermination(toAwait); } finally { if (idleStopListener != null) { getSparkSession().streams().removeListener(idleStopListener); @@ -220,14 +224,37 @@ private StreamingQuery startQuery( } } - private void awaitTermination(StreamingQuery query) { - try { - query.awaitTermination(); - } catch (StreamingQueryException e) { - LOG.error("Streaming query {} terminated with an exception.", query.id(), e); - // Make sure sibling queries do not keep running once one of them has failed. - stop(); - throw new RuntimeException(e); + /** + * Blocks until every query in {@code toAwait} has terminated, polling them round robin with a + * short timeout rather than awaiting them one after the other: a plain {@link + * StreamingQuery#awaitTermination()} on the first query would block for as long as that query + * keeps running and not surface a failure of a later query until then. Polling bounds the latency + * of surfacing any query's failure by roughly one round of poll timeouts, no matter which query + * fails. + * + *

On the first query that fails, {@link #stop()} makes sure sibling queries do not keep + * running, and the failure is rethrown. Deliberately not {@code + * StreamingQueryManager#awaitAnyTermination}: that relies on the session global {@code + * resetTerminated} bookkeeping and would interfere with concurrent tests sharing one session. + */ + private void awaitTermination(List toAwait) { + List active = new ArrayList<>(toAwait); + while (!active.isEmpty()) { + Iterator iterator = active.iterator(); + while (iterator.hasNext()) { + StreamingQuery query = iterator.next(); + try { + if (query.awaitTermination(AWAIT_POLL_TIMEOUT_MILLIS)) { + // Terminated without an exception, nothing left to await for this query. + iterator.remove(); + } + } catch (StreamingQueryException e) { + LOG.error("Streaming query {} terminated with an exception.", query.id(), e); + // Make sure sibling queries do not keep running once one of them has failed. + stop(); + throw new RuntimeException(e); + } + } } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java index d10c0730a232..335f12be3f71 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import java.io.File; @@ -35,6 +36,7 @@ import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.TimestampedValue; import org.joda.time.Duration; @@ -50,9 +52,10 @@ /** * {@link PipelineResult.State} transitions for a streaming pipeline: {@code RUNNING} immediately * after {@code run()}, {@code DONE} once a naturally idle pipeline's idle-stop listener stops it, - * and {@code CANCELLED} after an explicit {@code cancel()}. + * {@code CANCELLED} after an explicit {@code cancel()}, and {@code FAILED} once any leaf query + * fails. * - *

Both tests here set {@code testMode(false)} on top of {@link + *

The tests here set {@code testMode(false)} on top of {@link * StreamingTestUtils#streamingOptions}, unlike every other test in this package: {@code * SparkStructuredStreamingRunner#run()} calls {@code result.waitUntilFinish()} itself before * returning when {@code testMode} is {@code true} (see its implementation), which would make {@code @@ -195,6 +198,78 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception { } } + /** + * A failure in any leaf query must surface through {@code waitUntilFinish()} even while another + * leaf query is still running. The healthy leaf has idle-stop disabled, so it never terminates on + * its own; only the round robin await in {@code StreamingEvaluationContext} surfaces the poisoned + * leaf's failure promptly and stops the healthy sibling. With a sequential await this test would + * sit on the healthy query until the JUnit timeout whenever that query happens to be awaited + * first. + */ + @Test(timeout = 300_000) + public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-failed"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + options.setTestMode(false); + // Disabled so the healthy query only ever stops because the failure of its sibling is + // surfaced and stop() is called, not because it happened to go idle first. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply( + "ReadHealthy", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + pipeline + .apply( + "ReadPoisoned", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Throw", ParDo.of(new ThrowOnElementDoFn(5))); + + PipelineResult result = pipeline.run(); + + assertThrows(RuntimeException.class, result::waitUntilFinish); + assertEquals(PipelineResult.State.FAILED, result.getState()); + + // The failed pipeline has to have stopped the healthy sibling query too: with idle-stop + // disabled it would otherwise keep running until the JUnit timeout. + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "a sibling query was still active " + + QUERY_START_TIMEOUT_MILLIS + + "ms after the pipeline failed", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + /** Throws on one specific element, passes every other element through unchanged. */ + private static final class ThrowOnElementDoFn extends DoFn { + private final int poisonElement; + + ThrowOnElementDoFn(int poisonElement) { + this.poisonElement = poisonElement; + } + + @ProcessElement + public void processElement(@Element Integer element, OutputReceiver out) { + if (element == poisonElement) { + throw new IllegalStateException("poison element " + poisonElement); + } + out.output(element); + } + } + /** * When no {@code checkpointDir} is configured, {@code StreamingEvaluationContext} falls back to a * {@code beam-spark4-streaming-checkpoint*} temporary directory under {@code java.io.tmpdir}. From 8e72e679b2c750babfee1bc0b561454930cef2db Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 12:22:35 +0000 Subject: [PATCH 16/21] [Spark 4] Recover Beam sources durably from the checkpoint location Restarting a streaming pipeline previously lost all source progress, the source id was a fresh UUID per translation, checkpoint marks lived only in executor memory, and the split list was recomputed per run. The source id is now derived deterministically from the read transform's full name. The first run pins its split list under the checkpoint location, Beam sources do not guarantee deterministic splitting and the split index keys all per split state. Each partition reader persists its checkpoint mark per batch epoch, atomically via write and rename through the Hadoop filesystem of the checkpoint location, retaining the two newest epochs. A reader created without an in memory mark restores the newest durable mark at or before the epoch its batch starts at. The driver epoch counter fast forwards past every offset replayed from Spark's log so a restarted driver never emits an offset that goes backwards. Semantics are at least once, a mark is written when a batch finishes reading rather than transactionally with Spark's commit. --- .../io/streaming/BeamCheckpointFiles.java | 236 ++++++++++++++++++ .../io/streaming/BeamInputPartition.java | 14 ++ .../io/streaming/BeamMicroBatchStream.java | 62 ++++- .../io/streaming/BeamPartitionReader.java | 41 ++- .../io/streaming/BeamReaderCache.java | 35 ++- .../io/streaming/BeamStreamingSource.java | 6 +- .../io/streaming/UnboundedSourceDataset.java | 34 ++- .../streaming/ReadUnboundedTranslator.java | 7 +- .../streaming/BeamCheckpointRecoveryTest.java | 198 +++++++++++++++ .../streaming/BeamMicroBatchSourceTest.java | 3 +- 10 files changed, 614 insertions(+), 22 deletions(-) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java new file mode 100644 index 000000000000..069bb27e9f65 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java @@ -0,0 +1,236 @@ +/* + * 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.io.streaming; + +import java.io.IOException; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FSDataOutputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Durable state of the Beam micro-batch source, stored next to Spark's own streaming state under + * {@code /beam-source-/}. + * + *

Two kinds of state are kept per source. The pinned split list under {@code /splits} + * records the sub sources produced by the first run, Beam sources do not guarantee deterministic + * splitting and the split index is part of the reader cache key, so every later run must reuse the + * first run's splits. The checkpoint marks under {@code /marks//} record the + * read position of one split at the end of the micro-batch that ends at that epoch, which is also + * the position at the start of any batch whose start offset equals that epoch. + * + *

Every file is written to a {@code .tmp} sibling first and then renamed into place, so a + * partially written file is never observed under its final name. + * + *

The Hadoop {@link FileSystem} serving the checkpoint location is resolved from a default + * {@link Configuration}. On executors this means the Hadoop configuration comes from classpath + * defaults rather than from the Spark session, a known limitation of this helper. + */ +public final class BeamCheckpointFiles { + + private static final Logger LOG = LoggerFactory.getLogger(BeamCheckpointFiles.class); + + private static final String ROOT_PREFIX = "beam-source-"; + private static final String SPLITS_FILE = "splits"; + private static final String MARKS_DIR = "marks"; + private static final String TMP_SUFFIX = ".tmp"; + + /** Number of most recent mark files retained per split. */ + private static final int RETAINED_MARKS = 2; + + private BeamCheckpointFiles() {} + + /** + * Reads the pinned split list of {@code sourceId}, or returns {@code null} if no list has been + * pinned yet. + */ + public static @Nullable List readSplits(String checkpointLocation, String sourceId) + throws IOException { + Path path = new Path(root(checkpointLocation, sourceId), SPLITS_FILE); + FileSystem fs = fileSystem(path); + if (!fs.exists(path)) { + return null; + } + @SuppressWarnings("unchecked") + List splits = (List) deserialize(read(fs, path), "pinned split list " + path); + return splits; + } + + /** Pins the split list of {@code sourceId} so later runs reuse exactly these splits. */ + public static void writeSplits(String checkpointLocation, String sourceId, List splitsB64) + throws IOException { + Path path = new Path(root(checkpointLocation, sourceId), SPLITS_FILE); + FileSystem fs = fileSystem(path); + writeAtomically(fs, path, SerializableUtils.serializeToByteArray(new ArrayList<>(splitsB64))); + LOG.info("Pinned {} split(s) of Beam source {} at {}.", splitsB64.size(), sourceId, path); + } + + /** + * Writes the checkpoint mark of one split at the end of the batch ending at {@code endEpoch} and + * then, best effort, deletes mark files older than the two most recent epochs. + * + * @throws IOException if the mark is not {@link Serializable} or the write fails + */ + public static void writeMark( + String checkpointLocation, String sourceId, int splitId, long endEpoch, CheckpointMark mark) + throws IOException { + if (!(mark instanceof Serializable)) { + throw new IOException( + "Checkpoint mark " + + mark.getClass().getName() + + " is not Serializable, it cannot be persisted for durable recovery."); + } + Path dir = marksDir(checkpointLocation, sourceId, splitId); + FileSystem fs = fileSystem(dir); + writeAtomically( + fs, + new Path(dir, Long.toString(endEpoch)), + SerializableUtils.serializeToByteArray((Serializable) mark)); + deleteOldMarks(fs, dir); + } + + /** + * Restores the durable checkpoint mark of one split for a batch starting at {@code startEpoch}. + * + *

The mark written under exactly {@code startEpoch} is preferred, if it is absent the mark + * with the largest epoch not exceeding {@code startEpoch} is used. Returns {@code null}, meaning + * a fresh start, when no such mark exists or reading fails. + */ + public static @Nullable CheckpointMark readMark( + String checkpointLocation, String sourceId, int splitId, long startEpoch) { + Path dir = marksDir(checkpointLocation, sourceId, splitId); + try { + FileSystem fs = fileSystem(dir); + if (!fs.exists(dir)) { + return null; + } + long epoch = Long.MIN_VALUE; + if (fs.exists(new Path(dir, Long.toString(startEpoch)))) { + epoch = startEpoch; + } else { + for (FileStatus status : fs.listStatus(dir)) { + @Nullable Long candidate = parseEpoch(status.getPath().getName()); + if (candidate != null && candidate <= startEpoch && candidate > epoch) { + epoch = candidate; + } + } + } + if (epoch == Long.MIN_VALUE) { + return null; + } + Path path = new Path(dir, Long.toString(epoch)); + CheckpointMark mark = + (CheckpointMark) deserialize(read(fs, path), "durable checkpoint mark " + path); + LOG.info( + "Restored durable checkpoint mark of Beam source {} split {} at epoch {} " + + "(requested epoch {}).", + sourceId, + splitId, + epoch, + startEpoch); + return mark; + } catch (IOException e) { + LOG.warn( + "Failed to read a durable checkpoint mark of Beam source {} split {} at epoch {}, " + + "the reader starts without one.", + sourceId, + splitId, + startEpoch, + e); + return null; + } + } + + private static Path root(String checkpointLocation, String sourceId) { + return new Path(checkpointLocation, ROOT_PREFIX + sourceId); + } + + private static Path marksDir(String checkpointLocation, String sourceId, int splitId) { + return new Path( + new Path(root(checkpointLocation, sourceId), MARKS_DIR), String.valueOf(splitId)); + } + + private static FileSystem fileSystem(Path path) throws IOException { + return path.getFileSystem(new Configuration()); + } + + /** Writes {@code bytes} to a {@code .tmp} sibling of {@code target} and renames it into place. */ + private static void writeAtomically(FileSystem fs, Path target, byte[] bytes) throws IOException { + Path tmp = new Path(target.getParent(), target.getName() + TMP_SUFFIX); + try (FSDataOutputStream out = fs.create(tmp, true)) { + out.write(bytes); + } + // On HDFS a rename onto an existing target fails, a mark rewritten by a task retry is the + // same position, so an existing target counts as success. + if (!fs.rename(tmp, target) && !fs.exists(target)) { + throw new IOException("Failed to rename " + tmp + " to " + target); + } + if (fs.exists(tmp)) { + fs.delete(tmp, false); + } + } + + private static byte[] read(FileSystem fs, Path path) throws IOException { + try (FSDataInputStream in = fs.open(path)) { + return ByteStreams.toByteArray(in); + } + } + + private static Object deserialize(byte[] bytes, String description) { + return SerializableUtils.deserializeFromByteArray(bytes, description); + } + + /** Best effort deletion of mark files older than the {@value #RETAINED_MARKS} highest epochs. */ + private static void deleteOldMarks(FileSystem fs, Path dir) { + try { + List epochs = new ArrayList<>(); + for (FileStatus status : fs.listStatus(dir)) { + @Nullable Long epoch = parseEpoch(status.getPath().getName()); + if (epoch != null) { + epochs.add(epoch); + } + } + Collections.sort(epochs); + for (int i = 0; i < epochs.size() - RETAINED_MARKS; i++) { + fs.delete(new Path(dir, Long.toString(epochs.get(i))), false); + } + } catch (IOException e) { + LOG.warn("Failed to delete outdated checkpoint marks under {}.", dir, e); + } + } + + private static @Nullable Long parseEpoch(String fileName) { + try { + return Long.valueOf(fileName); + } catch (NumberFormatException e) { + return null; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java index 389ed3fec91e..905a5292b57e 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java @@ -35,6 +35,8 @@ public class BeamInputPartition implements InputPartition { private final String sourceId; private final int splitId; private final String checkpointLocation; + private final long startEpoch; + private final long endEpoch; private final int maxRecordsPerMicroBatch; private final long maxBatchDurationMillis; @@ -45,6 +47,8 @@ public class BeamInputPartition implements InputPartition { String sourceId, int splitId, String checkpointLocation, + long startEpoch, + long endEpoch, int maxRecordsPerMicroBatch, long maxBatchDurationMillis) { this.sourceB64 = sourceB64; @@ -53,6 +57,8 @@ public class BeamInputPartition implements InputPartition { this.sourceId = sourceId; this.splitId = splitId; this.checkpointLocation = checkpointLocation; + this.startEpoch = startEpoch; + this.endEpoch = endEpoch; this.maxRecordsPerMicroBatch = maxRecordsPerMicroBatch; this.maxBatchDurationMillis = maxBatchDurationMillis; } @@ -81,6 +87,14 @@ String checkpointLocation() { return checkpointLocation; } + long startEpoch() { + return startEpoch; + } + + long endEpoch() { + return endEpoch; + } + int maxRecordsPerMicroBatch() { return maxRecordsPerMicroBatch; } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java index 33d9198412d6..290faf012bbe 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.spark.structuredstreaming.io.streaming; +import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -42,8 +43,15 @@ * or an explicit {@code StreamingQuery.stop()}). * *

The wrapped source is split exactly once, on the driver, and the resulting sub sources are - * cached so every micro-batch plans the same, stable set of partitions. Splits must be stable - * across micro-batches because the executor side reader cache is keyed by split index. + * pinned to the checkpoint location so every micro-batch of every run plans the same, stable set of + * partitions. Splits must be stable across micro-batches and across restarts because the executor + * side reader cache and the durable checkpoint marks are keyed by split index, and Beam sources do + * not guarantee deterministic splitting. A restarted stream therefore loads the split list written + * by the first run instead of splitting again. + * + *

On a restart Spark replays offsets from its offset log through {@link #deserializeOffset} and + * {@link #planInputPartitions}. The epoch counter fast forwards past every epoch seen there, so + * {@link #latestOffset()} never emits an offset smaller than one already committed to the log. */ public class BeamMicroBatchStream implements MicroBatchStream { @@ -87,7 +95,9 @@ public synchronized Offset latestOffset() { @Override public Offset deserializeOffset(String json) { - return BeamOffset.fromJson(json); + BeamOffset offset = BeamOffset.fromJson(json); + fastForwardEpoch(offset.epoch()); + return offset; } @Override @@ -102,6 +112,9 @@ public void stop() { @Override public InputPartition[] planInputPartitions(Offset start, Offset end) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); List splits = splits(); InputPartition[] partitions = new InputPartition[splits.size()]; for (int i = 0; i < splits.size(); i++) { @@ -113,22 +126,58 @@ public InputPartition[] planInputPartitions(Offset start, Offset end) { sourceId, i, checkpointLocation, + startEpoch, + endEpoch, maxRecordsPerMicroBatch, maxBatchDurationMillis); } return partitions; } + /** + * Raises the epoch counter to {@code seen} if it is behind, keeping {@link #latestOffset()} ahead + * of every offset Spark already logged before a restart. + */ + private synchronized void fastForwardEpoch(long seen) { + if (seen > epoch) { + LOG.info( + "Fast forwarding the epoch of Beam source {} from {} to {} seen in Spark's offset log.", + sourceId, + epoch, + seen); + epoch = seen; + } + } + @Override public PartitionReaderFactory createReaderFactory() { return new BeamPartitionReaderFactory(); } - /** Splits the wrapped source once and memoizes the serialized sub sources. */ + /** + * Returns the pinned split list of this source, loading it from the checkpoint location if a + * previous run pinned one, and splitting the source and pinning the result otherwise. + */ private synchronized List splits() { if (splitsB64 != null) { return splitsB64; } + List pinned; + try { + pinned = BeamCheckpointFiles.readSplits(checkpointLocation, sourceId); + } catch (IOException e) { + throw new IllegalStateException( + "Failed to read the pinned split list of Beam source " + sourceId, e); + } + if (pinned != null) { + LOG.info( + "Restored {} pinned split(s) of Beam source {} from {}.", + pinned.size(), + sourceId, + checkpointLocation); + splitsB64 = pinned; + return pinned; + } UnboundedSource source = BeamStreamingSource.decode(sourceB64, "UnboundedSource"); SerializablePipelineOptions serializableOptions = BeamStreamingSource.decode(pipelineOptionsB64, "PipelineOptions"); @@ -152,6 +201,11 @@ private synchronized List splits() { sourceId, encoded.size(), desiredNumSplits); + try { + BeamCheckpointFiles.writeSplits(checkpointLocation, sourceId, encoded); + } catch (IOException e) { + throw new IllegalStateException("Failed to pin the split list of Beam source " + sourceId, e); + } splitsB64 = encoded; return encoded; } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java index ff6c319f01c8..bc7186584f83 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -62,6 +62,10 @@ public class BeamPartitionReader implements PartitionReader { private final String cacheKey; private final CachedReader cached; private final Coder> windowedValueCoder; + private final String checkpointLocation; + private final String sourceId; + private final int splitId; + private final long endEpoch; private final int maxRecordsPerMicroBatch; private final long maxBatchDurationMillis; @@ -76,12 +80,20 @@ public class BeamPartitionReader implements PartitionReader { BeamStreamingSource.decode(partition.coderB64(), "WindowedValue coder"); SerializablePipelineOptions options = BeamStreamingSource.decode(partition.pipelineOptionsB64(), "PipelineOptions"); + this.checkpointLocation = partition.checkpointLocation(); + this.sourceId = partition.sourceId(); + this.splitId = partition.splitId(); + this.endEpoch = partition.endEpoch(); this.maxRecordsPerMicroBatch = partition.maxRecordsPerMicroBatch(); this.maxBatchDurationMillis = partition.maxBatchDurationMillis(); - this.cacheKey = - BeamReaderCache.key( - partition.checkpointLocation(), partition.sourceId(), partition.splitId()); - this.cached = BeamReaderCache.getOrCreate(cacheKey, source, options.get()); + this.cacheKey = BeamReaderCache.key(checkpointLocation, sourceId, splitId); + long startEpoch = partition.startEpoch(); + this.cached = + BeamReaderCache.getOrCreate( + cacheKey, + source, + options.get(), + () -> BeamCheckpointFiles.readMark(checkpointLocation, sourceId, splitId, startEpoch)); } @Override @@ -119,7 +131,7 @@ public InternalRow get() { /** * Ends the micro-batch. The Beam reader deliberately stays open in {@link BeamReaderCache}, only - * its checkpoint mark is remembered and finalized. + * its checkpoint mark is remembered, persisted for durable recovery, and finalized. */ @Override public void close() { @@ -127,6 +139,7 @@ public void close() { try { UnboundedSource.CheckpointMark mark = cached.reader().getCheckpointMark(); BeamReaderCache.rememberCheckpointMark(cacheKey, mark); + persistMark(mark); mark.finalizeCheckpoint(); } catch (Exception e) { LOG.warn("Failed to finalize the checkpoint mark of Beam reader {}.", cacheKey, e); @@ -134,6 +147,24 @@ public void close() { LOG.debug("Beam reader {} emitted {} record(s) in this micro-batch.", cacheKey, recordsRead); } + /** + * Best effort persistence of {@code mark} under the checkpoint location. An IO failure only + * degrades recovery after a restart, the in memory path in {@link BeamReaderCache} still works, + * so the batch is never failed here. + */ + private void persistMark(UnboundedSource.CheckpointMark mark) { + try { + BeamCheckpointFiles.writeMark(checkpointLocation, sourceId, splitId, endEpoch, mark); + } catch (Exception e) { + LOG.warn( + "Failed to persist the checkpoint mark of Beam reader {} at epoch {}, recovery after a " + + "restart will fall back to an older mark or to a fresh start.", + cacheKey, + endEpoch, + e); + } + } + private InternalRow toRow() { Instant timestamp = cached.reader().getCurrentTimestamp(); WindowedValue windowedValue = diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index 9ccfc3388bd3..a7013186f75d 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -22,6 +22,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; @@ -44,12 +45,15 @@ * stopped. This mirrors what the legacy Spark runner does in {@code * org.apache.beam.runners.spark.io.MicrobatchSource}. * - *

POC scope, weak failure recovery. The {@code CheckpointMark} of every split is cached - * in executor memory only, it is never written to the Spark checkpoint location and never restored - * from it. If an executor JVM dies, or the cache entry expires, the reader is recreated from the - * last mark that happens to still be in memory, and from scratch if there is none. Data loss or - * duplication across executor failures is therefore possible. Making recovery durable means - * committing marks through the Spark offset log and is deliberately out of scope for the POC. + *

Durable recovery. The {@code CheckpointMark} of every split is remembered in executor + * memory after each micro-batch, and {@link BeamPartitionReader} additionally persists it under the + * checkpoint location, see {@link BeamCheckpointFiles}. When a reader has to be created and no mark + * is in memory, for example after an executor or driver restart or after the cache entry expired, + * the caller supplied fallback restores the newest durable mark at or before the epoch the batch + * starts at. Two caveats remain. The source is consumed with at least once semantics, a mark is + * written when a batch finished reading rather than transactionally with Spark's commit, so a crash + * between the two replays the last micro-batch. And persisting a mark is best effort, an IO failure + * only degrades recovery to an older mark or to a fresh start, it never fails the batch. */ public final class BeamReaderCache { @@ -92,15 +96,32 @@ public static String key(String checkpointLocation, String sourceId, int splitId * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if * there is none. */ - @SuppressWarnings({"unchecked", "nullness"}) // the mark type always matches the source public static CachedReader getOrCreate( String key, UnboundedSource source, PipelineOptions options) { + return getOrCreate(key, source, options, () -> null); + } + + /** + * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if + * there is none. The {@code durableMarkFallback} is consulted only when no mark is in memory + * either, it typically restores a mark persisted by {@link BeamCheckpointFiles} and may return + * {@code null} for a fresh start. + */ + @SuppressWarnings({"unchecked", "nullness"}) // the mark type always matches the source + public static CachedReader getOrCreate( + String key, + UnboundedSource source, + PipelineOptions options, + Supplier<@Nullable CheckpointMark> durableMarkFallback) { try { return (CachedReader) READERS.get( key, () -> { CheckpointMarkT mark = (CheckpointMarkT) MARKS.get(key); + if (mark == null) { + mark = (CheckpointMarkT) durableMarkFallback.get(); + } LOG.info( "No cached Beam reader for {}, creating one at checkpoint mark {}.", key, mark); return new CachedReader<>(source.createReader(options, mark)); diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java index d057d39bac55..ede523604c76 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java @@ -80,7 +80,11 @@ public class BeamStreamingSource implements TableProvider, DataSourceRegister { */ public static final String OPT_PIPELINE_OPTIONS = "beam.pipelineoptions"; - /** Identifier making the reader cache key unique per source instance. */ + /** + * Deterministic identifier of the source, derived from the full name of the read transform, see + * {@link UnboundedSourceDataset#sourceId}. It keys the reader cache and the durable checkpoint + * state, so it must stay stable across restarts against the same checkpoint location. + */ public static final String OPT_SOURCE_ID = "beam.sourceid"; /** Desired number of splits handed to {@code UnboundedSource.split}. */ diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java index 7e4af6687cfc..62a0f8d52f3a 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -17,14 +17,15 @@ */ package org.apache.beam.runners.spark.structuredstreaming.io.streaming; +import java.nio.charset.StandardCharsets; import java.util.HashMap; import java.util.Map; -import java.util.UUID; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; @@ -55,6 +56,9 @@ public final class UnboundedSourceDataset { /** Upper bound on the number of splits requested from a source, keeps the POC predictable. */ private static final int MAX_DESIRED_SPLITS = 8; + /** Upper bound on the sanitized transform name inside a source id, see {@link #sourceId}. */ + private static final int MAX_SANITIZED_NAME_LENGTH = 64; + private UnboundedSourceDataset() {} /** @@ -66,6 +70,8 @@ private UnboundedSourceDataset() {} * @param windowedValueCoder the coder used to encode the {@value #COL_PAYLOAD} column, normally a * {@code WindowedValues.FullWindowedValueCoder} * @param options the pipeline options, supplying the watermark delay and the micro-batch limits + * @param transformName the full name of the read transform, turned into the deterministic source + * id keying the durable checkpoint state, see {@link #sourceId} * @param the element type of the source * @param the checkpoint mark type of the source */ @@ -73,7 +79,8 @@ public static Datase SparkSession session, UnboundedSource source, Coder> windowedValueCoder, - SparkStructuredStreamingPipelineOptions options) { + SparkStructuredStreamingPipelineOptions options, + String transformName) { Map readerOptions = new HashMap<>(); readerOptions.put(BeamStreamingSource.OPT_SOURCE, BeamStreamingSource.encode(source)); @@ -82,7 +89,7 @@ public static Datase readerOptions.put( BeamStreamingSource.OPT_PIPELINE_OPTIONS, BeamStreamingSource.encode(new SerializablePipelineOptions(options))); - readerOptions.put(BeamStreamingSource.OPT_SOURCE_ID, UUID.randomUUID().toString()); + readerOptions.put(BeamStreamingSource.OPT_SOURCE_ID, sourceId(transformName)); readerOptions.put( BeamStreamingSource.OPT_NUM_SPLITS, Integer.toString(desiredNumSplits(session))); readerOptions.put( @@ -99,6 +106,27 @@ public static Datase return rows.withWatermark(COL_EVENT_TS, options.getWatermarkDelayMillis() + " milliseconds"); } + /** + * Derives the deterministic source id of a read transform from its full name. + * + *

The id keys the durable checkpoint state under the checkpoint location, so it must be + * filesystem safe and stable across JVMs. Characters outside {@code [A-Za-z0-9._-]} are replaced + * with {@code _}, the sanitized name is truncated to {@value #MAX_SANITIZED_NAME_LENGTH} + * characters, and an 8 hex character hash of the unsanitized name is appended to keep distinct + * transform names distinct. A pipeline restarted against the same checkpoint location must + * therefore keep the same transform names, the same requirement every Beam runner with durable + * state imposes. + */ + public static String sourceId(String transformName) { + String sanitized = transformName.replaceAll("[^A-Za-z0-9._-]", "_"); + if (sanitized.length() > MAX_SANITIZED_NAME_LENGTH) { + sanitized = sanitized.substring(0, MAX_SANITIZED_NAME_LENGTH); + } + String hash = + Hashing.murmur3_32_fixed().hashString(transformName, StandardCharsets.UTF_8).toString(); + return sanitized + "-" + hash; + } + private static int desiredNumSplits(SparkSession session) { int parallelism = session.sparkContext().defaultParallelism(); return Math.max(1, Math.min(MAX_DESIRED_SPLITS, parallelism)); diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java index 36cf305e3376..9f1368e19e8b 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java @@ -82,7 +82,12 @@ protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Co cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); Dataset rows = - UnboundedSourceDataset.of(cxt.getSparkSession(), source, payloadCoder, options); + UnboundedSourceDataset.of( + cxt.getSparkSession(), + source, + payloadCoder, + options, + cxt.getCurrentTransform().getFullName()); Encoder> encoder = cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java new file mode 100644 index 000000000000..4b7d7c55a00e --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java @@ -0,0 +1,198 @@ +/* + * 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.io.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.Serializable; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for the durable checkpoint recovery pieces of the Spark 4 micro-batch source, the {@link + * BeamCheckpointFiles} layout, the epoch fast forward of {@link BeamMicroBatchStream} and the + * deterministic source id of {@link UnboundedSourceDataset}. + */ +@RunWith(JUnit4.class) +public class BeamCheckpointRecoveryTest { + + private static final String SOURCE_ID = "read-source-cafe0123"; + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + @Test + public void testSplitsRoundTrip() throws Exception { + String checkpointLocation = temp.newFolder("splits").getAbsolutePath(); + List splits = Arrays.asList("c3BsaXQtMA==", "c3BsaXQtMQ==", "c3BsaXQtMg=="); + + assertNull( + "no splits pinned yet", BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); + + BeamCheckpointFiles.writeSplits(checkpointLocation, SOURCE_ID, splits); + assertEquals(splits, BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); + + File root = new File(checkpointLocation, "beam-source-" + SOURCE_ID); + assertEquals( + "only the final splits file may remain", + new HashSet<>(Arrays.asList("splits")), + fileNames(root)); + } + + @Test + public void testMarkRetentionAndRecovery() throws Exception { + String checkpointLocation = temp.newFolder("marks").getAbsolutePath(); + for (long epoch = 1; epoch <= 4; epoch++) { + BeamCheckpointFiles.writeMark( + checkpointLocation, SOURCE_ID, 0, epoch, new TestMark((int) epoch)); + } + + File marksDir = new File(new File(checkpointLocation, "beam-source-" + SOURCE_ID), "marks/0"); + assertEquals( + "retention must keep the two highest epochs only", + new HashSet<>(Arrays.asList("3", "4")), + fileNames(marksDir)); + + assertEquals(4, restoredPosition(checkpointLocation, 0, 4)); + assertEquals(3, restoredPosition(checkpointLocation, 0, 3)); + // No exact file for epoch 7, the largest epoch not exceeding it is 4. + assertEquals(4, restoredPosition(checkpointLocation, 0, 7)); + // Epochs 1 and 2 were deleted by retention, nothing at or below 2 is left. + assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 0, 2)); + // A different split has no marks at all. + assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 1, 4)); + } + + @Test + public void testEpochFastForwardsPastDeserializedOffset() throws Exception { + BeamMicroBatchStream stream = stream(temp.newFolder("ff-offset").getAbsolutePath()); + stream.deserializeOffset("{\"epoch\":7}"); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the replayed epoch 7, got " + next, next.epoch() > 7L); + } + + @Test + public void testEpochFastForwardsPastPlannedOffsets() throws Exception { + BeamMicroBatchStream stream = stream(temp.newFolder("ff-plan").getAbsolutePath()); + InputPartition[] partitions = + stream.planInputPartitions(new BeamOffset(3L), new BeamOffset(9L)); + assertTrue("at least one partition expected", partitions.length > 0); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the planned epoch 9, got " + next, next.epoch() > 9L); + } + + @Test + public void testSourceIdIsDeterministicAndFilesystemSafe() { + String name = "Read from PubSub/PubsubUnboundedSource (with: spaces & colons)"; + String id = UnboundedSourceDataset.sourceId(name); + + assertEquals("same name must give the same id", id, UnboundedSourceDataset.sourceId(name)); + assertNotEquals( + "different names must give different ids", + UnboundedSourceDataset.sourceId("Read A"), + UnboundedSourceDataset.sourceId("Read B")); + assertTrue("id must be filesystem safe: " + id, id.matches("[A-Za-z0-9._-]+")); + + String longName = String.join("/", java.util.Collections.nCopies(50, "NestedTransform")); + String longId = UnboundedSourceDataset.sourceId(longName); + assertTrue("id length must stay bounded: " + longId.length(), longId.length() <= 73); + assertNotEquals( + "truncated names must still differ through the hash", + longId, + UnboundedSourceDataset.sourceId(longName + "/Tail")); + } + + private static int restoredPosition(String checkpointLocation, int splitId, long startEpoch) { + CheckpointMark mark = + BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, splitId, startEpoch); + assertNotNull("expected a durable mark at or before epoch " + startEpoch, mark); + return ((TestMark) mark).position; + } + + /** + * Lists the visible files of {@code dir}, asserting no {@code .tmp} file was left behind and + * ignoring the hidden checksum sidecars of Hadoop's local filesystem. + */ + private static Set fileNames(File dir) { + String[] names = dir.list(); + assertNotNull("expected directory " + dir, names); + Set visible = new HashSet<>(); + for (String name : names) { + assertFalse("no temporary file may remain: " + name, name.endsWith(".tmp")); + if (!name.startsWith(".")) { + visible.add(name); + } + } + return visible; + } + + /** Builds a stream over a real source, driver side only, no Spark session required. */ + private static BeamMicroBatchStream stream(String checkpointLocation) { + Map options = new HashMap<>(); + options.put( + BeamStreamingSource.OPT_SOURCE, BeamStreamingSource.encode(CountingSource.unbounded())); + options.put( + BeamStreamingSource.OPT_CODER, + BeamStreamingSource.encode( + WindowedValues.getFullCoder( + org.apache.beam.sdk.coders.VarLongCoder.of(), GlobalWindow.Coder.INSTANCE))); + options.put( + BeamStreamingSource.OPT_PIPELINE_OPTIONS, + BeamStreamingSource.encode( + new SerializablePipelineOptions(PipelineOptionsFactory.create()))); + options.put(BeamStreamingSource.OPT_SOURCE_ID, SOURCE_ID); + options.put(BeamStreamingSource.OPT_NUM_SPLITS, "2"); + return new BeamMicroBatchStream(new CaseInsensitiveStringMap(options), checkpointLocation); + } + + /** A trivial serializable checkpoint mark carrying a read position. */ + private static class TestMark implements CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; + + private final int position; + + TestMark(int position) { + this.position = position; + } + + @Override + public void finalizeCheckpoint() {} + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java index 38b27db1f4b5..f9c0574c6874 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -229,7 +229,8 @@ private Dataset rows(int count, long watermarkDelayMillis) { options.setWatermarkDelayMillis(watermarkDelayMillis); options.setMaxRecordsPerMicroBatch(1000); options.setMaxBatchDurationMillis(200L); - return UnboundedSourceDataset.of(SESSION.getSession(), new ListSource(count), coder(), options); + return UnboundedSourceDataset.of( + SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); } private static Coder> coder() { From 715c3b6b0bcccbd28c69544aee3035ba68d7c7eb Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 12:43:05 +0000 Subject: [PATCH 17/21] [Spark 4] Prove durable restart recovery end to end Runs two pipelines with identical transform names against one checkpoint directory, clearing the in memory reader cache in between to simulate a fresh JVM. The second run must resume from the durable checkpoint marks, re reading the whole range would mean recovery regressed. ListBackedUnboundedSource gains a positional checkpoint mark so a resumed reader continues after the last emitted element. --- .../StreamingCheckpointRestartTest.java | 222 ++++++++++++++++++ .../streaming/StreamingTestUtils.java | 74 +++--- 2 files changed, 265 insertions(+), 31 deletions(-) create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java new file mode 100644 index 000000000000..6cef4c138593 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java @@ -0,0 +1,222 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.io.Serializable; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.Stream; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.TimestampedValue; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End to end proof that a Spark 4 structured streaming pipeline restarted against the same + * checkpoint location resumes an {@link org.apache.beam.sdk.io.UnboundedSource.UnboundedReader} + * from its durable checkpoint marks instead of re-reading the source from scratch. + * + *

Two independent {@link Pipeline}s are run one after the other against the same checkpoint + * directory, with identical read transform names so they derive the same deterministic {@code + * sourceId} (see {@link UnboundedSourceDataset#sourceId}). {@link BeamReaderCache#invalidateAll} is + * called between the two runs to simulate a fresh JVM: every in-memory reader and checkpoint mark + * is dropped, forcing the second run's readers to fall back to the durable marks persisted by + * {@link org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamCheckpointFiles} during + * the first run. + * + *

{@link StreamingTestUtils.ListBackedUnboundedSource} now carries a positional checkpoint mark, + * see {@link StreamingTestUtils.ListBackedUnboundedSource.Mark}, so a reader created from a + * restored mark resumes right after the last element it had emitted rather than always starting at + * the beginning of its split. + * + *

The second run's assertion is deliberately weak, per the at-least-once semantics documented on + * {@code BeamCheckpointFiles} and {@code BeamReaderCache}: Spark may replay the last micro-batch + * that was in flight when the query stopped, so a handful of already-seen elements MAY reappear. + * What must never happen is the second run re-reading the whole range, which is what a + * regression back to a fresh, non-durable source id or in-memory-only marks would look like. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingCheckpointRestartTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final int ELEMENT_COUNT = 10; + private static final String READ_TRANSFORM_NAME = "ReadUnbounded"; + + @Test(timeout = 300_000) + public void restartedPipelineResumesFromDurableCheckpointMarks() throws Exception { + String checkpointPath = checkpointDir.newFolder("checkpoint").getAbsolutePath(); + + String collectorA = StreamingTestUtils.newCollectorId("checkpoint-restart-a"); + String collectorB = StreamingTestUtils.newCollectorId("checkpoint-restart-b"); + StreamingTestUtils.clear(collectorA); + StreamingTestUtils.clear(collectorB); + + // First run: read the whole 0..9 range once, from a fresh checkpoint location. + PipelineResult first = runPipeline(checkpointPath, collectorA); + first.waitUntilFinish(); + + List collectedA = + new ArrayList<>(StreamingTestUtils.getCollected(collectorA)); + Collections.sort(collectedA); + assertEquals( + "first run (pipeline state=" + first.getState() + ") must read every element", + fullRange(), + collectedA); + + // The durable layout pinned by the first run must exist: the pinned split list and at least + // one persisted checkpoint mark, both keyed by the deterministic source id of "ReadUnbounded". + String sourceId = UnboundedSourceDataset.sourceId(READ_TRANSFORM_NAME); + File sourceRoot = findSourceRoot(new File(checkpointPath, "0"), sourceId); + assertNotNull( + "expected a beam-source-" + sourceId + " directory under " + checkpointPath + "/0", + sourceRoot); + File splitsFile = new File(sourceRoot, "splits"); + assertTrue("pinned splits file must exist: " + splitsFile, splitsFile.isFile()); + assertTrue( + "at least one split's marks directory must contain a persisted mark file", + hasAnyMarkFile(new File(sourceRoot, "marks"))); + + // Simulate a fresh JVM: drop every in-memory reader and checkpoint mark, forcing the next + // reader to fall back to the durable marks just asserted above. + BeamReaderCache.invalidateAll(); + + // Second run: identical transform name, identical checkpoint location, identical source + // content. A correct restart must resume from the durable marks rather than re-reading + // everything. + PipelineResult second = runPipeline(checkpointPath, collectorB); + second.waitUntilFinish(); + + List collectedB = + new ArrayList<>(StreamingTestUtils.getCollected(collectorB)); + Set seenB = new HashSet<>(collectedB); + assertFalse( + "a restarted run (pipeline state=" + + second.getState() + + ") must not re-read the whole range from scratch, saw " + + collectedB, + seenB.containsAll(fullRange())); + } + + private PipelineResult runPipeline(String checkpointPath, String collectorId) throws IOException { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.as(SparkStructuredStreamingPipelineOptions.class); + options.setRunner(SparkStructuredStreamingRunner.class); + options.setTestMode(true); + options.setStreaming(true); + options.setStreamingStopAfterIdleBatches(3); + options.setMaxBatchDurationMillis(200); + // One element per split per micro-batch: the 10 elements are spread across several + // micro-batches instead of a single one, so even a replayed in-flight batch can only ever + // carry a small tail of the range, never all of it. + options.setMaxRecordsPerMicroBatch(1); + options.setCheckpointDir(checkpointPath); + SESSION.configure(options); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + READ_TRANSFORM_NAME, + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>(elements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + return pipeline.run(); + } + + private static List> elements() { + List> elements = new ArrayList<>(ELEMENT_COUNT); + for (int i = 0; i < ELEMENT_COUNT; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + return elements; + } + + private static List fullRange() { + List range = new ArrayList<>(ELEMENT_COUNT); + for (int i = 0; i < ELEMENT_COUNT; i++) { + range.add(i); + } + return range; + } + + /** Recursively searches {@code root} for a directory named {@code beam-source-}. */ + private static @Nullable File findSourceRoot(File root, String sourceId) throws IOException { + if (!root.isDirectory()) { + return null; + } + String target = "beam-source-" + sourceId; + try (Stream paths = Files.walk(root.toPath())) { + return paths + .filter(Files::isDirectory) + .filter(path -> target.equals(path.getFileName().toString())) + .findFirst() + .map(Path::toFile) + .orElse(null); + } + } + + /** {@code true} if any split subdirectory of {@code marksRoot} contains a mark file. */ + private static boolean hasAnyMarkFile(File marksRoot) { + File[] splitDirs = marksRoot.listFiles(); + if (splitDirs == null) { + return false; + } + for (File splitDir : splitDirs) { + File[] markFiles = splitDir.listFiles(); + if (markFiles != null && markFiles.length > 0) { + return true; + } + } + return false; + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java index 2da0458dbc07..75c657ef39c6 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; import java.io.IOException; +import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -29,6 +30,7 @@ import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -209,9 +211,9 @@ public static SparkStructuredStreamingPipelineOptions streamingOptions( * terminate, see the {@link StreamingTestUtils} class javadoc). Supports explicit, possibly * out-of-order event timestamps so tests can inject late data. * - *

Uses {@link org.apache.beam.sdk.io.UnboundedSource.CheckpointMark#NOOP_CHECKPOINT_MARK}: - * this source never needs to resume a partially read split from a durable checkpoint, so there is - * no state worth persisting between micro-batches beyond the in-memory read position. + *

Checkpoint marks record the index of the next element to read, see {@link Mark}, so a reader + * created from a mark resumes right after the last element that mark's reader had emitted. A + * reader created with a {@code null} mark starts at the beginning of its split, as before. * *

Elements are stored pre-encoded (via {@code coder}) as {@code byte[]} plus a {@code long} * timestamp rather than kept as {@link TimestampedValue} objects, because this source (like any @@ -224,7 +226,7 @@ public static SparkStructuredStreamingPipelineOptions streamingOptions( * for correctness, round robin merely spreads elements across splits close to evenly. */ public static final class ListBackedUnboundedSource - extends UnboundedSource { + extends UnboundedSource { private final List encodedElements; private final List timestampsMillis; @@ -250,8 +252,8 @@ private ListBackedUnboundedSource( } @Override - public List> - split(int desiredNumSplits, PipelineOptions options) { + public List> split( + int desiredNumSplits, PipelineOptions options) { int numElements = encodedElements.size(); if (numElements == 0 || desiredNumSplits <= 1) { return Collections.singletonList(this); @@ -278,15 +280,14 @@ private ListBackedUnboundedSource( } @Override - public UnboundedReader createReader( - PipelineOptions options, - UnboundedSource.CheckpointMark.@Nullable NoopCheckpointMark checkpointMark) { - return new ListBackedUnboundedReader<>(this); + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark checkpointMark) { + return new ListBackedUnboundedReader<>( + this, checkpointMark == null ? 0 : checkpointMark.next); } @Override - public Coder getCheckpointMarkCoder() { - return new NoopCheckpointMarkCoder(); + public Coder getCheckpointMarkCoder() { + return SerializableCoder.of(Mark.class); } @Override @@ -312,11 +313,16 @@ private static T decode(Coder coder, byte[] bytes) { private static final class ListBackedUnboundedReader extends UnboundedReader { private final ListBackedUnboundedSource source; - private int index = -1; + private int index; private Instant maxTimestampSeen = BoundedWindow.TIMESTAMP_MIN_VALUE; - ListBackedUnboundedReader(ListBackedUnboundedSource source) { + /** + * @param nextIndex the index of the first element this reader has not yet emitted, {@code 0} + * for a fresh start or the value carried by a resumed {@link Mark} + */ + ListBackedUnboundedReader(ListBackedUnboundedSource source, int nextIndex) { this.source = source; + this.index = nextIndex - 1; } @Override @@ -368,8 +374,10 @@ public Instant getWatermark() { } @Override - public UnboundedSource.CheckpointMark.NoopCheckpointMark getCheckpointMark() { - return UnboundedSource.CheckpointMark.NOOP_CHECKPOINT_MARK; + public Mark getCheckpointMark() { + // index + 1 is the index of the next element to read, i.e. the position a resumed reader + // must start at. + return new Mark(index + 1); } @Override @@ -380,23 +388,27 @@ public UnboundedSource.CheckpointMark.NoopCheckpointMark getCheckpointMark() { @Override public void close() throws IOException {} } - } - /** - * A trivial coder for the stateless {@link UnboundedSource.CheckpointMark.NoopCheckpointMark}. - */ - private static final class NoopCheckpointMarkCoder - extends org.apache.beam.sdk.coders.AtomicCoder< - UnboundedSource.CheckpointMark.NoopCheckpointMark> { - @Override - public void encode( - UnboundedSource.CheckpointMark.NoopCheckpointMark value, java.io.OutputStream outStream) { - // Nothing to persist. - } + /** + * The read position of one {@link ListBackedUnboundedSource} split: the index of the next + * element a resumed reader must emit. {@code 0} means nothing was read yet. + */ + public static final class Mark implements UnboundedSource.CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; - @Override - public UnboundedSource.CheckpointMark.NoopCheckpointMark decode(java.io.InputStream inStream) { - return UnboundedSource.CheckpointMark.NOOP_CHECKPOINT_MARK; + private final int next; + + Mark(int next) { + this.next = next; + } + + @Override + public void finalizeCheckpoint() {} + + @Override + public String toString() { + return "Mark{next=" + next + '}'; + } } } } From be7e3ce51fb85a21137e6d6f148630786ae0769a Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 20:16:32 +0000 Subject: [PATCH 18/21] [Spark 4] Support PAssert in streaming via end-of-stream watermark advancement and streaming Impulse --- .../io/streaming/BeamPartitionReader.java | 28 +++ .../io/streaming/BeamReaderCache.java | 11 + .../PipelineTranslatorStreaming.java | 16 +- .../streaming/ImpulseStreamingTranslator.java | 237 ++++++++++++++++++ .../streaming/ReadUnboundedTranslator.java | 12 + .../StreamingTranslationHelpers.java | 13 +- .../streaming/PAssertStreamingTest.java | 128 ++++++++++ .../streaming/StreamingTestUtils.java | 24 +- 8 files changed, 459 insertions(+), 10 deletions(-) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java index bc7186584f83..a66a5bbe82c3 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -23,6 +23,8 @@ import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +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; @@ -116,6 +118,24 @@ public boolean next() throws IOException { current = toRow(); return true; } + if (recordsRead > 0) { + // Cut off the micro-batch at the first empty poll: once we have ingested data in this + // batch, + // finish it immediately so the data batch is processed and its natural event watermark + // declared + // before any end-of-stream sentinel row is emitted in a subsequent empty batch. + current = null; + return false; + } + if (!cached.hasEmittedSentinel()) { + Instant watermark = cached.reader().getWatermark(); + if (watermark != null && !watermark.isBefore(GlobalWindow.INSTANCE.maxTimestamp())) { + cached.markSentinelEmitted(); + recordsRead++; + current = sentinelRow(); + return true; + } + } Uninterruptibles.sleepUninterruptibly( Math.min(remaining, POLL_INTERVAL_MILLIS), java.util.concurrent.TimeUnit.MILLISECONDS); } @@ -178,4 +198,12 @@ private InternalRow toRow() { // Spark stores TimestampType as microseconds since the epoch. return new GenericInternalRow(new Object[] {payload, timestamp.getMillis() * 1000L}); } + + private InternalRow sentinelRow() { + // Spark stores TimestampType as microseconds since the epoch. The sentinel carries an empty + // payload that is filtered before reaching user transforms, and TIMESTAMP_MAX_VALUE + // which pushes Spark's data-driven watermark to infinity. + return new GenericInternalRow( + new Object[] {new byte[0], BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis() * 1000L}); + } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index a7013186f75d..60a8c76cba00 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -150,6 +150,7 @@ public static void invalidateAll() { public static final class CachedReader implements Closeable { private final UnboundedReader reader; private boolean started; + private boolean sentinelEmitted; CachedReader(UnboundedReader reader) { this.reader = reader; @@ -172,6 +173,16 @@ public synchronized boolean startOrAdvance() throws IOException { return reader.advance(); } + /** Returns whether an end-of-stream sentinel row has already been emitted for this reader. */ + public synchronized boolean hasEmittedSentinel() { + return sentinelEmitted; + } + + /** Marks that an end-of-stream sentinel row has been emitted for this reader. */ + public synchronized void markSentinelEmitted() { + this.sentinelEmitted = true; + } + @Override public void close() throws IOException { reader.close(); diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java index 7901fb9f4a31..c55801a4ccf6 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -21,11 +21,13 @@ import org.apache.beam.runners.spark.SparkCommonPipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.translation.batch.PipelineTranslatorBatch; import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.GroupByKeyStreamingTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ImpulseStreamingTranslator; 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.transforms.Combine; 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; @@ -42,10 +44,10 @@ * *

This extends {@link PipelineTranslatorBatch} purely for reuse: {@link * PipelineTranslatorBatch#getTransformTranslator} and its private registry are the only way to - * reach the (package-private) batch translators for {@code Impulse}, {@code Window.Assign}, {@code - * Flatten}, {@code Reshuffle}, the bounded read, and stateless {@code ParDo}, all of which are - * reused completely unchanged for streaming. This class only intercepts the handful of transforms - * that need genuinely different, streaming-aware handling before falling back to {@code super}. + * reach the (package-private) batch translators for {@code Window.Assign}, {@code Flatten}, {@code + * Reshuffle}, the bounded read, and stateless {@code ParDo}, all of which are reused completely + * unchanged for streaming. This class only intercepts the handful of transforms that need genuinely + * different, streaming-aware handling before falling back to {@code super}. */ @Internal public class PipelineTranslatorStreaming extends PipelineTranslatorBatch { @@ -61,6 +63,10 @@ TransformTranslator getTransformTranslator(TransformT tra return (TransformTranslator) new ReadUnboundedTranslator<>(); } + if (transform instanceof Impulse) { + return (TransformTranslator) new ImpulseStreamingTranslator(); + } + if (transform instanceof GroupByKey) { return (TransformTranslator) new GroupByKeyStreamingTranslator<>(); } @@ -82,7 +88,7 @@ TransformTranslator getTransformTranslator(TransformT tra // Stateless ParDo falls through to super, reusing ParDoTranslatorBatch unchanged. } - // Impulse, Window.Assign, Flatten, Reshuffle, the bounded read, and stateless ParDo: reused + // Window.Assign, Flatten, Reshuffle, the bounded read, and stateless ParDo: reused // unchanged from the batch registry. return super.getTransformTranslator(transform); } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java new file mode 100644 index 000000000000..30256b79bf2d --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java @@ -0,0 +1,237 @@ +/* + * 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 static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY; +import static org.apache.spark.sql.functions.col; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +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.AtomicCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.spark.api.java.function.FilterFunction; +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.Row; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Streaming translator for {@link Impulse}. + * + *

Emits a single empty byte array in a streaming micro-batch, advancing the watermark to + * infinity upon completion so that pipelines mixing Impulse-generated elements (e.g. {@code + * Create.of} in {@code PAssert}) with unbounded streams can safely union without schema or + * batch/streaming mismatch. + */ +public class ImpulseStreamingTranslator + extends TransformTranslator, Impulse> { + + public ImpulseStreamingTranslator() { + super(0.05f); + } + + @Override + protected void translate(Impulse transform, Context cxt) { + PCollection output = cxt.getOutput(); + Coder elementCoder = output.getCoder(); + + WindowedValues.FullWindowedValueCoder payloadCoder = + WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + SparkStructuredStreamingPipelineOptions options = + cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); + + Dataset rows = + UnboundedSourceDataset.of( + cxt.getSparkSession(), + new ImpulseSource(), + payloadCoder, + options, + cxt.getCurrentTransform().getFullName()); + + Encoder> encoder = + cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + Dataset> dataset = + rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) + .as(Encoders.BINARY()) + .filter(new FilterNonEmptyPayload()) + .map(new DecodePayload<>(payloadCoder), encoder); + + cxt.putDataset(output, dataset); + } + + private static final class FilterNonEmptyPayload implements FilterFunction { + @Override + public boolean call(byte[] payload) { + return payload != null && payload.length > 0; + } + } + + private static final class DecodePayload implements MapFunction> { + private final Coder> coder; + + DecodePayload(Coder> coder) { + this.coder = coder; + } + + @Override + public WindowedValue call(byte[] payload) { + return CoderHelpers.fromByteArray(payload, coder); + } + } + + /** An unbounded source that produces exactly one empty byte array and then finishes. */ + private static final class ImpulseSource + extends UnboundedSource { + + @Override + public List> split( + int desiredNumSplits, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public UnboundedReader createReader( + PipelineOptions options, @Nullable ImpulseCheckpointMark checkpointMark) { + return new ImpulseReader(this, checkpointMark != null && checkpointMark.done); + } + + @Override + public Coder getCheckpointMarkCoder() { + return ImpulseCheckpointMarkCoder.of(); + } + + @Override + public Coder getOutputCoder() { + return ByteArrayCoder.of(); + } + + static final class ImpulseCheckpointMark + implements UnboundedSource.CheckpointMark, Serializable { + final boolean done; + + ImpulseCheckpointMark(boolean done) { + this.done = done; + } + + @Override + public void finalizeCheckpoint() {} + } + + static final class ImpulseCheckpointMarkCoder extends AtomicCoder { + private static final ImpulseCheckpointMarkCoder INSTANCE = new ImpulseCheckpointMarkCoder(); + + public static ImpulseCheckpointMarkCoder of() { + return INSTANCE; + } + + @Override + public void encode(ImpulseCheckpointMark value, OutputStream outStream) throws IOException { + outStream.write(value.done ? 1 : 0); + } + + @Override + public ImpulseCheckpointMark decode(InputStream inStream) throws IOException { + return new ImpulseCheckpointMark(inStream.read() != 0); + } + } + + private static final class ImpulseReader extends UnboundedReader { + private final ImpulseSource source; + private boolean done; + private boolean started; + + ImpulseReader(ImpulseSource source, boolean done) { + this.source = source; + this.done = done; + } + + @Override + public boolean start() { + if (done) { + return false; + } + started = true; + done = true; + return true; + } + + @Override + public boolean advance() { + return false; + } + + @Override + public byte[] getCurrent() { + if (!started) { + throw new NoSuchElementException(); + } + return EMPTY_BYTE_ARRAY; + } + + @Override + public Instant getCurrentTimestamp() { + if (!started) { + throw new NoSuchElementException(); + } + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public void close() {} + + @Override + public Instant getWatermark() { + return done ? BoundedWindow.TIMESTAMP_MAX_VALUE : BoundedWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public CheckpointMark getCheckpointMark() { + return new ImpulseCheckpointMark(done); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java index 9f1368e19e8b..559318e33a58 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java @@ -31,6 +31,7 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.spark.api.java.function.FilterFunction; import org.apache.spark.api.java.function.MapFunction; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Encoder; @@ -95,11 +96,22 @@ protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Co Dataset> dataset = rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) .as(Encoders.BINARY()) + .filter(new FilterNonEmptyPayload()) .map(new DecodePayload<>(payloadCoder), encoder); cxt.putDataset(output, dataset); } + /** + * Filters out empty sentinel payloads emitted to advance the watermark upon stream exhaustion. + */ + private static final class FilterNonEmptyPayload implements FilterFunction { + @Override + public boolean call(byte[] payload) { + return payload != null && payload.length > 0; + } + } + /** Decodes the binary payload column back into a Beam {@code WindowedValue}. */ private static final class DecodePayload implements MapFunction> { private final Coder> coder; diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java index 6b5e4878613f..b9d3649e42ed 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java @@ -29,6 +29,7 @@ import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; import org.apache.beam.sdk.transforms.windowing.AfterWatermark; import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.transforms.windowing.Never; import org.apache.beam.sdk.transforms.windowing.Trigger; import org.apache.beam.sdk.transforms.windowing.WindowFn; import org.apache.beam.sdk.values.KV; @@ -87,9 +88,10 @@ static void checkSupportedWindowing(WindowingStrategy strategy, String ste } /** - * The default trigger, as Beam models it, is either {@link DefaultTrigger} itself or the - * equivalent {@code AfterWatermark.pastEndOfWindow()} without early or late firings, which is - * what {@code WindowingStrategy} may normalise it to. + * The default trigger, as Beam models it, is either {@link DefaultTrigger} itself, the equivalent + * {@code AfterWatermark.pastEndOfWindow()} without early or late firings, or {@link + * Never.NeverTrigger} used by {@code PAssert} on unbounded streams to gather all window contents + * before the end-of-stream watermark triggers final pane evaluation. */ private static boolean isDefaultTrigger(Trigger trigger) { if (trigger instanceof DefaultTrigger) { @@ -100,6 +102,11 @@ private static boolean isDefaultTrigger(Trigger trigger) { // configured Beam represents it as one of the AfterWatermarkEarlyAndLate subclasses instead. return true; } + if (trigger instanceof Never.NeverTrigger) { + // NeverTrigger is used by PAssert on unbounded PCollections to prevent intermediate firings. + // In this runner, end-of-stream watermark advancement fires the final pane. + return true; + } return false; } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java new file mode 100644 index 000000000000..6d0a5b6c1b3b --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java @@ -0,0 +1,128 @@ +/* + * 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 static org.junit.Assert.assertThrows; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End-to-end tests verifying {@link PAssert} evaluations against the Spark 4 Structured Streaming + * runner. + * + *

PAssert regroups elements and verifies expectations once a window closes upon watermark + * passage. In streaming, this relies on the end-of-stream sentinel row emitted when unbounded + * readers reach exhaustion to advance Spark's data-driven watermark past the global window end. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class PAssertStreamingTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + + private static class DoubleFn extends DoFn { + @ProcessElement + public void process(@Element Integer element, OutputReceiver out) { + out.output(element * 2); + } + } + + @Test(timeout = 300_000) + public void testPAssertInGlobalWindow() throws Exception { + List> elements = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("Double", ParDo.of(new DoubleFn())); + + PAssert.that(output).containsInAnyOrder(0, 2, 4, 6, 8, 10, 12, 14, 16, 18); + + pipeline.run().waitUntilFinish(); + } + + @Test(timeout = 300_000) + public void testPAssertFailureThrows() throws Exception { + List> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(1)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("Double", ParDo.of(new DoubleFn())); + + // Deliberately incorrect expectation: output is [2], expected is [999]. + PAssert.that(output).containsInAnyOrder(999); + + Exception e = assertThrows(Exception.class, () -> pipeline.run().waitUntilFinish()); + Throwable root = e; + while (root.getCause() != null && root.getCause() != root) { + root = root.getCause(); + } + org.junit.Assert.assertTrue( + "Expected AssertionError at root of cause chain, got: " + root, + root instanceof AssertionError); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java index 75c657ef39c6..52b8eb454ca0 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -231,9 +231,18 @@ public static final class ListBackedUnboundedSource private final List encodedElements; private final List timestampsMillis; private final Coder coder; + private final boolean advanceWatermarkToInfinityOnExhaustion; public ListBackedUnboundedSource(List> elements, Coder coder) { + this(elements, coder, false); + } + + public ListBackedUnboundedSource( + List> elements, + Coder coder, + boolean advanceWatermarkToInfinityOnExhaustion) { this.coder = Preconditions.checkNotNull(coder); + this.advanceWatermarkToInfinityOnExhaustion = advanceWatermarkToInfinityOnExhaustion; List encoded = new ArrayList<>(elements.size()); List timestamps = new ArrayList<>(elements.size()); for (TimestampedValue element : elements) { @@ -245,10 +254,14 @@ public ListBackedUnboundedSource(List> elements, Coder co } private ListBackedUnboundedSource( - List encodedElements, List timestampsMillis, Coder coder) { + List encodedElements, + List timestampsMillis, + Coder coder, + boolean advanceWatermarkToInfinityOnExhaustion) { this.encodedElements = encodedElements; this.timestampsMillis = timestampsMillis; this.coder = coder; + this.advanceWatermarkToInfinityOnExhaustion = advanceWatermarkToInfinityOnExhaustion; } @Override @@ -274,7 +287,10 @@ public List> split( for (int i = 0; i < numSplits; i++) { splits.add( new ListBackedUnboundedSource<>( - bucketedElements.get(i), bucketedTimestamps.get(i), coder)); + bucketedElements.get(i), + bucketedTimestamps.get(i), + coder, + advanceWatermarkToInfinityOnExhaustion)); } return splits; } @@ -368,6 +384,10 @@ public Instant getCurrentTimestamp() throws NoSuchElementException { @Override public Instant getWatermark() { + if (source.advanceWatermarkToInfinityOnExhaustion + && index + 1 >= source.encodedElements.size()) { + return BoundedWindow.TIMESTAMP_MAX_VALUE; + } // No idle advance, deliberately: once the list is exhausted this simply stops moving, // rather than jumping to +infinity. See the StreamingTestUtils class javadoc. return maxTimestampSeen; From e991c5d771222df3956b9d1a1706f59ba3490d88 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 20 Aug 2026 21:49:01 +0000 Subject: [PATCH 19/21] [Spark 4] Support PAssert in non global windows by clamping arrival side expiry PAssert on a non globally windowed collection rewindows into the global window and chains two GroupByKeys. The end of stream sentinel pushes the watermark to TIMESTAMP_MAX_VALUE so end of global window timers fire, but LateDataUtils truncates every garbage collection time to the end of the global window, one day earlier. Judged against the raw watermark the global window itself counts as expired, so the late data filter in front of the second GroupByKey dropped the entire pane flushed by the first and PAssert compared against an empty iterable. Arrival side expiry decisions now see the watermark clamped to the end of the global window through a delegating TimerInternals view installed in the step context. The timer firing path and the ReduceFnRunner keep the real watermark, clamping there would withhold the end of global window timers forever. Restores the fixed windows PAssert test and adds a negative variant proving a wrong expectation still fails on the rewindowed path. Late data drop semantics for finite windows are unchanged, the clamp only takes effect once the watermark has passed the end of the global window. --- .../state/BeamStatefulProcessor.java | 100 +++++++++++++++++- .../streaming/PAssertStreamingTest.java | 67 ++++++++++++ 2 files changed, 166 insertions(+), 1 deletion(-) 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 d61f4582a263..f4e38feee8c0 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 @@ -40,10 +40,12 @@ import org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.SerializableUtils; import org.apache.beam.sdk.util.WindowedValueMultiReceiver; @@ -103,6 +105,10 @@ * micro-batch after the one whose data crossed the end of the window. That is the same one batch * delay every micro-batch runner has. * + *

Once the end-of-stream sentinel has pushed the watermark beyond the end of the global window, + * arrival side expiry decisions see a clamped watermark, see {@link + * ArrivalExpiryClampedTimerInternals}. The timer firing path always sees the real watermark. + * *

Timers

* *

Beam timers fire only in {@code handleExpiredTimer}, never in {@code handleInputRows}. A timer @@ -245,6 +251,11 @@ private scala.collection.Iterator process( return ScalaInterop.scalaIterator(Collections.emptyList()); } + // The step context is what DoFnRunners consult for arrival side expiry decisions, the late + // data filter of the GROUP_ALSO_BY_WINDOW path and the expired window drop of the stateful + // ParDo path, so it hands out the clamped watermark view. The timer firing path keeps the + // real timerInternals, see ArrivalExpiryClampedTimerInternals. + TimerInternals arrivalTimerInternals = new ArrivalExpiryClampedTimerInternals(timerInternals); StepContext stepContext = new StepContext() { @Override @@ -254,7 +265,7 @@ public StateInternals stateInternals() { @Override public TimerInternals timerInternals() { - return timerInternals; + return arrivalTimerInternals; } }; @@ -396,6 +407,93 @@ private static T decode(Coder coder, byte[] bytes, String what) { } } + /** + * A view of {@link TwsTimerInternals} for arrival side expiry decisions whose input watermark is + * clamped to the end of the global window, every other method delegating unchanged. + * + *

The two sides of an expiry comparison live in different domains. {@code + * LateDataUtils.garbageCollectionTime} truncates every garbage collection time to {@code + * GlobalWindow.INSTANCE.maxTimestamp()}, so no window, not even the global one, ever expires + * later than that. The end-of-stream sentinel however pushes Spark's watermark all the way to + * {@link BoundedWindow#TIMESTAMP_MAX_VALUE}, one day further, so that end of global window timers + * can fire. Judged against that raw watermark the global window itself is expired, and every + * element that reaches a downstream stateful operator in the final micro-batch, such as the + * second GroupByKey inside a PAssert, is dropped as late data. A watermark beyond the end of the + * global window exists only to fire end of global window timers and must not be used to judge + * arrival side expiry, so this view reports at most {@code GlobalWindow.INSTANCE.maxTimestamp()} + * and is installed in the {@link StepContext} consulted by {@code LateDataDroppingDoFnRunner} and + * {@code StatefulDoFnRunner}. + * + *

It must never be used on the timer firing path. {@link + * TwsTimerInternals#removeTimersReadyToFire} releases a timer only once the watermark is strictly + * past it, so clamping there would withhold end of global window timers forever. Likewise the + * {@code ReduceFnRunner} behind the GROUP_ALSO_BY_WINDOW mode keeps the unclamped {@link + * TwsTimerInternals}, its triggers must see the watermark pass the end of the global window to + * emit the final pane. + */ + private static final class ArrivalExpiryClampedTimerInternals implements TimerInternals { + + private final TwsTimerInternals delegate; + + private ArrivalExpiryClampedTimerInternals(TwsTimerInternals delegate) { + this.delegate = delegate; + } + + @Override + public Instant currentInputWatermarkTime() { + Instant watermark = delegate.currentInputWatermarkTime(); + Instant endOfGlobalWindow = GlobalWindow.INSTANCE.maxTimestamp(); + return watermark.isAfter(endOfGlobalWindow) ? endOfGlobalWindow : watermark; + } + + @Override + public void setTimer( + StateNamespace namespace, + String timerId, + String timerFamilyId, + Instant target, + Instant outputTimestamp, + TimeDomain timeDomain) { + delegate.setTimer(namespace, timerId, timerFamilyId, target, outputTimestamp, timeDomain); + } + + @Override + public void setTimer(TimerInternals.TimerData timerData) { + delegate.setTimer(timerData); + } + + @Override + public void deleteTimer( + StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { + delegate.deleteTimer(namespace, timerId, timerFamilyId, timeDomain); + } + + @Override + public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) { + delegate.deleteTimer(namespace, timerId, timerFamilyId); + } + + @Override + public void deleteTimer(TimerInternals.TimerData timerKey) { + delegate.deleteTimer(timerKey); + } + + @Override + public Instant currentProcessingTime() { + return delegate.currentProcessingTime(); + } + + @Override + public @Nullable Instant currentSynchronizedProcessingTime() { + return delegate.currentSynchronizedProcessingTime(); + } + + @Override + public @Nullable Instant currentOutputWatermarkTime() { + return delegate.currentOutputWatermarkTime(); + } + } + /** Encodes every emitted element into the tagged output row layout and appends it. */ private final class EncodingReceiver implements WindowedValueMultiReceiver { private final List outputs; diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java index 6d0a5b6c1b3b..0a6311faf647 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java @@ -31,6 +31,8 @@ import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.TimestampedValue; import org.joda.time.Duration; @@ -94,6 +96,37 @@ public void testPAssertInGlobalWindow() throws Exception { pipeline.run().waitUntilFinish(); } + @Test(timeout = 300_000) + public void testPAssertInFixedWindows() throws Exception { + List> elements = new ArrayList<>(); + // Window 1: [0s, 5s) -> elements 0, 1, 2 (doubled: 0, 2, 4) + elements.add(TimestampedValue.of(0, BASE.plus(Duration.standardSeconds(1)))); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(2)))); + elements.add(TimestampedValue.of(2, BASE.plus(Duration.standardSeconds(4)))); + // Window 2: [5s, 10s) -> elements 3, 4 (doubled: 6, 8) + elements.add(TimestampedValue.of(3, BASE.plus(Duration.standardSeconds(6)))); + elements.add(TimestampedValue.of(4, BASE.plus(Duration.standardSeconds(8)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("FixedWindows", Window.into(FixedWindows.of(Duration.standardSeconds(5)))) + .apply("Double", ParDo.of(new DoubleFn())); + + PAssert.that(output).containsInAnyOrder(0, 2, 4, 6, 8); + + pipeline.run().waitUntilFinish(); + } + @Test(timeout = 300_000) public void testPAssertFailureThrows() throws Exception { List> elements = new ArrayList<>(); @@ -125,4 +158,38 @@ public void testPAssertFailureThrows() throws Exception { "Expected AssertionError at root of cause chain, got: " + root, root instanceof AssertionError); } + + @Test(timeout = 300_000) + public void testPAssertFailureThrowsInFixedWindows() throws Exception { + List> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(1)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("FixedWindows", Window.into(FixedWindows.of(Duration.standardSeconds(5)))) + .apply("Double", ParDo.of(new DoubleFn())); + + // Deliberately incorrect expectation: output is [2], expected is [999]. This proves the + // arrival side watermark clamp does not mask genuine failures in the rewindowed path. + PAssert.that(output).containsInAnyOrder(999); + + Exception e = assertThrows(Exception.class, () -> pipeline.run().waitUntilFinish()); + Throwable root = e; + while (root.getCause() != null && root.getCause() != root) { + root = root.getCause(); + } + org.junit.Assert.assertTrue( + "Expected AssertionError at root of cause chain, got: " + root, + root instanceof AssertionError); + } } From 2e365abf20a8ce4a5f9ef80c88596092829f5ae7 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 21 Aug 2026 12:18:15 +0000 Subject: [PATCH 20/21] [Spark 4] Support allowed lateness, late firings, and accumulating panes in streaming --- .../StreamingTranslationHelpers.java | 47 ++- .../WindowedGroupByKeyStreamingTest.java | 344 ++++++++++++++++++ 2 files changed, 373 insertions(+), 18 deletions(-) diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java index b9d3649e42ed..98d751d3abeb 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java @@ -27,6 +27,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.reflect.DoFnSignature; import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.transforms.windowing.AfterPane; import org.apache.beam.sdk.transforms.windowing.AfterWatermark; import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; import org.apache.beam.sdk.transforms.windowing.Never; @@ -35,7 +36,6 @@ import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowingStrategy; -import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; import org.apache.spark.api.java.function.FilterFunction; import org.apache.spark.api.java.function.MapFunction; @@ -59,7 +59,7 @@ private StreamingTranslationHelpers() {} /** * Rejects windowing strategies whose semantics the {@code transformWithState} bridge cannot - * reproduce: merging (session) windows, custom triggers and accumulating panes. + * reproduce: merging (session) windows and unsupported custom triggers. */ static void checkSupportedWindowing(WindowingStrategy strategy, String stepName) { WindowFn windowFn = strategy.getWindowFn(); @@ -72,41 +72,52 @@ static void checkSupportedWindowing(WindowingStrategy strategy, String ste + "Spark 4 streaming runner"); } Trigger trigger = strategy.getTrigger(); - if (!isDefaultTrigger(trigger)) { + if (!isSupportedTrigger(trigger)) { throw unsupported( stepName, "the custom trigger " + trigger + ". Only the default trigger (one on-time pane per window when the watermark passes " - + "its end) is implemented"); - } - if (strategy.getMode() != AccumulationMode.DISCARDING_FIRED_PANES) { - throw unsupported( - stepName, - "accumulating panes (" + strategy.getMode() + "). Only discarding panes are implemented"); + + "its end) and late firings with AfterPane are implemented"); } } /** - * The default trigger, as Beam models it, is either {@link DefaultTrigger} itself, the equivalent - * {@code AfterWatermark.pastEndOfWindow()} without early or late firings, or {@link - * Never.NeverTrigger} used by {@code PAssert} on unbounded streams to gather all window contents - * before the end-of-stream watermark triggers final pane evaluation. + * Identifies triggers supported by the streaming translation. + * + *

    + *
  • {@link DefaultTrigger} + *
  • {@link AfterWatermark.FromEndOfWindow} without early or late firings + *
  • {@link Never.NeverTrigger} used by {@code PAssert} on unbounded streams + *
  • {@link AfterWatermark.AfterWatermarkEarlyAndLate} with no early firings and late firings + * configured via {@link AfterPane} + *
*/ - private static boolean isDefaultTrigger(Trigger trigger) { + private static boolean isSupportedTrigger(Trigger trigger) { if (trigger instanceof DefaultTrigger) { return true; } if (trigger instanceof AfterWatermark.FromEndOfWindow) { - // FromEndOfWindow without early/late firings is exactly the default trigger; with firings - // configured Beam represents it as one of the AfterWatermarkEarlyAndLate subclasses instead. return true; } if (trigger instanceof Never.NeverTrigger) { - // NeverTrigger is used by PAssert on unbounded PCollections to prevent intermediate firings. - // In this runner, end-of-stream watermark advancement fires the final pane. return true; } + if (trigger instanceof AfterWatermark.AfterWatermarkEarlyAndLate) { + AfterWatermark.AfterWatermarkEarlyAndLate earlyAndLate = + (AfterWatermark.AfterWatermarkEarlyAndLate) trigger; + if (!(earlyAndLate.getEarlyTrigger() instanceof Never.NeverTrigger)) { + return false; + } + Trigger lateTrigger = earlyAndLate.getLateTrigger(); + if (lateTrigger == null || lateTrigger instanceof Never.NeverTrigger) { + return true; + } + if (lateTrigger instanceof AfterPane) { + return true; + } + return false; + } return false; } 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 32d0fd4544e6..f4c021f0b2aa 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 @@ -18,11 +18,14 @@ package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Objects; import org.apache.beam.runners.spark.StreamingTest; import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; @@ -32,12 +35,19 @@ import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.Read; import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.AfterPane; +import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime; +import org.apache.beam.sdk.transforms.windowing.AfterWatermark; import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Repeatedly; 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.TimestampedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.ClassRule; @@ -230,4 +240,338 @@ public void lateDataIsDropped() throws Exception { assertEquals( "pipeline state=" + result.getState(), "[a=1, sentinel=1]", collectedCounts(collectorId)); } + + /** + * Encapsulates the value and complete {@link PaneInfo} metadata of an emitted element for precise + * assertions in streaming tests. + */ + public static final class PaneRecord implements Serializable { + private final String key; + private final long value; + private final PaneInfo.Timing timing; + private final long index; + private final long onTimeIndex; + private final boolean isFirst; + private final boolean isLast; + + public PaneRecord( + String key, + long value, + PaneInfo.Timing timing, + long index, + long onTimeIndex, + boolean isFirst, + boolean isLast) { + this.key = key; + this.value = value; + this.timing = timing; + this.index = index; + this.onTimeIndex = onTimeIndex; + this.isFirst = isFirst; + this.isLast = isLast; + } + + public String getKey() { + return key; + } + + public long getValue() { + return value; + } + + public PaneInfo.Timing getTiming() { + return timing; + } + + public long getIndex() { + return index; + } + + public long getOnTimeIndex() { + return onTimeIndex; + } + + public boolean isFirst() { + return isFirst; + } + + public boolean isLast() { + return isLast; + } + + @Override + public String toString() { + return key + + "=" + + value + + ":" + + timing + + ":index=" + + index + + ":onTimeIndex=" + + onTimeIndex + + ":isFirst=" + + isFirst + + ":isLast=" + + isLast; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof PaneRecord)) { + return false; + } + PaneRecord that = (PaneRecord) o; + return value == that.value + && index == that.index + && onTimeIndex == that.onTimeIndex + && isFirst == that.isFirst + && isLast == that.isLast + && Objects.equals(key, that.key) + && timing == that.timing; + } + + @Override + public int hashCode() { + return Objects.hash(key, value, timing, index, onTimeIndex, isFirst, isLast); + } + } + + /** Converts {@code KV} into {@link PaneRecord} carrying full {@link PaneInfo}. */ + public static final class CollectPaneDoFn extends DoFn, PaneRecord> { + @ProcessElement + public void processElement( + @Element KV element, PaneInfo paneInfo, OutputReceiver out) { + PaneRecord record = + new PaneRecord( + element.getKey(), + element.getValue(), + paneInfo.getTiming(), + paneInfo.getIndex(), + paneInfo.getNonSpeculativeIndex(), + paneInfo.isFirst(), + paneInfo.isLast()); + out.output(record); + } + } + + private static List filterPanesForKey(String collectorId, String targetKey) { + List result = new ArrayList<>(); + for (PaneRecord record : StreamingTestUtils.getCollected(collectorId)) { + if (targetKey.equals(record.getKey())) { + result.add(record); + } + } + return result; + } + + private static List>> lateFiringsInputElements() { + List>> elements = new ArrayList<>(); + // Split 0: [a@1s, dummy@22s, a@2s, a@3s, dummy@36s, a@4s, dummy@62s] + // Split 1: [sentinel@20s, sentinel@25s, sentinel@30s, sentinel@35s, sentinel@45s, sentinel@60s, + // sentinel@70s] + + // Batch 1: a@1s (on-time in [0s, 10s)), sentinel@20s -> Watermark advances to 20s + elements.add(TimestampedValue.of(KV.of("a", "1"), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s1"), BASE.plus(Duration.standardSeconds(20)))); + + // Batch 2: dummy@22s, sentinel@25s -> Batch start watermark is 20s. On-time timer for 'a' fires + // here! + elements.add( + TimestampedValue.of(KV.of("dummy", "d1"), BASE.plus(Duration.standardSeconds(22)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s2"), BASE.plus(Duration.standardSeconds(25)))); + + // Batch 3: a@2s (late! watermark 25s < GC time 40s), sentinel@30s -> Late pane 1 fires + elements.add(TimestampedValue.of(KV.of("a", "2"), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s3"), BASE.plus(Duration.standardSeconds(30)))); + + // Batch 4: a@3s (late! watermark 30s < GC time 40s), sentinel@35s -> Late pane 2 fires + elements.add(TimestampedValue.of(KV.of("a", "3"), BASE.plus(Duration.standardSeconds(3)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s4"), BASE.plus(Duration.standardSeconds(35)))); + + // Batch 5: dummy@36s, sentinel@45s -> Watermark advances to 45s (past GC horizon 40s, window + // expires) + elements.add( + TimestampedValue.of(KV.of("dummy", "d2"), BASE.plus(Duration.standardSeconds(36)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s5"), BASE.plus(Duration.standardSeconds(45)))); + + // Batch 6: a@4s (late! arriving when start watermark is 45s > GC horizon 40s -> dropped), + // sentinel@60s + elements.add(TimestampedValue.of(KV.of("a", "4"), BASE.plus(Duration.standardSeconds(4)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s6"), BASE.plus(Duration.standardSeconds(60)))); + + // Batch 7: dummy@62s, sentinel@70s -> Trailing batch + elements.add( + TimestampedValue.of(KV.of("dummy", "d3"), BASE.plus(Duration.standardSeconds(62)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s7"), BASE.plus(Duration.standardSeconds(70)))); + + return elements; + } + + @Test(timeout = 300_000) + public void fixedWindowsWithLateFiringsDiscarding() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("late-firings-discarding"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerMicroBatch(1); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(1))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two late delta panes in discarding mode, with expired element dropped", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=1:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=1:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + @Test(timeout = 300_000) + public void fixedWindowsWithLateFiringsAccumulating() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("late-firings-accumulating"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerMicroBatch(1); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(1))) + .withAllowedLateness(Duration.standardSeconds(30)) + .accumulatingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two accumulating panes carrying full window counts, with expired element dropped", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=2:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=3:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + @Test(timeout = 300_000) + public void unsupportedEarlyFiringsTriggerThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + Repeatedly.forever( + AfterWatermark.pastEndOfWindow() + .withEarlyFirings(AfterPane.elementCountAtLeast(1)))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("withEarlyFirings")); + } + + @Test(timeout = 300_000) + public void unsupportedProcessingTimeTriggerThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering(Repeatedly.forever(AfterProcessingTime.pastFirstElementInPane())) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("AfterProcessingTime")); + } } From c30d6704f85bd931e86666f47ae97c4dbd30cf30 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 21 Aug 2026 15:44:31 +0000 Subject: [PATCH 21/21] [Spark 4] Tighten trigger guards, add DefaultTrigger accumulating and same-batch folding tests --- .../StreamingTranslationHelpers.java | 6 +- .../WindowedGroupByKeyStreamingTest.java | 155 +++++++++++++++++- 2 files changed, 154 insertions(+), 7 deletions(-) diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java index 98d751d3abeb..00aa839fcda4 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java @@ -78,7 +78,7 @@ static void checkSupportedWindowing(WindowingStrategy strategy, String ste "the custom trigger " + trigger + ". Only the default trigger (one on-time pane per window when the watermark passes " - + "its end) and late firings with AfterPane are implemented"); + + "its end) and late firings with AfterPane.elementCountAtLeast(1) are implemented"); } } @@ -90,7 +90,7 @@ static void checkSupportedWindowing(WindowingStrategy strategy, String ste *
  • {@link AfterWatermark.FromEndOfWindow} without early or late firings *
  • {@link Never.NeverTrigger} used by {@code PAssert} on unbounded streams *
  • {@link AfterWatermark.AfterWatermarkEarlyAndLate} with no early firings and late firings - * configured via {@link AfterPane} + * configured via {@link AfterPane#elementCountAtLeast(int)} with element count 1 * */ private static boolean isSupportedTrigger(Trigger trigger) { @@ -114,7 +114,7 @@ private static boolean isSupportedTrigger(Trigger trigger) { return true; } if (lateTrigger instanceof AfterPane) { - return true; + return ((AfterPane) lateTrigger).getElementCount() == 1; } return false; } 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 f4c021f0b2aa..51cd91b2ccab 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 @@ -512,6 +512,121 @@ public void fixedWindowsWithLateFiringsAccumulating() throws Exception { renderedPanes); } + @Test(timeout = 300_000) + public void fixedWindowsDefaultTriggerAccumulating() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("default-trigger-accumulating"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerMicroBatch(1); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .withAllowedLateness(Duration.standardSeconds(30)) + .accumulatingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two accumulating panes under default trigger", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=2:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=3:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + /** + * Under Spark Structured Streaming's micro-batch execution model, {@code transformWithState} + * processes input rows against the batch start watermark (the watermark established by the + * previous micro-batch). When an element timestamped in {@code [0s, 10s)} arrives in the same + * micro-batch that carries a watermark-advancing sentinel, the batch start watermark has not yet + * passed the window end. In Beam semantics, the window is still open at the moment the element is + * processed, so it is buffered into the open window accumulator. When the on-time timer + * subsequently fires at the end of the window, it emits a single {@code ON_TIME} pane containing + * both elements. + */ + @Test(timeout = 300_000) + public void fixedWindowsLateElementInSameBatchFoldedIntoOnTimePane() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("same-batch-merge"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // Split 0: [a@1s, a@2s, dummy@52s] + // Split 1: [sentinel@5s, sentinel@20s, sentinel@60s] + + // Batch 1: a@1s (in [0s, 10s)), sentinel@5s -> Watermark advances to 5s + elements.add(TimestampedValue.of(KV.of("a", "1"), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s1"), BASE.plus(Duration.standardSeconds(5)))); + + // Batch 2: a@2s (in [0s, 10s)), sentinel@20s -> Batch start watermark is 5s < 10s! + // Window [0s, 10s) is still open during input row processing in batch 2, so a@2s is folded in. + // Watermark advances to 20s at the end of batch 2. + elements.add(TimestampedValue.of(KV.of("a", "2"), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s2"), BASE.plus(Duration.standardSeconds(20)))); + + // Batch 3: dummy@52s, sentinel@60s -> Batch start watermark is 20s >= 10s. + // The on-time timer for [0s, 10s) fires here, emitting a single ON_TIME pane with count = 2. + elements.add(TimestampedValue.of(KV.of("dummy", "d"), BASE.plus(Duration.standardSeconds(52)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s3"), BASE.plus(Duration.standardSeconds(60)))); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerMicroBatch(1); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected single merged ON_TIME pane containing both elements", + List.of("a=2:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false"), + renderedPanes); + } + @Test(timeout = 300_000) public void unsupportedEarlyFiringsTriggerThrows() throws Exception { SparkStructuredStreamingPipelineOptions options = options(); @@ -530,9 +645,9 @@ public void unsupportedEarlyFiringsTriggerThrows() throws Exception { "FixedWindows", Window.>into(FixedWindows.of(WINDOW_SIZE)) .triggering( - Repeatedly.forever( - AfterWatermark.pastEndOfWindow() - .withEarlyFirings(AfterPane.elementCountAtLeast(1)))) + AfterWatermark.pastEndOfWindow() + .withEarlyFirings(AfterPane.elementCountAtLeast(1)) + .withLateFirings(AfterPane.elementCountAtLeast(1))) .withAllowedLateness(Duration.standardSeconds(30)) .discardingFiredPanes()) .apply("CountPerKey", Count.perKey()); @@ -546,7 +661,7 @@ public void unsupportedEarlyFiringsTriggerThrows() throws Exception { } @Test(timeout = 300_000) - public void unsupportedProcessingTimeTriggerThrows() throws Exception { + public void unsupportedRepeatedlyProcessingTimeTriggerThrows() throws Exception { SparkStructuredStreamingPipelineOptions options = options(); Pipeline pipeline = Pipeline.create(options); @@ -574,4 +689,36 @@ public void unsupportedProcessingTimeTriggerThrows() throws Exception { e.getMessage().contains("the custom trigger") && e.getMessage().contains("AfterProcessingTime")); } + + @Test(timeout = 300_000) + public void unsupportedAfterPaneCountGreaterThanOneThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(2))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("AfterPane.elementCountAtLeast(2)")); + } }