[Spark][#36841] Add the DataSourceV2 unbounded source for the Spark 4 streaming runner by tkaymak · Pull Request #39971 · apache/beam

tkaymak · 2026-09-02T11:48:23Z

Third slice of the Spark 4 Structured Streaming work split out of #39576, following the dispatch seam (#39906) and the Kryo registrations (#39939). Addresses #36841.

This adds the DataSourceV2 micro-batch source that exposes any Beam UnboundedSource as a Spark 4 streaming table. All 12 files are new, nothing existing changes.

Design notes:

Rows have a fixed two column schema, the element encoded with the Beam FullWindowedValueCoder as BINARY plus the event timestamp. No Catalyst encoder is generated for Beam types, payloads stay opaque until a downstream translator decodes them.
Offsets are opaque, strictly increasing epoch counters. latestOffset always advances so Spark keeps scheduling micro-batches, termination belongs to the lifecycle owner, not to the offsets.
Recovery is durable under the query's checkpoint location. The source id derives deterministically from the read transform's full name, the first run pins its split list because Beam sources do not guarantee deterministic splitting, and each split persists its CheckpointMark per epoch (atomic write via temp file and rename, retention two). On restart the epoch counter fast forwards past everything replayed from Spark's offset log and readers resume from the newest durable mark at or before the replayed epoch.
Executors cache live Beam readers between micro-batches, keyed by checkpoint location, source id and split, mirroring MicrobatchSource in the legacy runner.
Semantics are at least once, a mark is written when a batch finishes reading rather than transactionally with Spark's commit, so a crash between the two replays the last micro-batch.
The batch cutoff honors the maxRecordsPerBatch option from [Spark][#36841] Reuse the legacy maxRecordsPerBatch option in the Structured Streaming runner #39952, values below 1, including the default, mean no limit and the batch then ends on the maxBatchDurationMillis deadline.

Tests cover element delivery, watermark tracking through typed maps, the epoch offset round trip, the unlimited default, the checkpoint file layout with retention, epoch fast forward and the deterministic source id. The end of stream sentinel used by PAssert arrives with the translators slice, its hooks are deliberately absent here.

Remaining slices: the state and timer bridge on transformWithState, then the translators with the end to end tests. End to end evidence remains in draft #39576.

R: @Abacn

Exposes any Beam UnboundedSource as a Spark 4 DataSourceV2 streaming table with a fixed two column schema, encoded payload plus event timestamp. Offsets are opaque, strictly increasing epoch counters, so Spark keeps scheduling micro-batches and termination stays with the lifecycle owner. Recovery is durable under the query's checkpoint location: the source id derives deterministically from the read transform's full name, the first run pins its split list (Beam sources do not guarantee deterministic splitting), and every split persists its CheckpointMark per epoch with a retention of two, written atomically via temp file and rename. Executors cache live readers between micro-batches and fall back to the newest durable mark at or before the replayed epoch after a restart. Semantics are at least once, a crash between finishing a read and Spark's commit replays the last micro-batch. The batch cutoff honors maxRecordsPerBatch, values below 1, including the default, mean no limit and the batch ends on the duration deadline.

github-actions · 2026-09-02T13:16:31Z

Assigning reviewers:

R: @tvalentyn added as fallback since no labels match configuration

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

Abacn · 2026-09-02T15:09:49Z

+ *
+ * @param <T> the element type of the wrapped source
+ */
+@SuppressWarnings({


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.

Abacn · 2026-09-02T15:26:32Z

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


We need to clean up PoC hardcodes when checking in them into master branch

Consider make it in alignment with the Bounded source:

beam/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/io/BoundedDatasetFactory.java

Line 96 in 1040879

public static <T> Dataset<WindowedValue<T>> createDatasetFromRDD(

which uses session.sparkContext().defaultParallelism() or from pipeline options

Abacn · 2026-09-02T15:38:03Z

@@ -0,0 +1,236 @@
+/*


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.

Abacn · 2026-09-02T15:40:43Z

+  public static void writeMark(
+      String checkpointLocation, String sourceId, int splitId, long endEpoch, CheckpointMark mark)
+      throws IOException {
+    if (!(mark instanceof Serializable)) {


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

Abacn · 2026-09-02T15:44:30Z

+    InputPartition[] partitions = new InputPartition[splits.size()];
+    for (int i = 0; i < splits.size(); i++) {
+      partitions[i] =
+          new BeamInputPartition(


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

beam/runners/spark/src/main/java/org/apache/beam/runners/spark/io/MicrobatchSource.java

Line 101 in 7151bed

private static long[] splitNumRecords(final long numRecords, final int numSplits) {

, splitNumRecords(maxNumRecords, numSplits) evenly partitions the record quota across splits. BeamMicroBatchStream.planInputPartitions should use the same distribution logic.

Abacn · 2026-09-02T16:47:55Z

+  }
+
+  private static FileSystem fileSystem(Path path) throws IOException {
+    return path.getFileSystem(new Configuration());


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 /offsets and /commits.

Currently a default new Configuration() is subject to fail on a cloud based file system.

Abacn · 2026-09-02T16:49:55Z

+  private static final String TMP_SUFFIX = ".tmp";
+
+  /** Number of most recent mark files retained per split. */
+  private static final int RETAINED_MARKS = 2;


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).

Abacn · 2026-09-02T16:54:34Z

+          .build();
+
+  /** Last known checkpoint mark per key, used when a reader has to be recreated. */
+  private static final ConcurrentMap<String, CheckpointMark> MARKS = new ConcurrentHashMap<>();


Reference leak possible as MARKS holds references to CheckpointMark indefinitely, unless invalidateAll() gets called

Abacn · 2026-09-02T16:55:39Z

+
+  /** 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);


It may work by coincidence, doesn't sounds semantically correct way to extract an offset

Abacn · 2026-09-02T16:56:42Z

+        current = toRow();
+        return true;
+      }
+      Uninterruptibles.sleepUninterruptibly(


What is the consideration of sleepUninterruptibly? Consider using Beam's FluentBackoff

tkaymak · 2026-09-02T20:55:57Z

Thank you for the review @Abacn!
The TransformWithState was planned for the next slice, but given the points raised I will rework this slice first.

github-actions Bot added runners spark labels Sep 2, 2026

tkaymak mentioned this pull request Sep 2, 2026

[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState #39576

Draft

tkaymak changed the title ~~[#36841] Add the DataSourceV2 unbounded source for the Spark 4 streaming runner~~ [Spark][#36841] Add the DataSourceV2 unbounded source for the Spark 4 streaming runner Sep 2, 2026

github-actions Bot added the Next Action: Reviewers label Sep 2, 2026

Abacn reviewed Sep 2, 2026

View reviewed changes

Navigation Menu

Uh oh!

[Spark][#36841] Add the DataSourceV2 unbounded source for the Spark 4 streaming runner - #39971

[Spark][#36841] Add the DataSourceV2 unbounded source for the Spark 4 streaming runner#39971
tkaymak wants to merge 1 commit into
apache:masterfrom
tkaymak:spark4-streaming-slice3-dsv2-source

tkaymak commented Sep 2, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

Abacn Sep 2, 2026

Uh oh!

tkaymak commented Sep 2, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

tkaymak commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tkaymak commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tkaymak commented Sep 2, 2026 •

edited

Loading

tkaymak commented Sep 2, 2026 •

edited

Loading