tejasiyer-dev opened a new pull request, #40543:
URL: https://github.com/apache/beam/pull/40543
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 `FileResult`s 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
