[WriteFiles] Support direct writes of large streaming bundles to temp files (#40533)

[tejasiyer-dev]

Fixes #40533

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

[github-actions]

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]

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 160 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]

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]

R: @arunpandianp

[github-actions]

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]

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]

cc: @tomstepp

[tomstepp]

Agreed, it would be nice to move and reuse BundleLifter

[ahmedabu98]

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

[tejasiyer-dev]

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]

@ahmedabu98 PTAL.

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

[arunpandianp]

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