atognolag opened a new pull request, #40653:
URL: https://github.com/apache/beam/pull/40653
> Supersedes #38406 (which was auto-closed by the stale bot while in draft).
## Summary
Adds end-to-end support for sorting records on write according to an Iceberg
table's `SortOrder` (`.withSortedWrites()` / Managed IO `sort_on_write: true`)
as well as `DistributionMode.RANGE`
(`.withDistributionMode(DistributionMode.RANGE)` / Managed IO
`distribution_mode: "RANGE"`) across both Batch and Streaming pipelines:
1. **Order-Preserving Binary Key Encoding & External Spill-to-Disk Sorting
(`IcebergRowSorter`)**:
- Encodes composite Iceberg `SortOrder` keys into lexicographically
comparable `byte[]` arrays supporting `ASC`/`DESC`, `NULLS_FIRST`/`NULLS_LAST`,
all Iceberg primitive/logical types (`BOOLEAN`, `INTEGER`, `LONG`, `FLOAT`,
`DOUBLE`, `DECIMAL`, `STRING`, `UUID`, `FIXED`, `BINARY`, `DATE`, `TIME`,
`TIMESTAMP`, `TIMESTAMPTZ`, `EnumerationType`), and non-identity sort
transforms (`month`, `day`, `hour`, `year`, `bucket`, `truncate`).
- Uses `BufferedExternalSorter` (`NATIVE`) with bounded per-instance
memory (`64 MB` for single-partition grouped writes, `16 MB` per destination in
multi-destination bundle writes) to spill to local disk when a partition or
bundle exceeds memory.
2. **Zero-Shuffle Local Bundle Sorting for `DistributionMode.NONE` +
`withSortedWrites()` (`RecordWriterManager`)**:
- Buffers incoming rows per bundle in `IcebergRowSorter` keyed by
`(partitionPath, sortKey)` without opening any `RecordWriter`s during
`@ProcessElement`.
- At `@FinishBundle`, flushes sorted rows one partition at a time—keeping
only a single active `RecordWriter` open at any instant—eliminating any network
shuffle for `DistributionMode.NONE` in both Batch and Streaming.
3. **Dynamic Quantile Sharding for `DistributionMode.RANGE`
(`AssignDestinationsAndPartitions`)**:
- Supports both a custom
`withDistributionFunction(SerializableFunction<Row, Integer>)` and automatic
quantile sharding in Batch when `distributionFunction` is omitted.
- Uses **Efraimidis–Spirakis (`A-Res`) weighted reservoir sampling
without replacement** (`2,048` samples per `(destination, partition)`) combined
with estimated partition byte volume to compute balanced quantile split
boundaries sized to `DEFAULT_BYTES_PER_FILE` (`512 MiB`).
- Avoids materializing `byte[]` sort keys across the stage barrier for
un-sampled rows; only encodes sort keys during shard assignment for partitions
that split into multiple shards.
---
## Distribution & Sorting Modes Matrix
| Mode | Configuration | Network Shuffles (Batch) | Network Shuffles
(Streaming) | File Sort & Range Guarantees |
| :--- | :--- | :--- | :--- | :--- |
| **`NONE` (Unsorted)** | `.withDistributionMode(NONE)` | **0** (Direct
bundle write; 1 `GroupByKey` only if open writers overflow `MAX_NUM_WRITERS`) |
**0** (under `directWriteByteLimit`) or **1** (`GroupByKey` across `shardingKey
% 50`) | Files are written in arrival order. |
| **`NONE` (Sorted)** | `.withDistributionMode(NONE).withSortedWrites()` |
**0** (In-bundle external sort per `(dest, partition)`) | **0** (under
`directWriteByteLimit`) or **1** (`GroupByKey` across `shardingKey % 50`) |
Every file is internally sorted by `table.sortOrder()`; min/max key ranges
across different worker bundles may overlap. |
| **`HASH`** | `.withDistributionMode(HASH)` (+ optional
`.withSortedWrites()`, `.withAutosharding()`) | **1** (`GroupByKey` by `(dest,
partition)`) | **1** (`GroupIntoBatches` by `(dest, partition)`) | Consolidates
files per partition. With `.withSortedWrites()` (without autosharding), files
within each partition have strictly non-overlapping sort-key ranges. |
| **`RANGE` (Dynamic Quantiles)** | `.withDistributionMode(RANGE)` (no
`distributionFunction`, Batch) | **1** (`GroupByKey` by `(dest, partition,
quantileShard)` + lightweight quantile side-input `Combine.perKey`) |
*(Requires `withDistributionFunction` or `.withAutosharding()`)* | Dynamically
sizes shard count to `totalPartitionBytes / targetFileSize` using weighted
reservoir quantiles; files within each partition have non-overlapping sort-key
ranges. |
| **`RANGE` (Custom Function)** |
`.withDistributionMode(RANGE).withDistributionFunction(fn)` | **1**
(`GroupByKey` by `(dest, partition, fn(row))`) | **1** (`GroupIntoBatches` by
`(dest, partition, fn(row))`) | Every file is internally sorted; cross-file
ranges are non-overlapping if `fn(row)` is monotonic with respect to
`table.sortOrder()`. |
---
## Testing & Verification
- **Unit Tests (`265` tests in `:sdks:java:io:iceberg:test`)**:
- `IcebergRowSorterTest`: Order-preserving byte encoding across all
primitive/logical types, `ASC`/`DESC` inversion, `NULLS_FIRST`/`NULLS_LAST`
matrix, non-identity transforms (`month`, `hour`, `truncate`), and external
disk spill verification.
- `IcebergIOWriteTest`: End-to-end sorted writes across `NONE`, `HASH`,
`HASH + autosharding`, `RANGE` (custom function and dynamic quantile sharding),
post-sort file rolling without range overlap, and streaming
`directWriteByteLimit` bundle sorting.
- `AssignDestinationsAndPartitionsTest`: Side-input metadata resolution
and Efraimidis–Spirakis weighted reservoir sampling under extreme bundle skew.
- `IcebergIOSideInputTableCacheTest` &
`IcebergWriteSchemaTransformProviderTest`: Parameterized across `NONE`, `HASH`,
`HASH + autosharding`, and `RANGE`.
- **Live Dataflow Scale & Adversarial Skew Testing (Batch & Streaming)**:
- Verified on Dataflow across baseline and adversarial skew pipelines
(`17.32M` batch rows / `10.31M` streaming rows with 80% elephant partition
volume skew, 20% `null` sort keys, 30% duplicate heavy-hitter keys, 21x bundle
payload fanout skew, and 3 dynamic destinations $\times$ 34 partitions).
- Audited physical Parquet files and Dataflow worker logs: `0` internal
sort violations, `0` overlapping key ranges across dynamic `RANGE` shards (`23`
shards on `ELEPHANT_HOT_PARTITION`, `3` shards on medium partitions, `1` file
on each of `30` mouse partitions), `0` worker errors/OOMs, and `gc
thrashing=false` throughout external disk spills.
--
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]
