[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

ahmedabu98 · 2026-10-20T22:24:59Z

 * (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)).
+* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) for configuring AWS service clients such as S3 ([#40528](https://github.com/apache/beam/issues/40528)).


Remove irrelevant line? Maybe it's getting accidentally picked up from master

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?

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?

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?

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

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!

ahmedabu98 Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

ahmedabu98 Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

ahmedabu98 Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

ahmedabu98 Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Oct 20, 2026

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