Skip to content

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

Draft
Eliaaazzz wants to merge 1 commit into
apache:masterfrom
Eliaaazzz:spark4-portable-dataset-poc
Draft

Eliaaazzz wants to merge 1 commit into
apache:masterfrom
Eliaaazzz:spark4-portable-dataset-poc

Conversation

@Eliaaazzz

Copy link
Copy Markdown
Contributor

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 rejected at translation with a message pointing at the tracking issue, so nothing silently runs on a path that cannot support it yet.

The translation context attaches to the SparkContext that SparkContextFactory already created, so the job server code path, the metrics accumulator and result handling are unchanged. 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

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

  • SparkDatasetPortableExecutionTest runs Impulse, ParDo, Flatten, GroupByKey and PAssert end to end on the Dataset backend with the embedded SDK harness, and asserts that unbounded input and a stateful ParDo are rejected at translation with the expected message. Three tests, green on :runners:spark:4:test under JDK 17 and on :runners:spark:3:test under JDK 11.
  • Known coverage gaps in this PR, all of which the PVR run is meant to cover: multi-output stages, so the oneOfEncoder demux is only ever built with one alternative; Reshuffle; the Flatten re-encode branch for mismatched coders; side inputs, which the streaming suite also excludes via UsesSideInputs; and the timer half of the state and timer rejection.
  • Only the global window is covered by a test here. Sessions, sliding windows and non-default timestamp combiners are unverified on this backend.
  • 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 files suppress nullness at class level, matching the surrounding translator code.

Notes for review

Three things I would particularly like a second opinion on:

  • Setting pipelineOptions.setStreaming(false) in the new branch. MetricsAccumulator.init reads the option directly, not the mode the runner resolved, so without this it takes its checkpoint branch and seeds the accumulator from a checkpoint file this path never writes back.
  • 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)

Build python source distribution and wheels
Python tests
Java tests
Go tests

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

…e pipelines

Behind --useStructuredStreaming, SparkPipelineRunner translates the fused
Runner API pipeline into Spark Datasets. Impulse, Flatten, Reshuffle and
GroupByKey are Dataset operations, and executable stages run through the
existing Fn API bridge inside mapPartitions. Bounded pipelines run as batch
Datasets. Unbounded input, user state and timers are rejected at translation.

The branch clears the streaming option so the metrics accumulator and anything
else reading it agree with how the job runs.

The validatesPortableRunnerStructuredStreaming task runs the streaming
PortableValidatesRunner suite on this backend and is the exit gate for making
it the default for portable streaming pipelines.
@Eliaaazzz
Eliaaazzz force-pushed the spark4-portable-dataset-poc branch from d1b25b9 to 58719bd Compare September 15, 2026 12:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant