From d291ec8005a999634f56e334141d9289a764dfdc Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 2 Sep 2026 07:53:28 +0000 Subject: [PATCH 1/4] [Spark 4] Add DataSourceV2 micro-batch source for Beam unbounded sources Exposes any Beam UnboundedSource as a Spark 4 DataSourceV2 streaming table with a fixed two column schema, encoded payload plus event timestamp. Offsets are opaque, strictly increasing epoch counters, so Spark keeps scheduling micro-batches and termination stays with the lifecycle owner. Recovery is durable under the query's checkpoint location: the source id derives deterministically from the read transform's full name, the first run pins its split list (Beam sources do not guarantee deterministic splitting), and every split persists its CheckpointMark per epoch with a retention of two, written atomically via temp file and rename. Executors cache live readers between micro-batches and fall back to the newest durable mark at or before the replayed epoch after a restart. Semantics are at least once, a crash between finishing a read and Spark's commit replays the last micro-batch. The batch cutoff honors maxRecordsPerBatch, values below 1, including the default, mean no limit and the batch ends on the duration deadline. --- .../io/streaming/BeamCheckpointFiles.java | 236 ++++++++ .../io/streaming/BeamInputPartition.java | 110 ++++ .../io/streaming/BeamMicroBatchStream.java | 211 +++++++ .../io/streaming/BeamOffset.java | 79 +++ .../io/streaming/BeamPartitionReader.java | 181 ++++++ .../streaming/BeamPartitionReaderFactory.java | 34 ++ .../io/streaming/BeamReaderCache.java | 179 ++++++ .../io/streaming/BeamStreamingSource.java | 157 ++++++ .../io/streaming/BeamStreamingTable.java | 100 ++++ .../io/streaming/UnboundedSourceDataset.java | 133 +++++ .../streaming/BeamCheckpointRecoveryTest.java | 198 +++++++ .../streaming/BeamMicroBatchSourceTest.java | 520 ++++++++++++++++++ 12 files changed, 2138 insertions(+) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamOffset.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReaderFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingTable.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java new file mode 100644 index 000000000000..069bb27e9f65 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java @@ -0,0 +1,236 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FSDataOutputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Durable state of the Beam micro-batch source, stored next to Spark's own streaming state under + * {@code /beam-source-/}. + * + *

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

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

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

The mark written under exactly {@code startEpoch} is preferred, if it is absent the mark + * with the largest epoch not exceeding {@code startEpoch} is used. Returns {@code null}, meaning + * a fresh start, when no such mark exists or reading fails. + */ + public static @Nullable CheckpointMark readMark( + String checkpointLocation, String sourceId, int splitId, long startEpoch) { + Path dir = marksDir(checkpointLocation, sourceId, splitId); + try { + FileSystem fs = fileSystem(dir); + if (!fs.exists(dir)) { + return null; + } + long epoch = Long.MIN_VALUE; + if (fs.exists(new Path(dir, Long.toString(startEpoch)))) { + epoch = startEpoch; + } else { + for (FileStatus status : fs.listStatus(dir)) { + @Nullable Long candidate = parseEpoch(status.getPath().getName()); + if (candidate != null && candidate <= startEpoch && candidate > epoch) { + epoch = candidate; + } + } + } + if (epoch == Long.MIN_VALUE) { + return null; + } + Path path = new Path(dir, Long.toString(epoch)); + CheckpointMark mark = + (CheckpointMark) deserialize(read(fs, path), "durable checkpoint mark " + path); + LOG.info( + "Restored durable checkpoint mark of Beam source {} split {} at epoch {} " + + "(requested epoch {}).", + sourceId, + splitId, + epoch, + startEpoch); + return mark; + } catch (IOException e) { + LOG.warn( + "Failed to read a durable checkpoint mark of Beam source {} split {} at epoch {}, " + + "the reader starts without one.", + sourceId, + splitId, + startEpoch, + e); + return null; + } + } + + private static Path root(String checkpointLocation, String sourceId) { + return new Path(checkpointLocation, ROOT_PREFIX + sourceId); + } + + private static Path marksDir(String checkpointLocation, String sourceId, int splitId) { + return new Path( + new Path(root(checkpointLocation, sourceId), MARKS_DIR), String.valueOf(splitId)); + } + + private static FileSystem fileSystem(Path path) throws IOException { + return path.getFileSystem(new Configuration()); + } + + /** Writes {@code bytes} to a {@code .tmp} sibling of {@code target} and renames it into place. */ + private static void writeAtomically(FileSystem fs, Path target, byte[] bytes) throws IOException { + Path tmp = new Path(target.getParent(), target.getName() + TMP_SUFFIX); + try (FSDataOutputStream out = fs.create(tmp, true)) { + out.write(bytes); + } + // On HDFS a rename onto an existing target fails, a mark rewritten by a task retry is the + // same position, so an existing target counts as success. + if (!fs.rename(tmp, target) && !fs.exists(target)) { + throw new IOException("Failed to rename " + tmp + " to " + target); + } + if (fs.exists(tmp)) { + fs.delete(tmp, false); + } + } + + private static byte[] read(FileSystem fs, Path path) throws IOException { + try (FSDataInputStream in = fs.open(path)) { + return ByteStreams.toByteArray(in); + } + } + + private static Object deserialize(byte[] bytes, String description) { + return SerializableUtils.deserializeFromByteArray(bytes, description); + } + + /** Best effort deletion of mark files older than the {@value #RETAINED_MARKS} highest epochs. */ + private static void deleteOldMarks(FileSystem fs, Path dir) { + try { + List epochs = new ArrayList<>(); + for (FileStatus status : fs.listStatus(dir)) { + @Nullable Long epoch = parseEpoch(status.getPath().getName()); + if (epoch != null) { + epochs.add(epoch); + } + } + Collections.sort(epochs); + for (int i = 0; i < epochs.size() - RETAINED_MARKS; i++) { + fs.delete(new Path(dir, Long.toString(epochs.get(i))), false); + } + } catch (IOException e) { + LOG.warn("Failed to delete outdated checkpoint marks under {}.", dir, e); + } + } + + private static @Nullable Long parseEpoch(String fileName) { + try { + return Long.valueOf(fileName); + } catch (NumberFormatException e) { + return null; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java new file mode 100644 index 000000000000..8a1103dac46a --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java @@ -0,0 +1,110 @@ +/* + * 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 long startEpoch; + private final long endEpoch; + private final long maxRecordsPerBatch; + private final long maxBatchDurationMillis; + + BeamInputPartition( + String sourceB64, + String coderB64, + String pipelineOptionsB64, + String sourceId, + int splitId, + String checkpointLocation, + long startEpoch, + long endEpoch, + long maxRecordsPerBatch, + long maxBatchDurationMillis) { + this.sourceB64 = sourceB64; + this.coderB64 = coderB64; + this.pipelineOptionsB64 = pipelineOptionsB64; + this.sourceId = sourceId; + this.splitId = splitId; + this.checkpointLocation = checkpointLocation; + this.startEpoch = startEpoch; + this.endEpoch = endEpoch; + this.maxRecordsPerBatch = maxRecordsPerBatch; + 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; + } + + long startEpoch() { + return startEpoch; + } + + long endEpoch() { + return endEpoch; + } + + long maxRecordsPerBatch() { + return maxRecordsPerBatch; + } + + 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..f93fdc3113c4 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.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 + * pinned to the checkpoint location so every micro-batch of every run plans the same, stable set of + * partitions. Splits must be stable across micro-batches and across restarts because the executor + * side reader cache and the durable checkpoint marks are keyed by split index, and Beam sources do + * not guarantee deterministic splitting. A restarted stream therefore loads the split list written + * by the first run instead of splitting again. + * + *

On a restart Spark replays offsets from its offset log through {@link #deserializeOffset} and + * {@link #planInputPartitions}. The epoch counter fast forwards past every epoch seen there, so + * {@link #latestOffset()} never emits an offset smaller than one already committed to the log. + */ +public class BeamMicroBatchStream implements MicroBatchStream { + + 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 long maxRecordsPerBatch; + 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.maxRecordsPerBatch = options.getLong(BeamStreamingSource.OPT_MAX_RECORDS, -1L); + 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) { + BeamOffset offset = BeamOffset.fromJson(json); + fastForwardEpoch(offset.epoch()); + return offset; + } + + @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) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); + List splits = splits(); + InputPartition[] partitions = new InputPartition[splits.size()]; + for (int i = 0; i < splits.size(); i++) { + partitions[i] = + new BeamInputPartition( + splits.get(i), + coderB64, + pipelineOptionsB64, + sourceId, + i, + checkpointLocation, + startEpoch, + endEpoch, + maxRecordsPerBatch, + maxBatchDurationMillis); + } + return partitions; + } + + /** + * Raises the epoch counter to {@code seen} if it is behind, keeping {@link #latestOffset()} ahead + * of every offset Spark already logged before a restart. + */ + private synchronized void fastForwardEpoch(long seen) { + if (seen > epoch) { + LOG.info( + "Fast forwarding the epoch of Beam source {} from {} to {} seen in Spark's offset log.", + sourceId, + epoch, + seen); + epoch = seen; + } + } + + @Override + public PartitionReaderFactory createReaderFactory() { + return new BeamPartitionReaderFactory(); + } + + /** + * Returns the pinned split list of this source, loading it from the checkpoint location if a + * previous run pinned one, and splitting the source and pinning the result otherwise. + */ + private synchronized List splits() { + if (splitsB64 != null) { + return splitsB64; + } + List pinned; + try { + pinned = BeamCheckpointFiles.readSplits(checkpointLocation, sourceId); + } catch (IOException e) { + throw new IllegalStateException( + "Failed to read the pinned split list of Beam source " + sourceId, e); + } + if (pinned != null) { + LOG.info( + "Restored {} pinned split(s) of Beam source {} from {}.", + pinned.size(), + sourceId, + checkpointLocation); + splitsB64 = pinned; + return pinned; + } + UnboundedSource source = BeamStreamingSource.decode(sourceB64, "UnboundedSource"); + SerializablePipelineOptions serializableOptions = + BeamStreamingSource.decode(pipelineOptionsB64, "PipelineOptions"); + 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); + try { + BeamCheckpointFiles.writeSplits(checkpointLocation, sourceId, encoded); + } catch (IOException e) { + throw new IllegalStateException("Failed to pin the split list of Beam source " + sourceId, e); + } + splitsB64 = encoded; + return encoded; + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/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..a8dd28328b2e --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -0,0 +1,181 @@ +/* + * 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 maxRecordsPerBatch} elements were emitted (a limit + * below 1 means unlimited) 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 String checkpointLocation; + private final String sourceId; + private final int splitId; + private final long endEpoch; + private final long maxRecordsPerBatch; + 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.checkpointLocation = partition.checkpointLocation(); + this.sourceId = partition.sourceId(); + this.splitId = partition.splitId(); + this.endEpoch = partition.endEpoch(); + this.maxRecordsPerBatch = partition.maxRecordsPerBatch(); + this.maxBatchDurationMillis = partition.maxBatchDurationMillis(); + this.cacheKey = BeamReaderCache.key(checkpointLocation, sourceId, splitId); + long startEpoch = partition.startEpoch(); + this.cached = + BeamReaderCache.getOrCreate( + cacheKey, + source, + options.get(), + () -> BeamCheckpointFiles.readMark(checkpointLocation, sourceId, splitId, startEpoch)); + } + + @Override + public boolean next() throws IOException { + if (deadlineMillis < 0) { + deadlineMillis = System.currentTimeMillis() + maxBatchDurationMillis; + } + while (true) { + if (maxRecordsPerBatch > 0 && recordsRead >= maxRecordsPerBatch) { + 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, persisted for durable recovery, and finalized. + */ + @Override + public void close() { + current = null; + try { + UnboundedSource.CheckpointMark mark = cached.reader().getCheckpointMark(); + BeamReaderCache.rememberCheckpointMark(cacheKey, mark); + persistMark(mark); + mark.finalizeCheckpoint(); + } catch (Exception e) { + LOG.warn("Failed to finalize the checkpoint mark of Beam reader {}.", cacheKey, e); + } + LOG.debug("Beam reader {} emitted {} record(s) in this micro-batch.", cacheKey, recordsRead); + } + + /** + * Best effort persistence of {@code mark} under the checkpoint location. An IO failure only + * degrades recovery after a restart, the in memory path in {@link BeamReaderCache} still works, + * so the batch is never failed here. + */ + private void persistMark(UnboundedSource.CheckpointMark mark) { + try { + BeamCheckpointFiles.writeMark(checkpointLocation, sourceId, splitId, endEpoch, mark); + } catch (Exception e) { + LOG.warn( + "Failed to persist the checkpoint mark of Beam reader {} at epoch {}, recovery after a " + + "restart will fall back to an older mark or to a fresh start.", + cacheKey, + endEpoch, + e); + } + } + + private InternalRow toRow() { + Instant timestamp = cached.reader().getCurrentTimestamp(); + WindowedValue windowedValue = + 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..057883e92687 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -0,0 +1,179 @@ +/* + * 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 java.util.function.Supplier; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; +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 holds the read position. Keeping the reader alive + * between micro-batches lets the next batch continue where the previous one stopped, mirroring + * {@code org.apache.beam.runners.spark.io.MicrobatchSource} in the legacy runner. + * + *

Durable recovery. The {@code CheckpointMark} of every split is remembered in executor + * memory after each micro-batch, and {@link BeamPartitionReader} additionally persists it under the + * checkpoint location, see {@link BeamCheckpointFiles}. When a reader has to be created and no mark + * is in memory, for example after an executor or driver restart or after the cache entry expired, + * the caller supplied fallback restores the newest durable mark at or before the epoch the batch + * starts at. Two caveats remain. The source is consumed with at least once semantics, a mark is + * written when a batch finished reading rather than transactionally with Spark's commit, so a crash + * between the two replays the last micro-batch. And persisting a mark is best effort, an IO failure + * only degrades recovery to an older mark or to a fresh start, it never fails the batch. + */ +public final class BeamReaderCache { + + 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. + */ + public static CachedReader getOrCreate( + String key, UnboundedSource source, PipelineOptions options) { + return getOrCreate(key, source, options, () -> null); + } + + /** + * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if + * there is none. The {@code durableMarkFallback} is consulted only when no mark is in memory + * either, it typically restores a mark persisted by {@link BeamCheckpointFiles} and may return + * {@code null} for a fresh start. + */ + @SuppressWarnings({"unchecked", "nullness"}) // the mark type always matches the source + public static CachedReader getOrCreate( + String key, + UnboundedSource source, + PipelineOptions options, + Supplier<@Nullable CheckpointMark> durableMarkFallback) { + try { + return (CachedReader) + READERS.get( + key, + () -> { + CheckpointMarkT mark = (CheckpointMarkT) MARKS.get(key); + if (mark == null) { + mark = (CheckpointMarkT) durableMarkFallback.get(); + } + LOG.info( + "No cached Beam reader for {}, creating one at checkpoint mark {}.", key, mark); + return new CachedReader<>(source.createReader(options, mark)); + }); + } 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..e895d764ff1f --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java @@ -0,0 +1,157 @@ +/* + * 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: + * + *

    + *
  • {@code payload} of type {@code BINARY}, holding the element encoded with the Beam {@code + * WindowedValues.FullWindowedValueCoder} supplied by the translator. + *
  • {@code eventTimestamp} of type {@code TIMESTAMP}, holding the event timestamp reported by + * the Beam reader for that element. + *
+ * + *

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"; + + /** + * Deterministic identifier of the source, derived from the full name of the read transform, see + * {@link UnboundedSourceDataset#sourceId}. It keys the reader cache and the durable checkpoint + * state, so it must stay stable across restarts against the same checkpoint location. + */ + public static final String OPT_SOURCE_ID = "beam.sourceid"; + + /** Desired number of splits handed to {@code UnboundedSource.split}. */ + public static final String OPT_NUM_SPLITS = "beam.numsplits"; + + /** Maximum number of records read per split per micro-batch, values below 1 mean no limit. */ + 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. + * + *

The return type is inferred from the call site. Deserialization cannot check it, and the + * decoded types include generic ones such as {@code Coder>} that no {@code + * Class} token can express, so a checked variant is not available here. + */ + @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"}) + static T decode(String encoded, String description) { + return (T) + SerializableUtils.deserializeFromByteArray( + 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..e385542b50be --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -0,0 +1,133 @@ +/* + * 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.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.Map; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +/** + * 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; + + /** Upper bound on the sanitized transform name inside a source id, see {@link #sourceId}. */ + private static final int MAX_SANITIZED_NAME_LENGTH = 64; + + private UnboundedSourceDataset() {} + + /** + * 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 transformName the full name of the read transform, turned into the deterministic source + * id keying the durable checkpoint state, see {@link #sourceId} + * @param the element type of the source + * @param the checkpoint mark type of the source + */ + public static Dataset of( + SparkSession session, + UnboundedSource source, + Coder> windowedValueCoder, + SparkStructuredStreamingPipelineOptions options, + String transformName) { + + 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, sourceId(transformName)); + readerOptions.put( + BeamStreamingSource.OPT_NUM_SPLITS, Integer.toString(desiredNumSplits(session))); + readerOptions.put( + BeamStreamingSource.OPT_MAX_RECORDS, Long.toString(options.getMaxRecordsPerBatch())); + 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"); + } + + /** + * Derives the deterministic source id of a read transform from its full name. + * + *

The id keys the durable checkpoint state under the checkpoint location, so it must be + * filesystem safe and stable across JVMs. Characters outside {@code [A-Za-z0-9._-]} are replaced + * with {@code _}, the sanitized name is truncated to {@value #MAX_SANITIZED_NAME_LENGTH} + * characters, and an 8 hex character hash of the unsanitized name is appended to keep distinct + * transform names distinct. A pipeline restarted against the same checkpoint location must + * therefore keep the same transform names, the same requirement every Beam runner with durable + * state imposes. + */ + public static String sourceId(String transformName) { + String sanitized = transformName.replaceAll("[^A-Za-z0-9._-]", "_"); + if (sanitized.length() > MAX_SANITIZED_NAME_LENGTH) { + sanitized = sanitized.substring(0, MAX_SANITIZED_NAME_LENGTH); + } + String hash = + Hashing.murmur3_32_fixed().hashString(transformName, StandardCharsets.UTF_8).toString(); + return sanitized + "-" + hash; + } + + private static int desiredNumSplits(SparkSession session) { + int parallelism = session.sparkContext().defaultParallelism(); + return Math.max(1, Math.min(MAX_DESIRED_SPLITS, parallelism)); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java new file mode 100644 index 000000000000..4b7d7c55a00e --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.Serializable; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for the durable checkpoint recovery pieces of the Spark 4 micro-batch source, the {@link + * BeamCheckpointFiles} layout, the epoch fast forward of {@link BeamMicroBatchStream} and the + * deterministic source id of {@link UnboundedSourceDataset}. + */ +@RunWith(JUnit4.class) +public class BeamCheckpointRecoveryTest { + + private static final String SOURCE_ID = "read-source-cafe0123"; + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + @Test + public void testSplitsRoundTrip() throws Exception { + String checkpointLocation = temp.newFolder("splits").getAbsolutePath(); + List splits = Arrays.asList("c3BsaXQtMA==", "c3BsaXQtMQ==", "c3BsaXQtMg=="); + + assertNull( + "no splits pinned yet", BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); + + BeamCheckpointFiles.writeSplits(checkpointLocation, SOURCE_ID, splits); + assertEquals(splits, BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); + + File root = new File(checkpointLocation, "beam-source-" + SOURCE_ID); + assertEquals( + "only the final splits file may remain", + new HashSet<>(Arrays.asList("splits")), + fileNames(root)); + } + + @Test + public void testMarkRetentionAndRecovery() throws Exception { + String checkpointLocation = temp.newFolder("marks").getAbsolutePath(); + for (long epoch = 1; epoch <= 4; epoch++) { + BeamCheckpointFiles.writeMark( + checkpointLocation, SOURCE_ID, 0, epoch, new TestMark((int) epoch)); + } + + File marksDir = new File(new File(checkpointLocation, "beam-source-" + SOURCE_ID), "marks/0"); + assertEquals( + "retention must keep the two highest epochs only", + new HashSet<>(Arrays.asList("3", "4")), + fileNames(marksDir)); + + assertEquals(4, restoredPosition(checkpointLocation, 0, 4)); + assertEquals(3, restoredPosition(checkpointLocation, 0, 3)); + // No exact file for epoch 7, the largest epoch not exceeding it is 4. + assertEquals(4, restoredPosition(checkpointLocation, 0, 7)); + // Epochs 1 and 2 were deleted by retention, nothing at or below 2 is left. + assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 0, 2)); + // A different split has no marks at all. + assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 1, 4)); + } + + @Test + public void testEpochFastForwardsPastDeserializedOffset() throws Exception { + BeamMicroBatchStream stream = stream(temp.newFolder("ff-offset").getAbsolutePath()); + stream.deserializeOffset("{\"epoch\":7}"); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the replayed epoch 7, got " + next, next.epoch() > 7L); + } + + @Test + public void testEpochFastForwardsPastPlannedOffsets() throws Exception { + BeamMicroBatchStream stream = stream(temp.newFolder("ff-plan").getAbsolutePath()); + InputPartition[] partitions = + stream.planInputPartitions(new BeamOffset(3L), new BeamOffset(9L)); + assertTrue("at least one partition expected", partitions.length > 0); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the planned epoch 9, got " + next, next.epoch() > 9L); + } + + @Test + public void testSourceIdIsDeterministicAndFilesystemSafe() { + String name = "Read from PubSub/PubsubUnboundedSource (with: spaces & colons)"; + String id = UnboundedSourceDataset.sourceId(name); + + assertEquals("same name must give the same id", id, UnboundedSourceDataset.sourceId(name)); + assertNotEquals( + "different names must give different ids", + UnboundedSourceDataset.sourceId("Read A"), + UnboundedSourceDataset.sourceId("Read B")); + assertTrue("id must be filesystem safe: " + id, id.matches("[A-Za-z0-9._-]+")); + + String longName = String.join("/", java.util.Collections.nCopies(50, "NestedTransform")); + String longId = UnboundedSourceDataset.sourceId(longName); + assertTrue("id length must stay bounded: " + longId.length(), longId.length() <= 73); + assertNotEquals( + "truncated names must still differ through the hash", + longId, + UnboundedSourceDataset.sourceId(longName + "/Tail")); + } + + private static int restoredPosition(String checkpointLocation, int splitId, long startEpoch) { + CheckpointMark mark = + BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, splitId, startEpoch); + assertNotNull("expected a durable mark at or before epoch " + startEpoch, mark); + return ((TestMark) mark).position; + } + + /** + * Lists the visible files of {@code dir}, asserting no {@code .tmp} file was left behind and + * ignoring the hidden checksum sidecars of Hadoop's local filesystem. + */ + private static Set fileNames(File dir) { + String[] names = dir.list(); + assertNotNull("expected directory " + dir, names); + Set visible = new HashSet<>(); + for (String name : names) { + assertFalse("no temporary file may remain: " + name, name.endsWith(".tmp")); + if (!name.startsWith(".")) { + visible.add(name); + } + } + return visible; + } + + /** Builds a stream over a real source, driver side only, no Spark session required. */ + private static BeamMicroBatchStream stream(String checkpointLocation) { + Map options = new HashMap<>(); + options.put( + BeamStreamingSource.OPT_SOURCE, BeamStreamingSource.encode(CountingSource.unbounded())); + options.put( + BeamStreamingSource.OPT_CODER, + BeamStreamingSource.encode( + WindowedValues.getFullCoder( + org.apache.beam.sdk.coders.VarLongCoder.of(), GlobalWindow.Coder.INSTANCE))); + options.put( + BeamStreamingSource.OPT_PIPELINE_OPTIONS, + BeamStreamingSource.encode( + new SerializablePipelineOptions(PipelineOptionsFactory.create()))); + options.put(BeamStreamingSource.OPT_SOURCE_ID, SOURCE_ID); + options.put(BeamStreamingSource.OPT_NUM_SPLITS, "2"); + return new BeamMicroBatchStream(new CaseInsensitiveStringMap(options), checkpointLocation); + } + + /** A trivial serializable checkpoint mark carrying a read position. */ + private static class TestMark implements CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; + + private final int position; + + TestMark(int position) { + this.position = position; + } + + @Override + public void finalizeCheckpoint() {} + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java new file mode 100644 index 000000000000..0d1bee7a9dab --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -0,0 +1,520 @@ +/* + * 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 default record limit of -1 means unlimited, an available source drains in one batch. */ + @Test(timeout = 300_000) + public void testUnlimitedRecordsPerBatchByDefault() throws Exception { + int count = 2500; + SparkStructuredStreamingPipelineOptions options = + org.apache.beam.sdk.options.PipelineOptionsFactory.create() + .as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(0L); + options.setMaxBatchDurationMillis(5_000L); + Dataset rows = + UnboundedSourceDataset.of( + SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); + + String queryName = "beam_nolimit_" + QUERY_COUNTER.incrementAndGet(); + List batchSizes = Collections.synchronizedList(new ArrayList<>()); + StreamingQuery query = + rows.writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + long size = batch.count(); + if (size > 0) { + batchSizes.add(size); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(queryName).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + try { + long deadline = System.currentTimeMillis() + POLL_TIMEOUT_MILLIS; + while (System.currentTimeMillis() < deadline && batchSizes.isEmpty()) { + Thread.sleep(100L); + } + } finally { + stopQuietly(query); + } + assertEquals( + "without a limit all elements must arrive in the first non-empty micro-batch", + Collections.singletonList((long) count), + new ArrayList<>(batchSizes)); + } + + /** 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.setMaxRecordsPerBatch(1000L); + options.setMaxBatchDurationMillis(200L); + return UnboundedSourceDataset.of( + SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); + } + + 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 {} + } + } +} From 8af17d50a48ac7a9de48368357b964e6bc952d83 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 3 Sep 2026 06:28:15 +0000 Subject: [PATCH 2/4] [Spark 4] Align the micro-batch source with Spark's commit lifecycle Marks are no longer finalized when a partition reader closes. A reader finalizes the mark taken at its start offset when the next micro-batch for that split is scheduled, because Spark only starts a batch at the initial offset or at the end offset of a batch already in its commit log. A reader whose position does not match the start offset, or that moved without completing its batch (task retry, killed attempt, executor change, restart), is dropped without finalizing and recreated from the durable mark at that offset. Per source state lives under the checkpoint location Spark hands to toMicroBatchStream, written through CheckpointFileManager with the session Hadoop configuration broadcast to executors. Marks are coded with the source's checkpoint mark coder. commit(end) purges marks below end on a background thread, nothing is retained by a fixed count. The dataset is built like BoundedDatasetFactory, a Table holding the real objects wrapped in StreamingRelationV2, no string options, no Base64. Splits travel as objects in the InputPartition, options and Hadoop configuration as broadcasts. maxRecordsPerBatch is divided across splits like the legacy MicrobatchSource, defaultParallelism decides the split count, idle readers back off with FluentBackoff, offsets serialize as the bare epoch like LongOffset with the base class equality. A new option readerIdleTimeoutMillis bounds how long an executor keeps an idle reader. Tests drive the reader cache protocol directly and prove restart recovery, finalization only after commit, and mark purging against Spark's real offsets and commits logs. The JUnit per test timeout is removed from the streaming test, its throwaway thread group poisoned Spark's static pools for later batch tests in the same JVM. --- .../io/streaming/BeamCheckpointFiles.java | 236 ------- .../io/streaming/BeamInputPartition.java | 110 +-- .../io/streaming/BeamMicroBatchStream.java | 292 +++++--- .../io/streaming/BeamOffset.java | 36 +- .../io/streaming/BeamPartitionReader.java | 196 +++--- .../streaming/BeamPartitionReaderFactory.java | 8 +- .../io/streaming/BeamReaderCache.java | 286 +++++--- .../io/streaming/BeamSourceCheckpoint.java | 184 +++++ .../io/streaming/BeamSourceSpec.java | 100 +++ .../io/streaming/BeamStreamingSource.java | 157 ----- .../io/streaming/BeamStreamingTable.java | 46 +- .../io/streaming/UnboundedSourceDataset.java | 128 ++-- .../streaming/BeamCheckpointRecoveryTest.java | 198 ------ .../streaming/BeamMicroBatchSourceTest.java | 650 +++++++++++++++++- .../BeamReaderCacheProtocolTest.java | 478 +++++++++++++ .../streaming/BeamSourceCheckpointTest.java | 209 ++++++ ...arkStructuredStreamingPipelineOptions.java | 9 + 17 files changed, 2257 insertions(+), 1066 deletions(-) delete mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java delete mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java delete mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java deleted file mode 100644 index 069bb27e9f65..000000000000 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java +++ /dev/null @@ -1,236 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.runners.spark.structuredstreaming.io.streaming; - -import java.io.IOException; -import java.io.Serializable; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; -import org.apache.beam.sdk.util.SerializableUtils; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; -import org.apache.hadoop.conf.Configuration; -import org.apache.hadoop.fs.FSDataInputStream; -import org.apache.hadoop.fs.FSDataOutputStream; -import org.apache.hadoop.fs.FileStatus; -import org.apache.hadoop.fs.FileSystem; -import org.apache.hadoop.fs.Path; -import org.checkerframework.checker.nullness.qual.Nullable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Durable state of the Beam micro-batch source, stored next to Spark's own streaming state under - * {@code /beam-source-/}. - * - *

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

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

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

The mark written under exactly {@code startEpoch} is preferred, if it is absent the mark - * with the largest epoch not exceeding {@code startEpoch} is used. Returns {@code null}, meaning - * a fresh start, when no such mark exists or reading fails. - */ - public static @Nullable CheckpointMark readMark( - String checkpointLocation, String sourceId, int splitId, long startEpoch) { - Path dir = marksDir(checkpointLocation, sourceId, splitId); - try { - FileSystem fs = fileSystem(dir); - if (!fs.exists(dir)) { - return null; - } - long epoch = Long.MIN_VALUE; - if (fs.exists(new Path(dir, Long.toString(startEpoch)))) { - epoch = startEpoch; - } else { - for (FileStatus status : fs.listStatus(dir)) { - @Nullable Long candidate = parseEpoch(status.getPath().getName()); - if (candidate != null && candidate <= startEpoch && candidate > epoch) { - epoch = candidate; - } - } - } - if (epoch == Long.MIN_VALUE) { - return null; - } - Path path = new Path(dir, Long.toString(epoch)); - CheckpointMark mark = - (CheckpointMark) deserialize(read(fs, path), "durable checkpoint mark " + path); - LOG.info( - "Restored durable checkpoint mark of Beam source {} split {} at epoch {} " - + "(requested epoch {}).", - sourceId, - splitId, - epoch, - startEpoch); - return mark; - } catch (IOException e) { - LOG.warn( - "Failed to read a durable checkpoint mark of Beam source {} split {} at epoch {}, " - + "the reader starts without one.", - sourceId, - splitId, - startEpoch, - e); - return null; - } - } - - private static Path root(String checkpointLocation, String sourceId) { - return new Path(checkpointLocation, ROOT_PREFIX + sourceId); - } - - private static Path marksDir(String checkpointLocation, String sourceId, int splitId) { - return new Path( - new Path(root(checkpointLocation, sourceId), MARKS_DIR), String.valueOf(splitId)); - } - - private static FileSystem fileSystem(Path path) throws IOException { - return path.getFileSystem(new Configuration()); - } - - /** Writes {@code bytes} to a {@code .tmp} sibling of {@code target} and renames it into place. */ - private static void writeAtomically(FileSystem fs, Path target, byte[] bytes) throws IOException { - Path tmp = new Path(target.getParent(), target.getName() + TMP_SUFFIX); - try (FSDataOutputStream out = fs.create(tmp, true)) { - out.write(bytes); - } - // On HDFS a rename onto an existing target fails, a mark rewritten by a task retry is the - // same position, so an existing target counts as success. - if (!fs.rename(tmp, target) && !fs.exists(target)) { - throw new IOException("Failed to rename " + tmp + " to " + target); - } - if (fs.exists(tmp)) { - fs.delete(tmp, false); - } - } - - private static byte[] read(FileSystem fs, Path path) throws IOException { - try (FSDataInputStream in = fs.open(path)) { - return ByteStreams.toByteArray(in); - } - } - - private static Object deserialize(byte[] bytes, String description) { - return SerializableUtils.deserializeFromByteArray(bytes, description); - } - - /** Best effort deletion of mark files older than the {@value #RETAINED_MARKS} highest epochs. */ - private static void deleteOldMarks(FileSystem fs, Path dir) { - try { - List epochs = new ArrayList<>(); - for (FileStatus status : fs.listStatus(dir)) { - @Nullable Long epoch = parseEpoch(status.getPath().getName()); - if (epoch != null) { - epochs.add(epoch); - } - } - Collections.sort(epochs); - for (int i = 0; i < epochs.size() - RETAINED_MARKS; i++) { - fs.delete(new Path(dir, Long.toString(epochs.get(i))), false); - } - } catch (IOException e) { - LOG.warn("Failed to delete outdated checkpoint marks under {}.", dir, e); - } - } - - private static @Nullable Long parseEpoch(String fileName) { - try { - return Long.valueOf(fileName); - } catch (NumberFormatException e) { - return null; - } - } -} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java index 8a1103dac46a..39b562fe0f8f 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamInputPartition.java @@ -17,76 +17,84 @@ */ package org.apache.beam.runners.spark.structuredstreaming.io.streaming; +import java.util.Arrays; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +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.broadcast.Broadcast; import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.util.SerializableConfiguration; -/** - * 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 { +/** One split of a Beam unbounded source for one micro-batch, from epoch start to epoch end. */ +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 UnboundedSource split; + private final Coder> coder; + private final Broadcast options; + private final Broadcast hadoopConf; private final String checkpointLocation; + private final int splitId; private final long startEpoch; private final long endEpoch; - private final long maxRecordsPerBatch; + private final long maxRecords; private final long maxBatchDurationMillis; + private final long readerIdleTimeoutMillis; + private final String[] preferredLocations; BeamInputPartition( - String sourceB64, - String coderB64, - String pipelineOptionsB64, - String sourceId, - int splitId, + UnboundedSource split, + Coder> coder, + Broadcast options, + Broadcast hadoopConf, String checkpointLocation, + int splitId, long startEpoch, long endEpoch, - long maxRecordsPerBatch, - long maxBatchDurationMillis) { - this.sourceB64 = sourceB64; - this.coderB64 = coderB64; - this.pipelineOptionsB64 = pipelineOptionsB64; - this.sourceId = sourceId; - this.splitId = splitId; + long maxRecords, + long maxBatchDurationMillis, + long readerIdleTimeoutMillis, + String[] preferredLocations) { + this.split = split; + this.coder = coder; + this.options = options; + this.hadoopConf = hadoopConf; this.checkpointLocation = checkpointLocation; + this.splitId = splitId; this.startEpoch = startEpoch; this.endEpoch = endEpoch; - this.maxRecordsPerBatch = maxRecordsPerBatch; + this.maxRecords = maxRecords; this.maxBatchDurationMillis = maxBatchDurationMillis; + this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + this.preferredLocations = preferredLocations.clone(); } - String sourceB64() { - return sourceB64; + UnboundedSource split() { + return split; } - String coderB64() { - return coderB64; + Coder> coder() { + return coder; } - String pipelineOptionsB64() { - return pipelineOptionsB64; + Broadcast options() { + return options; } - String sourceId() { - return sourceId; - } - - int splitId() { - return splitId; + Broadcast hadoopConf() { + return hadoopConf; } String checkpointLocation() { return checkpointLocation; } + int splitId() { + return splitId; + } + long startEpoch() { return startEpoch; } @@ -95,16 +103,36 @@ long endEpoch() { return endEpoch; } - long maxRecordsPerBatch() { - return maxRecordsPerBatch; + /** Records this split may emit in this micro-batch, below 1 means unlimited. */ + long maxRecords() { + return maxRecords; } long maxBatchDurationMillis() { return maxBatchDurationMillis; } + long readerIdleTimeoutMillis() { + return readerIdleTimeoutMillis; + } + + @Override + public String[] preferredLocations() { + return preferredLocations.clone(); + } + @Override public String toString() { - return "BeamInputPartition{source=" + sourceId + ", split=" + splitId + "}"; + return "BeamInputPartition{checkpointLocation=" + + checkpointLocation + + ", split=" + + splitId + + ", epochs=" + + startEpoch + + ".." + + endEpoch + + ", locations=" + + Arrays.toString(preferredLocations) + + "}"; } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java index f93fdc3113c4..0ec8fb061f4e 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchStream.java @@ -19,67 +19,66 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; -import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.spark.SparkEnv; 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.apache.spark.storage.BlockManager; +import org.apache.spark.storage.BlockManagerId; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import scala.collection.Iterator; /** - * A {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * Driver side {@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()}). + *

Offsets are opaque epochs, {@link #latestOffset()} advances by one every trigger. The source + * is split once and the splits are pinned under the checkpoint location, every batch of every run + * plans the same partitions. {@link #commit} purges marks below the committed epoch on a background + * thread, Spark never asks for those again. * - *

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

On a restart Spark replays offsets from its offset log through {@link #deserializeOffset} and - * {@link #planInputPartitions}. The epoch counter fast forwards past every epoch seen there, so - * {@link #latestOffset()} never emits an offset smaller than one already committed to the log. + *

Each split prefers the executor rendezvous hashing assigns it, so an executor joining or + * leaving moves only the splits of that executor. Locality is a hint, the reader cache restores a + * split from its durable mark wherever it lands. */ -public class BeamMicroBatchStream implements MicroBatchStream { +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 BeamSourceSpec spec; private final String checkpointLocation; - private final int desiredNumSplits; - private final long maxRecordsPerBatch; - private final long maxBatchDurationMillis; + private final BeamSourceCheckpoint checkpoint; + + private final ExecutorService purger = + Executors.newSingleThreadExecutor( + runnable -> { + Thread thread = new Thread(runnable, "beam-source-mark-purge"); + thread.setDaemon(true); + return thread; + }); + private final AtomicBoolean purgeInFlight = new AtomicBoolean(); + private final AtomicLong purgeRequested = new AtomicLong(); 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); + private @Nullable List> splits; + + BeamMicroBatchStream(BeamSourceSpec spec, String checkpointLocation) { + this.spec = spec; this.checkpointLocation = checkpointLocation; - this.desiredNumSplits = Math.max(1, options.getInt(BeamStreamingSource.OPT_NUM_SPLITS, 1)); - this.maxRecordsPerBatch = options.getLong(BeamStreamingSource.OPT_MAX_RECORDS, -1L); - this.maxBatchDurationMillis = - Math.max(1L, options.getLong(BeamStreamingSource.OPT_MAX_BATCH_DURATION_MILLIS, 500L)); + this.checkpoint = + new BeamSourceCheckpoint(checkpointLocation, spec.hadoopConf().value().value()); } @Override @@ -99,113 +98,180 @@ public Offset deserializeOffset(String json) { return offset; } - @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) { long startEpoch = ((BeamOffset) start).epoch(); long endEpoch = ((BeamOffset) end).epoch(); fastForwardEpoch(endEpoch); - List splits = splits(); - InputPartition[] partitions = new InputPartition[splits.size()]; - for (int i = 0; i < splits.size(); i++) { + List> pinned = splits(); + long[] quotas = splitQuotas(spec.maxRecordsPerBatch(), pinned.size()); + List executors = sortedExecutors(); + InputPartition[] partitions = new InputPartition[pinned.size()]; + for (int i = 0; i < pinned.size(); i++) { + String[] locations = + executors.isEmpty() ? new String[0] : new String[] {assign(i, executors)}; partitions[i] = - new BeamInputPartition( - splits.get(i), - coderB64, - pipelineOptionsB64, - sourceId, - i, + new BeamInputPartition<>( + pinned.get(i), + spec.coder(), + spec.options(), + spec.hadoopConf(), checkpointLocation, + i, startEpoch, endEpoch, - maxRecordsPerBatch, - maxBatchDurationMillis); + quotas[i], + spec.maxBatchDurationMillis(), + spec.readerIdleTimeoutMillis(), + locations); } return partitions; } - /** - * Raises the epoch counter to {@code seen} if it is behind, keeping {@link #latestOffset()} ahead - * of every offset Spark already logged before a restart. - */ - private synchronized void fastForwardEpoch(long seen) { - if (seen > epoch) { - LOG.info( - "Fast forwarding the epoch of Beam source {} from {} to {} seen in Spark's offset log.", - sourceId, - epoch, - seen); - epoch = seen; + @Override + public PartitionReaderFactory createReaderFactory() { + return new BeamPartitionReaderFactory(); + } + + /** Purges marks below {@code end} off the stream thread, one purge runs at a time. */ + @Override + public void commit(Offset end) { + long endEpoch = ((BeamOffset) end).epoch(); + int numSplits = splits().size(); + purgeRequested.accumulateAndGet(endEpoch, Math::max); + if (purgeInFlight.compareAndSet(false, true)) { + purger.execute(() -> purgeRequested(numSplits)); } } + private void purgeRequested(int numSplits) { + long epoch; + do { + epoch = purgeRequested.get(); + try { + for (int i = 0; i < numSplits; i++) { + checkpoint.purgeMarksBelow(i, epoch); + } + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to purge marks below epoch {} at {}.", epoch, checkpointLocation, e); + } + purgeInFlight.set(false); + } while (purgeRequested.get() > epoch && purgeInFlight.compareAndSet(false, true)); + } + @Override - public PartitionReaderFactory createReaderFactory() { - return new BeamPartitionReaderFactory(); + public void stop() { + LOG.info( + "Stopping Beam micro-batch stream {} at {}.", spec.transformName(), checkpointLocation); + purger.shutdown(); } - /** - * Returns the pinned split list of this source, loading it from the checkpoint location if a - * previous run pinned one, and splitting the source and pinning the result otherwise. - */ - private synchronized List splits() { - if (splitsB64 != null) { - return splitsB64; + /** Keeps {@link #latestOffset()} ahead of every epoch Spark logged before a restart. */ + private synchronized void fastForwardEpoch(long seen) { + if (seen > epoch) { + LOG.info("Fast forwarding epoch of {} from {} to {}.", spec.transformName(), epoch, seen); + epoch = seen; + } + } + + private synchronized List> splits() { + if (splits != null) { + return splits; } - List pinned; + List> pinned; try { - pinned = BeamCheckpointFiles.readSplits(checkpointLocation, sourceId); + pinned = checkpoint.readSplits(); } catch (IOException e) { - throw new IllegalStateException( - "Failed to read the pinned split list of Beam source " + sourceId, e); + throw new IllegalStateException("Failed to read pinned splits at " + checkpointLocation, e); + } + if (pinned == null) { + pinned = new ArrayList<>(splitSource()); + try { + checkpoint.writeSplits(pinned); + } catch (IOException e) { + throw new IllegalStateException("Failed to pin splits at " + checkpointLocation, e); + } + } else { + LOG.info("Restored {} pinned split(s) from {}.", pinned.size(), checkpointLocation); } - if (pinned != null) { - LOG.info( - "Restored {} pinned split(s) of Beam source {} from {}.", - pinned.size(), - sourceId, - checkpointLocation); - splitsB64 = pinned; - return pinned; + List> typed = new ArrayList<>(pinned.size()); + for (UnboundedSource split : pinned) { + @SuppressWarnings("unchecked") // splits of this source share its element type + UnboundedSource cast = (UnboundedSource) split; + typed.add(cast); } - UnboundedSource source = BeamStreamingSource.decode(sourceB64, "UnboundedSource"); - SerializablePipelineOptions serializableOptions = - BeamStreamingSource.decode(pipelineOptionsB64, "PipelineOptions"); - PipelineOptions options = serializableOptions.get(); - List> split; + splits = typed; + return typed; + } + + private List> splitSource() { + UnboundedSource source = spec.source(); + PipelineOptions options = spec.options().value().get(); + List> result; try { - split = source.split(desiredNumSplits, options); + result = source.split(spec.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)); + if (result.isEmpty()) { + result = Collections.singletonList(source); } LOG.info( - "Split Beam source {} into {} partition(s) (desired {}).", - sourceId, - encoded.size(), - desiredNumSplits); + "Split {} into {} partition(s), desired {}.", + spec.transformName(), + result.size(), + spec.desiredNumSplits()); + return result; + } + + /** + * Divides the batch quota over the splits, the remainder goes to the first splits. Below 1 means + * unlimited for every split. A quota below the split count gives every split one record. + */ + static long[] splitQuotas(long maxRecordsPerBatch, int numSplits) { + long[] quotas = new long[numSplits]; + if (maxRecordsPerBatch < 1) { + Arrays.fill(quotas, maxRecordsPerBatch); + return quotas; + } + long base = maxRecordsPerBatch / numSplits; + long remainder = maxRecordsPerBatch % numSplits; + for (int i = 0; i < numSplits; i++) { + quotas[i] = Math.max(1L, base + (i < remainder ? 1 : 0)); + } + return quotas; + } + + /** Rendezvous hashing of a split over the executors, stable under membership changes. */ + static String assign(int splitId, List executors) { + String best = executors.get(0); + int bestHash = Integer.MIN_VALUE; + for (String executor : executors) { + int hash = Hashing.murmur3_32_fixed().hashUnencodedChars(executor + '#' + splitId).asInt(); + if (hash > bestHash) { + bestHash = hash; + best = executor; + } + } + return best; + } + + /** Sorted {@code executor__} locations, empty in local mode or on any failure. */ + private static List sortedExecutors() { try { - BeamCheckpointFiles.writeSplits(checkpointLocation, sourceId, encoded); - } catch (IOException e) { - throw new IllegalStateException("Failed to pin the split list of Beam source " + sourceId, e); + BlockManager bm = SparkEnv.get().blockManager(); + Iterator peers = bm.master().getPeers(bm.blockManagerId()).iterator(); + List executors = new ArrayList<>(); + while (peers.hasNext()) { + BlockManagerId peer = peers.next(); + executors.add("executor_" + peer.host() + "_" + peer.executorId()); + } + Collections.sort(executors); + return executors; + } catch (RuntimeException e) { + LOG.debug("No executor list available for preferred locations.", e); + return Collections.emptyList(); } - 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 index f7e5e0e56027..ba5e9825e27a 100644 --- 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 @@ -17,59 +17,39 @@ */ 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. + * Opaque epoch counter used as the Spark {@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}. + *

The read position lives in Beam checkpoint marks on the executors, see {@link + * BeamSourceCheckpoint}. Equality is the base class comparison of {@link #json()}. */ 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 + "}"; + return Long.toString(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); + try { + return new BeamOffset(Long.parseLong(json.trim())); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Not a valid BeamOffset: " + json, e); } - 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 diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java index a8dd28328b2e..d500aec44644 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -18,82 +18,80 @@ 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.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; 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.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.util.BackOff; +import org.apache.beam.sdk.util.BackOffUtils; +import org.apache.beam.sdk.util.FluentBackoff; +import org.apache.beam.sdk.util.Sleeper; 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.hadoop.conf.Configuration; +import org.apache.spark.TaskContext; 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.Duration; 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. + * Reads one split of a Beam {@link UnboundedSource} for one micro-batch. * - *

The batch ends as soon as either {@code maxRecordsPerBatch} elements were emitted (a limit - * below 1 means unlimited) 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 batch ends at the record quota or at the deadline. The reader then writes its checkpoint + * mark durably at the end epoch and stays in {@link BeamReaderCache} for the next batch. A failed + * mark write fails the task. An attempt Spark killed or failed writes nothing and its reader is + * dropped, the retry restores from the durable mark at the start epoch. * - *

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 + * @param the element type of the split */ -@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 static final Duration INITIAL_BACKOFF = Duration.millis(10); - private final String cacheKey; + private final String key; + private final UnboundedSource split; + private final Coder> coder; + private final BeamSourceCheckpoint checkpoint; private final CachedReader cached; - private final Coder> windowedValueCoder; - private final String checkpointLocation; - private final String sourceId; private final int splitId; private final long endEpoch; - private final long maxRecordsPerBatch; + private final long maxRecords; private final long maxBatchDurationMillis; private long recordsRead; private long deadlineMillis = -1L; + private boolean batchEnded; 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.checkpointLocation = partition.checkpointLocation(); - this.sourceId = partition.sourceId(); + BeamPartitionReader(BeamInputPartition partition) throws IOException { + this.split = partition.split(); + this.coder = partition.coder(); this.splitId = partition.splitId(); this.endEpoch = partition.endEpoch(); - this.maxRecordsPerBatch = partition.maxRecordsPerBatch(); + this.maxRecords = partition.maxRecords(); this.maxBatchDurationMillis = partition.maxBatchDurationMillis(); - this.cacheKey = BeamReaderCache.key(checkpointLocation, sourceId, splitId); + PipelineOptions options = partition.options().value().get(); + Configuration conf = partition.hadoopConf().value().value(); + this.checkpoint = new BeamSourceCheckpoint(partition.checkpointLocation(), conf); + this.key = BeamReaderCache.key(partition.checkpointLocation(), splitId); long startEpoch = partition.startEpoch(); this.cached = - BeamReaderCache.getOrCreate( - cacheKey, - source, - options.get(), - () -> BeamCheckpointFiles.readMark(checkpointLocation, sourceId, splitId, startEpoch)); + BeamReaderCache.acquire( + key, + startEpoch, + split, + options, + partition.readerIdleTimeoutMillis(), + () -> checkpoint.readMark(splitId, startEpoch)); } @Override @@ -101,81 +99,111 @@ public boolean next() throws IOException { if (deadlineMillis < 0) { deadlineMillis = System.currentTimeMillis() + maxBatchDurationMillis; } + BackOff backOff = null; while (true) { - if (maxRecordsPerBatch > 0 && recordsRead >= maxRecordsPerBatch) { - current = null; - return false; + if (maxRecords > 0 && recordsRead >= maxRecords) { + return endOfBatch(false); } long remaining = deadlineMillis - System.currentTimeMillis(); if (remaining <= 0) { - current = null; - return false; + return endOfBatch(false); } if (cached.startOrAdvance()) { recordsRead++; current = toRow(); return true; } - Uninterruptibles.sleepUninterruptibly( - Math.min(remaining, POLL_INTERVAL_MILLIS), java.util.concurrent.TimeUnit.MILLISECONDS); + if (backOff == null) { + backOff = backOff(remaining); + } + try { + if (!BackOffUtils.next(Sleeper.DEFAULT, backOff)) { + return endOfBatch(false); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return endOfBatch(true); + } } } @Override public InternalRow get() { - if (current == null) { + InternalRow row = current; + if (row == null) { throw new IllegalStateException("No current row, next() did not return true."); } - return current; + return row; } - /** - * Ends the micro-batch. The Beam reader deliberately stays open in {@link BeamReaderCache}, only - * its checkpoint mark is remembered, persisted for durable recovery, and finalized. - */ @Override - public void close() { + public void close() throws IOException { + endBatch(attemptDiscarded()); current = null; - try { - UnboundedSource.CheckpointMark mark = cached.reader().getCheckpointMark(); - BeamReaderCache.rememberCheckpointMark(cacheKey, mark); - persistMark(mark); - mark.finalizeCheckpoint(); - } catch (Exception e) { - LOG.warn("Failed to finalize the checkpoint mark of Beam reader {}.", cacheKey, e); - } - LOG.debug("Beam reader {} emitted {} record(s) in this micro-batch.", cacheKey, recordsRead); + } + + private boolean endOfBatch(boolean discarded) throws IOException { + endBatch(discarded); + current = null; + return false; } /** - * Best effort persistence of {@code mark} under the checkpoint location. An IO failure only - * degrades recovery after a restart, the in memory path in {@link BeamReaderCache} still works, - * so the batch is never failed here. + * Ends the batch once. A discarded attempt drops the reader and writes nothing. A reader that was + * never started has not moved, its start mark is written forward, an empty file standing for a + * fresh start. */ - private void persistMark(UnboundedSource.CheckpointMark mark) { - try { - BeamCheckpointFiles.writeMark(checkpointLocation, sourceId, splitId, endEpoch, mark); - } catch (Exception e) { - LOG.warn( - "Failed to persist the checkpoint mark of Beam reader {} at epoch {}, recovery after a " - + "restart will fall back to an older mark or to a fresh start.", - cacheKey, - endEpoch, - e); + private void endBatch(boolean discarded) throws IOException { + if (batchEnded) { + return; } + batchEnded = true; + if (discarded) { + LOG.info("Attempt for Beam reader {} was discarded, dropping the reader.", key); + BeamReaderCache.invalidate(key); + return; + } + if (!cached.started()) { + byte[] startMark = cached.positionMark(); + byte[] codedMark = startMark == null ? new byte[0] : startMark; + checkpoint.writeMark(splitId, endEpoch, codedMark); + cached.endBatch(endEpoch, null, codedMark); + return; + } + CheckpointMark mark = cached.reader().getCheckpointMark(); + byte[] codedMark = encodeMark(split, mark); + checkpoint.writeMark(splitId, endEpoch, codedMark); + cached.endBatch(endEpoch, mark, codedMark); + LOG.debug("Beam reader {} read {} record(s) up to epoch {}.", key, recordsRead, endEpoch); + } + + private static boolean attemptDiscarded() { + TaskContext context = TaskContext.get(); + return context != null && (context.isInterrupted() || context.isFailed()); + } + + private static byte[] encodeMark( + UnboundedSource source, CheckpointMark mark) { + @SuppressWarnings("unchecked") // getCheckpointMark returns the source's own mark type + MarkT typed = (MarkT) mark; + return CoderHelpers.toByteArray(typed, source.getCheckpointMarkCoder()); + } + + private static BackOff backOff(long remainingMillis) { + Duration remaining = Duration.millis(remainingMillis); + return FluentBackoff.DEFAULT + .withInitialBackoff(INITIAL_BACKOFF) + .withMaxBackoff(remaining) + .withMaxCumulativeBackoff(remaining) + .backoff(); } private InternalRow toRow() { Instant timestamp = cached.reader().getCurrentTimestamp(); - WindowedValue windowedValue = + WindowedValue value = 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. + byte[] payload = CoderHelpers.toByteArray(value, coder); + // Spark stores TimestampType as microseconds. 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 index db2573eff315..a360cf895b7b 100644 --- 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 @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.spark.structuredstreaming.io.streaming; +import java.io.IOException; +import java.io.UncheckedIOException; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.connector.read.InputPartition; import org.apache.spark.sql.connector.read.PartitionReader; @@ -29,6 +31,10 @@ public class BeamPartitionReaderFactory implements PartitionReaderFactory { @Override public PartitionReader createReader(InputPartition partition) { - return new BeamPartitionReader<>((BeamInputPartition) partition); + try { + return new BeamPartitionReader<>((BeamInputPartition) partition); + } catch (IOException e) { + throw new UncheckedIOException("Failed to open Beam reader for " + partition, e); + } } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index 057883e92687..cbc697412b18 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -19,151 +19,188 @@ import java.io.Closeable; import java.io.IOException; +import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.TimeUnit; -import java.util.function.Supplier; +import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; 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). + * Executor side cache of live Beam {@link UnboundedReader}s keyed by checkpoint location and split. * - *

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

An entry records the epoch its reader is positioned at and the mark taken there. A batch + * starting at that epoch reuses the reader and finalizes the pending mark, the start epoch of a + * batch is always committed by Spark. Any other start epoch, or a reader that moved without + * completing its batch, closes the entry without finalizing and restores the reader from the + * durable mark at the start epoch. * - *

Durable recovery. The {@code CheckpointMark} of every split is remembered in executor - * memory after each micro-batch, and {@link BeamPartitionReader} additionally persists it under the - * checkpoint location, see {@link BeamCheckpointFiles}. When a reader has to be created and no mark - * is in memory, for example after an executor or driver restart or after the cache entry expired, - * the caller supplied fallback restores the newest durable mark at or before the epoch the batch - * starts at. Two caveats remain. The source is consumed with at least once semantics, a mark is - * written when a batch finished reading rather than transactionally with Spark's commit, so a crash - * between the two replays the last micro-batch. And persisting a mark is best effort, an IO failure - * only degrades recovery to an older mark or to a fresh start, it never fails the batch. + *

A reader idle for longer than its idle timeout is closed without finalizing its mark, the + * source redelivers. The timeout must exceed the longest gap between two micro-batches of one + * split. Speculative execution can leave a losing attempt's mark finalized on another executor, + * this source is not safe under {@code spark.speculation} with sources whose reads are not + * deterministic. */ 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(); + private static final ConcurrentMap> READERS = new ConcurrentHashMap<>(); - /** Last known checkpoint mark per key, used when a reader has to be recreated. */ - private static final ConcurrentMap MARKS = new ConcurrentHashMap<>(); + /** One monitor per key, acquire serializes per split, not across splits. */ + private static final ConcurrentMap LOCKS = 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; + public static String key(String checkpointLocation, int splitId) { + return checkpointLocation + '|' + splitId; } - /** - * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if - * there is none. - */ - public static CachedReader getOrCreate( - String key, UnboundedSource source, PipelineOptions options) { - return getOrCreate(key, source, options, () -> null); + /** Supplies the durable coded mark at the start epoch of a batch, null if there is none. */ + @FunctionalInterface + interface MarkRestorer { + byte @Nullable [] restore() throws IOException; } /** - * Returns the cached reader for {@code key}, creating it from the last cached checkpoint mark if - * there is none. The {@code durableMarkFallback} is consulted only when no mark is in memory - * either, it typically restores a mark persisted by {@link BeamCheckpointFiles} and may return - * {@code null} for a fresh start. + * Returns the reader for {@code key} positioned at {@code startEpoch}, reusing the cached one if + * it is there, restoring from the durable mark otherwise. A zero length durable mark means a + * fresh start. + * + * @throws IllegalStateException if {@code startEpoch > 0} and no durable mark exists */ - @SuppressWarnings({"unchecked", "nullness"}) // the mark type always matches the source - public static CachedReader getOrCreate( + public static CachedReader acquire( String key, - UnboundedSource source, + long startEpoch, + UnboundedSource source, PipelineOptions options, - Supplier<@Nullable CheckpointMark> durableMarkFallback) { - try { - return (CachedReader) - READERS.get( - key, - () -> { - CheckpointMarkT mark = (CheckpointMarkT) MARKS.get(key); - if (mark == null) { - mark = (CheckpointMarkT) durableMarkFallback.get(); - } - LOG.info( - "No cached Beam reader for {}, creating one at checkpoint mark {}.", key, mark); - return new CachedReader<>(source.createReader(options, mark)); - }); - } catch (Exception e) { - throw new IllegalStateException("Failed to get or create Beam unbounded reader " + key, e); + long idleTimeoutMillis, + MarkRestorer restorer) + throws IOException { + closeIdle(); + synchronized (lock(key)) { + CachedReader existing = READERS.get(key); + if (existing != null) { + if (existing.beginBatch(startEpoch)) { + existing.finalizePendingMark(key); + @SuppressWarnings("unchecked") // one source per key, its element type never changes + CachedReader reused = (CachedReader) existing; + return reused; + } + LOG.info( + "Cached Beam reader {} is at epoch {}, batch starts at {}, restoring from the durable" + + " mark.", + key, + existing.positionEpoch(), + startEpoch); + invalidate(key); + } + byte[] codedMark = restorer.restore(); + if (codedMark == null && startEpoch > 0) { + throw new IllegalStateException( + "No durable checkpoint mark for Beam reader " + key + " at epoch " + startEpoch); + } + if (codedMark != null && codedMark.length == 0) { + codedMark = null; + } + LOG.info( + "Creating Beam reader {} at epoch {} ({} mark).", + key, + startEpoch, + codedMark == null ? "no" : "restored"); + CachedReader created = + new CachedReader<>( + createReader(source, options, codedMark), startEpoch, codedMark, idleTimeoutMillis); + created.beginBatch(startEpoch); + READERS.put(key, created); + return created; } } - /** 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); + private static UnboundedReader createReader( + UnboundedSource source, PipelineOptions options, byte @Nullable [] codedMark) + throws IOException { + MarkT mark = + codedMark == null + ? null + : CoderHelpers.fromByteArray(codedMark, source.getCheckpointMarkCoder()); + return source.createReader(options, mark); + } + + /** Closes and forgets the reader of {@code key}, nothing is finalized. */ + public static void invalidate(String key) { + synchronized (lock(key)) { + CachedReader removed = READERS.remove(key); + if (removed != null) { + close(key, removed); + } } } - /** Closes and forgets every cached reader, intended for tests and for query shutdown. */ - @VisibleForTesting + /** Closes and forgets every cached reader. */ public static void invalidateAll() { - READERS.invalidateAll(); - READERS.cleanUp(); - MARKS.clear(); + for (String key : READERS.keySet()) { + invalidate(key); + } + } + + private static void closeIdle() { + long now = System.currentTimeMillis(); + for (Map.Entry> entry : READERS.entrySet()) { + if (entry.getValue().isIdleSince(now)) { + LOG.info("Closing idle Beam reader {}.", entry.getKey()); + invalidate(entry.getKey()); + } + } + } + + private static void close(String key, CachedReader reader) { + try { + reader.close(); + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to close Beam reader {}.", key, e); + } + } + + private static Object lock(String key) { + return LOCKS.computeIfAbsent(key, k -> new Object()); } - /** A cached {@link UnboundedReader} that remembers whether it has been started already. */ + /** A live reader with the epoch it is positioned at and the coded mark taken there. */ public static final class CachedReader implements Closeable { private final UnboundedReader reader; + private final long idleTimeoutMillis; private boolean started; + private boolean inBatch; + private boolean moved; + private long positionEpoch; + private byte @Nullable [] positionMark; + private @Nullable CheckpointMark pendingMark; + private long lastUsedMillis; - CachedReader(UnboundedReader reader) { + CachedReader( + UnboundedReader reader, + long positionEpoch, + byte @Nullable [] positionMark, + long idleTimeoutMillis) { this.reader = reader; + this.positionEpoch = positionEpoch; + this.positionMark = positionMark; + this.idleTimeoutMillis = idleTimeoutMillis; + this.lastUsedMillis = System.currentTimeMillis(); } - /** 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 { + moved = true; if (!started) { started = true; return reader.start(); @@ -171,6 +208,67 @@ public synchronized boolean startOrAdvance() throws IOException { return reader.advance(); } + /** + * Whether {@link #startOrAdvance()} was called at least once, only then may a mark be taken. + */ + public synchronized boolean started() { + return started; + } + + public synchronized long positionEpoch() { + return positionEpoch; + } + + /** The coded mark of the current position, null for a fresh start. */ + synchronized byte @Nullable [] positionMark() { + return positionMark; + } + + /** + * Claims the reader for a batch starting at {@code epoch}, false if it cannot continue there. + */ + synchronized boolean beginBatch(long epoch) { + if (positionEpoch != epoch || moved) { + return false; + } + inBatch = true; + lastUsedMillis = System.currentTimeMillis(); + return true; + } + + /** Records a completed batch, the reader is positioned at {@code endEpoch} from now on. */ + synchronized void endBatch(long endEpoch, @Nullable CheckpointMark mark, byte[] codedMark) { + positionEpoch = endEpoch; + positionMark = codedMark; + pendingMark = mark; + moved = false; + inBatch = false; + lastUsedMillis = System.currentTimeMillis(); + } + + synchronized boolean isIdleSince(long nowMillis) { + return !inBatch && nowMillis - lastUsedMillis > idleTimeoutMillis; + } + + /** Finalizes the pending mark if any, a failure is logged. */ + synchronized void finalizePendingMark(String key) { + CheckpointMark mark = pendingMark; + pendingMark = null; + if (mark == null) { + return; + } + LOG.debug("Finalizing checkpoint mark of Beam reader {} at epoch {}.", key, positionEpoch); + try { + mark.finalizeCheckpoint(); + } catch (IOException | RuntimeException e) { + LOG.warn( + "Failed to finalize checkpoint mark of Beam reader {} at epoch {}.", + key, + positionEpoch, + e); + } + } + @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/BeamSourceCheckpoint.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java new file mode 100644 index 000000000000..74b4bd101a1d --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java @@ -0,0 +1,184 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FileStatus; +import org.apache.hadoop.fs.Path; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager; +import org.apache.spark.sql.execution.streaming.CheckpointFileManager.CancellableFSDataOutputStream; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Durable state of one Beam unbounded source under the per source checkpoint location Spark hands + * to {@code toMicroBatchStream}. + * + *

{@code /splits} pins the split list, written once by the driver. {@code + * /marks//} holds the coded checkpoint mark of a split at the end of the + * batch ending at that epoch. All IO goes through Spark's {@link CheckpointFileManager}, writes are + * atomic renames. + */ +public final class BeamSourceCheckpoint { + + private static final Logger LOG = LoggerFactory.getLogger(BeamSourceCheckpoint.class); + + private static final String SPLITS_FILE = "splits"; + private static final String MARKS_DIR = "marks"; + + /** A purge further than this above the last one lists the directory instead of probing epochs. */ + private static final long MAX_BLIND_PURGE_RANGE = 1_000L; + + private final String location; + private final CheckpointFileManager fm; + private final Path root; + private final Path splitsPath; + private final Path marksRoot; + + /** Splits whose mark directory this instance already created. */ + private final Set preparedSplits = ConcurrentHashMap.newKeySet(); + + /** Per split, every mark epoch strictly below the value is known to be deleted. */ + private final ConcurrentMap purgeFloors = new ConcurrentHashMap<>(); + + public BeamSourceCheckpoint(String checkpointLocation, Configuration hadoopConf) { + this.location = checkpointLocation; + this.root = new Path(checkpointLocation); + this.fm = CheckpointFileManager.create(root, hadoopConf); + this.splitsPath = new Path(root, SPLITS_FILE); + this.marksRoot = new Path(root, MARKS_DIR); + } + + public String location() { + return location; + } + + /** The pinned split list, or null if none was pinned yet. */ + public @Nullable List> readSplits() throws IOException { + if (!fm.exists(splitsPath)) { + return null; + } + @SuppressWarnings("unchecked") // written by writeSplits as an ArrayList of sources + List> splits = + (List>) + SerializableUtils.deserializeFromByteArray(read(splitsPath), "splits at " + splitsPath); + return splits; + } + + /** Pins the split list, fails if one is pinned already. */ + public void writeSplits(List> splits) throws IOException { + fm.mkdirs(root); + if (fm.exists(splitsPath)) { + throw new IOException("Split list already pinned at " + splitsPath); + } + write(splitsPath, SerializableUtils.serializeToByteArray(new ArrayList<>(splits)), false); + LOG.info("Pinned {} split(s) at {}.", splits.size(), splitsPath); + } + + public void writeMark(int splitId, long epoch, byte[] codedMark) throws IOException { + if (preparedSplits.add(splitId)) { + fm.mkdirs(marksDir(splitId)); + } + write(markPath(splitId, epoch), codedMark, true); + } + + /** The coded mark of a split at an epoch, or null if absent. */ + public byte @Nullable [] readMark(int splitId, long epoch) throws IOException { + Path path = markPath(splitId, epoch); + if (!fm.exists(path)) { + return null; + } + return read(path); + } + + /** + * Deletes every mark of a split with an epoch strictly below {@code epoch}. Lists the directory + * once per split, later calls delete the range above the previous floor only. Idempotent. + */ + public void purgeMarksBelow(int splitId, long epoch) throws IOException { + Long floor = purgeFloors.get(splitId); + if (floor != null && epoch - floor > MAX_BLIND_PURGE_RANGE) { + floor = null; + } + if (floor == null) { + Path dir = marksDir(splitId); + if (!fm.exists(dir)) { + return; + } + for (FileStatus status : fm.list(dir)) { + long existing = parseEpoch(status.getPath().getName()); + if (existing >= 0 && existing < epoch) { + fm.delete(status.getPath()); + } + } + purgeFloors.put(splitId, epoch); + return; + } + for (long e = floor; e < epoch; e++) { + fm.delete(markPath(splitId, e)); + } + if (epoch > floor) { + purgeFloors.put(splitId, epoch); + } + } + + private Path marksDir(int splitId) { + return new Path(marksRoot, Integer.toString(splitId)); + } + + private Path markPath(int splitId, long epoch) { + return new Path(marksDir(splitId), Long.toString(epoch)); + } + + private byte[] read(Path path) throws IOException { + try (FSDataInputStream in = fm.open(path)) { + return ByteStreams.toByteArray(in); + } + } + + private void write(Path path, byte[] bytes, boolean overwrite) throws IOException { + CancellableFSDataOutputStream out = fm.createAtomic(path, overwrite); + try { + out.write(bytes); + out.close(); + } catch (IOException | RuntimeException e) { + out.cancel(); + throw e; + } + } + + /** The epoch encoded in a mark file name, or -1 for anything else. */ + private static long parseEpoch(String name) { + try { + return Long.parseLong(name); + } catch (NumberFormatException e) { + return -1L; + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.java new file mode 100644 index 000000000000..e024550d29a3 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceSpec.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.io.Serializable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +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.broadcast.Broadcast; +import org.apache.spark.util.SerializableConfiguration; + +/** Everything the driver needs to plan micro-batches of one Beam {@link UnboundedSource}. */ +final class BeamSourceSpec implements Serializable { + + private static final long serialVersionUID = 1L; + + private final UnboundedSource source; + private final Coder> coder; + private final Broadcast options; + private final Broadcast hadoopConf; + private final int desiredNumSplits; + private final long maxRecordsPerBatch; + private final long maxBatchDurationMillis; + private final long readerIdleTimeoutMillis; + private final String transformName; + + BeamSourceSpec( + UnboundedSource source, + Coder> coder, + Broadcast options, + Broadcast hadoopConf, + int desiredNumSplits, + long maxRecordsPerBatch, + long maxBatchDurationMillis, + long readerIdleTimeoutMillis, + String transformName) { + this.source = source; + this.coder = coder; + this.options = options; + this.hadoopConf = hadoopConf; + this.desiredNumSplits = desiredNumSplits; + this.maxRecordsPerBatch = maxRecordsPerBatch; + this.maxBatchDurationMillis = maxBatchDurationMillis; + this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + this.transformName = transformName; + } + + UnboundedSource source() { + return source; + } + + Coder> coder() { + return coder; + } + + Broadcast options() { + return options; + } + + Broadcast hadoopConf() { + return hadoopConf; + } + + int desiredNumSplits() { + return desiredNumSplits; + } + + /** Records per micro-batch across all splits, below 1 means unlimited. */ + long maxRecordsPerBatch() { + return maxRecordsPerBatch; + } + + long maxBatchDurationMillis() { + return maxBatchDurationMillis; + } + + long readerIdleTimeoutMillis() { + return readerIdleTimeoutMillis; + } + + String transformName() { + return transformName; + } +} 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 deleted file mode 100644 index e895d764ff1f..000000000000 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamStreamingSource.java +++ /dev/null @@ -1,157 +0,0 @@ -/* - * 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: - * - *

    - *
  • {@code payload} of type {@code BINARY}, holding the element encoded with the Beam {@code - * WindowedValues.FullWindowedValueCoder} supplied by the translator. - *
  • {@code eventTimestamp} of type {@code TIMESTAMP}, holding the event timestamp reported by - * the Beam reader for that element. - *
- * - *

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"; - - /** - * Deterministic identifier of the source, derived from the full name of the read transform, see - * {@link UnboundedSourceDataset#sourceId}. It keys the reader cache and the durable checkpoint - * state, so it must stay stable across restarts against the same checkpoint location. - */ - public static final String OPT_SOURCE_ID = "beam.sourceid"; - - /** Desired number of splits handed to {@code UnboundedSource.split}. */ - public static final String OPT_NUM_SPLITS = "beam.numsplits"; - - /** Maximum number of records read per split per micro-batch, values below 1 mean no limit. */ - 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. - * - *

The return type is inferred from the call site. Deserialization cannot check it, and the - * decoded types include generic ones such as {@code Coder>} that no {@code - * Class} token can express, so a checked variant is not available here. - */ - @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"}) - static T decode(String encoded, String description) { - return (T) - SerializableUtils.deserializeFromByteArray( - 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 index 4ba6c24a5cef..d6fc2ada4dab 100644 --- 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 @@ -28,28 +28,23 @@ 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. - */ +/** DataSourceV2 {@link Table} over a Beam unbounded source, micro-batch reads only. */ public class BeamStreamingTable implements Table, SupportsRead { - private final CaseInsensitiveStringMap options; + private final BeamSourceSpec spec; - BeamStreamingTable(CaseInsensitiveStringMap options) { - this.options = options; + BeamStreamingTable(BeamSourceSpec spec) { + this.spec = spec; } @Override public String name() { - return "BeamUnboundedSource[" - + options.getOrDefault(BeamStreamingSource.OPT_SOURCE_ID, "?") - + "]"; + return "BeamUnboundedSource[" + spec.transformName() + "]"; } @Override public StructType schema() { - return BeamStreamingSource.SCHEMA; + return UnboundedSourceDataset.SCHEMA; } @Override @@ -58,43 +53,30 @@ public Set capabilities() { } @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); + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap ignored) { + return () -> new BeamScan(spec); } - /** The {@link Scan} of a Beam unbounded source, micro-batch only. */ private static class BeamScan implements Scan { - private final CaseInsensitiveStringMap options; + private final BeamSourceSpec spec; - BeamScan(CaseInsensitiveStringMap options) { - this.options = options; + BeamScan(BeamSourceSpec spec) { + this.spec = spec; } @Override public StructType readSchema() { - return BeamStreamingSource.SCHEMA; + return UnboundedSourceDataset.SCHEMA; } @Override public String description() { - return "BeamUnboundedSource[" - + options.getOrDefault(BeamStreamingSource.OPT_SOURCE_ID, "?") - + "]"; + return "BeamUnboundedSource[" + spec.transformName() + "]"; } @Override public MicroBatchStream toMicroBatchStream(String checkpointLocation) { - return new BeamMicroBatchStream(options, checkpointLocation); + return new BeamMicroBatchStream<>(spec, 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 index e385542b50be..87a0fa0aeff6 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -17,61 +17,61 @@ */ package org.apache.beam.runners.spark.structuredstreaming.io.streaming; -import java.nio.charset.StandardCharsets; -import java.util.HashMap; -import java.util.Map; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.values.WindowedValue; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.broadcast.Broadcast; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; +import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2; +import org.apache.spark.sql.catalyst.types.DataTypeUtils; +import org.apache.spark.sql.classic.Dataset$; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.apache.spark.util.SerializableConfiguration; +import scala.Option; +import scala.reflect.ClassTag; /** * 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 dataset has two columns, {@value #COL_PAYLOAD} of type {@code BINARY} holding the element + * encoded with the supplied {@code WindowedValue} coder, and {@value #COL_EVENT_TS} of type {@code + * TIMESTAMP} holding 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. + *

The event time watermark is declared here and only here. Spark 4 rejects a second {@code + * withWatermark} further down the plan, so downstream translators must never call it again. */ 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; + public static final StructType SCHEMA = + new StructType() + .add(COL_PAYLOAD, DataTypes.BinaryType, false) + .add(COL_EVENT_TS, DataTypes.TimestampType, false); - /** Upper bound on the sanitized transform name inside a source id, see {@link #sourceId}. */ - private static final int MAX_SANITIZED_NAME_LENGTH = 64; + private static final String SOURCE_NAME = "beam-unbounded"; private UnboundedSourceDataset() {} /** - * Builds the streaming {@link Dataset} for {@code source}, with the event time watermark already - * applied. + * Builds the streaming {@link Dataset} for {@code source} with the event time watermark 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 windowedValueCoder the coder of the {@value #COL_PAYLOAD} column * @param options the pipeline options, supplying the watermark delay and the micro-batch limits - * @param transformName the full name of the read transform, turned into the deterministic source - * id keying the durable checkpoint state, see {@link #sourceId} + * @param transformName the full name of the read transform, used for naming only * @param the element type of the source * @param the checkpoint mark type of the source */ @@ -81,53 +81,41 @@ public static Datase Coder> windowedValueCoder, SparkStructuredStreamingPipelineOptions options, String transformName) { - - 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, sourceId(transformName)); - readerOptions.put( - BeamStreamingSource.OPT_NUM_SPLITS, Integer.toString(desiredNumSplits(session))); - readerOptions.put( - BeamStreamingSource.OPT_MAX_RECORDS, Long.toString(options.getMaxRecordsPerBatch())); - 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. + org.apache.spark.sql.classic.SparkSession classic = + (org.apache.spark.sql.classic.SparkSession) session; + Configuration hadoopConf = classic.sessionState().newHadoopConf(); + BeamSourceSpec spec = + new BeamSourceSpec<>( + source, + windowedValueCoder, + broadcast( + session, + new SerializablePipelineOptions(options), + SerializablePipelineOptions.class), + broadcast( + session, + new SerializableConfiguration(hadoopConf), + SerializableConfiguration.class), + session.sparkContext().defaultParallelism(), + options.getMaxRecordsPerBatch(), + Math.max(1L, options.getMaxBatchDurationMillis()), + options.getReaderIdleTimeoutMillis(), + transformName); + LogicalPlan plan = + new StreamingRelationV2( + Option.empty(), + SOURCE_NAME, + new BeamStreamingTable(spec), + CaseInsensitiveStringMap.empty(), + DataTypeUtils.toAttributes(SCHEMA), + Option.empty(), + Option.empty(), + Option.empty()); + Dataset rows = Dataset$.MODULE$.ofRows(classic, plan); return rows.withWatermark(COL_EVENT_TS, options.getWatermarkDelayMillis() + " milliseconds"); } - /** - * Derives the deterministic source id of a read transform from its full name. - * - *

The id keys the durable checkpoint state under the checkpoint location, so it must be - * filesystem safe and stable across JVMs. Characters outside {@code [A-Za-z0-9._-]} are replaced - * with {@code _}, the sanitized name is truncated to {@value #MAX_SANITIZED_NAME_LENGTH} - * characters, and an 8 hex character hash of the unsanitized name is appended to keep distinct - * transform names distinct. A pipeline restarted against the same checkpoint location must - * therefore keep the same transform names, the same requirement every Beam runner with durable - * state imposes. - */ - public static String sourceId(String transformName) { - String sanitized = transformName.replaceAll("[^A-Za-z0-9._-]", "_"); - if (sanitized.length() > MAX_SANITIZED_NAME_LENGTH) { - sanitized = sanitized.substring(0, MAX_SANITIZED_NAME_LENGTH); - } - String hash = - Hashing.murmur3_32_fixed().hashString(transformName, StandardCharsets.UTF_8).toString(); - return sanitized + "-" + hash; - } - - private static int desiredNumSplits(SparkSession session) { - int parallelism = session.sparkContext().defaultParallelism(); - return Math.max(1, Math.min(MAX_DESIRED_SPLITS, parallelism)); + private static Broadcast broadcast(SparkSession session, T value, Class type) { + return session.sparkContext().broadcast(value, ClassTag.apply(type)); } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java deleted file mode 100644 index 4b7d7c55a00e..000000000000 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointRecoveryTest.java +++ /dev/null @@ -1,198 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.runners.spark.structuredstreaming.io.streaming; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotEquals; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertNull; -import static org.junit.Assert.assertTrue; - -import java.io.File; -import java.io.Serializable; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.apache.beam.runners.core.construction.SerializablePipelineOptions; -import org.apache.beam.sdk.io.CountingSource; -import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.transforms.windowing.GlobalWindow; -import org.apache.beam.sdk.values.WindowedValues; -import org.apache.spark.sql.connector.read.InputPartition; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; -import org.junit.Rule; -import org.junit.Test; -import org.junit.rules.TemporaryFolder; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Tests for the durable checkpoint recovery pieces of the Spark 4 micro-batch source, the {@link - * BeamCheckpointFiles} layout, the epoch fast forward of {@link BeamMicroBatchStream} and the - * deterministic source id of {@link UnboundedSourceDataset}. - */ -@RunWith(JUnit4.class) -public class BeamCheckpointRecoveryTest { - - private static final String SOURCE_ID = "read-source-cafe0123"; - - @Rule public TemporaryFolder temp = new TemporaryFolder(); - - @Test - public void testSplitsRoundTrip() throws Exception { - String checkpointLocation = temp.newFolder("splits").getAbsolutePath(); - List splits = Arrays.asList("c3BsaXQtMA==", "c3BsaXQtMQ==", "c3BsaXQtMg=="); - - assertNull( - "no splits pinned yet", BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); - - BeamCheckpointFiles.writeSplits(checkpointLocation, SOURCE_ID, splits); - assertEquals(splits, BeamCheckpointFiles.readSplits(checkpointLocation, SOURCE_ID)); - - File root = new File(checkpointLocation, "beam-source-" + SOURCE_ID); - assertEquals( - "only the final splits file may remain", - new HashSet<>(Arrays.asList("splits")), - fileNames(root)); - } - - @Test - public void testMarkRetentionAndRecovery() throws Exception { - String checkpointLocation = temp.newFolder("marks").getAbsolutePath(); - for (long epoch = 1; epoch <= 4; epoch++) { - BeamCheckpointFiles.writeMark( - checkpointLocation, SOURCE_ID, 0, epoch, new TestMark((int) epoch)); - } - - File marksDir = new File(new File(checkpointLocation, "beam-source-" + SOURCE_ID), "marks/0"); - assertEquals( - "retention must keep the two highest epochs only", - new HashSet<>(Arrays.asList("3", "4")), - fileNames(marksDir)); - - assertEquals(4, restoredPosition(checkpointLocation, 0, 4)); - assertEquals(3, restoredPosition(checkpointLocation, 0, 3)); - // No exact file for epoch 7, the largest epoch not exceeding it is 4. - assertEquals(4, restoredPosition(checkpointLocation, 0, 7)); - // Epochs 1 and 2 were deleted by retention, nothing at or below 2 is left. - assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 0, 2)); - // A different split has no marks at all. - assertNull(BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, 1, 4)); - } - - @Test - public void testEpochFastForwardsPastDeserializedOffset() throws Exception { - BeamMicroBatchStream stream = stream(temp.newFolder("ff-offset").getAbsolutePath()); - stream.deserializeOffset("{\"epoch\":7}"); - BeamOffset next = (BeamOffset) stream.latestOffset(); - assertTrue("latestOffset must move past the replayed epoch 7, got " + next, next.epoch() > 7L); - } - - @Test - public void testEpochFastForwardsPastPlannedOffsets() throws Exception { - BeamMicroBatchStream stream = stream(temp.newFolder("ff-plan").getAbsolutePath()); - InputPartition[] partitions = - stream.planInputPartitions(new BeamOffset(3L), new BeamOffset(9L)); - assertTrue("at least one partition expected", partitions.length > 0); - BeamOffset next = (BeamOffset) stream.latestOffset(); - assertTrue("latestOffset must move past the planned epoch 9, got " + next, next.epoch() > 9L); - } - - @Test - public void testSourceIdIsDeterministicAndFilesystemSafe() { - String name = "Read from PubSub/PubsubUnboundedSource (with: spaces & colons)"; - String id = UnboundedSourceDataset.sourceId(name); - - assertEquals("same name must give the same id", id, UnboundedSourceDataset.sourceId(name)); - assertNotEquals( - "different names must give different ids", - UnboundedSourceDataset.sourceId("Read A"), - UnboundedSourceDataset.sourceId("Read B")); - assertTrue("id must be filesystem safe: " + id, id.matches("[A-Za-z0-9._-]+")); - - String longName = String.join("/", java.util.Collections.nCopies(50, "NestedTransform")); - String longId = UnboundedSourceDataset.sourceId(longName); - assertTrue("id length must stay bounded: " + longId.length(), longId.length() <= 73); - assertNotEquals( - "truncated names must still differ through the hash", - longId, - UnboundedSourceDataset.sourceId(longName + "/Tail")); - } - - private static int restoredPosition(String checkpointLocation, int splitId, long startEpoch) { - CheckpointMark mark = - BeamCheckpointFiles.readMark(checkpointLocation, SOURCE_ID, splitId, startEpoch); - assertNotNull("expected a durable mark at or before epoch " + startEpoch, mark); - return ((TestMark) mark).position; - } - - /** - * Lists the visible files of {@code dir}, asserting no {@code .tmp} file was left behind and - * ignoring the hidden checksum sidecars of Hadoop's local filesystem. - */ - private static Set fileNames(File dir) { - String[] names = dir.list(); - assertNotNull("expected directory " + dir, names); - Set visible = new HashSet<>(); - for (String name : names) { - assertFalse("no temporary file may remain: " + name, name.endsWith(".tmp")); - if (!name.startsWith(".")) { - visible.add(name); - } - } - return visible; - } - - /** Builds a stream over a real source, driver side only, no Spark session required. */ - private static BeamMicroBatchStream stream(String checkpointLocation) { - Map options = new HashMap<>(); - options.put( - BeamStreamingSource.OPT_SOURCE, BeamStreamingSource.encode(CountingSource.unbounded())); - options.put( - BeamStreamingSource.OPT_CODER, - BeamStreamingSource.encode( - WindowedValues.getFullCoder( - org.apache.beam.sdk.coders.VarLongCoder.of(), GlobalWindow.Coder.INSTANCE))); - options.put( - BeamStreamingSource.OPT_PIPELINE_OPTIONS, - BeamStreamingSource.encode( - new SerializablePipelineOptions(PipelineOptionsFactory.create()))); - options.put(BeamStreamingSource.OPT_SOURCE_ID, SOURCE_ID); - options.put(BeamStreamingSource.OPT_NUM_SPLITS, "2"); - return new BeamMicroBatchStream(new CaseInsensitiveStringMap(options), checkpointLocation); - } - - /** A trivial serializable checkpoint mark carrying a read position. */ - private static class TestMark implements CheckpointMark, Serializable { - private static final long serialVersionUID = 1L; - - private final int position; - - TestMark(int position) { - this.position = position; - } - - @Override - public void finalizeCheckpoint() {} - } -} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java index 0d1bee7a9dab..cc7007267601 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -19,46 +19,69 @@ 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.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import java.io.File; import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; import java.sql.Timestamp; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.Set; +import java.util.TreeSet; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; import javax.annotation.Nullable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; 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.CustomCoder; import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; 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.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.hadoop.conf.Configuration; import org.apache.spark.api.java.function.MapFunction; import org.apache.spark.api.java.function.VoidFunction2; +import org.apache.spark.broadcast.Broadcast; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Encoders; import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark; import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan; +import org.apache.spark.sql.connector.read.InputPartition; +import org.apache.spark.sql.connector.read.streaming.MicroBatchStream; import org.apache.spark.sql.streaming.StreamingQuery; import org.apache.spark.sql.streaming.StreamingQueryProgress; import org.apache.spark.sql.streaming.Trigger; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.apache.spark.util.SerializableConfiguration; import org.joda.time.Instant; import org.junit.After; import org.junit.ClassRule; @@ -68,6 +91,7 @@ import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import scala.reflect.ClassTag; /** * Tests for the Spark 4 DataSourceV2 micro-batch source wrapping a Beam {@link UnboundedSource}. @@ -90,6 +114,9 @@ public class BeamMicroBatchSourceTest implements Serializable { /** Rows collected per query name by {@link #startCollecting}, driver side. */ private static final Map> COLLECTED = new ConcurrentHashMap<>(); + /** Rows per micro-batch per query name, see {@link #startCollectingBatches}, driver side. */ + private static final Map>> BATCHES = 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; @@ -101,6 +128,8 @@ public class BeamMicroBatchSourceTest implements Serializable { public void tearDown() { BeamReaderCache.invalidateAll(); COLLECTED.clear(); + BATCHES.clear(); + ShardedListSource.FINALIZED.clear(); } /** @@ -120,7 +149,7 @@ public void tearDown() { * TimeMode.EventTime} does, are unaffected, see {@link * #testWatermarkIsTrackedAtRuntimeAfterTypedMap}. */ - @Test(timeout = 300_000) + @Test public void testEventTimeWatermarkSurvivesTypedMap() { Dataset rows = rows(4, 1_000L); assertTrue("source dataset must be streaming", rows.isStreaming()); @@ -148,7 +177,7 @@ public void testEventTimeWatermarkSurvivesTypedMap() { * 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) + @Test public void testWatermarkIsTrackedAtRuntimeAfterTypedMap() throws Exception { Dataset rows = rows(8, 0L); Dataset typed = @@ -168,7 +197,7 @@ public void testWatermarkIsTrackedAtRuntimeAfterTypedMap() throws Exception { } /** Reads a finite set of elements through the DSv2 source and checks payloads and timestamps. */ - @Test(timeout = 300_000) + @Test public void testReadsElementsFromUnboundedSource() throws Exception { int count = 8; Dataset rows = rows(count, 0L); @@ -210,12 +239,11 @@ public void testReadsElementsFromUnboundedSource() throws Exception { } /** The default record limit of -1 means unlimited, an available source drains in one batch. */ - @Test(timeout = 300_000) + @Test public void testUnlimitedRecordsPerBatchByDefault() throws Exception { int count = 2500; SparkStructuredStreamingPipelineOptions options = - org.apache.beam.sdk.options.PipelineOptionsFactory.create() - .as(SparkStructuredStreamingPipelineOptions.class); + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); options.setWatermarkDelayMillis(0L); options.setMaxBatchDurationMillis(5_000L); Dataset rows = @@ -253,13 +281,268 @@ public void testUnlimitedRecordsPerBatchByDefault() throws Exception { new ArrayList<>(batchSizes)); } - /** The offset is an opaque, strictly increasing epoch counter. */ - @Test(timeout = 300_000) + /** The offset is an opaque, strictly increasing epoch counter, its JSON is the bare number. */ + @Test public void testEpochOffsetRoundTrip() { BeamOffset offset = new BeamOffset(42L); - assertEquals("{\"epoch\":42}", offset.json()); - assertEquals(offset, BeamOffset.fromJson(offset.json())); + assertEquals("42", offset.json()); + assertEquals(42L, BeamOffset.fromJson("42").epoch()); assertEquals(0L, BeamOffset.ZERO.epoch()); + assertEquals(new BeamOffset(7L), new BeamOffset(7L)); + assertThrows(IllegalArgumentException.class, () -> BeamOffset.fromJson("x")); + } + + @Test + public void testEpochFastForwardsPastDeserializedOffset() throws Exception { + BeamMicroBatchStream stream = newStream(temp.newFolder("ff-offset").getAbsolutePath()); + stream.deserializeOffset("7"); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the replayed epoch 7, got " + next, next.epoch() > 7L); + } + + @Test + public void testEpochFastForwardsPastPlannedOffsets() throws Exception { + BeamMicroBatchStream stream = newStream(temp.newFolder("ff-plan").getAbsolutePath()); + InputPartition[] partitions = + stream.planInputPartitions(new BeamOffset(3L), new BeamOffset(9L)); + assertTrue("at least one partition expected", partitions.length > 0); + BeamOffset next = (BeamOffset) stream.latestOffset(); + assertTrue("latestOffset must move past the planned epoch 9, got " + next, next.epoch() > 9L); + } + + /** The batch quota is divided over the splits, remainder first, every split gets at least one. */ + @Test + public void testSplitQuotas() { + assertArrayEquals(new long[] {1, 1, 1, 1, 1, 1, 1, 1}, BeamMicroBatchStream.splitQuotas(3, 8)); + assertArrayEquals(new long[] {4, 3, 3}, BeamMicroBatchStream.splitQuotas(10, 3)); + assertArrayEquals(new long[] {5, 5}, BeamMicroBatchStream.splitQuotas(10, 2)); + assertArrayEquals(new long[] {0, 0, 0}, BeamMicroBatchStream.splitQuotas(0, 3)); + assertArrayEquals(new long[] {-1, -1, -1}, BeamMicroBatchStream.splitQuotas(-1, 3)); + long[] many = BeamMicroBatchStream.splitQuotas(1, 200); + assertEquals(200, many.length); + for (long quota : many) { + assertEquals(1L, quota); + } + } + + /** Rendezvous hashing: a joining executor only takes splits for itself, order does not matter. */ + @Test + public void testSplitAssignmentIsStableWhenExecutorJoins() { + List three = Arrays.asList("executor_a_1", "executor_b_2", "executor_c_3"); + List shuffled = Arrays.asList("executor_c_3", "executor_a_1", "executor_b_2"); + List four = new ArrayList<>(three); + four.add("executor_d_4"); + int splits = 200; + int moved = 0; + for (int split = 0; split < splits; split++) { + String before = BeamMicroBatchStream.assign(split, three); + String after = BeamMicroBatchStream.assign(split, four); + assertTrue(three.contains(before)); + assertEquals(before, BeamMicroBatchStream.assign(split, shuffled)); + if (!after.equals(before)) { + assertEquals("split " + split + " moved to an old executor", "executor_d_4", after); + moved++; + } + } + assertTrue("the new executor took no split", moved > 0); + assertTrue("the new executor took every split", moved < splits); + assertEquals( + "executor_a_1", BeamMicroBatchStream.assign(7, Collections.singletonList("executor_a_1"))); + } + + /** The record limit of a micro-batch is a total over all splits, not a per split allowance. */ + @Test + public void testMaxRecordsPerBatchIsSharedAcrossSplits() throws Exception { + int shards = 2; + int count = 30; + long limit = 10L; + String queryName = "beam_shared_limit_" + QUERY_COUNTER.incrementAndGet(); + Dataset rows = shardedRows(queryName, shards, count, limit); + StreamingQuery query = startCollectingBatches(rows, queryName, temp.newFolder(queryName)); + try { + await("all rows", () -> values(batches(queryName)).size() >= count); + } finally { + stopQuietly(query); + } + + List sizes = new ArrayList<>(); + for (List batch : batches(queryName)) { + if (!batch.isEmpty()) { + sizes.add(batch.size()); + } + } + assertFalse("no rows arrived", sizes.isEmpty()); + assertTrue("first batch exceeds the shared limit: " + sizes, sizes.get(0) <= limit); + for (int size : sizes) { + assertTrue("batch exceeds the shared limit: " + sizes, size <= limit); + } + List values = values(batches(queryName)); + assertEquals(count, values.size()); + assertEquals(ShardedListSource.elements(shards, count), new HashSet<>(values)); + } + + /** + * A restart resumes every split from the durable mark of the last committed batch, replaying at + * most the one uncommitted batch. + */ + @Test + public void testRestartResumesFromCommittedMark() throws Exception { + int shards = 2; + int count = 80; + long limit = 4L; + File checkpointDir = temp.newFolder("restart"); + Set all = ShardedListSource.elements(shards, count); + + String first = "beam_restart_a_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(first, shards, count, limit), first, checkpointDir); + try { + await("two commits", () -> committedBatchIds(checkpointDir).size() >= 2); + } finally { + stopQuietly(query); + } + BeamReaderCache.invalidateAll(); + List firstValues = values(batches(first)); + + String second = "beam_restart_b_" + QUERY_COUNTER.incrementAndGet(); + query = + startCollectingBatches(shardedRows(second, shards, count, limit), second, checkpointDir); + try { + await( + "union of both runs", + () -> { + Set union = new HashSet<>(firstValues); + union.addAll(values(batches(second))); + return union.containsAll(all); + }); + } finally { + stopQuietly(query); + } + List secondValues = values(batches(second)); + + Set union = new HashSet<>(firstValues); + union.addAll(secondValues); + assertEquals(all, union); + assertTrue( + "more than the uncommitted batch replayed: " + + firstValues.size() + + " + " + + secondValues.size() + + " rows", + firstValues.size() + secondValues.size() <= count + limit); + for (int shard = 0; shard < shards; shard++) { + int min = Integer.MAX_VALUE; + for (String value : secondValues) { + if (ShardedListSource.shardOf(value) == shard) { + min = Math.min(min, ShardedListSource.indexOf(value)); + } + } + assertTrue("run 2 delivered nothing for shard " + shard, min < Integer.MAX_VALUE); + assertTrue("run 2 restarted shard " + shard + " from element 0", min > 0); + } + } + + /** + * Every finalized position of a split is at most the position in that split's mark at the end + * epoch of the highest committed batch, so no mark is finalized before Spark commits its batch. + */ + @Test + public void testMarksAreFinalizedOnlyAfterSparkCommit() throws Exception { + int shards = 2; + // Never exhausted while the test runs, positions strictly increase with the epoch. + int count = 4_000; + File checkpointDir = temp.newFolder("finalize"); + String queryName = "beam_finalize_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(queryName, shards, count, 4L), queryName, checkpointDir); + try { + await("three commits", () -> committedBatchIds(checkpointDir).size() >= 3); + } finally { + stopQuietly(query); + } + + long committedEpoch = + endEpoch(checkpointDir, Collections.max(committedBatchIds(checkpointDir))); + BeamSourceCheckpoint files = sourceCheckpoint(checkpointDir); + int finalizations = 0; + for (int shard = 0; shard < shards; shard++) { + byte[] coded = files.readMark(shard, committedEpoch); + assertNotNull("no mark at committed epoch " + committedEpoch + " for split " + shard, coded); + int committedPosition = + CoderUtils.decodeFromByteArray(ShardedListSource.MARK_CODER, coded).next; + List finalized = ShardedListSource.finalized(queryName, shard); + for (int position : finalized) { + assertTrue( + "split " + + shard + + " finalized position " + + position + + " beyond committed position " + + committedPosition, + position <= committedPosition); + } + finalizations += finalized.size(); + } + assertTrue("no mark was finalized", finalizations > 0); + } + + /** + * Spark reports the commit of batch N-1 when it constructs batch N, so after a run the surviving + * marks of every split are at or above the end epoch of the batch before the last constructed + * one, the mark at the highest committed end epoch exists, and the mark of batch 0 is gone. + */ + @Test + public void testMarksBelowCommittedOffsetArePurged() throws Exception { + int shards = 2; + int count = 4_000; + File checkpointDir = temp.newFolder("purge"); + String queryName = "beam_purge_" + QUERY_COUNTER.incrementAndGet(); + StreamingQuery query = + startCollectingBatches(shardedRows(queryName, shards, count, 4L), queryName, checkpointDir); + try { + await("three commits", () -> committedBatchIds(checkpointDir).size() >= 3); + } finally { + stopQuietly(query); + } + + long lastConstructed = Collections.max(batchIds(new File(checkpointDir, "offsets"))); + long purgeFloor = endEpoch(checkpointDir, lastConstructed - 1); + long committedEpoch = + endEpoch(checkpointDir, Collections.max(committedBatchIds(checkpointDir))); + long firstEpoch = endEpoch(checkpointDir, 0); + assertTrue(committedEpoch >= purgeFloor); + assertTrue(purgeFloor > firstEpoch); + File sourceDir = new File(checkpointDir, "sources/0"); + // The last purge runs asynchronously on the driver and may still be in flight after stop. + awaitQuietly( + 10_000L, + () -> { + for (int shard = 0; shard < shards; shard++) { + for (long epoch : batchIds(new File(sourceDir, "marks/" + shard))) { + if (epoch < purgeFloor) { + return false; + } + } + } + return true; + }); + for (int shard = 0; shard < shards; shard++) { + Set remaining = batchIds(new File(sourceDir, "marks/" + shard)); + assertTrue( + "split " + + shard + + " lost the mark at committed epoch " + + committedEpoch + + ": " + + remaining, + remaining.contains(committedEpoch)); + for (long epoch : remaining) { + assertTrue( + "split " + shard + " kept mark " + epoch + " below purge floor " + purgeFloor, + epoch >= purgeFloor); + } + assertFalse("split " + shard + " kept the mark of batch 0", remaining.contains(firstEpoch)); + } } // --------------------------------------------------------------------------------------------- @@ -268,8 +551,7 @@ public void testEpochOffsetRoundTrip() { private Dataset rows(int count, long watermarkDelayMillis) { SparkStructuredStreamingPipelineOptions options = - org.apache.beam.sdk.options.PipelineOptionsFactory.create() - .as(SparkStructuredStreamingPipelineOptions.class); + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); options.setWatermarkDelayMillis(watermarkDelayMillis); options.setMaxRecordsPerBatch(1000L); options.setMaxBatchDurationMillis(200L); @@ -277,6 +559,158 @@ private Dataset rows(int count, long watermarkDelayMillis) { SESSION.getSession(), new ListSource(count), coder(), options, "Read(ListSource)"); } + /** Builds the driver side stream through the table, with real broadcasts from the session. */ + private static BeamMicroBatchStream newStream(String checkpointLocation) { + SparkSession session = SESSION.getSession(); + Broadcast options = + session + .sparkContext() + .broadcast( + new SerializablePipelineOptions(PipelineOptionsFactory.create()), + ClassTag.apply(SerializablePipelineOptions.class)); + Broadcast hadoopConf = + session + .sparkContext() + .broadcast( + new SerializableConfiguration(new Configuration()), + ClassTag.apply(SerializableConfiguration.class)); + BeamSourceSpec spec = + new BeamSourceSpec<>( + new ListSource(4), + coder(), + options, + hadoopConf, + 2, + -1L, + 200L, + 600_000L, + "Read(ListSource)"); + MicroBatchStream stream = + new BeamStreamingTable(spec) + .newScanBuilder(CaseInsensitiveStringMap.empty()) + .build() + .toMicroBatchStream(checkpointLocation); + return (BeamMicroBatchStream) stream; + } + + private Dataset shardedRows(String tag, int shards, int count, long maxRecordsPerBatch) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(0L); + options.setMaxRecordsPerBatch(maxRecordsPerBatch); + options.setMaxBatchDurationMillis(1_000L); + return UnboundedSourceDataset.of( + SESSION.getSession(), + new ShardedListSource(tag, shards, count), + coder(), + options, + "Read(ShardedListSource)"); + } + + /** Starts a query collecting every micro-batch as one list into {@link #BATCHES}. */ + private static StreamingQuery startCollectingBatches( + Dataset dataset, String queryName, File checkpointDir) throws Exception { + BATCHES.put(queryName, Collections.synchronizedList(new ArrayList<>())); + return dataset + .writeStream() + .foreachBatch( + (VoidFunction2, Long>) + (batch, batchId) -> { + List> target = BATCHES.get(queryName); + if (target != null) { + target.add(batch.collectAsList()); + } + }) + .queryName(queryName) + .outputMode("append") + .option("checkpointLocation", checkpointDir.getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + } + + private static List> batches(String queryName) { + List> batches = BATCHES.getOrDefault(queryName, Collections.emptyList()); + synchronized (batches) { + return new ArrayList<>(batches); + } + } + + private static List values(List> batches) { + Coder> coder = coder(); + List values = new ArrayList<>(); + for (List batch : batches) { + for (Row row : batch) { + byte[] payload = row.getAs(COL_PAYLOAD); + try { + values.add(CoderUtils.decodeFromByteArray(coder, payload).getValue()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + } + return values; + } + + private static void await(String what, BooleanSupplier condition) throws Exception { + if (!awaitQuietly(POLL_TIMEOUT_MILLIS, condition)) { + throw new AssertionError("timed out waiting for " + what); + } + } + + private static boolean awaitQuietly(long timeoutMillis, BooleanSupplier condition) + throws Exception { + long deadline = System.currentTimeMillis() + timeoutMillis; + while (System.currentTimeMillis() < deadline) { + if (condition.getAsBoolean()) { + return true; + } + Thread.sleep(50L); + } + return condition.getAsBoolean(); + } + + /** Numeric file names in a Spark log directory, temp and hidden files excluded. */ + private static Set batchIds(File dir) { + Set ids = new TreeSet<>(); + String[] names = dir.list(); + if (names == null) { + return ids; + } + for (String name : names) { + if (!name.startsWith(".") && !name.endsWith(".tmp")) { + try { + ids.add(Long.parseLong(name)); + } catch (NumberFormatException e) { + // not a log entry + } + } + } + return ids; + } + + private static Set committedBatchIds(File checkpointDir) { + return batchIds(new File(checkpointDir, "commits")); + } + + /** The end epoch of a batch, the offset line of the single source in {@code offsets/}. */ + private static long endEpoch(File checkpointDir, long batchId) throws IOException { + File file = new File(new File(checkpointDir, "offsets"), Long.toString(batchId)); + List lines = new ArrayList<>(); + for (String line : Files.readAllLines(file.toPath(), StandardCharsets.UTF_8)) { + if (!line.trim().isEmpty()) { + lines.add(line.trim()); + } + } + assertTrue("offset log entry too short: " + lines, lines.size() >= 3); + assertEquals("one source expected in " + lines, 3, lines.size()); + return BeamOffset.fromJson(lines.get(2)).epoch(); + } + + private static BeamSourceCheckpoint sourceCheckpoint(File checkpointDir) { + return new BeamSourceCheckpoint( + new File(checkpointDir, "sources/0").getAbsolutePath(), new Configuration()); + } + private static Coder> coder() { return WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); } @@ -517,4 +951,196 @@ public CheckpointMark getCheckpointMark() { public void close() throws IOException {} } } + + // --------------------------------------------------------------------------------------------- + // a multi split in-memory UnboundedSource with finalization counting marks + // --------------------------------------------------------------------------------------------- + + /** + * Splits into one sub source per shard, each over {@code count / shards} elements named {@code + * shard-N-element-M} with evenly spaced timestamps. Marks are not Java serializable, they carry + * the shard's read position and record it under {@code /} when finalized. + */ + static class ShardedListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + static final Coder MARK_CODER = new MarkCoder(); + + /** Finalized positions keyed by {@code /}. */ + static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); + + private final String tag; + private final int shard; + private final int shards; + private final int perShard; + + ShardedListSource(String tag, int shards, int count) { + this(tag, -1, shards, count / shards); + } + + private ShardedListSource(String tag, int shard, int shards, int perShard) { + this.tag = tag; + this.shard = shard; + this.shards = shards; + this.perShard = perShard; + } + + static Set elements(int shards, int count) { + Set elements = new HashSet<>(); + for (int shard = 0; shard < shards; shard++) { + for (int index = 0; index < count / shards; index++) { + elements.add(element(shard, index)); + } + } + return elements; + } + + static String element(int shard, int index) { + return "shard-" + shard + "-element-" + index; + } + + static int shardOf(String element) { + return Integer.parseInt(element.substring("shard-".length(), element.indexOf("-element-"))); + } + + static int indexOf(String element) { + return Integer.parseInt(element.substring(element.lastIndexOf('-') + 1)); + } + + static List finalized(String tag, int shard) { + List positions = FINALIZED.get(key(tag, shard)); + if (positions == null) { + return Collections.emptyList(); + } + synchronized (positions) { + return new ArrayList<>(positions); + } + } + + private static String key(String tag, int shard) { + return tag + "/" + shard; + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + if (shard >= 0) { + return Collections.singletonList(this); + } + List splits = new ArrayList<>(); + for (int i = 0; i < shards; i++) { + splits.add(new ShardedListSource(tag, i, shards, perShard)); + } + return splits; + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable ShardMark mark) { + if (shard < 0) { + throw new IllegalStateException("split before reading"); + } + return new ShardReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return MARK_CODER; + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + /** Position of the next element of a shard, deliberately not {@link Serializable}. */ + static final class ShardMark implements UnboundedSource.CheckpointMark { + private final String key; + final int next; + + ShardMark(String key, int next) { + this.key = key; + this.next = next; + } + + @Override + public void finalizeCheckpoint() { + FINALIZED + .computeIfAbsent(key, k -> Collections.synchronizedList(new ArrayList<>())) + .add(next); + } + } + + private static final class MarkCoder extends CustomCoder { + private static final long serialVersionUID = 1L; + + @Override + public void encode(ShardMark mark, OutputStream out) throws IOException { + StringUtf8Coder.of().encode(mark.key, out); + VarIntCoder.of().encode(mark.next, out); + } + + @Override + public ShardMark decode(InputStream in) throws IOException { + return new ShardMark(StringUtf8Coder.of().decode(in), VarIntCoder.of().decode(in)); + } + } + + private static final class ShardReader extends UnboundedReader { + private final ShardedListSource source; + private int next; + private int current = -1; + + ShardReader(ShardedListSource source, int next) { + this.source = source; + this.next = next; + } + + @Override + public boolean start() { + return advance(); + } + + @Override + public boolean advance() { + if (next < source.perShard) { + current = next++; + return true; + } + return false; + } + + @Override + public String getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return element(source.shard, current); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(timestampMillis(source.shard * source.perShard + current)); + } + + @Override + public Instant getWatermark() { + return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new ShardMark(key(source.tag, source.shard), next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() {} + } + } } diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java new file mode 100644 index 000000000000..138d44acdab8 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCacheProtocolTest.java @@ -0,0 +1,478 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +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.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.hadoop.conf.Configuration; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.util.SerializableConfiguration; +import org.joda.time.Instant; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +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; +import scala.reflect.ClassTag; + +/** + * Drives {@link BeamPartitionReader} directly over hand built partitions to prove the reader cache + * protocol. The session only supplies the broadcasts, no query runs. + */ +@Category(StreamingTest.class) +@RunWith(JUnit4.class) +public class BeamReaderCacheProtocolTest { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + private static final AtomicInteger TAGS = new AtomicInteger(); + + private static final Coder> CODER = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + + private static final long MAX_RECORDS = 3L; + + private static final long MAX_BATCH_MILLIS = 30_000L; + + private static final long READER_IDLE_MILLIS = 600_000L; + + private static Broadcast options; + private static Broadcast hadoopConf; + private static Configuration conf; + + private String tag; + private IntListSource source; + private String location; + + @BeforeClass + public static void broadcastOnce() { + SparkSession session = SESSION.getSession(); + conf = ((org.apache.spark.sql.classic.SparkSession) session).sessionState().newHadoopConf(); + options = + session + .sparkContext() + .broadcast( + new SerializablePipelineOptions(PipelineOptionsFactory.create()), + ClassTag.apply(SerializablePipelineOptions.class)); + hadoopConf = + session + .sparkContext() + .broadcast( + new SerializableConfiguration(conf), + ClassTag.apply(SerializableConfiguration.class)); + } + + @Before + public void setUp() throws IOException { + tag = "protocol-" + TAGS.incrementAndGet(); + source = new IntListSource(tag, 100); + location = temp.newFolder("source").getAbsolutePath(); + } + + @After + public void tearDown() { + BeamReaderCache.invalidateAll(); + IntListSource.forget(tag); + } + + /** + * A batch starting where the cached reader stopped reuses it and finalizes the mark taken there. + */ + @Test + public void testContinuationReusesReaderAndFinalizesOnce() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals("nothing is finalized before the next batch opens", noPositions(), finalized()); + + BeamPartitionReader second = open(1, 2); + assertEquals("the epoch 1 mark is finalized on open", positions(3), finalized()); + assertEquals(Arrays.asList(3, 4, 5), drain(second)); + assertEquals(positions(3), finalized()); + assertEquals("one reader serves both batches", 1, IntListSource.created(tag)); + } + + /** A retried batch restarts from the durable mark at its start and finalizes nothing. */ + @Test + public void testRetryRecreatesReaderWithoutFinalizing() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(noPositions(), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** After the cache is lost the durable mark at the start epoch restores the position. */ + @Test + public void testExecutorHopRestoresFromDurableMark() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals(Arrays.asList(3, 4, 5), readBatch(1, 2)); + BeamReaderCache.invalidateAll(); + + assertEquals(Arrays.asList(6, 7, 8), readBatch(2, 3)); + assertEquals( + "only the epoch 1 mark was live when its batch was committed", positions(3), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** + * A start epoch above zero without a durable mark is an invariant violation, not a fresh start. + */ + @Test + public void testMissingMarkThrows() { + assertThrows(IllegalStateException.class, () -> new BeamPartitionReader<>(partition(5, 6))); + assertThrows( + IllegalStateException.class, + () -> new BeamPartitionReaderFactory().createReader(partition(5, 6))); + assertEquals(0, IntListSource.created(tag)); + } + + /** + * A reader closed before it started writes its start mark forward at E and stays reusable from E, + * a later batch from S is a fresh start. + */ + @Test + public void testNeverStartedReaderWritesStartMarkForward() throws Exception { + BeamPartitionReader idle = open(0, 1); + idle.close(); + assertTrue(new File(location, "marks/0/1").exists()); + assertEquals(1, IntListSource.created(tag)); + + assertEquals(Arrays.asList(0, 1, 2), readBatch(1, 2)); + assertEquals("the idle reader is reused", 1, IntListSource.created(tag)); + + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + assertEquals("a batch from epoch 0 is a fresh start", 2, IntListSource.created(tag)); + assertEquals(noPositions(), finalized()); + } + + /** A failed mark write fails the batch instead of advancing the reader silently. */ + @Test + public void testMarkWriteFailureFailsBatch() throws Exception { + String file = temp.newFile("not-a-directory").getAbsolutePath(); + BeamPartitionReader reader = new BeamPartitionReader<>(partition(file, 0, 1)); + List values = new ArrayList<>(); + assertThrows(IOException.class, () -> drainInto(reader, values)); + assertEquals(Arrays.asList(0, 1, 2), values); + } + + /** A retry after a failed mark write recreates the reader from the durable mark at S. */ + @Test + public void testRetryAfterFailedMarkWriteRecreatesReader() throws Exception { + String file = temp.newFile("not-a-directory").getAbsolutePath(); + assertEquals( + Arrays.asList(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals( + Arrays.asList(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals(noPositions(), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + /** A retry after a failed mark encode above epoch 0 restores from the durable mark at S. */ + @Test + public void testRetryAfterFailedMarkEncodeRestoresFromDurableMark() throws Exception { + assertEquals(Arrays.asList(0, 1, 2), readBatch(0, 1)); + IntListSource.failEncoding(tag, true); + assertEquals( + Arrays.asList(3, 4, 5), drainUntilFailure(partition(1, 2), IllegalStateException.class)); + IntListSource.failEncoding(tag, false); + + assertEquals(Arrays.asList(3, 4, 5), readBatch(1, 2)); + assertEquals( + "the epoch 1 mark was finalized by the first open only", positions(3), finalized()); + assertEquals(2, IntListSource.created(tag)); + } + + // --------------------------------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------------------------------- + + private BeamInputPartition partition(long start, long end) { + return partition(location, start, end); + } + + private BeamInputPartition partition(String checkpointLocation, long start, long end) { + return new BeamInputPartition<>( + source, + CODER, + options, + hadoopConf, + checkpointLocation, + 0, + start, + end, + MAX_RECORDS, + MAX_BATCH_MILLIS, + READER_IDLE_MILLIS, + new String[0]); + } + + private BeamPartitionReader open(long start, long end) throws IOException { + return new BeamPartitionReader<>(partition(start, end)); + } + + private List readBatch(long start, long end) throws IOException { + return drain(open(start, end)); + } + + private static List drain(BeamPartitionReader reader) throws IOException { + List values = new ArrayList<>(); + drainInto(reader, values); + return values; + } + + private static void drainInto(BeamPartitionReader reader, List values) + throws IOException { + while (reader.next()) { + InternalRow row = reader.get(); + values.add(CoderUtils.decodeFromByteArray(CODER, row.getBinary(0)).getValue()); + } + reader.close(); + } + + /** Opens and drains a batch expected to fail, returns what it delivered before failing. */ + private static List drainUntilFailure( + BeamInputPartition partition, Class failure) + throws IOException { + BeamPartitionReader reader = new BeamPartitionReader<>(partition); + List values = new ArrayList<>(); + assertThrows(failure, () -> drainInto(reader, values)); + return values; + } + + private List finalized() { + return IntListSource.finalized(tag); + } + + private static List positions(Integer... positions) { + return Arrays.asList(positions); + } + + private static List noPositions() { + return Collections.emptyList(); + } + + // --------------------------------------------------------------------------------------------- + // a list source whose marks count their finalizations + // --------------------------------------------------------------------------------------------- + + /** + * Single split source over the integers {@code 0..count-1}. Marks are not Java serializable and + * record the position they finalize under the source's tag, readers are counted per tag. + */ + static class IntListSource extends UnboundedSource { + private static final long serialVersionUID = 1L; + + static final Coder MARK_CODER = new MarkCoder(); + + private static final ConcurrentMap> FINALIZED = new ConcurrentHashMap<>(); + private static final ConcurrentMap CREATED = new ConcurrentHashMap<>(); + private static final Set FAILING_ENCODE = ConcurrentHashMap.newKeySet(); + + private final String tag; + private final int count; + + IntListSource(String tag, int count) { + this.tag = tag; + this.count = count; + } + + static List finalized(String tag) { + List positions = FINALIZED.get(tag); + if (positions == null) { + return Collections.emptyList(); + } + synchronized (positions) { + return new ArrayList<>(positions); + } + } + + static int created(String tag) { + AtomicInteger created = CREATED.get(tag); + return created == null ? 0 : created.get(); + } + + static void failEncoding(String tag, boolean fail) { + if (fail) { + FAILING_ENCODE.add(tag); + } else { + FAILING_ENCODE.remove(tag); + } + } + + static void forget(String tag) { + FINALIZED.remove(tag); + CREATED.remove(tag); + FAILING_ENCODE.remove(tag); + } + + @Override + public List split(int desiredNumSplits, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark mark) { + CREATED.computeIfAbsent(tag, t -> new AtomicInteger()).incrementAndGet(); + return new IntListReader(this, mark == null ? 0 : mark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return MARK_CODER; + } + + @Override + public Coder getOutputCoder() { + return VarIntCoder.of(); + } + + /** Position of the next element, deliberately not {@link java.io.Serializable}. */ + static final class Mark implements UnboundedSource.CheckpointMark { + private final String tag; + private final int next; + + Mark(String tag, int next) { + this.tag = tag; + this.next = next; + } + + @Override + public void finalizeCheckpoint() { + FINALIZED + .computeIfAbsent(tag, t -> Collections.synchronizedList(new ArrayList<>())) + .add(next); + } + } + + private static final class MarkCoder extends CustomCoder { + private static final long serialVersionUID = 1L; + + @Override + public void encode(Mark mark, OutputStream out) throws IOException { + if (FAILING_ENCODE.contains(mark.tag)) { + throw new CoderException("injected mark encode failure for " + mark.tag); + } + StringUtf8Coder.of().encode(mark.tag, out); + VarIntCoder.of().encode(mark.next, out); + } + + @Override + public Mark decode(InputStream in) throws IOException { + return new Mark(StringUtf8Coder.of().decode(in), VarIntCoder.of().decode(in)); + } + } + + private static final class IntListReader extends UnboundedReader { + private final IntListSource source; + private int next; + private int current = -1; + + IntListReader(IntListSource 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 Integer getCurrent() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return current; + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant(1_700_000_000_000L + current * 1_000L); + } + + @Override + public Instant getWatermark() { + return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new Mark(source.tag, next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() {} + } + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java new file mode 100644 index 000000000000..e65918f219cf --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpointTest.java @@ -0,0 +1,209 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.io.streaming; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.fail; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.TreeSet; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.hadoop.conf.Configuration; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link BeamSourceCheckpoint} on the local file system, no Spark session. */ +@RunWith(JUnit4.class) +public class BeamSourceCheckpointTest { + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + private File root; + private BeamSourceCheckpoint checkpoint; + + @Before + public void setUp() { + root = new File(temp.getRoot(), "sources/0"); + checkpoint = fresh(); + } + + /** Splits round trip once, a second pin fails, a fresh location has none. */ + @Test + public void testSplitsArePinnedOnce() throws Exception { + assertNull("fresh location must have no splits", checkpoint.readSplits()); + + List> splits = + CountingSource.unbounded().split(2, PipelineOptionsFactory.create()); + assertEquals(2, splits.size()); + checkpoint.writeSplits(splits); + + List> restored = fresh().readSplits(); + assertNotNull(restored); + assertEquals(splits.size(), restored.size()); + for (int i = 0; i < splits.size(); i++) { + assertEquals(splits.get(i).getClass(), restored.get(i).getClass()); + } + + assertThrows(IOException.class, () -> checkpoint.writeSplits(splits)); + assertThrows(IOException.class, () -> fresh().writeSplits(splits)); + assertEquals(2, fresh().readSplits().size()); + assertCleanTree(); + } + + /** Marks round trip byte for byte, the latest write at an epoch wins, absent marks are null. */ + @Test + public void testMarkRoundTrip() throws Exception { + byte[] first = {0, 1, -1, 127, -128, 42}; + byte[] second = {7}; + + assertNull(checkpoint.readMark(0, 1)); + checkpoint.writeMark(0, 1, first); + assertArrayEquals(first, checkpoint.readMark(0, 1)); + assertArrayEquals(first, fresh().readMark(0, 1)); + + checkpoint.writeMark(0, 1, second); + assertArrayEquals(second, checkpoint.readMark(0, 1)); + assertArrayEquals(second, fresh().readMark(0, 1)); + + checkpoint.writeMark(0, 2, new byte[0]); + assertArrayEquals(new byte[0], checkpoint.readMark(0, 2)); + + assertNull("missing epoch", checkpoint.readMark(0, 3)); + assertNull("missing split", checkpoint.readMark(1, 1)); + assertNull("missing split, fresh instance", fresh().readMark(1, 1)); + assertCleanTree(); + } + + /** Purging keeps epochs at or above the floor, is idempotent and ignores absent splits. */ + @Test + public void testPurgeMarksBelow() throws Exception { + for (long epoch = 1; epoch <= 5; epoch++) { + checkpoint.writeMark(0, epoch, bytes(epoch)); + } + for (long epoch = 1; epoch <= 3; epoch++) { + checkpoint.writeMark(1, epoch, bytes(epoch)); + } + + checkpoint.purgeMarksBelow(0, 3); + assertEquals(epochs(3, 4, 5), markEpochs(0)); + assertEquals("other split untouched", epochs(1, 2, 3), markEpochs(1)); + + checkpoint.purgeMarksBelow(0, 3); + assertEquals("idempotent", epochs(3, 4, 5), markEpochs(0)); + + checkpoint.purgeMarksBelow(0, 1); + assertEquals("lower floor deletes nothing", epochs(3, 4, 5), markEpochs(0)); + + checkpoint.purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + assertArrayEquals(bytes(5), checkpoint.readMark(0, 5)); + + checkpoint.purgeMarksBelow(0, 6); + assertEquals(epochs(), markEpochs(0)); + + checkpoint.purgeMarksBelow(7, 3); + checkpoint.purgeMarksBelow(7, 3); + assertEquals(epochs(), markEpochs(7)); + assertCleanTree(); + } + + /** A fresh instance without an in memory floor purges the same files as the writer. */ + @Test + public void testPurgeWithoutInMemoryFloor() throws Exception { + for (long epoch = 1; epoch <= 5; epoch++) { + checkpoint.writeMark(0, epoch, bytes(epoch)); + } + checkpoint.purgeMarksBelow(0, 2); + assertEquals(epochs(2, 3, 4, 5), markEpochs(0)); + + BeamSourceCheckpoint other = fresh(); + other.purgeMarksBelow(0, 4); + assertEquals(epochs(4, 5), markEpochs(0)); + other.purgeMarksBelow(0, 4); + assertEquals(epochs(4, 5), markEpochs(0)); + + // The writer's floor is 2, epochs 2 and 3 are already gone. + checkpoint.purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + + fresh().purgeMarksBelow(0, 5); + assertEquals(epochs(5), markEpochs(0)); + assertCleanTree(); + } + + private BeamSourceCheckpoint fresh() { + return new BeamSourceCheckpoint(root.getAbsolutePath(), new Configuration()); + } + + private static byte[] bytes(long epoch) { + return new byte[] {(byte) epoch, (byte) (epoch * 3)}; + } + + private static Set epochs(long... epochs) { + return Arrays.stream(epochs).boxed().collect(Collectors.toCollection(TreeSet::new)); + } + + /** Numeric file names under {@code marks/}, empty if the directory is absent. */ + private Set markEpochs(int split) { + File dir = new File(root, "marks/" + split); + String[] names = dir.list(); + Set epochs = new TreeSet<>(); + if (names == null) { + return epochs; + } + for (String name : names) { + if (!name.startsWith(".")) { + epochs.add(Long.parseLong(name)); + } + } + return epochs; + } + + /** No temp files and no hidden files other than Hadoop's checksum sidecars remain. */ + private void assertCleanTree() throws IOException { + List offenders = new ArrayList<>(); + try (Stream paths = Files.walk(temp.getRoot().toPath())) { + paths + .map(path -> path.getFileName().toString()) + .filter(name -> name.endsWith(".tmp") || (name.startsWith(".") && !name.endsWith(".crc"))) + .forEach(offenders::add); + } + if (!offenders.isEmpty()) { + fail("unexpected files under " + temp.getRoot() + ": " + offenders); + } + } +} 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 391350fd348e..fb0192dba868 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 @@ -63,6 +63,15 @@ public interface SparkStructuredStreamingPipelineOptions extends SparkCommonPipe void setMaxBatchDurationMillis(long value); + @Description( + "Idle time in milliseconds after which an executor closes a cached unbounded reader. Must " + + "exceed the longest gap between two micro-batches, a closed reader's last checkpoint " + + "mark is not finalized and the source redelivers (streaming mode only).") + @Default.Long(600_000) + long getReaderIdleTimeoutMillis(); + + void setReaderIdleTimeoutMillis(long value); + @Description( "Test-oriented: gracefully stop streaming queries after this many consecutive empty " + "micro-batches. Disabled if negative (streaming mode only).") From d12a500236a146f9799be2cb17ec2d6899b437f6 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 3 Sep 2026 11:44:29 +0000 Subject: [PATCH 3/4] Trigger CI after the setup-gradle allowlist fix From 800c9bdb0b8df8847c700acf394e781495a15f45 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Thu, 3 Sep 2026 12:31:18 +0000 Subject: [PATCH 4/4] [Spark 4] Rebase the Structured Streaming POC on the reworked micro-batch source Brings the final file states of the spark4-streaming-poc branch onto the head of the slice 3 rework (#39971) as one commit: the streaming pipeline translator and evaluation context, the Read, Impulse, GroupByKey and stateful ParDo translators, the transformWithState state and timer bridge, and the end to end streaming tests. The io/streaming package of the rework is kept as is, the POC's own version of it is dropped. The end of stream sentinel the POC had added to the old source is re-applied on the reworked BeamPartitionReader and BeamReaderCache: a batch that holds data ends at the first empty poll so its watermark is declared first, and an exhausted reader whose watermark reached the end of the global window emits one empty payload row at the maximum timestamp once per cached reader. The translators filter that row. Callers of the removed int maxRecordsPerMicroBatch option now use the long maxRecordsPerBatch option of master, whose per batch quota is split across the splits with at least one record per split, which is what the tests relied on. StreamingCheckpointRestartTest asserts the reworked checkpoint layout, splits and marks under the per source location Spark hands the stream, instead of the old beam-source- directory found by a recursive search. The JUnit method timeouts are removed from every streaming test. Its timeout thread group leaks into Spark's static pools and breaks later tests in the same JVM. StreamingTestUtils gains run and waitUntilFinish helpers with a five minute deadline that cancel the pipeline and fail the test instead. SparkSessionFactory, build.gradle, the pipeline options, result, runner, evaluation context and pipeline translator of the shared base need no change, the merged slices already carry the POC's deltas including the RocksDB state store default and the Kryo registrations. --- .../io/streaming/BeamPartitionReader.java | 27 +- .../io/streaming/BeamReaderCache.java | 10 + .../PipelineTranslatorFactory.java | 40 + .../PipelineTranslatorStreaming.java | 103 +++ .../StreamingEvaluationContext.java | 377 +++++++++ .../GroupByKeyStreamingTranslator.java | 120 +++ .../streaming/ImpulseStreamingTranslator.java | 237 ++++++ .../streaming/ReadUnboundedTranslator.java | 128 ++++ .../StatefulParDoStreamingTranslator.java | 201 +++++ .../StreamingTranslationHelpers.java | 255 +++++++ .../streaming/TwsTransformFactory.java | 198 +++++ .../state/BeamStatefulProcessor.java | 526 +++++++++++++ .../state/BeamStatefulProcessorConfig.java | 242 ++++++ .../translation/streaming/state/BytesKV.java | 106 +++ .../streaming/state/TwsStateInternals.java | 577 ++++++++++++++ .../streaming/state/TwsTimerInternals.java | 390 ++++++++++ .../ChainedStatefulStreamingEvidenceTest.java | 392 ++++++++++ .../ChainedStatefulStreamingTest.java | 158 ++++ .../streaming/PAssertStreamingTest.java | 195 +++++ .../streaming/StatefulParDoStreamingTest.java | 153 ++++ .../StatelessParDoStreamingTest.java | 113 +++ .../StreamingCheckpointRestartTest.java | 194 +++++ .../StreamingPipelineLifecycleTest.java | 321 ++++++++ .../streaming/StreamingTestUtils.java | 478 ++++++++++++ .../WindowedGroupByKeyStreamingTest.java | 718 ++++++++++++++++++ .../state/BeamStatefulProcessorTest.java | 440 +++++++++++ .../state/TwsStateInternalsTest.java | 359 +++++++++ .../state/TwsTimerInternalsTest.java | 418 ++++++++++ 28 files changed, 7474 insertions(+), 2 deletions(-) create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/GroupByKeyStreamingTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/TwsTransformFactory.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorConfig.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BytesKV.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java create mode 100644 runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternals.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessorTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternalsTest.java create mode 100644 runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsTimerInternalsTest.java diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java index d500aec44644..79f155c4e797 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamPartitionReader.java @@ -24,6 +24,8 @@ import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; 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.BackOff; import org.apache.beam.sdk.util.BackOffUtils; import org.apache.beam.sdk.util.FluentBackoff; @@ -44,8 +46,10 @@ /** * Reads one split of a Beam {@link UnboundedSource} for one micro-batch. * - *

The batch ends at the record quota or at the deadline. The reader then writes its checkpoint - * mark durably at the end epoch and stays in {@link BeamReaderCache} for the next batch. A failed + *

The batch ends at the record quota, at the deadline, or at the first empty poll once it holds + * data. The reader then writes its checkpoint mark durably at the end epoch and stays in {@link + * BeamReaderCache} for the next batch. An exhausted source whose watermark reached the end of the + * global window yields one sentinel row with an empty payload, the translators filter it. A failed * mark write fails the task. An attempt Spark killed or failed writes nothing and its reader is * dropped, the retry restores from the durable mark at the start epoch. * @@ -113,6 +117,19 @@ public boolean next() throws IOException { current = toRow(); return true; } + if (recordsRead > 0) { + // The data batch declares its watermark before any sentinel arrives in a later batch. + return endOfBatch(false); + } + if (!cached.hasEmittedSentinel()) { + Instant watermark = cached.reader().getWatermark(); + if (watermark != null && !watermark.isBefore(GlobalWindow.INSTANCE.maxTimestamp())) { + cached.markSentinelEmitted(); + recordsRead++; + current = sentinelRow(); + return true; + } + } if (backOff == null) { backOff = backOff(remaining); } @@ -206,4 +223,10 @@ private InternalRow toRow() { // Spark stores TimestampType as microseconds. return new GenericInternalRow(new Object[] {payload, timestamp.getMillis() * 1000L}); } + + /** Empty payload at the end of time, the cross package end of stream contract. */ + private static InternalRow sentinelRow() { + return new GenericInternalRow( + new Object[] {new byte[0], BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis() * 1000L}); + } } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java index cbc697412b18..a7ac2d4b45e1 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -178,6 +178,7 @@ public static final class CachedReader implements Closeable { private boolean started; private boolean inBatch; private boolean moved; + private boolean sentinelEmitted; private long positionEpoch; private byte @Nullable [] positionMark; private @Nullable CheckpointMark pendingMark; @@ -219,6 +220,15 @@ public synchronized long positionEpoch() { return positionEpoch; } + /** Whether the end of stream sentinel row was emitted for this reader. */ + public synchronized boolean hasEmittedSentinel() { + return sentinelEmitted; + } + + public synchronized void markSentinelEmitted() { + sentinelEmitted = true; + } + /** The coded mark of the current position, null for a fresh start. */ synchronized byte @Nullable [] positionMark() { return positionMark; 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..c55801a4ccf6 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/PipelineTranslatorStreaming.java @@ -0,0 +1,103 @@ +/* + * 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.ImpulseStreamingTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.ReadUnboundedTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.streaming.StatefulParDoStreamingTranslator; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.transforms.Combine; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +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 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 Impulse) { + return (TransformTranslator) new ImpulseStreamingTranslator(); + } + + 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. + } + + // 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..51f9bfaa6310 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -0,0 +1,377 @@ +/* + * 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.nio.file.Path; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.annotations.Internal; +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.checkerframework.checker.nullness.qual.Nullable; +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: + * + *

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

Draining on {@link #stop()}

+ * + *

{@link #stop()} does not attempt a full {@code Trigger.AvailableNow()} drain pass that + * processes every already-buffered offset before halting: doing so would require stopping the query + * and restarting it with a different trigger against the same checkpoint, which is more machinery + * than this POC's lifecycle warrants and risks checkpoint-compatibility bugs of its own. Instead, + * {@link #stop()} relies on {@link StreamingQuery#stop()}'s own graceful behaviour, which lets a + * micro-batch that is already in flight finish normally instead of interrupting it, and only then + * halts the query. Data that was not yet pulled into an in-flight micro-batch at the moment {@link + * #stop()} is called is simply left unprocessed. This is a documented limitation of the POC, not an + * oversight. + */ +@Internal +public class StreamingEvaluationContext extends EvaluationContext { + private static final Logger LOG = LoggerFactory.getLogger(StreamingEvaluationContext.class); + + // How long one awaitTermination poll blocks on a single query before moving on to the next one, + // see awaitTermination(List). Short enough to surface a failure promptly, long enough to not + // busy-spin while every query is healthy. + private static final long AWAIT_POLL_TIMEOUT_MILLIS = 100; + + private final SparkStructuredStreamingPipelineOptions options; + + // Guards both `queries` and `stopped` so evaluate() (which appends to `queries` as it starts + // 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; + + // Set by checkpointBaseDir() when no checkpointDir was configured, so a temporary directory was + // created as a fallback. Written once in evaluate() before any query starts and read back in the + // finally block of that same method on the same thread, so unlike `queries` and `stopped` it does + // not need `lock`. Never set when the user configured a checkpointDir, which evaluate() must + // never delete. + private @Nullable Path tempCheckpointDir; + + StreamingEvaluationContext( + Collection> leaves, + SparkSession session, + 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); + } + awaitTermination(toAwait); + } finally { + if (idleStopListener != null) { + getSparkSession().streams().removeListener(idleStopListener); + } + if (tempCheckpointDir != null) { + deleteTempCheckpointDir(tempCheckpointDir); + } + } + } + + /** + * 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); + } + } + + /** + * Blocks until every query in {@code toAwait} has terminated, polling them round robin with a + * short timeout rather than awaiting them one after the other: a plain {@link + * StreamingQuery#awaitTermination()} on the first query would block for as long as that query + * keeps running and not surface a failure of a later query until then. Polling bounds the latency + * of surfacing any query's failure by roughly one round of poll timeouts, no matter which query + * fails. + * + *

On the first query that fails, {@link #stop()} makes sure sibling queries do not keep + * running, and the failure is rethrown. Deliberately not {@code + * StreamingQueryManager#awaitAnyTermination}: that relies on the session global {@code + * resetTerminated} bookkeeping and would interfere with concurrent tests sharing one session. + */ + private void awaitTermination(List toAwait) { + List active = new ArrayList<>(toAwait); + while (!active.isEmpty()) { + Iterator iterator = active.iterator(); + while (iterator.hasNext()) { + StreamingQuery query = iterator.next(); + try { + if (query.awaitTermination(AWAIT_POLL_TIMEOUT_MILLIS)) { + // Terminated without an exception, nothing left to await for this query. + iterator.remove(); + } + } catch (StreamingQueryException e) { + LOG.error("Streaming query {} terminated with an exception.", query.id(), e); + // Make sure sibling queries do not keep running once one of them has failed. + stop(); + throw new RuntimeException(e); + } + } + } + } + + /** + * 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 String checkpointBaseDir(SparkCommonPipelineOptions options) { + String dir = options.getCheckpointDir(); + if (dir == null || dir.isEmpty()) { + try { + tempCheckpointDir = Files.createTempDirectory("beam-spark4-streaming-checkpoint"); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + dir = tempCheckpointDir.toString(); + LOG.warn("No checkpoint directory configured, falling back to temporary directory {}.", dir); + } + return dir; + } + + /** + * Best-effort, recursive delete of the fallback temporary checkpoint directory created by {@link + * #checkpointBaseDir}. RocksDB and Spark write nested state and log files under it, so plain + * {@link Files#delete} is not enough. Failures are logged rather than thrown: cleanup is a + * courtesy, not something that should fail a pipeline that otherwise ran to completion. + */ + private static void deleteTempCheckpointDir(Path dir) { + try (Stream paths = Files.walk(dir)) { + paths + .sorted(Comparator.reverseOrder()) + .forEach( + path -> { + try { + Files.delete(path); + } catch (IOException e) { + LOG.warn( + "Failed to delete temporary checkpoint path {}: {}", + path, + String.valueOf(e.getMessage())); + } + }); + } catch (IOException e) { + LOG.warn( + "Failed to clean up temporary checkpoint directory {}: {}", + dir, + String.valueOf(e.getMessage())); + } + } + + /** + * Counts, per query, the number of consecutive micro-batches with zero input rows, and gracefully + * stops a query once its count reaches {@code threshold}. The count for a query resets to zero as + * 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/ImpulseStreamingTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java new file mode 100644 index 000000000000..30256b79bf2d --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ImpulseStreamingTranslator.java @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils.EMPTY_BYTE_ARRAY; +import static org.apache.spark.sql.functions.col; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.CoderHelpers; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.spark.api.java.function.FilterFunction; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Streaming translator for {@link Impulse}. + * + *

Emits a single empty byte array in a streaming micro-batch, advancing the watermark to + * infinity upon completion so that pipelines mixing Impulse-generated elements (e.g. {@code + * Create.of} in {@code PAssert}) with unbounded streams can safely union without schema or + * batch/streaming mismatch. + */ +public class ImpulseStreamingTranslator + extends TransformTranslator, Impulse> { + + public ImpulseStreamingTranslator() { + super(0.05f); + } + + @Override + protected void translate(Impulse transform, Context cxt) { + PCollection output = cxt.getOutput(); + Coder elementCoder = output.getCoder(); + + WindowedValues.FullWindowedValueCoder payloadCoder = + WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + SparkStructuredStreamingPipelineOptions options = + cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); + + Dataset rows = + UnboundedSourceDataset.of( + cxt.getSparkSession(), + new ImpulseSource(), + payloadCoder, + options, + cxt.getCurrentTransform().getFullName()); + + Encoder> encoder = + cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + Dataset> dataset = + rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) + .as(Encoders.BINARY()) + .filter(new FilterNonEmptyPayload()) + .map(new DecodePayload<>(payloadCoder), encoder); + + cxt.putDataset(output, dataset); + } + + private static final class FilterNonEmptyPayload implements FilterFunction { + @Override + public boolean call(byte[] payload) { + return payload != null && payload.length > 0; + } + } + + private static final class DecodePayload implements MapFunction> { + private final Coder> coder; + + DecodePayload(Coder> coder) { + this.coder = coder; + } + + @Override + public WindowedValue call(byte[] payload) { + return CoderHelpers.fromByteArray(payload, coder); + } + } + + /** An unbounded source that produces exactly one empty byte array and then finishes. */ + private static final class ImpulseSource + extends UnboundedSource { + + @Override + public List> split( + int desiredNumSplits, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public UnboundedReader createReader( + PipelineOptions options, @Nullable ImpulseCheckpointMark checkpointMark) { + return new ImpulseReader(this, checkpointMark != null && checkpointMark.done); + } + + @Override + public Coder getCheckpointMarkCoder() { + return ImpulseCheckpointMarkCoder.of(); + } + + @Override + public Coder getOutputCoder() { + return ByteArrayCoder.of(); + } + + static final class ImpulseCheckpointMark + implements UnboundedSource.CheckpointMark, Serializable { + final boolean done; + + ImpulseCheckpointMark(boolean done) { + this.done = done; + } + + @Override + public void finalizeCheckpoint() {} + } + + static final class ImpulseCheckpointMarkCoder extends AtomicCoder { + private static final ImpulseCheckpointMarkCoder INSTANCE = new ImpulseCheckpointMarkCoder(); + + public static ImpulseCheckpointMarkCoder of() { + return INSTANCE; + } + + @Override + public void encode(ImpulseCheckpointMark value, OutputStream outStream) throws IOException { + outStream.write(value.done ? 1 : 0); + } + + @Override + public ImpulseCheckpointMark decode(InputStream inStream) throws IOException { + return new ImpulseCheckpointMark(inStream.read() != 0); + } + } + + private static final class ImpulseReader extends UnboundedReader { + private final ImpulseSource source; + private boolean done; + private boolean started; + + ImpulseReader(ImpulseSource source, boolean done) { + this.source = source; + this.done = done; + } + + @Override + public boolean start() { + if (done) { + return false; + } + started = true; + done = true; + return true; + } + + @Override + public boolean advance() { + return false; + } + + @Override + public byte[] getCurrent() { + if (!started) { + throw new NoSuchElementException(); + } + return EMPTY_BYTE_ARRAY; + } + + @Override + public Instant getCurrentTimestamp() { + if (!started) { + throw new NoSuchElementException(); + } + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public void close() {} + + @Override + public Instant getWatermark() { + return done ? BoundedWindow.TIMESTAMP_MAX_VALUE : BoundedWindow.TIMESTAMP_MIN_VALUE; + } + + @Override + public CheckpointMark getCheckpointMark() { + return new ImpulseCheckpointMark(done); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + } + } +} diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java new file mode 100644 index 000000000000..559318e33a58 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ReadUnboundedTranslator.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.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.FilterFunction; +import org.apache.spark.api.java.function.MapFunction; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; + +/** + * Translator for {@link SplittableParDo.PrimitiveUnboundedRead}, the streaming counterpart of the + * batch {@code ReadSourceTranslatorBatch}. + * + *

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

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

    + *
  • The watermark is already declared by {@link UnboundedSourceDataset} and must never + * be re-declared, Spark rejects a second {@code withWatermark} in the same plan. The {@code + * EventTimeWatermark} plan node sits below the projection and the typed map applied here and + * survives both, and a {@code transformWithState} operator downstream reads the query wide + * watermark rather than a column, so dropping the timestamp column here is safe. + *
  • Elements arrive in the global window, timestamped with the reader's record + * timestamp. A windowed pipeline therefore still needs its {@code Window.Assign}, which is + * translated by the reused batch translator. + *
+ */ +public class ReadUnboundedTranslator + extends TransformTranslator, SplittableParDo.PrimitiveUnboundedRead> { + + public ReadUnboundedTranslator() { + super(0.05f); + } + + @Override + protected void translate(SplittableParDo.PrimitiveUnboundedRead transform, Context cxt) { + PCollection output = cxt.getOutput(); + UnboundedSource source = transform.getSource(); + Coder elementCoder = output.getCoder(); + + // Matches what the partition readers emit: a value in the global window, timestamped with the + // record's own event timestamp. + WindowedValues.FullWindowedValueCoder payloadCoder = + WindowedValues.getFullCoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + SparkStructuredStreamingPipelineOptions options = + cxt.getOptions().as(SparkStructuredStreamingPipelineOptions.class); + + Dataset rows = + UnboundedSourceDataset.of( + cxt.getSparkSession(), + source, + payloadCoder, + options, + cxt.getCurrentTransform().getFullName()); + + Encoder> encoder = + cxt.windowedEncoder(elementCoder, GlobalWindow.Coder.INSTANCE); + + Dataset> dataset = + rows.select(col(UnboundedSourceDataset.COL_PAYLOAD)) + .as(Encoders.BINARY()) + .filter(new FilterNonEmptyPayload()) + .map(new DecodePayload<>(payloadCoder), encoder); + + cxt.putDataset(output, dataset); + } + + /** + * Filters out empty sentinel payloads emitted to advance the watermark upon stream exhaustion. + */ + private static final class FilterNonEmptyPayload implements FilterFunction { + @Override + public boolean call(byte[] payload) { + return payload != null && payload.length > 0; + } + } + + /** Decodes the binary payload column back into a Beam {@code WindowedValue}. */ + private static final class DecodePayload implements MapFunction> { + private final Coder> coder; + + 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..ef6df4777c47 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTranslator.java @@ -0,0 +1,201 @@ +/* + * 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); + } + + /** + * Unlike the batch translators, this override never returns {@code false}: {@link + * #rejectUnsupported} throws with the offending feature named instead. Silently declining here + * would make {@code PipelineTranslator#getSupportedTranslator} fall back to the batch {@code + * ParDo} translator, which would run the stateful {@code DoFn} without the streaming state and + * timer semantics and produce quietly wrong results. + */ + @Override + protected boolean canTranslate(ParDo.MultiOutput, OutputT> transform) { + rejectUnsupported(transform); + return true; + } + + /** + * Throws when the stateful {@code ParDo} uses a feature this translator does not implement, + * naming that feature; returns normally otherwise. + */ + private void rejectUnsupported(ParDo.MultiOutput, OutputT> transform) { + DoFn, OutputT> doFn = transform.getFn(); + String stepName = doFn.getClass().getName(); + DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn); + + 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"); + } + } + + @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..00aa839fcda4 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTranslationHelpers.java @@ -0,0 +1,255 @@ +/* + * 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.AfterPane; +import org.apache.beam.sdk.transforms.windowing.AfterWatermark; +import org.apache.beam.sdk.transforms.windowing.DefaultTrigger; +import org.apache.beam.sdk.transforms.windowing.Never; +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.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 and unsupported custom triggers. + */ + 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 (!isSupportedTrigger(trigger)) { + throw unsupported( + stepName, + "the custom trigger " + + trigger + + ". Only the default trigger (one on-time pane per window when the watermark passes " + + "its end) and late firings with AfterPane.elementCountAtLeast(1) are implemented"); + } + } + + /** + * Identifies triggers supported by the streaming translation. + * + *

    + *
  • {@link DefaultTrigger} + *
  • {@link AfterWatermark.FromEndOfWindow} without early or late firings + *
  • {@link Never.NeverTrigger} used by {@code PAssert} on unbounded streams + *
  • {@link AfterWatermark.AfterWatermarkEarlyAndLate} with no early firings and late firings + * configured via {@link AfterPane#elementCountAtLeast(int)} with element count 1 + *
+ */ + private static boolean isSupportedTrigger(Trigger trigger) { + if (trigger instanceof DefaultTrigger) { + return true; + } + if (trigger instanceof AfterWatermark.FromEndOfWindow) { + return true; + } + if (trigger instanceof Never.NeverTrigger) { + return true; + } + if (trigger instanceof AfterWatermark.AfterWatermarkEarlyAndLate) { + AfterWatermark.AfterWatermarkEarlyAndLate earlyAndLate = + (AfterWatermark.AfterWatermarkEarlyAndLate) trigger; + if (!(earlyAndLate.getEarlyTrigger() instanceof Never.NeverTrigger)) { + return false; + } + Trigger lateTrigger = earlyAndLate.getLateTrigger(); + if (lateTrigger == null || lateTrigger instanceof Never.NeverTrigger) { + return true; + } + if (lateTrigger instanceof AfterPane) { + return ((AfterPane) lateTrigger).getElementCount() == 1; + } + return false; + } + 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..f4e38feee8c0 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/BeamStatefulProcessor.java @@ -0,0 +1,526 @@ +/* + * 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.state.TimeDomain; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.reflect.DoFnInvoker; +import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.SerializableUtils; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +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. + * + *

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

Timers

+ * + *

Beam timers fire only in {@code handleExpiredTimer}, never in {@code handleInputRows}. A timer + * 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(); + + 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()); + } + + // The step context is what DoFnRunners consult for arrival side expiry decisions, the late + // data filter of the GROUP_ALSO_BY_WINDOW path and the expired window drop of the stateful + // ParDo path, so it hands out the clamped watermark view. The timer firing path keeps the + // real timerInternals, see ArrivalExpiryClampedTimerInternals. + TimerInternals arrivalTimerInternals = new ArrivalExpiryClampedTimerInternals(timerInternals); + StepContext stepContext = + new StepContext() { + @Override + public StateInternals stateInternals() { + return stateInternals; + } + + @Override + public TimerInternals timerInternals() { + return arrivalTimerInternals; + } + }; + + 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); + } + } + + /** + * A view of {@link TwsTimerInternals} for arrival side expiry decisions whose input watermark is + * clamped to the end of the global window, every other method delegating unchanged. + * + *

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

It must never be used on the timer firing path. {@link + * TwsTimerInternals#removeTimersReadyToFire} releases a timer only once the watermark is strictly + * past it, so clamping there would withhold end of global window timers forever. Likewise the + * {@code ReduceFnRunner} behind the GROUP_ALSO_BY_WINDOW mode keeps the unclamped {@link + * TwsTimerInternals}, its triggers must see the watermark pass the end of the global window to + * emit the final pane. + */ + private static final class ArrivalExpiryClampedTimerInternals implements TimerInternals { + + private final TwsTimerInternals delegate; + + private ArrivalExpiryClampedTimerInternals(TwsTimerInternals delegate) { + this.delegate = delegate; + } + + @Override + public Instant currentInputWatermarkTime() { + Instant watermark = delegate.currentInputWatermarkTime(); + Instant endOfGlobalWindow = GlobalWindow.INSTANCE.maxTimestamp(); + return watermark.isAfter(endOfGlobalWindow) ? endOfGlobalWindow : watermark; + } + + @Override + public void setTimer( + StateNamespace namespace, + String timerId, + String timerFamilyId, + Instant target, + Instant outputTimestamp, + TimeDomain timeDomain) { + delegate.setTimer(namespace, timerId, timerFamilyId, target, outputTimestamp, timeDomain); + } + + @Override + public void setTimer(TimerInternals.TimerData timerData) { + delegate.setTimer(timerData); + } + + @Override + public void deleteTimer( + StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { + delegate.deleteTimer(namespace, timerId, timerFamilyId, timeDomain); + } + + @Override + public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) { + delegate.deleteTimer(namespace, timerId, timerFamilyId); + } + + @Override + public void deleteTimer(TimerInternals.TimerData timerKey) { + delegate.deleteTimer(timerKey); + } + + @Override + public Instant currentProcessingTime() { + return delegate.currentProcessingTime(); + } + + @Override + public @Nullable Instant currentSynchronizedProcessingTime() { + return delegate.currentSynchronizedProcessingTime(); + } + + @Override + public @Nullable Instant currentOutputWatermarkTime() { + return delegate.currentOutputWatermarkTime(); + } + } + + /** Encodes every emitted element into the tagged output row layout and appends it. */ + private final class EncodingReceiver implements WindowedValueMultiReceiver { + private final List outputs; + + 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..e0a18a27840e --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/state/TwsStateInternals.java @@ -0,0 +1,577 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming.state; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.coders.MapCoder; +import org.apache.beam.sdk.coders.SetCoder; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.CombiningState; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.MultimapState; +import org.apache.beam.sdk.state.OrderedListState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.state.ReadableStates; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.state.State; +import org.apache.beam.sdk.state.StateBinder; +import org.apache.beam.sdk.state.StateContext; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.CombineWithContext; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.CombineFnUtil; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Beam {@link StateInternals} on top of Spark 4 {@code transformWithState}, a port of the legacy + * {@code org.apache.beam.runners.spark.stateful.SparkStateInternals} with its Guava {@code + * Table} replaced by the {@link BytesKV} SPI. + * + *

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Must be called exactly once, at the end of the {@code handleInputRows} or {@code + * handleExpiredTimer} invocation this instance belongs to. + */ + public void flush() { + if (flushed) { + throw new IllegalStateException("TwsTimerInternals.flush() called more than once"); + } + flushed = true; + + for (Map.Entry entry : current.entrySet()) { + TimerData before = loaded.get(entry.getKey()); + if (before == null || !before.equals(entry.getValue())) { + store.put(entry.getKey(), encode(entry.getValue())); + } + } + for (String key : loaded.keySet()) { + if (!current.containsKey(key)) { + store.remove(key); + } + } + + Set desired = new HashSet<>(); + for (TimerData timer : current.values()) { + desired.add(wakeupFor(timer.getTimestamp().getMillis())); + } + Set alreadyRegistered = registry.registered(); + for (Long expiry : desired) { + if (!alreadyRegistered.contains(expiry)) { + registry.register(expiry); + } + } + for (Long expiry : alreadyRegistered) { + // Spark removes the expiry it is currently firing on its own once the callback completes. + if (!desired.contains(expiry) && !expiry.equals(firedExpiryMs)) { + registry.delete(expiry); + } + } + } + + /** Maps a Beam timer timestamp onto the Spark wake-up millisecond to register for it. */ + private long wakeupFor(long timestampMs) { + if (firedExpiryMs != null && timestampMs <= firedExpiryMs) { + return firedExpiryMs + 1; + } + return timestampMs; + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/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..e4be10d06d04 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingEvidenceTest.java @@ -0,0 +1,392 @@ +/* + * 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.setMaxRecordsPerBatch(1L); + 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 + 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 = StreamingTestUtils.run(pipeline); + } 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 + 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 = StreamingTestUtils.run(pipeline); + } 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..bd619c201a5c --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/ChainedStatefulStreamingTest.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.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 + 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 = StreamingTestUtils.run(pipeline); + + // 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/PAssertStreamingTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java new file mode 100644 index 000000000000..d4df7cb424aa --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/PAssertStreamingTest.java @@ -0,0 +1,195 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertThrows; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End-to-end tests verifying {@link PAssert} evaluations against the Spark 4 Structured Streaming + * runner. + * + *

PAssert regroups elements and verifies expectations once a window closes upon watermark + * passage. In streaming, this relies on the end-of-stream sentinel row emitted when unbounded + * readers reach exhaustion to advance Spark's data-driven watermark past the global window end. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class PAssertStreamingTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + + private static class DoubleFn extends DoFn { + @ProcessElement + public void process(@Element Integer element, OutputReceiver out) { + out.output(element * 2); + } + } + + @Test + public void testPAssertInGlobalWindow() throws Exception { + List> elements = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("Double", ParDo.of(new DoubleFn())); + + PAssert.that(output).containsInAnyOrder(0, 2, 4, 6, 8, 10, 12, 14, 16, 18); + + StreamingTestUtils.run(pipeline); + } + + @Test + public void testPAssertInFixedWindows() throws Exception { + List> elements = new ArrayList<>(); + // Window 1: [0s, 5s) -> elements 0, 1, 2 (doubled: 0, 2, 4) + elements.add(TimestampedValue.of(0, BASE.plus(Duration.standardSeconds(1)))); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(2)))); + elements.add(TimestampedValue.of(2, BASE.plus(Duration.standardSeconds(4)))); + // Window 2: [5s, 10s) -> elements 3, 4 (doubled: 6, 8) + elements.add(TimestampedValue.of(3, BASE.plus(Duration.standardSeconds(6)))); + elements.add(TimestampedValue.of(4, BASE.plus(Duration.standardSeconds(8)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("FixedWindows", Window.into(FixedWindows.of(Duration.standardSeconds(5)))) + .apply("Double", ParDo.of(new DoubleFn())); + + PAssert.that(output).containsInAnyOrder(0, 2, 4, 6, 8); + + StreamingTestUtils.run(pipeline); + } + + @Test + public void testPAssertFailureThrows() throws Exception { + List> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(1)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("Double", ParDo.of(new DoubleFn())); + + // Deliberately incorrect expectation: output is [2], expected is [999]. + PAssert.that(output).containsInAnyOrder(999); + + Exception e = assertThrows(Exception.class, () -> StreamingTestUtils.run(pipeline)); + Throwable root = e; + while (root.getCause() != null && root.getCause() != root) { + root = root.getCause(); + } + org.junit.Assert.assertTrue( + "Expected AssertionError at root of cause chain, got: " + root, + root instanceof AssertionError); + } + + @Test + public void testPAssertFailureThrowsInFixedWindows() throws Exception { + List> elements = new ArrayList<>(); + elements.add(TimestampedValue.of(1, BASE.plus(Duration.standardSeconds(1)))); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + Pipeline pipeline = Pipeline.create(options); + + PCollection output = + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, VarIntCoder.of(), true))) + .apply("FixedWindows", Window.into(FixedWindows.of(Duration.standardSeconds(5)))) + .apply("Double", ParDo.of(new DoubleFn())); + + // Deliberately incorrect expectation: output is [2], expected is [999]. This proves the + // arrival side watermark clamp does not mask genuine failures in the rewindowed path. + PAssert.that(output).containsInAnyOrder(999); + + Exception e = assertThrows(Exception.class, () -> StreamingTestUtils.run(pipeline)); + Throwable root = e; + while (root.getCause() != null && root.getCause() != root) { + root = root.getCause(); + } + org.junit.Assert.assertTrue( + "Expected AssertionError at root of cause chain, got: " + root, + root instanceof AssertionError); + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/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..78a270b951d7 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatefulParDoStreamingTest.java @@ -0,0 +1,153 @@ +/* + * 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 + 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 = StreamingTestUtils.run(pipeline); + + // 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..d90ba1d4998b --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StatelessParDoStreamingTest.java @@ -0,0 +1,113 @@ +/* + * 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 + 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 = StreamingTestUtils.run(pipeline); + + // 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/StreamingCheckpointRestartTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java new file mode 100644 index 000000000000..7d8f79088141 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingCheckpointRestartTest.java @@ -0,0 +1,194 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming.translation.streaming; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingRunner; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamReaderCache; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.BeamSourceCheckpoint; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End to end proof that a Spark 4 structured streaming pipeline restarted against the same + * checkpoint location resumes an {@link org.apache.beam.sdk.io.UnboundedSource.UnboundedReader} + * from its durable checkpoint marks instead of re-reading the source from scratch. + * + *

Two independent {@link Pipeline}s are run one after the other against the same checkpoint + * directory. Spark hands the source the per source location {@code /sources/0}, + * where {@link BeamSourceCheckpoint} pins the split list and writes one mark file per split and + * epoch. {@link BeamReaderCache#invalidateAll} is called between the two runs to simulate a fresh + * JVM: every in-memory reader and pending mark is dropped, forcing the second run's readers to + * restore from the durable marks the first run wrote. + * + *

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

The second run's assertion is deliberately weak, per the at-least-once semantics documented on + * {@link BeamSourceCheckpoint} and {@link BeamReaderCache}: Spark may replay the last micro-batch + * that was in flight when the query stopped, so a handful of already-seen elements MAY reappear. + * What must never happen is the second run re-reading the whole range, which is what a + * regression to in-memory-only marks would look like. + */ +@RunWith(JUnit4.class) +@Category(StreamingTest.class) +public class StreamingCheckpointRestartTest implements Serializable { + + @ClassRule public static final SparkSessionRule SESSION = new SparkSessionRule(); + + @Rule public transient TemporaryFolder checkpointDir = new TemporaryFolder(); + + private static final Instant BASE = new Instant(0); + private static final int ELEMENT_COUNT = 10; + + @Test + public void restartedPipelineResumesFromDurableCheckpointMarks() throws Exception { + String checkpointPath = checkpointDir.newFolder("checkpoint").getAbsolutePath(); + + String collectorA = StreamingTestUtils.newCollectorId("checkpoint-restart-a"); + String collectorB = StreamingTestUtils.newCollectorId("checkpoint-restart-b"); + StreamingTestUtils.clear(collectorA); + StreamingTestUtils.clear(collectorB); + + // First run: read the whole 0..9 range once, from a fresh checkpoint location. + PipelineResult first = runPipeline(checkpointPath, collectorA); + + List collectedA = + new ArrayList<>(StreamingTestUtils.getCollected(collectorA)); + Collections.sort(collectedA); + assertEquals( + "first run (pipeline state=" + first.getState() + ") must read every element", + fullRange(), + collectedA); + + // The durable layout of the first run must exist under the location Spark hands the source, + // //sources/: the pinned split list and at least one + // persisted checkpoint mark. + File sourceRoot = new File(new File(checkpointPath, "0"), "sources/0"); + assertTrue("expected the source checkpoint directory " + sourceRoot, sourceRoot.isDirectory()); + File splitsFile = new File(sourceRoot, "splits"); + assertTrue("pinned splits file must exist: " + splitsFile, splitsFile.isFile()); + assertTrue( + "at least one split's marks directory must contain a persisted mark file", + hasAnyMarkFile(new File(sourceRoot, "marks"))); + + // Simulate a fresh JVM: drop every in-memory reader and checkpoint mark, forcing the next + // reader to fall back to the durable marks just asserted above. + BeamReaderCache.invalidateAll(); + + // Second run: identical transform name, identical checkpoint location, identical source + // content. A correct restart must resume from the durable marks rather than re-reading + // everything. + PipelineResult second = runPipeline(checkpointPath, collectorB); + + List collectedB = + new ArrayList<>(StreamingTestUtils.getCollected(collectorB)); + Set seenB = new HashSet<>(collectedB); + assertFalse( + "a restarted run (pipeline state=" + + second.getState() + + ") must not re-read the whole range from scratch, saw " + + collectedB, + seenB.containsAll(fullRange())); + } + + private PipelineResult runPipeline(String checkpointPath, String collectorId) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.as(SparkStructuredStreamingPipelineOptions.class); + options.setRunner(SparkStructuredStreamingRunner.class); + options.setTestMode(true); + options.setStreaming(true); + options.setStreamingStopAfterIdleBatches(3); + options.setMaxBatchDurationMillis(200); + // One element per split per micro-batch (a batch limit below the split count still gives every + // split one record): the 10 elements are spread across several micro-batches instead of a + // single one, so even a replayed in-flight batch can only ever carry a small tail of the range. + options.setMaxRecordsPerBatch(1L); + options.setCheckpointDir(checkpointPath); + SESSION.configure(options); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>(elements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + return StreamingTestUtils.run(pipeline); + } + + private static List> elements() { + List> elements = new ArrayList<>(ELEMENT_COUNT); + for (int i = 0; i < ELEMENT_COUNT; i++) { + elements.add(TimestampedValue.of(i, BASE.plus(Duration.standardSeconds(i)))); + } + return elements; + } + + private static List fullRange() { + List range = new ArrayList<>(ELEMENT_COUNT); + for (int i = 0; i < ELEMENT_COUNT; i++) { + range.add(i); + } + return range; + } + + /** {@code true} if any split subdirectory of {@code marksRoot} contains a mark file. */ + private static boolean hasAnyMarkFile(File marksRoot) { + File[] splitDirs = marksRoot.listFiles(); + if (splitDirs == null) { + return false; + } + for (File splitDir : splitDirs) { + File[] markFiles = splitDir.listFiles(); + if (markFiles != null && markFiles.length > 0) { + return true; + } + } + return false; + } +} diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/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..c34fc0f28c53 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java @@ -0,0 +1,321 @@ +/* + * 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.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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; + +/** + * {@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, + * {@code CANCELLED} after an explicit {@code cancel()}, and {@code FAILED} once any leaf query + * fails. + * + *

The tests here set {@code testMode(false)} on top of {@link + * StreamingTestUtils#streamingOptions}, unlike every other test in this package: {@code + * SparkStructuredStreamingRunner#run()} calls {@code result.waitUntilFinish()} itself before + * returning when {@code testMode} is {@code true} (see its implementation), which would make {@code + * 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 + 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 = StreamingTestUtils.waitUntilFinish(result); + 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 + 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 forever. + 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); + } + } + + /** + * A failure in any leaf query must surface through {@code waitUntilFinish()} even while another + * leaf query is still running. The healthy leaf has idle-stop disabled, so it never terminates on + * its own; only the round robin await in {@code StreamingEvaluationContext} surfaces the poisoned + * leaf's failure promptly and stops the healthy sibling. With a sequential await this test would + * sit on the healthy query until the JUnit timeout whenever that query happens to be awaited + * first. + */ + @Test + public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-failed"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + options.setTestMode(false); + // Disabled so the healthy query only ever stops because the failure of its sibling is + // surfaced and stop() is called, not because it happened to go idle first. + options.setStreamingStopAfterIdleBatches(-1); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply( + "ReadHealthy", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + pipeline + .apply( + "ReadPoisoned", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Throw", ParDo.of(new ThrowOnElementDoFn(5))); + + PipelineResult result = pipeline.run(); + + assertThrows(RuntimeException.class, () -> StreamingTestUtils.waitUntilFinish(result)); + assertEquals(PipelineResult.State.FAILED, result.getState()); + + // The failed pipeline has to have stopped the healthy sibling query too: with idle-stop + // disabled it would otherwise keep running forever. + long deadline = System.currentTimeMillis() + QUERY_START_TIMEOUT_MILLIS; + while (SESSION.getSession().streams().active().length > 0) { + assertTrue( + "a sibling query was still active " + + QUERY_START_TIMEOUT_MILLIS + + "ms after the pipeline failed", + System.currentTimeMillis() < deadline); + Thread.sleep(50L); + } + } + + /** Throws on one specific element, passes every other element through unchanged. */ + private static final class ThrowOnElementDoFn extends DoFn { + private final int poisonElement; + + ThrowOnElementDoFn(int poisonElement) { + this.poisonElement = poisonElement; + } + + @ProcessElement + public void processElement(@Element Integer element, OutputReceiver out) { + if (element == poisonElement) { + throw new IllegalStateException("poison element " + poisonElement); + } + out.output(element); + } + } + + /** + * When no {@code checkpointDir} is configured, {@code StreamingEvaluationContext} falls back to a + * {@code beam-spark4-streaming-checkpoint*} temporary directory under {@code java.io.tmpdir}. + * That fallback directory must not survive the pipeline: it is cleaned up in {@code evaluate()}'s + * {@code finally} block once every query has reached a terminal state. + */ + @Test + public void tempCheckpointDirIsCleanedUpWhenNoneConfigured() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("lifecycle-tempdir"); + StreamingTestUtils.clear(collectorId); + + SparkStructuredStreamingPipelineOptions options = + StreamingTestUtils.streamingOptions(checkpointDir); + SESSION.configure(options); + options.setTestMode(false); + // Unset the checkpointDir that streamingOptions() configured, to exercise the fallback path. + options.setCheckpointDir(null); + Pipeline pipeline = Pipeline.create(options); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + tenElements(), VarIntCoder.of()))) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + Set tempDirsBefore = tempCheckpointDirNames(); + + PipelineResult result = pipeline.run(); + PipelineResult.State finalState = StreamingTestUtils.waitUntilFinish(result); + assertEquals(PipelineResult.State.DONE, finalState); + + Set survivors = tempCheckpointDirNames(); + survivors.removeAll(tempDirsBefore); + assertTrue("leftover temporary checkpoint directories: " + survivors, survivors.isEmpty()); + } + + /** + * Names of {@code beam-spark4-streaming-checkpoint*} entries currently under {@code + * java.io.tmpdir}. + */ + private static Set tempCheckpointDirNames() { + File tmpDir = new File(System.getProperty("java.io.tmpdir")); + String[] names = + tmpDir.list((dir, name) -> name.startsWith("beam-spark4-streaming-checkpoint")); + return names == null ? new HashSet<>() : new HashSet<>(Arrays.asList(names)); + } +} 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..6d03cc9a43a7 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingTestUtils.java @@ -0,0 +1,478 @@ +/* + * 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.io.Serializable; +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.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.SerializableCoder; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +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.Duration; +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. Run it through {@link #run(Pipeline)}, which waits at most {@link #FINISH_TIMEOUT} for the + * pipeline to finish, and only then assert against the {@link #getCollected} static + * collector, never against the {@code PCollection} itself. No test in this package uses a + * JUnit method timeout, its timeout thread group leaks into Spark's static thread pools and + * breaks later tests in the same JVM. + *
+ * + *

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(); + } + + // --------------------------------------------------------------------------------------------- + // Running a pipeline with a deadline. + // --------------------------------------------------------------------------------------------- + + /** Upper bound on the wall clock time one streaming pipeline in this package may take. */ + public static final Duration FINISH_TIMEOUT = Duration.standardMinutes(5); + + /** + * Runs {@code pipeline} without the runner's own unbounded wait ({@code testMode} is switched + * off) and returns once it reached a terminal state. Cancels the pipeline and fails the test + * after {@link #FINISH_TIMEOUT}. + */ + public static PipelineResult run(Pipeline pipeline) { + pipeline.getOptions().as(SparkStructuredStreamingPipelineOptions.class).setTestMode(false); + PipelineResult result = pipeline.run(); + waitUntilFinish(result); + return result; + } + + /** + * Waits at most {@link #FINISH_TIMEOUT} for {@code result} to reach a terminal state. A failed + * pipeline rethrows its failure, a pipeline still running at the deadline is cancelled and the + * test fails. + */ + public static PipelineResult.State waitUntilFinish(PipelineResult result) { + PipelineResult.State state = result.waitUntilFinish(FINISH_TIMEOUT); + if (state == null || !state.isTerminal()) { + try { + result.cancel(); + } catch (IOException | RuntimeException e) { + // Best effort, the AssertionError below is what matters. + } + throw new AssertionError( + "pipeline did not finish within " + FINISH_TIMEOUT + ", last state " + state); + } + return state; + } + + // --------------------------------------------------------------------------------------------- + // 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. + * + *

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

Elements are stored pre-encoded (via {@code coder}) as {@code byte[]} plus a {@code long} + * timestamp rather than kept as {@link TimestampedValue} objects, because this source (like any + * {@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; + private final boolean advanceWatermarkToInfinityOnExhaustion; + + public ListBackedUnboundedSource(List> elements, Coder coder) { + this(elements, coder, false); + } + + public ListBackedUnboundedSource( + List> elements, + Coder coder, + boolean advanceWatermarkToInfinityOnExhaustion) { + this.coder = Preconditions.checkNotNull(coder); + this.advanceWatermarkToInfinityOnExhaustion = advanceWatermarkToInfinityOnExhaustion; + List encoded = new ArrayList<>(elements.size()); + List timestamps = new ArrayList<>(elements.size()); + for (TimestampedValue element : elements) { + 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, + boolean advanceWatermarkToInfinityOnExhaustion) { + this.encodedElements = encodedElements; + this.timestampsMillis = timestampsMillis; + this.coder = coder; + this.advanceWatermarkToInfinityOnExhaustion = advanceWatermarkToInfinityOnExhaustion; + } + + @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, + advanceWatermarkToInfinityOnExhaustion)); + } + return splits; + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark checkpointMark) { + return new ListBackedUnboundedReader<>( + this, checkpointMark == null ? 0 : checkpointMark.next); + } + + @Override + public Coder getCheckpointMarkCoder() { + return SerializableCoder.of(Mark.class); + } + + @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; + private Instant maxTimestampSeen = BoundedWindow.TIMESTAMP_MIN_VALUE; + + /** + * @param nextIndex the index of the first element this reader has not yet emitted, {@code 0} + * for a fresh start or the value carried by a resumed {@link Mark} + */ + ListBackedUnboundedReader(ListBackedUnboundedSource source, int nextIndex) { + this.source = source; + this.index = nextIndex - 1; + } + + @Override + 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() { + if (source.advanceWatermarkToInfinityOnExhaustion + && index + 1 >= source.encodedElements.size()) { + return BoundedWindow.TIMESTAMP_MAX_VALUE; + } + // No idle advance, deliberately: once the list is exhausted this simply stops moving, + // rather than jumping to +infinity. See the StreamingTestUtils class javadoc. + return maxTimestampSeen; + } + + @Override + public Mark getCheckpointMark() { + // index + 1 is the index of the next element to read, i.e. the position a resumed reader + // must start at. + return new Mark(index + 1); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() throws IOException {} + } + + /** + * The read position of one {@link ListBackedUnboundedSource} split: the index of the next + * element a resumed reader must emit. {@code 0} means nothing was read yet. + */ + public static final class Mark implements UnboundedSource.CheckpointMark, Serializable { + private static final long serialVersionUID = 1L; + + private final int next; + + Mark(int next) { + this.next = next; + } + + @Override + public void finalizeCheckpoint() {} + + @Override + public String toString() { + return "Mark{next=" + next + '}'; + } + } + } +} 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..0c7dea8e47fc --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/WindowedGroupByKeyStreamingTest.java @@ -0,0 +1,718 @@ +/* + * 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.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.apache.beam.runners.spark.StreamingTest; +import org.apache.beam.runners.spark.structuredstreaming.SparkSessionRule; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.AfterPane; +import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime; +import org.apache.beam.sdk.transforms.windowing.AfterWatermark; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Repeatedly; +import org.apache.beam.sdk.transforms.windowing.SlidingWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TimestampedValue; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.ClassRule; +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 + 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 = StreamingTestUtils.run(pipeline); + + // 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 + 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 = StreamingTestUtils.run(pipeline); + + // 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 + 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, a batch limit below the split count still gives + // every split one record. 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.setMaxRecordsPerBatch(1L); + + // 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 defaultParallelism, 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 = StreamingTestUtils.run(pipeline); + + // 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)); + } + + /** + * Encapsulates the value and complete {@link PaneInfo} metadata of an emitted element for precise + * assertions in streaming tests. + */ + public static final class PaneRecord implements Serializable { + private final String key; + private final long value; + private final PaneInfo.Timing timing; + private final long index; + private final long onTimeIndex; + private final boolean isFirst; + private final boolean isLast; + + public PaneRecord( + String key, + long value, + PaneInfo.Timing timing, + long index, + long onTimeIndex, + boolean isFirst, + boolean isLast) { + this.key = key; + this.value = value; + this.timing = timing; + this.index = index; + this.onTimeIndex = onTimeIndex; + this.isFirst = isFirst; + this.isLast = isLast; + } + + public String getKey() { + return key; + } + + public long getValue() { + return value; + } + + public PaneInfo.Timing getTiming() { + return timing; + } + + public long getIndex() { + return index; + } + + public long getOnTimeIndex() { + return onTimeIndex; + } + + public boolean isFirst() { + return isFirst; + } + + public boolean isLast() { + return isLast; + } + + @Override + public String toString() { + return key + + "=" + + value + + ":" + + timing + + ":index=" + + index + + ":onTimeIndex=" + + onTimeIndex + + ":isFirst=" + + isFirst + + ":isLast=" + + isLast; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof PaneRecord)) { + return false; + } + PaneRecord that = (PaneRecord) o; + return value == that.value + && index == that.index + && onTimeIndex == that.onTimeIndex + && isFirst == that.isFirst + && isLast == that.isLast + && Objects.equals(key, that.key) + && timing == that.timing; + } + + @Override + public int hashCode() { + return Objects.hash(key, value, timing, index, onTimeIndex, isFirst, isLast); + } + } + + /** Converts {@code KV} into {@link PaneRecord} carrying full {@link PaneInfo}. */ + public static final class CollectPaneDoFn extends DoFn, PaneRecord> { + @ProcessElement + public void processElement( + @Element KV element, PaneInfo paneInfo, OutputReceiver out) { + PaneRecord record = + new PaneRecord( + element.getKey(), + element.getValue(), + paneInfo.getTiming(), + paneInfo.getIndex(), + paneInfo.getNonSpeculativeIndex(), + paneInfo.isFirst(), + paneInfo.isLast()); + out.output(record); + } + } + + private static List filterPanesForKey(String collectorId, String targetKey) { + List result = new ArrayList<>(); + for (PaneRecord record : StreamingTestUtils.getCollected(collectorId)) { + if (targetKey.equals(record.getKey())) { + result.add(record); + } + } + return result; + } + + private static List>> lateFiringsInputElements() { + List>> elements = new ArrayList<>(); + // Split 0: [a@1s, dummy@22s, a@2s, a@3s, dummy@36s, a@4s, dummy@62s] + // Split 1: [sentinel@20s, sentinel@25s, sentinel@30s, sentinel@35s, sentinel@45s, sentinel@60s, + // sentinel@70s] + + // Batch 1: a@1s (on-time in [0s, 10s)), sentinel@20s -> Watermark advances to 20s + elements.add(TimestampedValue.of(KV.of("a", "1"), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s1"), BASE.plus(Duration.standardSeconds(20)))); + + // Batch 2: dummy@22s, sentinel@25s -> Batch start watermark is 20s. On-time timer for 'a' fires + // here! + elements.add( + TimestampedValue.of(KV.of("dummy", "d1"), BASE.plus(Duration.standardSeconds(22)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s2"), BASE.plus(Duration.standardSeconds(25)))); + + // Batch 3: a@2s (late! watermark 25s < GC time 40s), sentinel@30s -> Late pane 1 fires + elements.add(TimestampedValue.of(KV.of("a", "2"), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s3"), BASE.plus(Duration.standardSeconds(30)))); + + // Batch 4: a@3s (late! watermark 30s < GC time 40s), sentinel@35s -> Late pane 2 fires + elements.add(TimestampedValue.of(KV.of("a", "3"), BASE.plus(Duration.standardSeconds(3)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s4"), BASE.plus(Duration.standardSeconds(35)))); + + // Batch 5: dummy@36s, sentinel@45s -> Watermark advances to 45s (past GC horizon 40s, window + // expires) + elements.add( + TimestampedValue.of(KV.of("dummy", "d2"), BASE.plus(Duration.standardSeconds(36)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s5"), BASE.plus(Duration.standardSeconds(45)))); + + // Batch 6: a@4s (late! arriving when start watermark is 45s > GC horizon 40s -> dropped), + // sentinel@60s + elements.add(TimestampedValue.of(KV.of("a", "4"), BASE.plus(Duration.standardSeconds(4)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s6"), BASE.plus(Duration.standardSeconds(60)))); + + // Batch 7: dummy@62s, sentinel@70s -> Trailing batch + elements.add( + TimestampedValue.of(KV.of("dummy", "d3"), BASE.plus(Duration.standardSeconds(62)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s7"), BASE.plus(Duration.standardSeconds(70)))); + + return elements; + } + + @Test + public void fixedWindowsWithLateFiringsDiscarding() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("late-firings-discarding"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerBatch(1L); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(1))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + StreamingTestUtils.run(pipeline); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two late delta panes in discarding mode, with expired element dropped", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=1:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=1:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + @Test + public void fixedWindowsWithLateFiringsAccumulating() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("late-firings-accumulating"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerBatch(1L); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(1))) + .withAllowedLateness(Duration.standardSeconds(30)) + .accumulatingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + StreamingTestUtils.run(pipeline); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two accumulating panes carrying full window counts, with expired element dropped", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=2:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=3:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + @Test + public void fixedWindowsDefaultTriggerAccumulating() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("default-trigger-accumulating"); + StreamingTestUtils.clear(collectorId); + + List>> elements = lateFiringsInputElements(); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerBatch(1L); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .withAllowedLateness(Duration.standardSeconds(30)) + .accumulatingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + StreamingTestUtils.run(pipeline); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected on-time pane and two accumulating panes under default trigger", + List.of( + "a=1:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false", + "a=2:LATE:index=1:onTimeIndex=1:isFirst=false:isLast=false", + "a=3:LATE:index=2:onTimeIndex=2:isFirst=false:isLast=false"), + renderedPanes); + } + + /** + * Under Spark Structured Streaming's micro-batch execution model, {@code transformWithState} + * processes input rows against the batch start watermark (the watermark established by the + * previous micro-batch). When an element timestamped in {@code [0s, 10s)} arrives in the same + * micro-batch that carries a watermark-advancing sentinel, the batch start watermark has not yet + * passed the window end. In Beam semantics, the window is still open at the moment the element is + * processed, so it is buffered into the open window accumulator. When the on-time timer + * subsequently fires at the end of the window, it emits a single {@code ON_TIME} pane containing + * both elements. + */ + @Test + public void fixedWindowsLateElementInSameBatchFoldedIntoOnTimePane() throws Exception { + String collectorId = StreamingTestUtils.newCollectorId("same-batch-merge"); + StreamingTestUtils.clear(collectorId); + + List>> elements = new ArrayList<>(); + // Split 0: [a@1s, a@2s, dummy@52s] + // Split 1: [sentinel@5s, sentinel@20s, sentinel@60s] + + // Batch 1: a@1s (in [0s, 10s)), sentinel@5s -> Watermark advances to 5s + elements.add(TimestampedValue.of(KV.of("a", "1"), BASE.plus(Duration.standardSeconds(1)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s1"), BASE.plus(Duration.standardSeconds(5)))); + + // Batch 2: a@2s (in [0s, 10s)), sentinel@20s -> Batch start watermark is 5s < 10s! + // Window [0s, 10s) is still open during input row processing in batch 2, so a@2s is folded in. + // Watermark advances to 20s at the end of batch 2. + elements.add(TimestampedValue.of(KV.of("a", "2"), BASE.plus(Duration.standardSeconds(2)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s2"), BASE.plus(Duration.standardSeconds(20)))); + + // Batch 3: dummy@52s, sentinel@60s -> Batch start watermark is 20s >= 10s. + // The on-time timer for [0s, 10s) fires here, emitting a single ON_TIME pane with count = 2. + elements.add(TimestampedValue.of(KV.of("dummy", "d"), BASE.plus(Duration.standardSeconds(52)))); + elements.add( + TimestampedValue.of(KV.of("sentinel", "s3"), BASE.plus(Duration.standardSeconds(60)))); + + SparkStructuredStreamingPipelineOptions options = options(); + options.setMaxRecordsPerBatch(1L); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()) + .apply("CollectPane", ParDo.of(new CollectPaneDoFn())) + .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId))); + + StreamingTestUtils.run(pipeline); + + List aPanes = filterPanesForKey(collectorId, "a"); + List renderedPanes = new ArrayList<>(); + for (PaneRecord r : aPanes) { + renderedPanes.add(r.toString()); + } + + assertEquals( + "expected single merged ON_TIME pane containing both elements", + List.of("a=2:ON_TIME:index=0:onTimeIndex=0:isFirst=true:isLast=false"), + renderedPanes); + } + + @Test + public void unsupportedEarlyFiringsTriggerThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withEarlyFirings(AfterPane.elementCountAtLeast(1)) + .withLateFirings(AfterPane.elementCountAtLeast(1))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("withEarlyFirings")); + } + + @Test + public void unsupportedRepeatedlyProcessingTimeTriggerThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering(Repeatedly.forever(AfterProcessingTime.pastFirstElementInPane())) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("AfterProcessingTime")); + } + + @Test + public void unsupportedAfterPaneCountGreaterThanOneThrows() throws Exception { + SparkStructuredStreamingPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + + List>> elements = + List.of(TimestampedValue.of(KV.of("a", "1"), BASE)); + + pipeline + .apply( + "ReadUnbounded", + Read.from( + new StreamingTestUtils.ListBackedUnboundedSource<>( + elements, KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))) + .apply( + "FixedWindows", + Window.>into(FixedWindows.of(WINDOW_SIZE)) + .triggering( + AfterWatermark.pastEndOfWindow() + .withLateFirings(AfterPane.elementCountAtLeast(2))) + .withAllowedLateness(Duration.standardSeconds(30)) + .discardingFiredPanes()) + .apply("CountPerKey", Count.perKey()); + + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, pipeline::run); + assertTrue( + "Expected unsupported trigger message, got: " + e.getMessage(), + e.getMessage().contains("the custom trigger") + && e.getMessage().contains("AfterPane.elementCountAtLeast(2)")); + } +} 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..4a41f3259162 --- /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 + 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 + 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 + 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 + 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 + 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..edda652e4e17 --- /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 + public void testKeyIsExposed() { + assertEquals("key", internals(store()).getKey()); + } + + @Test + public void testStoreKeyLayout() { + assertEquals(NS_A.stringKey() + "+" + "tag", TwsStateInternals.storeKey(NS_A, "tag")); + } + + @Test + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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..bcd89c14595e --- /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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + public void testDeleteTimerWithoutTimeDomainIsRejected() { + TwsTimerInternals timers = internals(new InMemoryBytesKV(), new RecordingRegistry(), null); + assertThrows(UnsupportedOperationException.class, () -> timers.deleteTimer(NS, "t", "")); + } + + @Test + 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 + public void testFlushIsSingleUse() { + TwsTimerInternals timers = internals(new InMemoryBytesKV(), new RecordingRegistry(), null); + timers.flush(); + assertThrows(IllegalStateException.class, timers::flush); + } +}