Abacn commented on code in PR #39971: URL: https://github.com/apache/beam/pull/39971#discussion_r3915592536
########## 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. + * + * <p>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. + * + * <p>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 <T> the element type of the wrapped source + */ +@SuppressWarnings({ Review Comment: Avoid SuppressWarnings in new codes. If there is a compelling reason (I see there are comments here), put this in specific chunk, not whole class. This applies to the other few places. ########## runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/streaming/BeamCheckpointFiles.java: ########## @@ -0,0 +1,236 @@ +/* Review Comment: In #39576 it was noted to support streaming with TransformWithState API, which only exists in Spark 4. However, it appears current change does not yet involve TransformWithState API. Put it in `spark/4/` sounds fine as we only aim to support streaming for Spark 4. However, it may be more straightforward to re-use existing code if we work inside `spark/src` as long as it doesn't involve TransformWithState, as this addition-only change suggests there may be duplicated codes that should be shared with common/batch paths, see below. ########## 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. + * + * <p>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. + * + * <p>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 <T> 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<T> implements PartitionReader<InternalRow> { + + 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<T> cached; + private final Coder<WindowedValue<T>> 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<T, ?> 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<T> windowedValue = + WindowedValues.timestampedValueInGlobalWindow(cached.reader().getCurrent(), timestamp); + byte[] payload; + try { Review Comment: This duplicates https://github.com/apache/beam/blob/7151bed2bbdc1431b575c19d4f597d29493c1b2e/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/CoderHelpers.java#L52 ########## 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. + * + * <p>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; Review Comment: The encoding/decoding handling here lazily uses a Base64. Beam Source is Java serializable. We should be able to just Java-serialize them. In fact we already do it here but wrap it into a base64 str causing double serialization ########## 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. + * + * <p>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. + * + * <p><b>The watermark is declared here and only here.</b> 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. */ Review Comment: We need to clean up PoC hardcodes when checking in them into master branch Consider make it in alignment with the Bounded source: https://github.com/apache/beam/blob/10408793b66c4fdde67249b6ec2debdee2aec7f6/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/BoundedDatasetFactory.java#L96 which uses `session.sparkContext().defaultParallelism()` or from pipeline options ########## 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 <checkpointLocation>/beam-source-<sourceId>/}. + * + * <p>Two kinds of state are kept per source. The pinned split list under {@code <root>/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 <root>/marks/<splitId>/<epoch>} 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. + * + * <p>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. + * + * <p>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<String> 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<String> splits = (List<String>) 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<String> 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)) { Review Comment: This restriction doesn't sounds right as Beam CheckpointMark doesn't require Serializable. We should use provided coders via `source.getCheckpointMarkCoder()` to serialize Beam checkpoints ########## 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}. + * + * <p>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()}). + * + * <p>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. + * + * <p>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<String> 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<String> splits = splits(); + InputPartition[] partitions = new InputPartition[splits.size()]; + for (int i = 0; i < splits.size(); i++) { + partitions[i] = + new BeamInputPartition( Review Comment: maxRecordsPerBatch is passed to every BeamInputPartition, results in each micro-batch actually pulling (`maxRecordsPerBatch * splits`) records, inconsistent with `SparkStructuredStreamingPipelineOptions.getMaxRecordsPerBatch() ("Max records per micro-batch").` definition. In https://github.com/apache/beam/blob/7151bed2bbdc1431b575c19d4f597d29493c1b2e/runners/spark/src/main/java/org/apache/beam/runners/spark/io/MicrobatchSource.java#L101, splitNumRecords(maxNumRecords, numSplits) evenly partitions the record quota across splits. BeamMicroBatchStream.planInputPartitions should use the same distribution logic. ########## 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 <checkpointLocation>/beam-source-<sourceId>/}. + * + * <p>Two kinds of state are kept per source. The pinned split list under {@code <root>/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 <root>/marks/<splitId>/<epoch>} 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. + * + * <p>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. + * + * <p>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<String> 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<String> splits = (List<String>) 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<String> 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}. + * + * <p>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()); Review Comment: If we decide to stay with hadoop-file-based checkpointing, a recommendation is to use CheckpointFileManager: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointFileManager.scala Spark uses this for its own <checkpointLocation>/offsets and <checkpointLocation>/commits. Currently a default `new Configuration()` is subject to fail on a cloud based file system. ########## 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 <checkpointLocation>/beam-source-<sourceId>/}. + * + * <p>Two kinds of state are kept per source. The pinned split list under {@code <root>/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 <root>/marks/<splitId>/<epoch>} 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. + * + * <p>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. + * + * <p>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 { Review Comment: This implementation of checkpointing directly interact with low-level architectures (Hadoop filesystem) and completely bypassing Spark's offset tracking and commit lifecycle. This sounds brittle to me. From the PoC PR (#39576) description IIUC it's expected to use transformWithState (that replaces DStream mapWithState which was used in classic Spark runner for streaming to handle checkpoint). Arguably checkpoint is also a state. ########## 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. + * + * <p>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. + * + * <p>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 <T> 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<T> implements PartitionReader<InternalRow> { + + 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<T> cached; + private final Coder<WindowedValue<T>> 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<T, ?> 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( Review Comment: What is the consideration of sleepUninterruptibly? Consider using Beam's FluentBackoff ########## 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). + * + * <p>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. + * + * <p><b>Durable recovery.</b> 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<String, CachedReader<?>> 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<String, CachedReader<?>> 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<String, CheckpointMark> MARKS = new ConcurrentHashMap<>(); Review Comment: Reference leak possible as MARKS holds references to CheckpointMark indefinitely, unless `invalidateAll()` gets called ########## 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. + * + * <p>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); Review Comment: It may work by coincidence, doesn't sounds semantically correct way to extract an offset ########## 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. + * + * <p>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. + * + * <p>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 <T> 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<T> implements PartitionReader<InternalRow> { + + 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<T> cached; + private final Coder<WindowedValue<T>> 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<T, ?> 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(); Review Comment: Premature Checkpoint Finalization: `PartitionReader.close()` is called on the executor as soon as the task finishes reading its micro-batch partition. At this point, downstream operators have not processed the batch, sinks have not written data, and Spark has not committed the micro-batch to its WAL. Checkpoint marks must only be finalized when Spark invokes `MicroBatchStream.commit(Offset)` ########## 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 <checkpointLocation>/beam-source-<sourceId>/}. + * + * <p>Two kinds of state are kept per source. The pinned split list under {@code <root>/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 <root>/marks/<splitId>/<epoch>} 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. + * + * <p>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. + * + * <p>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; Review Comment: We should revisit this hard coded number and throughout the PR. RETAINED_MARKS = 2 is dangerous because Spark's offset log keeps 100 batches by default (`spark.sql.streaming.minBatchesToRetain`). If a restarted query rewinds 3 batches, the mark is missing and the stream replays from scratch. Keep at least minBatchesToRetain marks, or delete old marks only upon MicroBatchStream.commit(Offset). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
