From d291ec8005a999634f56e334141d9289a764dfdc Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Wed, 2 Sep 2026 07:53:28 +0000 Subject: [PATCH 1/3] [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/3] [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/3] Trigger CI after the setup-gradle allowlist fix