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..1eb31e766e78 --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamReaderCache.java @@ -0,0 +1,352 @@ +/* + * 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.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.LongSupplier; +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.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 and split. + * + *

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

A sweeper thread closes readers idle for longer than their idle timeout, finalizing marks + * whose epoch Spark committed, see {@link BeamSourceCheckpoint#readSparkCommittedEpoch()}, and + * dropping the others, the source redelivers. The timeout must exceed the longest gap between two + * micro-batches of one split. Under {@code spark.sql.streaming.asyncProgressTrackingEnabled} the + * commit log lags, idle readers then drop their marks. 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); + + private static final ConcurrentMap> READERS = new ConcurrentHashMap<>(); + + /** One monitor per key, acquire serializes per split, not across splits. */ + private static final ConcurrentMap LOCKS = new ConcurrentHashMap<>(); + + private static final long SWEEP_INTERVAL_MILLIS = 10_000L; + + private static final AtomicBoolean SWEEPER_STARTED = new AtomicBoolean(); + + private BeamReaderCache() {} + + public static String key(String checkpointLocation, int splitId) { + return checkpointLocation + '|' + splitId; + } + + /** 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 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. {@code committedEpoch} supplies the epoch Spark last committed, -1 if unknown. + * + * @throws IllegalStateException if {@code startEpoch > 0} and no durable mark exists + */ + public static CachedReader acquire( + String key, + long startEpoch, + UnboundedSource source, + PipelineOptions options, + long idleTimeoutMillis, + LongSupplier committedEpoch, + MarkRestorer restorer) + throws IOException { + startSweeper(); + 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, + committedEpoch); + created.beginBatch(startEpoch); + READERS.put(key, created); + return created; + } + } + + 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. */ + public static void invalidateAll() { + for (String key : READERS.keySet()) { + invalidate(key); + } + } + + @SuppressWarnings("FutureReturnValueIgnored") // the sweep runs until the JVM exits + private static void startSweeper() { + if (!SWEEPER_STARTED.compareAndSet(false, true)) { + return; + } + ScheduledExecutorService sweeper = + Executors.newSingleThreadScheduledExecutor( + runnable -> { + Thread thread = new Thread(runnable, "beam-reader-idle-sweep"); + thread.setDaemon(true); + return thread; + }); + sweeper.scheduleWithFixedDelay( + BeamReaderCache::sweep, + SWEEP_INTERVAL_MILLIS, + SWEEP_INTERVAL_MILLIS, + TimeUnit.MILLISECONDS); + } + + private static void sweep() { + try { + closeIdle(System.currentTimeMillis()); + } catch (RuntimeException e) { + LOG.warn("Idle sweep of Beam readers failed.", e); + } + } + + private static void closeIdle() { + closeIdle(System.currentTimeMillis()); + } + + /** Closes every reader idle at {@code nowMillis}, finalizing marks of committed epochs. */ + static void closeIdle(long nowMillis) { + for (Map.Entry> entry : READERS.entrySet()) { + String key = entry.getKey(); + CachedReader reader = entry.getValue(); + synchronized (lock(key)) { + if (!reader.isIdleSince(nowMillis) || !READERS.remove(key, reader)) { + continue; + } + LOG.info("Closing idle Beam reader {}.", key); + reader.finalizeIfCommitted(key); + close(key, reader); + } + } + } + + 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 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 final LongSupplier committedEpoch; + 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, + long positionEpoch, + byte @Nullable [] positionMark, + long idleTimeoutMillis, + LongSupplier committedEpoch) { + this.reader = reader; + this.positionEpoch = positionEpoch; + this.positionMark = positionMark; + this.idleTimeoutMillis = idleTimeoutMillis; + this.committedEpoch = committedEpoch; + this.lastUsedMillis = System.currentTimeMillis(); + } + + public UnboundedReader reader() { + return reader; + } + + public synchronized boolean startOrAdvance() throws IOException { + moved = true; + if (!started) { + started = true; + return reader.start(); + } + 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); + } + } + + /** Finalizes the pending mark if Spark committed its epoch, drops it otherwise. */ + synchronized void finalizeIfCommitted(String key) { + if (pendingMark == null) { + return; + } + long committed = committedEpoch.getAsLong(); + if (positionEpoch <= committed) { + finalizePendingMark(key); + return; + } + LOG.info( + "Dropping mark of Beam reader {} at epoch {}, Spark committed up to {}.", + key, + positionEpoch, + committed); + pendingMark = null; + } + + @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..d667976edcec --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamSourceCheckpoint.java @@ -0,0 +1,227 @@ +/* + * 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.FileNotFoundException; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +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. The epoch Spark last committed is read from Spark's own {@code + * commits} and {@code offsets} logs two levels up. 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"; + private static final String SPARK_COMMITS_DIR = "commits"; + private static final String SPARK_OFFSETS_DIR = "offsets"; + private static final String SERIALIZED_VOID_OFFSET = "-"; + + /** 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; + + /** Every mark epoch strictly below the value is known to be deleted, -1 for unknown. */ + private volatile long purgeFloor = -1L; + + 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); + } + + /** Creates the mark directory of {@code epoch}, the driver calls this once per batch. */ + public void prepareEpoch(long epoch) throws IOException { + fm.mkdirs(epochDir(epoch)); + } + + /** + * Writes the mark, creating the epoch directory if a manager without parent creation needs it. + */ + public void writeMark(int splitId, long epoch, byte[] codedMark) throws IOException { + Path path = markPath(splitId, epoch); + try { + write(path, codedMark, true); + } catch (FileNotFoundException e) { + fm.mkdirs(epochDir(epoch)); + write(path, 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); + } + + /** + * The end epoch of this source in the last batch Spark committed, or -1 if there is none or the + * logs cannot be read. The location is {@code /sources/}, the batch id is the + * highest entry of {@code /commits} and its epoch is line {@code index} after the version + * and metadata lines of {@code /offsets/}. + */ + public long readSparkCommittedEpoch() { + try { + Path sparkRoot = root.getParent().getParent(); + int sourceIndex = Integer.parseInt(root.getName()); + Path commits = new Path(sparkRoot, SPARK_COMMITS_DIR); + if (!fm.exists(commits)) { + return -1L; + } + long batchId = -1L; + for (FileStatus status : fm.list(commits)) { + batchId = Math.max(batchId, parseEpoch(status.getPath().getName())); + } + if (batchId < 0) { + return -1L; + } + Path offsets = new Path(new Path(sparkRoot, SPARK_OFFSETS_DIR), Long.toString(batchId)); + List lines = + Arrays.asList(new String(read(offsets), StandardCharsets.UTF_8).split("\n", -1)); + String line = lines.get(2 + sourceIndex).trim(); + return line.equals(SERIALIZED_VOID_OFFSET) ? -1L : Long.parseLong(line); + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to read the epoch Spark committed for {}.", location, e); + return -1L; + } + } + + /** + * Deletes the marks of every epoch strictly below {@code epoch}, one recursive delete per epoch + * directory. Lists the marks directory once, later calls delete the range above the previous + * floor only. Idempotent. + */ + public void purgeMarksBelow(long epoch) throws IOException { + long floor = purgeFloor; + if (floor >= 0 && epoch - floor > MAX_BLIND_PURGE_RANGE) { + floor = -1L; + } + if (floor < 0) { + if (fm.exists(marksRoot)) { + for (FileStatus status : fm.list(marksRoot)) { + long existing = parseEpoch(status.getPath().getName()); + if (existing >= 0 && existing < epoch) { + fm.delete(status.getPath()); + } + } + } + purgeFloor = epoch; + return; + } + for (long e = floor; e < epoch; e++) { + fm.delete(epochDir(e)); + } + if (epoch > floor) { + purgeFloor = epoch; + } + } + + private Path epochDir(long epoch) { + return new Path(marksRoot, Long.toString(epoch)); + } + + private Path markPath(int splitId, long epoch) { + return new Path(epochDir(epoch), Integer.toString(splitId)); + } + + 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 directory 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/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..0368f95ba6db --- /dev/null +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/UnboundedSourceDataset.java @@ -0,0 +1,721 @@ +/* + * 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.UncheckedIOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Set; +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.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions; +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.io.UnboundedSource; +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.collect.ImmutableSet; +import org.apache.hadoop.conf.Configuration; +import org.apache.spark.TaskContext; +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.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +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.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.InputPartition; +import org.apache.spark.sql.connector.read.PartitionReader; +import org.apache.spark.sql.connector.read.PartitionReaderFactory; +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.connector.read.streaming.Offset; +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 org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import scala.Option; +import scala.reflect.ClassTag; + +/** + * Translator facing entry point turning a Beam {@link UnboundedSource} into a streaming Spark + * {@link Dataset} of rows, with the DataSourceV2 micro-batch glue as nested classes. + * + *

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 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 { + + private static final Logger LOG = LoggerFactory.getLogger(UnboundedSourceDataset.class); + + public static final String COL_PAYLOAD = "payload"; + + public static final String COL_EVENT_TS = "eventTimestamp"; + + public static final StructType SCHEMA = + new StructType() + .add(COL_PAYLOAD, DataTypes.BinaryType, false) + .add(COL_EVENT_TS, DataTypes.TimestampType, false); + + private static final String SOURCE_NAME = "beam-unbounded"; + + private UnboundedSourceDataset() {} + + /** + * 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 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, used for naming only + * @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) { + org.apache.spark.sql.classic.SparkSession classic = + (org.apache.spark.sql.classic.SparkSession) session; + Configuration hadoopConf = classic.sessionState().newHadoopConf(); + BeamTable table = + new BeamTable<>( + 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, + table, + 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"); + } + + private static Broadcast broadcast(SparkSession session, T value, Class type) { + return session.sparkContext().broadcast(value, ClassTag.apply(type)); + } + + /** + * Opaque epoch counter used as the Spark {@link Offset} of a Beam unbounded source. + * + *

The read position lives in Beam checkpoint marks on the executors, see {@link + * BeamSourceCheckpoint}. Equality is the base class comparison of {@link #json()}. + */ + static class BeamOffset extends Offset { + + public static final BeamOffset ZERO = new BeamOffset(0L); + + private final long epoch; + + public BeamOffset(long epoch) { + this.epoch = epoch; + } + + public long epoch() { + return epoch; + } + + @Override + public String json() { + return Long.toString(epoch); + } + + public static BeamOffset fromJson(String json) { + try { + return new BeamOffset(Long.parseLong(json.trim())); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Not a valid BeamOffset: " + json, e); + } + } + + @Override + public String toString() { + return json(); + } + } + + /** DataSourceV2 {@link Table} over one Beam unbounded source, micro-batch reads only. */ + static final class BeamTable implements Table, SupportsRead { + final UnboundedSource source; + final Coder> coder; + final Broadcast options; + final Broadcast hadoopConf; + final int desiredNumSplits; + + /** Records per micro-batch across all splits, below 1 means unlimited. */ + final long maxRecordsPerBatch; + + final long maxBatchDurationMillis; + final long readerIdleTimeoutMillis; + final String transformName; + + BeamTable( + 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; + } + + @Override + public String name() { + return "BeamUnboundedSource[" + transformName + "]"; + } + + @Override + public StructType schema() { + return SCHEMA; + } + + @Override + public Set capabilities() { + return ImmutableSet.of(TableCapability.MICRO_BATCH_READ); + } + + @Override + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap ignored) { + return () -> new BeamScan<>(this); + } + } + + private static final class BeamScan implements Scan { + private final BeamTable table; + + BeamScan(BeamTable table) { + this.table = table; + } + + @Override + public StructType readSchema() { + return SCHEMA; + } + + @Override + public String description() { + return table.name(); + } + + @Override + public MicroBatchStream toMicroBatchStream(String checkpointLocation) { + return new BeamMicroBatchStream<>(table, checkpointLocation); + } + } + + /** + * Driver side {@link MicroBatchStream} over a Beam {@link UnboundedSource}. + * + *

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. Partitions carry no locality hint, the + * reader cache restores a split from its durable mark wherever it lands. + */ + static class BeamMicroBatchStream implements MicroBatchStream { + + private final BeamTable table; + private final String checkpointLocation; + 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> splits; + + BeamMicroBatchStream(BeamTable table, String checkpointLocation) { + this.table = table; + this.checkpointLocation = checkpointLocation; + this.checkpoint = + new BeamSourceCheckpoint(checkpointLocation, table.hadoopConf.value().value()); + } + + @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 InputPartition[] planInputPartitions(Offset start, Offset end) { + long startEpoch = ((BeamOffset) start).epoch(); + long endEpoch = ((BeamOffset) end).epoch(); + fastForwardEpoch(endEpoch); + List> pinned = splits(); + try { + checkpoint.prepareEpoch(endEpoch); + } catch (IOException e) { + LOG.warn( + "Failed to prepare mark directory of epoch {} at {}.", endEpoch, checkpointLocation, e); + } + long[] quotas = splitQuotas(table.maxRecordsPerBatch, pinned.size(), endEpoch); + InputPartition[] partitions = new InputPartition[pinned.size()]; + for (int i = 0; i < pinned.size(); i++) { + partitions[i] = + new BeamInputPartition<>( + pinned.get(i), + table.coder, + table.options, + table.hadoopConf, + checkpointLocation, + i, + startEpoch, + endEpoch, + quotas[i], + table.maxBatchDurationMillis, + table.readerIdleTimeoutMillis); + } + return partitions; + } + + @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) { + purgeRequested.accumulateAndGet(((BeamOffset) end).epoch(), Math::max); + if (purgeInFlight.compareAndSet(false, true)) { + purger.execute(this::purgeRequested); + } + } + + private void purgeRequested() { + long epoch; + do { + epoch = purgeRequested.get(); + try { + checkpoint.purgeMarksBelow(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 void stop() { + LOG.info( + "Stopping Beam micro-batch stream {} at {}.", table.transformName, checkpointLocation); + purger.shutdown(); + } + + /** 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 {}.", table.transformName, epoch, seen); + epoch = seen; + } + } + + private synchronized List> splits() { + if (splits != null) { + return splits; + } + List> pinned; + try { + pinned = checkpoint.readSplits(); + } catch (IOException 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); + } + 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); + } + splits = typed; + return typed; + } + + private List> splitSource() { + UnboundedSource source = table.source; + PipelineOptions options = table.options.value().get(); + List> result; + try { + result = source.split(table.desiredNumSplits, options); + } catch (Exception e) { + throw new IllegalStateException( + "Failed to split UnboundedSource " + source.getClass().getCanonicalName(), e); + } + if (result.isEmpty()) { + result = Collections.singletonList(source); + } + LOG.info( + "Split {} into {} partition(s), desired {}.", + table.transformName, + result.size(), + table.desiredNumSplits); + return result; + } + + /** + * Divides the batch limit over the splits. A limit below 1 means unlimited and yields -1 for + * every split. Otherwise the remainder rotates with the epoch, so a limit below the split count + * gives one record to a rotating subset of splits per batch and 0 to the others. + */ + static long[] splitQuotas(long maxRecordsPerBatch, int numSplits, long epoch) { + long[] quotas = new long[numSplits]; + if (maxRecordsPerBatch < 1) { + Arrays.fill(quotas, -1L); + return quotas; + } + long base = maxRecordsPerBatch / numSplits; + long remainder = maxRecordsPerBatch % numSplits; + for (int i = 0; i < numSplits; i++) { + quotas[i] = base + ((i + epoch) % numSplits < remainder ? 1 : 0); + } + return quotas; + } + } + + /** One split of a Beam unbounded source for one micro-batch, from epoch start to epoch end. */ + static final class BeamInputPartition implements InputPartition { + + private static final long serialVersionUID = 1L; + + final UnboundedSource split; + final Coder> coder; + final Broadcast options; + final Broadcast hadoopConf; + final String checkpointLocation; + final int splitId; + final long startEpoch; + final long endEpoch; + + /** Records this split may emit in this micro-batch, below 0 means unlimited, 0 means none. */ + final long maxRecords; + + final long maxBatchDurationMillis; + final long readerIdleTimeoutMillis; + + BeamInputPartition( + UnboundedSource split, + Coder> coder, + Broadcast options, + Broadcast hadoopConf, + String checkpointLocation, + int splitId, + long startEpoch, + long endEpoch, + long maxRecords, + long maxBatchDurationMillis, + long readerIdleTimeoutMillis) { + 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.maxRecords = maxRecords; + this.maxBatchDurationMillis = maxBatchDurationMillis; + this.readerIdleTimeoutMillis = readerIdleTimeoutMillis; + } + + @Override + public String toString() { + return "BeamInputPartition{checkpointLocation=" + + checkpointLocation + + ", split=" + + splitId + + ", epochs=" + + startEpoch + + ".." + + endEpoch + + "}"; + } + } + + /** Creates a {@link BeamPartitionReader} for a {@link BeamInputPartition} on the executor. */ + static final class BeamPartitionReaderFactory implements PartitionReaderFactory { + + private static final long serialVersionUID = 1L; + + @Override + public PartitionReader createReader(InputPartition partition) { + try { + return new BeamPartitionReader<>((BeamInputPartition) partition); + } catch (IOException e) { + throw new UncheckedIOException("Failed to open Beam reader for " + partition, e); + } + } + } + + /** + * Reads one split of a Beam {@link UnboundedSource} for one micro-batch. + * + *

The batch ends at the record quota or at the deadline. The reader then writes its checkpoint + * mark durably at the end epoch and stays in {@link BeamReaderCache} for the next batch. A failed + * 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. + * + * @param the element type of the split + */ + static final class BeamPartitionReader implements PartitionReader { + + private static final Duration INITIAL_BACKOFF = Duration.millis(10); + + private final String key; + private final UnboundedSource split; + private final Coder> coder; + private final BeamSourceCheckpoint checkpoint; + private final CachedReader cached; + private final int splitId; + private final long endEpoch; + 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) throws IOException { + this.split = partition.split; + this.coder = partition.coder; + this.splitId = partition.splitId; + this.endEpoch = partition.endEpoch; + this.maxRecords = partition.maxRecords; + this.maxBatchDurationMillis = partition.maxBatchDurationMillis; + PipelineOptions options = partition.options.value().get(); + Configuration conf = partition.hadoopConf.value().value(); + BeamSourceCheckpoint checkpoint = + new BeamSourceCheckpoint(partition.checkpointLocation, conf); + this.checkpoint = checkpoint; + this.key = BeamReaderCache.key(partition.checkpointLocation, splitId); + long startEpoch = partition.startEpoch; + int splitId = this.splitId; + this.cached = + BeamReaderCache.acquire( + key, + startEpoch, + split, + options, + partition.readerIdleTimeoutMillis, + checkpoint::readSparkCommittedEpoch, + () -> checkpoint.readMark(splitId, startEpoch)); + } + + @Override + public boolean next() throws IOException { + if (deadlineMillis < 0) { + deadlineMillis = System.currentTimeMillis() + maxBatchDurationMillis; + } + BackOff backOff = null; + while (true) { + if (maxRecords >= 0 && recordsRead >= maxRecords) { + return endOfBatch(false); + } + long remaining = deadlineMillis - System.currentTimeMillis(); + if (remaining <= 0) { + return endOfBatch(false); + } + if (cached.startOrAdvance()) { + recordsRead++; + current = toRow(); + return true; + } + 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() { + InternalRow row = current; + if (row == null) { + throw new IllegalStateException("No current row, next() did not return true."); + } + return row; + } + + @Override + public void close() throws IOException { + endBatch(attemptDiscarded()); + current = null; + } + + private boolean endOfBatch(boolean discarded) throws IOException { + endBatch(discarded); + current = null; + return false; + } + + /** + * 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 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 value = + WindowedValues.timestampedValueInGlobalWindow(cached.reader().getCurrent(), timestamp); + 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/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..4616c7563a20 --- /dev/null +++ b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamMicroBatchSourceTest.java @@ -0,0 +1,1028 @@ +/* + * 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.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.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamInputPartition; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamMicroBatchStream; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamOffset; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamPartitionReader; +import org.apache.beam.runners.spark.structuredstreaming.io.streaming.UnboundedSourceDataset.BeamTable; +import org.apache.beam.sdk.coders.Coder; +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.coders.VarLongCoder; +import org.apache.beam.sdk.io.CountingSource; +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.InternalRow; +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.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.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; + +/** + * Tests for the Spark 4 DataSourceV2 micro-batch source wrapping a Beam {@link UnboundedSource}. + * + *

The epoch offsets of this source never settle, so {@code processAllAvailable()} would block + * forever. Every query runs with {@code Trigger.ProcessingTime(100)} and is stopped 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 TAGS = new AtomicInteger(); + + /** Rows per micro-batch per query name, driver side. */ + private static final Map>> BATCHES = new ConcurrentHashMap<>(); + + private static final Coder> CODER = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + + /** 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; + + private static Broadcast optionsBroadcast; + private static Broadcast hadoopConfBroadcast; + + private String tag; + + @BeforeClass + public static void broadcastOnce() { + SparkSession session = SESSION.getSession(); + Configuration conf = + ((org.apache.spark.sql.classic.SparkSession) session).sessionState().newHadoopConf(); + optionsBroadcast = + session + .sparkContext() + .broadcast( + new SerializablePipelineOptions(PipelineOptionsFactory.create()), + ClassTag.apply(SerializablePipelineOptions.class)); + hadoopConfBroadcast = + session + .sparkContext() + .broadcast( + new SerializableConfiguration(conf), + ClassTag.apply(SerializableConfiguration.class)); + } + + @Before + public void setUp() { + tag = "src" + TAGS.incrementAndGet(); + } + + @After + public void tearDown() { + BeamReaderCache.invalidateAll(); + BATCHES.clear(); + TestSource.forget(tag); + } + + /** The {@code EventTimeWatermark} node survives typed maps in the logical and analyzed plan. */ + @Test + public void testEventTimeWatermarkSurvivesTypedMap() { + Dataset rows = rows(1, 4, limited(1_000L, 200L)); + assertTrue("source dataset must be streaming", rows.isStreaming()); + assertWatermark("directly after withWatermark, logical plan", logical(rows)); + assertWatermark("directly after withWatermark, analyzed plan", analyzed(rows)); + + 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)); + + 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)); + } + + /** A running query tracks the event time watermark past a typed map. */ + @Test + public void testWatermarkIsTrackedAtRuntimeAfterTypedMap() throws Exception { + Dataset typed = + rows(1, 8, limited(1_000L, 200L)) + .map((MapFunction) row -> row.getAs(COL_PAYLOAD), Encoders.BINARY()); + StreamingQuery query = + typed + .writeStream() + .format("noop") + .queryName(tag) + .outputMode("append") + .option("checkpointLocation", temp.newFolder(tag).getAbsolutePath()) + .trigger(Trigger.ProcessingTime(100)) + .start(); + try { + String watermark = awaitWatermark(query); + assertNotNull("query never reported an event time watermark", watermark); + assertFalse("watermark stuck at the epoch: " + watermark, watermark.startsWith("1970-")); + } finally { + stopQuietly(query); + } + } + + /** Payloads decode to the source elements and the timestamp column matches the element. */ + @Test + public void testReadsElementsFromUnboundedSource() throws Exception { + int count = 8; + StreamingQuery query = start(rows(1, count, limited(1_000L, 200L)), tag, temp.newFolder(tag)); + try { + await("all rows", () -> values(batches(tag)).size() >= count); + } finally { + stopQuietly(query); + } + List values = new ArrayList<>(); + for (List batch : batches(tag)) { + for (Row row : batch) { + WindowedValue value = CoderUtils.decodeFromByteArray(CODER, row.getAs(COL_PAYLOAD)); + values.add(value.getValue()); + assertEquals( + value.getTimestamp().getMillis(), row.getAs(COL_EVENT_TS).getTime()); + assertEquals( + Collections.singletonList(GlobalWindow.INSTANCE), new ArrayList<>(value.getWindows())); + assertEquals( + BASE_MILLIS + TestSource.indexOf(value.getValue()) * INTERVAL_MILLIS, + value.getTimestamp().getMillis()); + } + } + assertEquals(count, values.size()); + assertEquals(TestSource.elements(tag, 1, count), new HashSet<>(values)); + } + + /** The default record limit is unlimited, an available source drains in one micro-batch. */ + @Test + public void testUnlimitedRecordsPerBatchByDefault() throws Exception { + int count = 2500; + StreamingQuery query = start(rows(1, count, options(5_000L)), tag, temp.newFolder(tag)); + try { + await("a non empty batch", () -> !nonEmptySizes(batches(tag)).isEmpty()); + } finally { + stopQuietly(query); + } + assertEquals(Collections.singletonList(count), nonEmptySizes(batches(tag))); + } + + /** The offset is an opaque epoch counter whose JSON is the bare number. */ + @Test + public void testEpochOffsetRoundTrip() { + BeamOffset offset = new BeamOffset(42L); + 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")); + } + + /** A deserialized offset moves the epoch counter past itself. */ + @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); + } + + /** A planned end offset moves the epoch counter past itself. */ + @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 and the remainder rotates with the epoch. */ + @Test + public void testSplitQuotas() { + assertArrayEquals( + new long[] {1, 1, 1, 0, 0, 0, 0, 0}, BeamMicroBatchStream.splitQuotas(3, 8, 0)); + assertArrayEquals( + new long[] {0, 0, 0, 0, 0, 1, 1, 1}, BeamMicroBatchStream.splitQuotas(3, 8, 3)); + assertArrayEquals(new long[] {4, 3, 3}, BeamMicroBatchStream.splitQuotas(10, 3, 0)); + assertArrayEquals(new long[] {3, 4, 3}, BeamMicroBatchStream.splitQuotas(10, 3, 2)); + assertArrayEquals(new long[] {5, 5}, BeamMicroBatchStream.splitQuotas(10, 2, 0)); + assertArrayEquals(new long[] {-1, -1, -1}, BeamMicroBatchStream.splitQuotas(0, 3, 0)); + assertArrayEquals(new long[] {-1, -1, -1}, BeamMicroBatchStream.splitQuotas(-1, 3, 0)); + long[] many = BeamMicroBatchStream.splitQuotas(1, 200, 0); + assertEquals(200, many.length); + assertEquals(1L, many[0]); + assertEquals(1L, Arrays.stream(many).sum()); + } + + /** The record limit of a micro-batch is a total over all splits. */ + @Test + public void testMaxRecordsPerBatchIsSharedAcrossSplits() throws Exception { + int shards = 2; + int count = 30; + StreamingQuery query = + start(rows(shards, count, limited(10L, 1_000L)), tag, temp.newFolder(tag)); + try { + await("all rows", () -> values(batches(tag)).size() >= count); + } finally { + stopQuietly(query); + } + List sizes = nonEmptySizes(batches(tag)); + assertFalse("no rows arrived", sizes.isEmpty()); + assertTrue("batch exceeds the shared limit: " + sizes, Collections.max(sizes) <= 10); + assertEquals(TestSource.elements(tag, shards, count), new HashSet<>(values(batches(tag)))); + } + + /** + * A limit below the split count emits at most the limit per batch and rotates over the splits. + */ + @Test + public void testQuotaBelowSplitCountRotates() throws Exception { + int shards = 4; + StreamingQuery query = start(rows(shards, 40, limited(1L, 1_000L)), tag, temp.newFolder(tag)); + try { + await("every shard", () -> shardsOf(values(batches(tag))).size() == shards); + } finally { + stopQuietly(query); + } + List sizes = nonEmptySizes(batches(tag)); + assertTrue("batch exceeds the limit of 1: " + sizes, Collections.max(sizes) <= 1); + } + + /** A restart resumes every split from the last committed mark, replaying at most one batch. */ + @Test + public void testRestartResumesFromCommittedMark() throws Exception { + int shards = 2; + int count = 80; + long limit = 4L; + File checkpointDir = temp.newFolder("restart"); + String first = tag + "_a"; + String second = tag + "_b"; + + StreamingQuery query = start(rows(shards, count, limited(limit, 1_000L)), first, checkpointDir); + try { + await("two commits", () -> committedBatchIds(checkpointDir).size() >= 2); + } finally { + stopQuietly(query); + } + BeamReaderCache.invalidateAll(); + List firstValues = values(batches(first)); + + Set all = TestSource.elements(tag, shards, count); + query = start(rows(shards, count, limited(limit, 1_000L)), 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 one batch replayed: " + firstValues.size() + " + " + secondValues.size(), + firstValues.size() + secondValues.size() <= count + limit); + for (int shard = 0; shard < shards; shard++) { + int min = Integer.MAX_VALUE; + for (String value : secondValues) { + if (TestSource.shardOf(value) == shard) { + min = Math.min(min, TestSource.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); + } + } + + /** No split finalizes a position beyond its mark at the end epoch of the last committed batch. */ + @Test + public void testMarksAreFinalizedOnlyAfterSparkCommit() throws Exception { + File checkpointDir = temp.newFolder("finalize"); + runUntilCommits(checkpointDir, 3); + + int finalizations = 0; + for (int shard = 0; shard < 2; shard++) { + int committed = committedPosition(checkpointDir, shard); + List finalized = TestSource.finalized(tag, shard); + assertTrue( + "shard " + shard + " finalized " + finalized + " beyond committed " + committed, + finalized.isEmpty() || Collections.max(finalized) <= committed); + finalizations += finalized.size(); + } + assertTrue("no mark was finalized", finalizations > 0); + } + + /** + * After a run the mark at the last committed epoch exists, every surviving mark is at or above + * the end epoch of the batch before the last constructed one, and the mark of batch 0 is gone. + */ + @Test + public void testMarksBelowCommittedOffsetArePurged() throws Exception { + File checkpointDir = temp.newFolder("purge"); + runUntilCommits(checkpointDir, 3); + + long lastConstructed = Collections.max(batchIds(new File(checkpointDir, "offsets"))); + long purgeFloor = endEpoch(checkpointDir, lastConstructed - 1); + long committedEpoch = committedEpoch(checkpointDir); + long firstEpoch = endEpoch(checkpointDir, 0); + assertTrue(committedEpoch >= purgeFloor); + assertTrue(purgeFloor > firstEpoch); + File sourceDir = sourceDir(checkpointDir); + awaitQuietly( + 10_000L, + () -> { + for (int shard = 0; shard < 2; shard++) { + TreeSet epochs = markEpochs(sourceDir, shard); + if (epochs.isEmpty() || epochs.first() < purgeFloor) { + return false; + } + } + return true; + }); + for (int shard = 0; shard < 2; shard++) { + TreeSet remaining = markEpochs(sourceDir, shard); + assertTrue( + "shard " + + shard + + " lost the mark at committed epoch " + + committedEpoch + + ": " + + remaining, + remaining.contains(committedEpoch)); + assertTrue( + "shard " + shard + " kept marks below purge floor " + purgeFloor + ": " + remaining, + remaining.first() >= purgeFloor); + assertFalse("shard " + shard + " kept the mark of batch 0", remaining.contains(firstEpoch)); + } + } + + /** + * After a stop an idle sweep finalizes exactly the marks of the last committed epoch, {@link + * BeamReaderCache#closeIdle(long)} is the one white box hook these tests use. + */ + @Test + public void testStoppedQueryFinalizesLastCommittedMarks() throws Exception { + File checkpointDir = temp.newFolder("stopped"); + runUntilCommits(checkpointDir, 3); + BeamReaderCache.closeIdle(Long.MAX_VALUE); + + for (int shard = 0; shard < 2; shard++) { + int committed = committedPosition(checkpointDir, shard); + List finalized = TestSource.finalized(tag, shard); + assertTrue( + "shard " + shard + " finalized " + finalized + ", committed " + committed, + finalized.contains(committed) && Collections.max(finalized) == committed); + } + } + + /** A retried batch restarts from the durable mark at its start and finalizes nothing. */ + @Test + public void testRetriedBatchRestartsFromDurableMark() throws Exception { + String location = sourceDir(temp.newFolder("protocol")).getAbsolutePath(); + assertEquals(shardZero(0, 1, 2), readBatch(partition(location, 0, 1))); + assertEquals(shardZero(0, 1, 2), readBatch(partition(location, 0, 1))); + assertEquals(Collections.emptyList(), TestSource.finalized(tag, 0)); + assertEquals(2, TestSource.created(tag)); + } + + /** A start epoch above zero without a durable mark is an invariant violation. */ + @Test + public void testMissingMarkThrows() throws Exception { + String location = sourceDir(temp.newFolder("protocol")).getAbsolutePath(); + assertThrows( + IllegalStateException.class, () -> new BeamPartitionReader<>(partition(location, 5, 6))); + assertEquals(0, TestSource.created(tag)); + } + + /** A failed mark write fails the batch after its rows, the retry recreates the reader. */ + @Test + public void testRetryAfterFailedMarkWriteRecreatesReader() throws Exception { + File location = sourceDir(temp.newFolder("protocol")); + assertTrue(location.getParentFile().mkdirs() && location.createNewFile()); + String file = location.getAbsolutePath(); + assertEquals(shardZero(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals(shardZero(0, 1, 2), drainUntilFailure(partition(file, 0, 1), IOException.class)); + assertEquals(Collections.emptyList(), TestSource.finalized(tag, 0)); + assertEquals(2, TestSource.created(tag)); + } + + // --------------------------------------------------------------------------------------------- + // query helpers + // --------------------------------------------------------------------------------------------- + + private static SparkStructuredStreamingPipelineOptions options(long maxBatchDurationMillis) { + SparkStructuredStreamingPipelineOptions options = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + options.setWatermarkDelayMillis(0L); + options.setMaxBatchDurationMillis(maxBatchDurationMillis); + return options; + } + + private static SparkStructuredStreamingPipelineOptions limited( + long maxRecordsPerBatch, long maxBatchDurationMillis) { + SparkStructuredStreamingPipelineOptions options = options(maxBatchDurationMillis); + options.setMaxRecordsPerBatch(maxRecordsPerBatch); + return options; + } + + private Dataset rows( + int shards, int count, SparkStructuredStreamingPipelineOptions options) { + return UnboundedSourceDataset.of( + SESSION.getSession(), + new TestSource(tag, shards, count), + CODER, + options, + "Read(TestSource)"); + } + + /** Builds the driver side stream through the table, with the session's broadcasts. */ + private static BeamMicroBatchStream newStream(String checkpointLocation) { + BeamTable table = + new BeamTable<>( + CountingSource.unbounded(), + WindowedValues.getFullCoder(VarLongCoder.of(), GlobalWindow.Coder.INSTANCE), + optionsBroadcast, + hadoopConfBroadcast, + 2, + -1L, + 200L, + 600_000L, + "Read(CountingSource)"); + return (BeamMicroBatchStream) + table + .newScanBuilder(CaseInsensitiveStringMap.empty()) + .build() + .toMicroBatchStream(checkpointLocation); + } + + /** Starts a query collecting every micro-batch as one list into {@link #BATCHES}. */ + private static StreamingQuery start(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(); + } + + /** Runs two shards with a limit of 4 over a source that never drains until Spark committed. */ + private void runUntilCommits(File checkpointDir, int commits) throws Exception { + StreamingQuery query = start(rows(2, 4_000, limited(4L, 1_000L)), tag, checkpointDir); + try { + await(commits + " commits", () -> committedBatchIds(checkpointDir).size() >= commits); + } finally { + stopQuietly(query); + } + } + + private static void stopQuietly(StreamingQuery query) { + try { + query.stop(); + } catch (Exception e) { + // Nothing useful to do while tearing a test query down. + } + } + + 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) { + List values = new ArrayList<>(); + for (List batch : batches) { + for (Row row : batch) { + values.add(decode(row.getAs(COL_PAYLOAD))); + } + } + return values; + } + + private static String decode(byte[] payload) { + try { + return CoderUtils.decodeFromByteArray(CODER, payload).getValue(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + private static List nonEmptySizes(List> batches) { + List sizes = new ArrayList<>(); + for (List batch : batches) { + if (!batch.isEmpty()) { + sizes.add(batch.size()); + } + } + return sizes; + } + + private static Set shardsOf(List values) { + Set shards = new HashSet<>(); + for (String value : values) { + shards.add(TestSource.shardOf(value)); + } + return shards; + } + + 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(); + } + + /** 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()) { + String watermark = progress.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; + } + + // --------------------------------------------------------------------------------------------- + // checkpoint helpers + // --------------------------------------------------------------------------------------------- + + private static File sourceDir(File checkpointDir) { + return new File(checkpointDir, "sources/0"); + } + + /** Numeric file names in a Spark log directory, temp and hidden files excluded. */ + private static TreeSet batchIds(File dir) { + TreeSet 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; + } + + /** Epochs under {@code marks//} holding a mark file of {@code shard}. */ + private static TreeSet markEpochs(File sourceDir, int shard) { + TreeSet epochs = new TreeSet<>(); + for (long epoch : batchIds(new File(sourceDir, "marks"))) { + if (new File(sourceDir, "marks/" + epoch + "/" + shard).exists()) { + epochs.add(epoch); + } + } + return epochs; + } + + private static TreeSet 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()); + } + } + assertEquals("one source expected in " + lines, 3, lines.size()); + return BeamOffset.fromJson(lines.get(2)).epoch(); + } + + private static long committedEpoch(File checkpointDir) throws IOException { + return endEpoch(checkpointDir, committedBatchIds(checkpointDir).last()); + } + + /** The position in the mark of {@code shard} at the end epoch of the last committed batch. */ + private static int committedPosition(File checkpointDir, int shard) throws IOException { + long epoch = committedEpoch(checkpointDir); + BeamSourceCheckpoint checkpoint = + new BeamSourceCheckpoint(sourceDir(checkpointDir).getAbsolutePath(), new Configuration()); + byte[] coded = checkpoint.readMark(shard, epoch); + assertNotNull("no mark at committed epoch " + epoch + " for shard " + shard, coded); + return CoderUtils.decodeFromByteArray(TestSource.MARK_CODER, coded).next; + } + + // --------------------------------------------------------------------------------------------- + // hand built partition helpers + // --------------------------------------------------------------------------------------------- + + /** Split 0 of a single shard source of 100 elements from epoch {@code start} to {@code end}. */ + private BeamInputPartition partition(String location, long start, long end) { + TestSource split = new TestSource(tag, 1, 100).split(1, PipelineOptionsFactory.create()).get(0); + return new BeamInputPartition<>( + split, + CODER, + optionsBroadcast, + hadoopConfBroadcast, + location, + 0, + start, + end, + 3L, + 30_000L, + 600_000L); + } + + private static List readBatch(BeamInputPartition partition) throws IOException { + List values = new ArrayList<>(); + drainInto(new BeamPartitionReader<>(partition), values); + return values; + } + + private static void drainInto(BeamPartitionReader reader, List values) + throws IOException { + while (reader.next()) { + InternalRow row = reader.get(); + values.add(decode(row.getBinary(0))); + } + 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 shardZero(int... indexes) { + List elements = new ArrayList<>(); + for (int index : indexes) { + elements.add(TestSource.element(tag, 0, index)); + } + return elements; + } + + // --------------------------------------------------------------------------------------------- + // the shared in memory UnboundedSource + // --------------------------------------------------------------------------------------------- + + /** + * Splits into one sub source per shard, each over {@code count / shards} elements named {@code + * --} with evenly spaced timestamps. Marks are not Java serializable, they + * record the position they finalize under {@code /}, readers are counted per tag. + */ + static final class TestSource 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 final String tag; + private final int shard; + private final int shards; + private final int perShard; + + TestSource(String tag, int shards, int count) { + this(tag, -1, shards, count / shards); + } + + private TestSource(String tag, int shard, int shards, int perShard) { + this.tag = tag; + this.shard = shard; + this.shards = shards; + this.perShard = perShard; + } + + static Set elements(String tag, 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(tag, shard, index)); + } + } + return elements; + } + + static String element(String tag, int shard, int index) { + return tag + "-" + shard + "-" + index; + } + + static int shardOf(String element) { + String head = element.substring(0, element.lastIndexOf('-')); + return Integer.parseInt(head.substring(head.lastIndexOf('-') + 1)); + } + + 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); + } + } + + static int created(String tag) { + AtomicInteger created = CREATED.get(tag); + return created == null ? 0 : created.get(); + } + + static void forget(String tag) { + FINALIZED.keySet().removeIf(key -> key.startsWith(tag + "/")); + CREATED.remove(tag); + } + + 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 TestSource(tag, i, shards, perShard)); + } + return splits; + } + + @Override + public UnboundedReader createReader(PipelineOptions options, @Nullable Mark mark) { + if (shard < 0) { + throw new IllegalStateException("split before reading"); + } + CREATED.computeIfAbsent(tag, t -> new AtomicInteger()).incrementAndGet(); + return new Reader(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 Mark implements UnboundedSource.CheckpointMark { + private final String tag; + private final int shard; + final int next; + + Mark(String tag, int shard, int next) { + this.tag = tag; + this.shard = shard; + this.next = next; + } + + @Override + public void finalizeCheckpoint() { + FINALIZED + .computeIfAbsent(key(tag, shard), k -> 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 { + StringUtf8Coder.of().encode(mark.tag, out); + VarIntCoder.of().encode(mark.shard, 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), + VarIntCoder.of().decode(in)); + } + } + + private static final class Reader extends UnboundedReader { + private final TestSource source; + private int next; + private int current = -1; + + Reader(TestSource 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.tag, source.shard, current); + } + + @Override + public Instant getCurrentTimestamp() throws NoSuchElementException { + if (current < 0) { + throw new NoSuchElementException(); + } + return new Instant( + BASE_MILLIS + (source.shard * source.perShard + current) * INTERVAL_MILLIS); + } + + @Override + public Instant getWatermark() { + return current < 0 ? BoundedWindow.TIMESTAMP_MIN_VALUE : getCurrentTimestamp(); + } + + @Override + public CheckpointMark getCheckpointMark() { + return new Mark(source.tag, source.shard, next); + } + + @Override + public UnboundedSource getCurrentSource() { + return source; + } + + @Override + public void close() {} + } + } +} 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).")