[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. The only change outside the new package is 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 under its end epoch. commit(end) purges marks below end.
• 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.
• Delivery is at least once. A crash between the end of a batch and its commit replays that batch, a mark whose split is never scheduled again is not finalized and the source redelivers. spark.speculation is not supported for sources with non deterministic reads.
• maxRecordsPerBatch is a per batch total divided across splits like the legacy MicrobatchSource, values below 1 mean no limit and the batch ends on the maxBatchDurationMillis deadline. defaultParallelism decides the desired split count. Idle readers back off with FluentBackoff. The new option readerIdleTimeoutMillis bounds how long an executor keeps an idle reader, its last mark is not finalized when it is closed.

Tests cover element delivery, watermark tracking through typed maps, the offset round trip, the quota division, the reader cache protocol driven directly (continuation, retry, executor change, missing mark, never started reader, failed mark write), the checkpoint layout, and, against Spark's real offsets and commits logs, restart recovery with at most one replayed batch, finalization only after commit, and mark purging.

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.