From 4ca982ea435affc059ddac68e0f06e1b103b8281 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 31 Jul 2026 11:14:00 +0000 Subject: [PATCH 01/10] [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/10] [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/10] [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/10] [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/10] [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/10] [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/10] [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/10] [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/10] [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/10] [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