Skip to content

[IcebergIO] Support sorted writes and RANGE distribution mode - #40653

Open
atognolag wants to merge 3 commits into
apache:masterfrom
atognolag:feature/sorted-icebergIO
Open

atognolag wants to merge 3 commits into
apache:masterfrom
atognolag:feature/sorted-icebergIO

Conversation

@atognolag

Copy link
Copy Markdown
Contributor

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

@codecov

codecov Bot commented Oct 19, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 65.07%. Comparing base (76f6e6d) to head (b19c382).
⚠️ Report is 224 commits behind head on master.

Additional details and impacted files
@@             Coverage Diff              @@
##             master   #40653      +/-   ##
============================================
+ Coverage     56.16%   65.07%   +8.90%     
- Complexity     2288    13993   +11705     
============================================
  Files          1124     1729     +605     
  Lines        178336   108515   -69821     
  Branches       1489    11313    +9824     
============================================
- Hits         100170    70611   -29559     
+ Misses        75645    33694   -41951     
- Partials       2521     4210    +1689     
Flag Coverage Δ
go ?
java 65.07% <ø> (-10.72%) ⬇️
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

Copy link
Copy Markdown
Contributor

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

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

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

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant