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

[tkaymak]

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. Three files, UnboundedSourceDataset with the DataSourceV2 glue as nested classes like BoundedDatasetFactory, BeamReaderCache and BeamSourceCheckpoint, plus one option on SparkStructuredStreamingPipelineOptions.

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. The event time watermark is declared once, on this dataset.
• The dataset is built like BoundedDatasetFactory, a Table holding the source, coder and broadcasts wrapped in StreamingRelationV2. No string options, splits travel as objects inside the InputPartition, pipeline options and the session Hadoop configuration as broadcasts.
• Offsets are opaque, strictly increasing epoch counters serialized as the bare number like LongOffset. latestOffset always advances so Spark keeps scheduling micro-batches, termination belongs to the lifecycle owner.
• Per source state lives under the location Spark passes to toMicroBatchStream, written through CheckpointFileManager. The first run pins its split list because Beam sources do not guarantee deterministic splitting. Each split writes its CheckpointMark, coded with getCheckpointMarkCoder(), at the end of every micro-batch to marks/<epoch>/<split>. commit(end) purges the epoch directories below end, one recursive delete each, off the stream thread.
• A batch reuses the cached reader only if it is positioned at the batch's start offset and finalizes the mark taken there at that moment. Spark only starts a batch at the initial offset or at the end offset of a batch already in its commit log, so a mark is never finalized before Spark committed the batch it closes. Any other case (task retry, killed attempt, executor change, restart) drops the reader without finalizing and restores it from the durable mark at the start offset. A missing mark at a start offset above zero fails the task instead of silently restarting the source.
• Finalization stays on the executor because it must run on the live mark instance, PubsubCheckpoint throws on a restored checkpoint and KafkaCheckpointMark is a no-op without its reader, and DSv2 has no executor side commit callback.
• Idle readers are closed by a sweeper thread after readerIdleTimeoutMillis. A pending mark is finalized then if Spark's commit log covers its epoch, read from the query's commits and offsets entries, otherwise dropped. This also finalizes the last batch of a stopped query while the executor lives.
• Delivery is at least once. A crash between the end of a batch and its commit replays that batch, an executor that exits before the sweep leaves its last mark unfinalized and the source redelivers. spark.speculation is not supported for sources with non deterministic reads, and async progress tracking lags the commit log.
• maxRecordsPerBatch is a per batch total divided across splits like the legacy MicrobatchSource, a zero quota split emits nothing, values below 1 mean no limit and the batch ends on the maxBatchDurationMillis deadline. The remainder rotates by epoch so a limit below the split count does not starve the same splits. defaultParallelism decides the desired split count, Spark places the tasks. Readers waiting for data back off with FluentBackoff.

Tests, one class over one test source, cover element delivery, watermark tracking through typed maps, the offset round trip, the quota division and its rotation, and, against Spark's real offsets and commits logs, restart recovery with at most one replayed batch, finalization only after commit, mark purging and finalization of a stopped query's last batch. Three cases a live query cannot provoke are driven directly: a retried batch, a missing mark, a failed mark write.

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

[github-actions]

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

[tkaymak]

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.

[tkaymak]

@Abacn thanks for the thorough review, the failing tests seem to be unrelated (the setup environment step fails, because GitHub's org policy rejects the pinned SHA), will have a look at that later if I can.
I tried to address all points in 8af17d5:

• Checkpointing now follows Spark's lifecycle. Per source state lives under the location Spark passes to toMicroBatchStream, written through CheckpointFileManager with the session Hadoop conf broadcast to executors, marks coded with getCheckpointMarkCoder(). commit(end) purges marks below end, there is no fixed retention any more.

• Finalization moved out of PartitionReader.close(). A reader finalizes the mark taken at its start offset when the next batch for that split is scheduled. Spark only starts a batch at the initial offset or at the end offset of a batch already in the commit log (MicroBatchExecution.populateStartOffsets), so that point is always after the commit. Any mismatch (retry, killed attempt, executor change, restart) drops the reader without finalizing and restores it from the durable mark at the start offset. Finalizing decoded marks on the driver in commit() is not an option, PubsubCheckpoint.finalizeCheckpoint throws on a restored checkpoint and KafkaCheckpointMark is a no-op without the live reader, and DSv2 has no executor side commit hook. The remaining limit is spark.speculation with non deterministic sources, a losing attempt that completed before its kill arrived can be reused, DSv2 gives a task no way to learn it lost.

• On TransformWithState: it carries Beam state and timers in the next slices. Reading inside a stateful operator would be the legacy mapWithState design and it cannot own the watermark declaration the DSv2 source provides. On spark/4: Dataset.ofRows and StreamingRelationV2 moved packages in 4.0, the module already forks BoundedDatasetFactory for the first, and the only consumer is the Spark 4 streaming translator. CheckpointFileManager moves to a checkpointing package in later 4.x, one import to adjust in the future.

• Smaller items: the dataset is built like BoundedDatasetFactory (real objects in a Table wrapped in StreamingRelationV2, no string options, no Base64), splits are plain fields of the InputPartition, maxRecordsPerBatch is divided across splits like MicrobatchSource.splitNumRecords, defaultParallelism decides the split count, CoderHelpers is reused, FluentBackoff replaces the sleep, the offset serializes as the bare epoch like LongOffset with the base class equality, no class wide SuppressWarnings, the mark map is gone. One new option, readerIdleTimeoutMillis, replaces a hardcoded reader expiry.

• New tests drive the reader cache protocol directly and prove restart recovery, finalization only after commit and mark purging against Spark's real offsets and commits logs. The per test JUnit timeout is removed from the streaming test, its throwaway thread group poisoned Spark's static pools for later batch tests in the same forked JVM, which explains the flaky batch failures seen locally.

[Abacn]

Thanks for addressing the Round 1 comments and migrating to CheckpointFileManager and StreamingRelationV2.

However, the PR diff has expanded from ~1,000 to over 3,300 lines (with tests now comprising over 55% of the code). Much of this inflation comes from duplicate test fixtures, over-engineered scheduling heuristics, and white-box protocol testing:

• Rendezvous Hashing & Private Spark API: sortedExecutors() calls SparkEnv.get().blockManager().master().getPeers(). This is a private Spark internal API that is brittle across cluster managers (K8s, YARN, standalone) and dynamic allocation. Spark's scheduler already handles task placement, and fallback recovery from durable marks is already required and implemented. We can simply remove the rendezvous hashing and murmur3 logic and return empty preferredLocations().

• Mock Source Duplication: We currently have 3 separate synthetic unbounded sources (ListSource, ShardedListSource, and IntListSource) spanning ~500 lines across test files. ShardedListSource and IntListSource are practically identical (custom non-serializable marks, static maps for tracking finalizations, custom coders). Using existing Beam test sources is preferred, or consolidate them into a single reusable test source.

• Prune White-Box Unit Tests: BeamReaderCacheProtocolTest (478 lines) and BeamSourceCheckpointTest (209 lines) test internal bookkeeping without Spark. Meanwhile, BeamMicroBatchSourceTest already tests recovery, finalization, and purging against a live Spark query. We can significantly cut down maintenance overhead by pruning these unit tests and relying on integration tests.

• Consolidate POJO Classes: Compare with BoundedDatasetFactory.java, which contains the entire bounded source implementation, having 10 separate files for streaming (BeamSourceSpec, BeamPartitionReaderFactory, BeamStreamingTable, etc.) likely introduced unnecessary boilerplate.

Other comments --- need to verify if it's factually correct

• Deferred Finalization Lifecycle Edge Cases: In BeamReaderCache, deferring mark finalization until the next batch's acquire() means:

  • The final micro-batch on graceful query shutdown is never finalized.
  • An idle reader eviction (closeIdle()) drops the reader without finalizing the pending mark. We should ensure marks are finalized on clean query termination and idle close.

• splitQuotas Distribution Bug: In BeamMicroBatchStream.splitQuotas, using Math.max(1L, ...) causes queries where maxRecordsPerBatch < numSplits to emit more records than the configured limit (e.g. 20 records for limit of 10). Please adopt the quota splitting logic from MicrobatchSource.splitNumRecords where partitions with 0 quota emit 0 records (using < 0 for unlimited).

[Abacn]

It issues individual synchronous delete() RPCs for every epoch. On cloud object stores like GCS/S3, this can result in hundreds of sequential HTTP calls on every commit. Use directory listing or batch deletions instead.

[Abacn]

(from AI reivew) sortedExecutors uses .master().getPeers(...) and assumes naming format of Spark internals, both are unsupported and fragile in managed environments. We already support restoring from durable marks when a partition runs on another executor, so we should let Spark handle partition locality.

[tkaymak]

Thank you, will have a look!

[tkaymak]

@Abacn thanks, all addressed in 7376415.

Size. Three main files now, UnboundedSourceDataset with the DataSourceV2 glue as nested classes like BoundedDatasetFactory, BeamReaderCache, BeamSourceCheckpoint. One test class over one test source. The PR is 5 files and 2337 lines against master, tests 44 percent.

Locality. Rendezvous hashing and the getPeers call are gone, preferredLocations() is not overridden.

Tests. The three synthetic sources are one TestSource. BeamSourceCheckpointTest is deleted. BeamReaderCacheProtocolTest is deleted except three cases folded into the live test class because a live query cannot provoke them: a retried batch restarts from the durable mark and finalizes nothing, a start offset above zero without a mark fails instead of restarting the source, a failed mark write fails the batch and the retry recreates the reader.

Quotas. splitQuotas follows MicrobatchSource.splitNumRecords, a zero quota split emits nothing, below zero is unlimited. One addition, the remainder rotates by epoch so a limit below the split count does not starve the same splits forever. Verified live with a limit of 1 over 4 splits.

Purge. Layout is marks/<epoch>/<split>, purging below the committed offset is one recursive directory delete per epoch, independent of the split count, still off the stream thread. The driver creates the epoch directory when it plans the batch.

Finalization, both observations were correct. Idle close now finalizes the pending mark when Spark's commit log covers its epoch. The executor reads <root>/commits and this source's line in the matching <root>/offsets/<id>, which has no lag because Spark writes commits/N in markMicroBatchEnd while commit(end) reaches the source only when N+1 is constructed. A sweeper thread runs the idle check every 10 s, so after query.stop() the last completed batch is finalized within the idle timeout. Still unfinalized: an executor JVM that exits before the sweep, a batch Spark aborted mid flight, and async progress tracking where the commit log itself lags. DSv2 has no executor side stop or commit hook, in those cases the source redelivers. A wrong or racing read can only yield a lower epoch, so the failure direction is redelivery, never an extra finalization.

[tkaymak]

Run Java PreCommit