[Spark][#36841] Add an opt-in Dataset-based backend for portable pipelines

[Eliaaazzz]

Proof of concept for the Dataset-based portable Spark 4 backend discussed in the "[DISCUSS] Spark 4 portable runner direction" thread on dev@. Opening this as a draft so the discussion has concrete code to work from. Addresses #36841.

What this adds

An opt-in --useStructuredStreaming flag on the portable Spark runner. With the flag set, SparkPipelineRunner translates the fused Runner API pipeline into Spark Datasets through a new SparkDatasetPortablePipelineTranslator:

• Impulse, Flatten, Reshuffle and GroupByKey become Dataset operations. GroupByKey groups on the encoded key and runs SparkGroupAlsoByWindowViaOutputBufferFn per key. That is the general GroupAlsoByWindow the RDD batch translator falls back to when a windowing strategy is not eligible for the GroupNonMergingWindowsFunctions fast paths. The Dataset path always takes the general branch, so it gives up the non-merging memory optimization and keeps the ReduceFnRunner semantics.
• Executable stages run inside Dataset.mapPartitions through the existing Fn API bridge, SparkExecutableStageFunction, which is unchanged. Outputs are demultiplexed with EncoderHelpers.oneOfEncoder, following ParDoTranslatorBatch. Side inputs are collected during translation and broadcast.
• Fully bounded pipelines run as batch Datasets. Unbounded input, user state and timers are not supported yet. They fail at translation with a message pointing at the tracking issue, so nothing silently runs on a path that cannot support it.

The translation context attaches to the SparkContext that SparkContextFactory already created, so SparkJobInvoker and result handling are unchanged. The runner selects the backend by the flag alone and leaves the streaming option as submitted. MetricsAccumulator.init gained an explicit useCheckpoint argument, which the portable runner passes as the streaming option restricted to the DStream path, because that path is the only one using the metrics checkpoint. The translator lives in the shared runners/spark base and compiles for both the Spark 3 and Spark 4 modules.

The flag name matches the terminology in the dev@ thread and is provisional. Today it only selects the Dataset translator. No Structured Streaming query is built yet.

The exit gate

Following Yi's suggestion on the thread, validatesPortableRunnerStructuredStreaming runs the streaming PortableValidatesRunner suite with the flag on, keeping exactly the same categories and test filters as validatesPortableRunnerStreaming, and a PostCommit Java PVR Spark4 StructuredStreaming workflow runs it. That suite is the exit gate for flipping the default for portable streaming pipelines.

What that gate currently measures is narrower than its name suggests. --streaming does not by itself mark any PCollection unbounded, and the streaming exclusion list already removes the categories that would produce one. ValidatesRunner extends NeedsRunner, so tests annotated with NeedsRunner and UsesUnboundedPCollections are outside the suite, and the remaining ValidatesRunner tests carrying UsesUnboundedPCollections are each excluded by UsesTimersInParDo, UsesTestStream, UsesStrictTimerOrdering or UsesOnWindowExpiration. The first run therefore exercises bounded Dataset execution under the streaming exclusion list. The risk it actually probes is windowing, triggers and GroupByKey behaviour on this backend. This matches the point on the thread that the suite is mostly batch pipelines after the exclusions.

The new task is deliberately not part of the aggregate validatesPortableRunner task while the backend is a preview.

Out of scope, tracked separately

• Unbounded splittable DoFn residuals across micro-batches, Spark portable runner: support SDF #19468.
• Stateful stages, Add support for state in Spark portable streaming  #20396.
• Timers, Add support for timers in Spark portable streaming #20397.
• The Structured Streaming query itself for pipelines with unbounded input, [Feature Request]: Support Apache Spark 4.x in Beam Spark Runner #36841.

This does not modify the in-process Structured Streaming runner, including the work merged in #39906 and #39939 and the open #39576. It does reuse two of its classes, EncoderHelpers and EvaluationContext, so the portable path also picks up the Spark 4 EncoderHelpers override that came with the Spark 4 runner in #38255.

Validation

• SparkDatasetPortablePipelineTranslatorTest covers one piece at a time. Flatten with mismatched input coders, GroupByKey with fixed windows and Reshuffle are translated on injected Datasets. Impulse, a multi-output stage and a side input go through the translator on a small pipeline and run in the embedded SDK harness. The unbounded, state and timer rejections are asserted on the translator directly. The gate excludes UsesSideInputs, so the side-input test here is the only side-input coverage. SparkDatasetTranslationContextTest covers the persist decisions and that only unconsumed Datasets are evaluated. SparkDatasetPortableExecutionTest keeps one end to end run through the job invoker with --streaming and --useStructuredStreaming, as the gate task sets them, and checks that the streaming option is left untouched. Fifteen tests, green on :runners:spark:4:test under JDK 17 and on :runners:spark:3:test under JDK 11.
• Sessions, sliding windows and non-default timestamp combiners are unverified on this backend, and the streaming exclusion list also drops several of those ValidatesRunner tests, so the first gate run will not close that gap.
• I have not run the full PVR suite locally, so the first run of the new workflow is the real signal.
• The Spark 3 compile runs with the nullness checker enabled. The two new main-source files suppress nullness at class level, matching the surrounding translator code.

Notes for review

Three things I would particularly like a second opinion on:

• The MetricsAccumulator.init overload. The alternative was to read useStructuredStreaming inside MetricsAccumulator, which would tie a shared class to a portable-only flag.
• Persisted Datasets are never unpersisted and the SparkSession is never cleared. That is fine when PortableBatchMode.stop() stops the context, and it accumulates when the context is provided or reused.
• Side inputs are collected eagerly during translation, which makes translation itself trigger Spark actions.

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
• Update CHANGES.md with noteworthy changes. Holding this until the flag is more than a preview.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

[github-actions]

Assigning reviewers:

R: @Abacn for label build.

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]

We shouldn't overwrite pipelineOptions.setStreaming here as it doesn't sounds semantically correct. We should prevent the logic ends up detecting "isStreaming" to avoid expansion involving DStream. The correct flag is "useStructuredStreaming"

"The Dataset backend rejects unbounded input at translation" is current supporting status and isn't a spec

[Eliaaazzz]

Fixed in f4a2d39. The runner no longer touches the streaming option. With useStructuredStreaming set it picks the Dataset translator and skips the DStream translator and streaming context. MetricsAccumulator.init got an explicit useCheckpoint argument, so the Dataset path never opens the DStream metrics checkpoint and the DStream path keeps its old condition. I also reworded the comments, option description and README to describe the unbounded rejection as current status.

[Abacn]

We can have a few full pipeline test but the majority of runner unit tests are targeted, testing component but not a pipeline. The latter is covered by validate runner tests.

[Eliaaazzz]

Fixed in f4a2d39. I kept one end to end test through the job invoker, with --streaming and --useStructuredStreaming like the gate task, and moved the rest into component tests. SparkDatasetPortablePipelineTranslatorTest translates single transforms on injected Datasets and runs stages in the embedded harness, SparkDatasetTranslationContextTest checks the persist and evaluation logic, and the rejection cases are asserted on the translator directly.