Skip to content

[Spark] Support SDF self-checkpointing and bundle finalization on the portable runner#39331

Draft
Eliaaazzz wants to merge 1 commit into
apache:masterfrom
Eliaaazzz:spark-portable-sdf
Draft

[Spark] Support SDF self-checkpointing and bundle finalization on the portable runner#39331
Eliaaazzz wants to merge 1 commit into
apache:masterfrom
Eliaaazzz:spark-portable-sdf

Conversation

@Eliaaazzz

@Eliaaazzz Eliaaazzz commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

Reading a splittable DoFn that self-checkpoints has never worked on the portable Spark runner: no bundle checkpoint handler was registered, so the SDK's first residual raised UnsupportedOperationException: The ActiveBundle does not have a registered bundle checkpoint handler and the read died on its first bundle. This is why the SDF tests are excluded and skipped for Spark.

fixes #19468

Also addresses #19517 (bundle finalization).

Batch

SparkExecutableStageFunction registers a checkpoint handler that collects residual roots, then re-feeds them in fresh bundles until the SDK returns none, resuming each residual at its own requested time. A bounded restriction always runs out, so the loop terminates. An unbounded residual would loop forever in a batch task, so it is rejected with a clear error rather than drained, which keeps the previous fail-fast behaviour for that case.

The stage also registers BundleFinalizationHandlers.inMemoryFinalizer, so bundle finalization callbacks now run (#19517).

Streaming

Residuals leave the stage on a reserved union tag instead of being drained in place, and SdfResidualRelay relays them on the driver: each micro-batch's residuals are collected, held until their requested resume time, and fed back as stage input in a later micro-batch, so the read advances across micro-batches instead of occupying one task forever. The relay derives the stage's watermark from the residuals' output watermarks, which is what lets downstream event-time windows fire, and bounds it by the stage's upstream sources so it can never outrun live input. Relays are keyed per job so concurrent jobs in one job server stay isolated.

Known limitation of the streaming relay

The relay reports a watermark every micro-batch, while an impulse reports once with a fixed synchronized processing time. A GroupByKey whose inputs are flattened from both therefore sees two sources with different synchronized processing times, which SparkTimerInternals.forStreamFromSources rejects with Synchronized time is expected to keep synchronized across sources. Grouping an SDF's output on its own is unaffected, since the relay replaces the stage's stream sources, and that is the case the tests and demos below cover. This is the main reason UsesUnboundedSplittableParDo stays excluded for streaming.

I tried making an impulse report every micro-batch to reconcile the two, and it regressed 69 streaming VR tests, so it is not in this PR. The single shared timestamp an impulse reports is what keeps all impulses in a job mutually consistent; making each report per batch lets them drift apart. Reconciling the two source kinds properly looks like its own change.

Testing

  • Removes UsesBoundedSplittableParDo and UsesBundleFinalizer from the batch portable VR excludes; beam_PostCommit_Java_PVR_Spark_Batch passes with them enabled on this PR.
  • Unskips the SDF and bundle finalization tests in spark_runner_test.py; 9 pass, 0 fail locally.
  • New unit tests cover the streaming residual emission and the unbounded-in-batch rejection.
  • Manually: a Python UnboundedSource read streams continuously on a local job server and on a standalone cluster with a separate worker process; FixedWindows(10s) + GroupByKey fires with exact per-window counts; a 34 minute soak held steady throughput with no memory growth.

I could not run the Java validates-runner suites locally, since they do not currently run on Windows (#39332), so I am relying on the CI triggers in this PR for that signal.

beam_PostCommit_Python_ValidatesRunner_Spark is red here, but it is also red on master, failing in test_assert_that with IllegalAccessError: class org.apache.spark.storage.StorageUtils$ cannot access class sun.nio.ch.DirectBuffer, which looks like Spark needing --add-opens=java.base/sun.nio.ch=ALL-UNNAMED on newer JDKs. It does not appear related to this change. Happy to file that separately if it is not already tracked.

Open questions for reviewers

  • test_unbounded_source_read stays skipped, but no longer for SDF reasons: the source self-terminates while a streaming pipeline here runs until its streaming timeout, so the pipeline never reports completion. Worth a separate issue?
  • Residuals taken for a micro-batch are retired when that batch's collection job reports back, rather than after every output for the batch has succeeded. Since a failed job stops the streaming context and there is no driver recovery on this path, I judged this not worth the extra machinery, but I would like a second opinion.
  • Streaming still excludes UsesUnboundedSplittableParDo, since the tests assert through assert_that, which flattens an impulse-derived stream into the group and hits the synchronized-time limitation above. Enabling it needs that resolved first.
  • Happy to split this into a batch-only PR plus a streaming follow-up if that is easier to review. The batch half is the part CI already covers end to end.

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

@codecov

codecov Bot commented Jul 15, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 54.76%. Comparing base (de21b1c) to head (b1428e7).

Additional details and impacted files
@@            Coverage Diff            @@
##             master   #39331   +/-   ##
=========================================
  Coverage     54.76%   54.76%           
  Complexity     1716     1716           
=========================================
  Files          1066     1066           
  Lines        169075   169075           
  Branches       1255     1255           
=========================================
+ Hits          92589    92591    +2     
+ Misses        74269    74267    -2     
  Partials       2217     2217           
Flag Coverage Δ
python 79.71% <ø> (+<0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@Eliaaazzz Eliaaazzz force-pushed the spark-portable-sdf branch from b1428e7 to 17df8a3 Compare July 15, 2026 05:04
… portable runner

The portable Spark runner never registered a bundle checkpoint handler, so a
splittable DoFn that self-checkpoints failed on its first bundle and unbounded
reads could not run at all.

In batch, the executable stage function now collects residual roots and re-feeds
them in fresh bundles until the SDK returns none, resuming each residual at its
requested time. Unbounded residuals are rejected, since draining them would never
terminate. The stage also registers an in-memory bundle finalizer.

In streaming, residuals leave the stage on a reserved union tag and are relayed
on the driver: they are held until their resume time, fed back as stage input in
a later micro-batch, and their output watermarks advance the stage's watermark so
downstream event-time windows fire. The relay bounds its watermark by the stage's
upstream sources, and is keyed per job so concurrent jobs stay isolated.

Unskips the splittable DoFn and bundle finalization tests for the Spark runner.
@Eliaaazzz Eliaaazzz force-pushed the spark-portable-sdf branch from 17df8a3 to d99c843 Compare July 15, 2026 14:25
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.

Spark portable runner: support SDF

1 participant