Skip to content

[Java] FirestoreIO: ramp up writes from a pipeline-wide start time - #40391

Open
nbali wants to merge 4 commits into
apache:masterfrom
nbali:firestore-rampup-pipeline-start
Open

nbali wants to merge 4 commits into
apache:masterfrom
nbali:firestore-rampup-pipeline-start

Conversation

@nbali

@nbali nbali commented Oct 2, 2026 •

Copy link
Copy Markdown
Contributor

Addresses #27170 (alleviates it, does not fix it).

What

FirestoreIO batch writes now grow their ramp-up budget from one pipeline-wide start time, instead of from each DoFn instance's first write.

  • FirestoreV1.BatchWriteWithSummary and BatchWriteWithDeadLetterQueue create a singleton side input holding clock.instant(), computed once per pipeline (Create → MapElements → View.asSingleton()), and pass it to the write ParDo.
  • FirestoreV1WriteFn.BaseBatchWriteFn.processElement reads the side input and hands it to RpcQos.setRampUpStart(Instant).
  • RpcQosImpl.WriteRampUp uses that instant as its ramp-up start. If it's never set, behavior is unchanged: the ramp-up starts lazily at the first write.

Why

#27170 describes how slow the write ramp-up is for small jobs. On top of the low starting budget (500 / hintMaxNumWorkers, 1 write/s with the default hint), each RpcQos instance (one per DoFn instance, created in @Setup) starts its own ramp-up clock at its first write. Every DoFn instance created later starts again at the base budget. This happens on autoscaled workers, when concurrency grows, and after a failed bundle discards its DoFn instances. Time the pipeline spends before its first write doesn't count either.

With a shared start time:

  • DoFn instances created later join at the budget the pipeline has already grown to, instead of starting over at 1 write/s;
  • time spent before the first write (for example a long read or transform stage) counts toward the ramp-up.

Trade-off, the same one DatastoreIO accepts: if the first write happens long after the pipeline starts, the budget has already grown, so those writes get less ramp-up protection.

Where it comes from

This copies the approach DatastoreIO already uses. DatastoreV1.Mutate.expand creates a "Generate start timestamp" side input (Create.of("side input") → MapElements.via(s -> Instant.now()) → View.asSingleton()) and passes it to RampupThrottlingFn. Its processElement reads c.sideInput(firstInstantSideInput) and computes the budget relative to that instant. This PR applies the same pattern to the Firestore write path. The Firestore ramp-up lives inside RpcQosImpl rather than in a separate DoFn, so the instant is passed into it.

Pipeline update compatibility

The write transforms gain the side-input transforms, which changes the pipeline graph. Following the existing pattern (for example Reshuffle), the previous expansion is kept unchanged in expand_2_78_0, and expand delegates to it when --updateCompatibilityVersion is 2.78.0 or earlier. Pipelines that need to update a running job from an older SDK can set it and keep the previous shape and behavior. The cutoff is 2.79.0 because 2.78.0 was cut without this change.

Why this alleviates rather than fixes the referenced issue

The rate formula is unchanged. A small job that writes from the start still begins at 500 / hintMaxNumWorkers writes per second and grows by 50% every 5 minutes, so with the default hint it is still slow for the first 10–15 minutes. This PR removes the extra slowdown from per-instance restarts and makes Firestore behave like Datastore. Changing the default growth is left for a follow-up.

Testing

  • RpcQosTest.rampUp_growsFromProvidedStart: with a start 20 minutes before the first write, the first batch is sized from the grown budget (3), not the base budget (1).
  • FirestoreV1FnBatchWriteWithSummaryTest.rampUpStartIsReadFromSideInput: the DoFn reads the side input and passes it to RpcQos.setRampUpStart.
  • FirestoreV1BatchWriteRampUpStartTest: both write transforms add the side input to their batchWrite ParDo; with updateCompatibilityVersion=2.78.0 they keep the previous shape, and with 2.79.0 they add it.

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 for a list of workflows and how to trigger them.

Batch writes now grow their 500/50/5 ramp-up budget from one start time
shared by the whole pipeline, provided as a singleton side input, instead
of from each DoFn instance's first write. This follows DatastoreV1, so
DoFn instances created later (autoscaling, retried bundles) no longer
restart the ramp-up. The side input is omitted when updateCompatibilityVersion
is below 2.78.0.

Addresses apache#27170.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BofSrVTSEGJvEzsSNXfU2o
@github-actions

github-actions Bot commented Oct 2, 2026

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @ahmedabu98 for label java.

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

@github-actions

Copy link
Copy Markdown
Contributor

Reminder, please take a look at this pr: @ahmedabu98

@github-actions

Copy link
Copy Markdown
Contributor

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @kennknowles for label java.
Backup reviewer: damccorm (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

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)

*/
private static @Nullable PCollectionView<Instant> createRampUpStartView(
PCollection<?> input, JodaClock clock) {
if (StreamingOptions.updateCompatibilityVersionLessThan(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The way we should use this flag is that we should clone the prior expand method to expand_2_77(...).

  • The expand_2_77(...) method will always have an unchanged topology.
  • The new expand delegates to expand_2_77 in its first line(s) and after that has the new, totally separate, topology.
  • This eliminate the "action at a distance" of this return value being null causing the graph topology to be altered latter.

Check other uses of this pipeline option for examples. Particularly a simple case I can think of is Reshuffle.

They can share sub-transforms or helper functions, but the forking based on update compatibility should happen at the highest level. This is a case where redundancy is a benefit, and pushing the cyclomatic complexity to the earliest possible point is the way to do it.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, done in 53a7ca9. Each batch write transform keeps the previous expansion unchanged in expand_2_78_0, and expand delegates to it in its first lines when updateCompatibilityVersion is 2.78.0 or earlier, as Reshuffle does. The new path no longer has a nullable side input.

I named it _2_78_0 rather than _2_77 because 2.78.0 was cut without this change, so the new shape starts with 2.79.0. I also merged master, which now includes #40390.

@github-actions

Copy link
Copy Markdown
Contributor

Reminder, please take a look at this pr: @kennknowles @ahmedabu98

Per review, the previous batch write expansion is kept unchanged in
expand_2_78_0, and expand delegates to it when updateCompatibilityVersion
is 2.78.0 or earlier, instead of threading a nullable side input through
the new expansion. The cutoff moves to 2.79.0 since 2.78.0 was cut
without this change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BofSrVTSEGJvEzsSNXfU2o
@github-actions github-actions Bot added the size/m M: 50-249 lines changed label Nov 1, 2026
@nbali

nbali commented Nov 1, 2026

Copy link
Copy Markdown
Contributor Author

The failing job seems to be a flaky one, unrelated to this PR: it's beam_PreCommit_Java_GCP_IO_Direct_BigQuery, which only runs the BigQuery integration tests. This PR only touches the Firestore package, and the same job failed on master on Oct 30 (https://github.com/apache/beam/actions/runs/41065748867).

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.

3 participants