[WriteFiles] Support direct writes of large streaming bundles to temp files (#40533) by tejasiyer-dev · Pull Request #40543 · apache/beam

tejasiyer-dev · 2026-10-13T00:03:28Z

Summary

In streaming pipelines, WriteFiles shuffles every element either via GroupByKey (when a shard count or runner-determined sharding is used) or GroupIntoBatches.withShardedKey() (when auto-sharded) before writing temporary files. When upstream bundles are already large enough to produce well-sized temporary files on their own, shuffling all elements incurs unnecessary network and state overhead.

This PR adds an optional direct-write byte threshold that allows sufficiently large streaming bundles to write directly to temporary files, while routing smaller bundles through the existing grouped write path. Similar to
IcebergIO direct write (#36720).

Changes

WriteFiles:
- Added withDirectWriteByteThreshold(long) to configure the byte threshold for direct bundle writes on unbounded inputs.
- Added WriteShardedBundlesWithDirectWrite, which buffers incoming elements per bundle up to thresholdBytes using the input coder's byte-size estimation:
  - Large bundles (>= thresholdBytes): Written directly to temporary files via DirectWriteLargeBundlesFn without shuffling. If a bundle touches more than maxNumWritersPerBundle distinct (window, destination) writers, excess records spill over to the grouped path.
  - Small bundles (< thresholdBytes) and spillover records: Routed through WriteShardedBundlesToTempFiles (when a shard count is configured) or GroupIntoBatches.withShardedKey() (when unsharded / auto-sharded), then merged with the direct-write FileResults before FinalizeTempFileBundles.
TextIO & FileIO:
- Added .withDirectWriteByteThreshold(long) pass-through methods to TextIO.Write, TextIO.TypedWrite, and FileIO.Write.
DataflowRunner:
- Updated StreamingShardedWriteFactory to preserve getDirectWriteByteThreshold() when constructing the replacement WriteFiles transform on Dataflow Runner v1.
Metrics:
- Added counters under WriteFiles: directWriteBundles and groupedWriteBundles.

Testing

Added pipeline tests in WriteFilesTest:
- testDirectWriteAllBundlesLarge (large bundles bypass shuffle and write directly)
- testDirectWriteAllBundlesSmall (small bundles fall back to the grouped write path)
- testDirectWriteUnshardedMixedBundles (exercises both direct-write and auto-sharded fallback paths in the same pipeline)
- testDirectWriteNonDefaultTriggerLargeBundles & testDirectWriteNonDefaultTriggerSmallBundles (verifies non-default user triggers across both paths)
- testDirectWriteDisabledWithSideInputs (verifies fallback when dynamic destinations use side inputs)
Added Dataflow runner override tests in DataflowRunnerTest:
- testStreamingWriteWithDirectWriteByteThreshold
- testStreamingWriteDirectWriteThresholdDefaultNull
Added configuration pass-through tests in TextIOWriteTest and FileIOTest.

Benchmark

30-minute Dataflow streaming run (5x n1-standard-4, ~80 MB/s input, 1-minute FixedWindows, TextIO.write().withWindowedWrites()):

Before (direct-write-before): 2026-10-12_12_51_21-9229133822682131587
After (direct-write-after): 2026-10-12_12_51_50-15846649412145844281

Metric	Before (Default)	After (`directWriteByteThreshold = 1 MB`)
Data shuffled in `WriteFiles`	125.18 GB (`32.5M` records)	516.84 KB (`4,040` file paths)
System latency	~2.5 min (`2.37 GB` backlog)	3–5 sec (`0` backlog)
Output files in GCS	`300` files (`445 MB` avg)	`4,040` files (`35.2 MB` avg)

Before vs. After Shuffle Step Input (125.18 GB vs. 516.84 KB):
https://screenshot-v2.corp.google.com/image/504bv5gm3tlro.png
https://screenshot-v2.corp.google.com/image/3ndvsk8951v6g.png

Before vs. After System Latency (~2.5 min vs. 3–5s):
https://screenshot-v2.corp.google.com/image/4mv9vipuam300.png
https://screenshot-v2.corp.google.com/image/6507lv1t9mj68.png

… files

…ctWriteByteThreshold

github-actions · 2026-10-13T01:39:53Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

codecov · 2026-10-13T21:31:55Z

Codecov Report

❌ Patch coverage is 4.52675% with 232 lines in your changes missing coverage. Please review.
✅ Project coverage is 64.98%. Comparing base (76f6e6d) to head (3e064fd).
⚠️ Report is 319 commits behind head on master.

Files with missing lines	Patch %	Lines
...c/main/java/org/apache/beam/sdk/io/WriteFiles.java	1.33%	221 Missing and 1 partial ⚠️
...g/apache/beam/runners/dataflow/DataflowRunner.java	55.55%	3 Missing and 1 partial ⚠️
...e/src/main/java/org/apache/beam/sdk/io/FileIO.java	25.00%	2 Missing and 1 partial ⚠️
...e/src/main/java/org/apache/beam/sdk/io/TextIO.java	40.00%	1 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #40543      +/-   ##
============================================
+ Coverage     56.16%   64.98%   +8.81%     
- Complexity     2288    13977   +11689     
============================================
  Files          1124     1729     +605     
  Lines        178336   108641   -69695     
  Branches       1489    11332    +9843     
============================================
- Hits         100170    70595   -29575     
+ Misses        75645    33838   -41807     
- Partials       2521     4208    +1687

Flag	Coverage Δ
go	`?`
java	`64.98% <4.52%> (-10.81%)`	⬇️
python	`?`

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.

github-actions · 2026-10-13T21:38:00Z

Assigning reviewers:

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

This pull request likely touches a core component ("core" label). Please review with scrutiny.

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

tejasiyer-dev · 2026-10-13T22:24:43Z

R: @arunpandianp

github-actions · 2026-10-13T22:25:28Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

arunpandianp · 2026-10-13T22:31:07Z

Can we move sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BundleLifter.java into sdks/java/io/common and reuse it?

arunpandianp · 2026-10-13T22:31:13Z

cc: @tomstepp

tomstepp · 2026-10-13T22:36:26Z

Agreed, it would be nice to move and reuse BundleLifter

ahmedabu98 · 2026-10-13T22:41:50Z

+1 to that, but please make sure it doesn't break streaming update compatibility for the iceberg sink

tejasiyer-dev · 2026-10-13T23:32:51Z

Can we move sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BundleLifter.java into sdks/java/io/common and reuse it?

I think moving BundleLifter to sdks/java/io/common would create a circular dependency. sdks/java/io/common already depends on sdks/java/core and WriteFiles is in sdks/java/core, so to use BundleLifter from io/common, core would also have to depend on io/common.

Would moving BundleLifter to core be better? Iceberg already depends on Core so both Iceberg and WriteFiles could share it with no new dependencies

chamikaramj · 2026-10-19T16:47:08Z

@ahmedabu98 PTAL.

Also, seems like BundleLifter move is unrelated and probably should be a separate PR ?

arunpandianp · 2026-10-19T18:57:41Z

Would moving BundleLifter to core be better?

I think that is okay. We could move it to sdks/java/core/src/main/java/org/apache/beam/sdk/util and mark it with @org.apache.beam.sdk.annotations.Internal

ahmedabu98

Generally LGTM, just a nit about modularization and small question about the experiment flag

ahmedabu98 · 2026-10-20T22:21:46Z

+                      TupleTagList.of(ImmutableList.of(smallBundlesTag, BAD_RECORD_TAG))));
+      addErrorCollection(routed);
+
+      PCollection<UserT> smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder());


Is there a way to reuse the existing pipeline graph (built in WriteFilesResult's expand) and pass smallBundles through to it? Instead of attempting to recreate the same steps here?

I see you did a good amount for that using writeAutoShardedRecordsToTempFiles(...). Wondering if other parts can be similarly modularized

Partly. The small bundles now reuse the existing write steps, and both paths share the grouping steps and batch-size defaults through new helpers.

The small bundles can't go through the entire existing graph, though. If they did, the direct-write files and the small-bundle files would each go through a separate finalize step, and their file names would collide. Instead, both sets of temp files are merged before finalize, so a single finalize step assigns unique shard numbers to every file in the window.

ahmedabu98 · 2026-10-20T22:27:01Z

 * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)).
 * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)).
 * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)).
+* (Java) `TextIO`, `FileIO`, and `WriteFiles`: streaming writes can skip the shuffle for large bundles by writing them directly to temp files (`withDirectWriteByteThreshold` or experiment `streaming_file_direct_write_threshold_bytes`). Off by default ([#40533](https://github.com/apache/beam/issues/40533)).


Not seeing thestreaming_file_direct_write_threshold_bytes flag in the code here. is it in the Dataflow service?

No, it's not in the Dataflow service. I removed that experiment so direct write is configured only through withDirectWriteByteThreshold (matching Iceberg), but forgot to update CHANGES.md. Fixed now, thanks!

arunpandianp · 2026-10-20T22:47:59Z

+      for (BufferedRecord<UserT> record : buffer) {
+        c.output(smallBundlesTag, record.value, record.timestamp, record.window);
+      }
+      for (BufferedRecord<UserT> record : spilled) {
+        c.output(smallBundlesTag, record.value, record.timestamp, record.window);


seems like this will reorder outputs. Can we preserve the order?

The two loops are no longer back to back and each one now sits inside its own branch, so only one runs per bundle.

Direct mode: the buffer is flushed when we switch, so only spilled is emitted.
Grouped mode: nothing is written directly, so only buffer is emitted.

Each list is emitted in arrival order, so nothing gets reordered.

arunpandianp · 2026-10-20T22:50:22Z

+    public void finishBundle(FinishBundleContext c) throws Exception {
+      if (directMode) {
+        directBundles.inc();
+      } else if (!buffer.isEmpty()) {


since directMode, !buffer.isEmpty() can both be true, do we need the else here?

Yes, the else is needed, because directMode and !buffer.isEmpty() can't both be true: the buffer is flushed and cleared the moment we switch to direct mode.

…sharded paths

into gcs-direct-write

tejasiyer-dev · 2026-11-03T22:01:15Z

@ahmedabu98 PTAL.

Also, seems like BundleLifter move is unrelated and probably should be a separate PR ?

@chamikaramj @arunpandianp @ahmedabu98, I've moved BundleLifter to sdk/util with @internal in a separate PR: #40978. Making WriteFiles reuse BundleLifter will be a follow-up once both are merged.

Could you take another look and approve the workflows so CI can run on the latest changes? Thanks!

arunpandianp · 2026-11-05T07:56:04Z

+        return;
+      }
+      DestinationT destination = maybeDestination.destination;
+      WriterKey<DestinationT> key = new WriterKey<>(record.window, record.pane, destination);


since the goal is to write large files when possible, should we be grouping records by WriterKey and see if the grouped records meet the threshold, instead of applying the threshold to the whole bundle across all windows and destinations?

…reached. Evaluate the direct-write threshold per WriterKey instead of per bundle.

tejasiyer-dev · 2026-11-11T03:36:38Z

#40978 (moving BundleLifter from Iceberg to sdk/util) has merged, so I've updated this PR to use it.

Per-WriterKey threshold: BundleLifter now takes an optional key function, and WriteFiles uses it to check the threshold per window/pane/destination (WriterKey). Without a key function, BundleLifter still checks the threshold against the whole bundle.

Iceberg can also check the threshold per table/partition instead of per bundle by passing a key function but I've left its threshold behavior unchanged for now. Maybe it can be incorporated in a follow up PR.

Streaming: Once a key reaches the threshold, its buffered elements are flushed to the large output, and later elements for that key are output directly without sizing or buffering.

@arunpandianp @ahmedabu98 @scwhittle PTAL, and could someone also approve the workflows and merge afterwards? Thanks!

arunpandianp · 2026-11-13T03:59:24Z

+    Long directWriteThreshold = getDirectWriteByteThreshold();
+    if (directWriteThreshold != null
+        && getSideInputs().isEmpty()
+        && input.isBounded() == IsBounded.UNBOUNDED) {


We should disable direct write when the window is merging like in https://github.com/tejasiyer-dev/beam/blob/b1fe445521135fe623ef2d69ae1dd11361dbf1e4/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java#L573

when directWriteThreshold is present and if we are not enabling it, should we throw an error or fallback? When falling back adding a log saying direct write is disabled with a reason would be nice to have.

Good catch! I disabled direct write for merging windows using the same needsMerge() check.

When directWriteThreshold is set but direct write can't be used, I think we should fall back to the default grouped path and log a warning with the reason. The pipeline still writes all the same records to the same files, it just doesn't get the optimization.

…with log warning if directWriteThreshold is set but direct write can't be used. Also evict BundleLifter keys in one pass when over the memory cap.

tejasiyer-dev added 2 commits October 12, 2026 18:55

[WriteFiles] Support direct writes of large streaming bundles to temp…

0a49326

… files

Remove experiment flag and configure direct write solely via withDire…

d9d8af4

…ctWriteByteThreshold

github-actions Bot added java runners dataflow core labels Oct 13, 2026

Merge branch 'master' into gcs-direct-write

9485325

Merge branch 'master' into gcs-direct-write

3e064fd

github-actions Bot requested a review from ahmedabu98 October 13, 2026 21:38

github-actions Bot added the Next Action: Reviewers label Oct 13, 2026

ahmedabu98 approved these changes Oct 20, 2026

View reviewed changes

arunpandianp reviewed Oct 20, 2026

View reviewed changes

Clarify finishBundle branches and remove stale lines from CHANGES.md

2f21a80

github-actions Bot added the size/xl XL: >=1000 lines changed label Nov 3, 2026

tejasiyer-dev added 3 commits November 2, 2026 17:24

Merge branch 'master' into gcs-direct-write

8dab8bc

Share grouping and batch-default logic between direct-write and auto-…

14904e4

…sharded paths

Merge branch 'gcs-direct-write' of https://github.com/tejasiyer-dev/beam

dcb9079

into gcs-direct-write

This was referenced Nov 3, 2026

[Task]: Move BundleLifter to sdk/util and reuse it in WriteFiles #40976

Open

Move BundleLifter from Iceberg to sdk/util #40978

Merged

arunpandianp reviewed Nov 5, 2026

View reviewed changes

tejasiyer-dev added 2 commits November 10, 2026 19:18

Merge remote-tracking branch 'upstream/master' into gcs-direct-write

de6452b

Stream elements to the large output once BundleLifter's threshold is …

b1fe445

…reached. Evaluate the direct-write threshold per WriterKey instead of per bundle.

arunpandianp reviewed Nov 13, 2026

View reviewed changes

Disabled Direct Write for merging windows. Fall back to default path …

fec8a2e

…with log warning if directWriteThreshold is set but direct write can't be used. Also evict BundleLifter keys in one pass when over the memory cap.

Conversation

tejasiyer-dev commented Oct 13, 2026

Summary

Changes

Testing

Benchmark

Uh oh!

github-actions Bot commented Oct 13, 2026

Uh oh!

codecov Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Oct 13, 2026

Uh oh!

tejasiyer-dev commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026

Uh oh!

arunpandianp commented Oct 13, 2026

Uh oh!

arunpandianp commented Oct 13, 2026

Uh oh!

tomstepp commented Oct 13, 2026

Uh oh!

ahmedabu98 commented Oct 13, 2026

Uh oh!

tejasiyer-dev commented Oct 13, 2026

Uh oh!

chamikaramj commented Oct 19, 2026

Uh oh!

arunpandianp commented Oct 19, 2026

Uh oh!

ahmedabu98 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tejasiyer-dev commented Nov 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tejasiyer-dev commented Nov 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codecov Bot commented Oct 13, 2026 •

edited

Loading