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/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/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..7901fb9f4a31 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.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; + +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; +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 { + + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + @Nullable + protected > + TransformTranslator getTransformTranslator(TransformT transform) { + + if (transform instanceof SplittableParDo.PrimitiveUnboundedRead) { + return (TransformTranslator) new ReadUnboundedTranslator<>(); + } + + if (transform instanceof GroupByKey) { + 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 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()); + if (signature.usesState() || signature.usesTimers()) { + return (TransformTranslator) new StatefulParDoStreamingTranslator<>(); + } + // 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); + } +} 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..c45ec0de30ab --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -0,0 +1,306 @@ +/* + * 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}. 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

+ * + *

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: + * + *

+ * + *

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()); + } + } +} 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. + * + *

+ */ +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; + } + } +} 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 + * + *
    + *
  • {@code keyBytes} is the Beam key {@code K} encoded with {@code config.keyCoder()}. It is + * the Spark grouping key, so it must be a deterministic encoding, two elements of the same + * Beam key must produce byte-identical {@code keyBytes}. + *
  • {@code windowedValueBytes} is the {@code WindowedValue} of the value side only, + * encoded with {@code config.inputValueCoder()}, that is {@code + * WindowedValues.getFullCoder(config.valueCoder(), config.windowCoder())}. The key is not + * repeated inside the payload, the operator re-attaches it. + *
+ * + *

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[]}: + * + *

    + *
  • {@code beamState} holds all Beam user and system state, addressed by {@code namespace + + * tag}, see {@link TwsStateInternals}. + *
  • {@code beamTimers} holds the encoded {@link TimerInternals.TimerData}, see {@link + * TwsTimerInternals}. + *
+ * + *

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/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..38b27db1f4b5 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -0,0 +1,475 @@ +/* + * 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: 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 + .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 {} + } + } +} 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/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 new file mode 100644 index 000000000000..dd3ec0a9f782 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.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.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.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.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * 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. + * + *

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) +public class ChainedStatefulStreamingTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final 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) + 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); + Pipeline pipeline = Pipeline.create(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(); + + // 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 new file mode 100644 index 000000000000..1bd925b4c78d --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.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 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.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.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.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}). + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StatefulParDoStreamingTest implements Serializable { + + /** + * 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(); + + 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) + 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); + Pipeline pipeline = Pipeline.create(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(); + + // 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 new file mode 100644 index 000000000000..a69409681eda --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -0,0 +1,114 @@ +/* + * 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.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.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.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 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 { + + /** + * 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(); + + @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) + 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); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(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(); + + // 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 new file mode 100644 index 000000000000..f58a809d591b --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -0,0 +1,193 @@ +/* + * 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.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.transforms.ParDo; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * {@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}. 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 + * (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 { + + /** + * 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 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(); + + @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++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + 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) + 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); + Pipeline pipeline = Pipeline.create(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()); + + // 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) + 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); + Pipeline pipeline = Pipeline.create(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()); + + 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 new file mode 100644 index 000000000000..2da0458dbc07 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -0,0 +1,402 @@ +/* + * 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.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 a factory for + * the {@link SparkStructuredStreamingPipelineOptions} every streaming test in this package needs. + * + *

Why streaming tests in this suite look the way they do

+ * + *

Three 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. + *
  • {@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: + * + *

    + *
  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 factory. + // --------------------------------------------------------------------------------------------- + + /** + * 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 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( + 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; + } + + // --------------------------------------------------------------------------------------------- + // 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..32d0fd4544e6 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java @@ -0,0 +1,233 @@ +/* + * 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.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.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.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}. + * + *

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) +public class WindowedGroupByKeyStreamingTest implements Serializable { + + /** + * 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(); + + @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; + } + + /** 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) + 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)))); + + Pipeline pipeline = Pipeline.create(options()); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply("FixedWindows", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("CountPerKey", Count.perKey()) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // 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) + public void slidingWindowsCountPerKey() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("sliding-windows"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // 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)))); + + Pipeline pipeline = Pipeline.create(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(); + + // 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) + 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)))); + + 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", + 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(); + + // 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)); + } +} 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..bdfee6cf2113 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java @@ -0,0 +1,440 @@ +/* + * 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 { + + /** + * 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(); + + @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); + } +} 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..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 @@ -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..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 @@ -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; @@ -175,6 +176,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); } @@ -229,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); @@ -284,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); }