Skip to content

feat: Sort-aware Iceberg reads in Comet via a per-partition streaming merge - #5331

Draft
parthchandra wants to merge 3 commits into
apache:mainfrom
parthchandra:stream-merge
Draft

feat: Sort-aware Iceberg reads in Comet via a per-partition streaming merge#5331
parthchandra wants to merge 3 commits into
apache:mainfrom
parthchandra:stream-merge

Conversation

@parthchandra

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5323.

Rationale for this change

Currently when Comet reads a sorted Iceberg table, the scan throws away the ordering to achieve parallelism. As a result Spark can't tell the data is already sorted, and it re-sorts on every read — in joins, aggregates, windows, and order-by queries — even though the work was already done at write time.

This PR modifies the native scan to preserve and report that ordering, so Spark can drop the redundant sorts. It builds on Iceberg's own SupportsReportOrdering (apache/iceberg#14948): when Iceberg reports a sort
order, merge the sorted files per partition, and tell Spark the result is sorted.

The scan also reports Iceberg's key-group partitioning. (For Spark to eliminate shuffle in SMJ, the scan must also report how the data is grouped by the join key (storage-partitioned join)).

What changes are included in this PR?

For every Spark partition, the scan now reads each sorted file as its own stream and k-way-merges them into one sorted stream using DataFusion's SortPreservingMergeExec.

Summary of changes -

  • Proto: a new table_sort_orders field on IcebergScanCommon carries the reported sort order to the native side.
  • Scala (serde + scan exec): report the sort order and the key-grouped partitioning to Spark. Ordering reporting is limited to the safe case for v1 — identity-transform sort fields on top-level columns that are
    actually in the projection. Anything else (transforms, a sort key that isn't selected) falls back to today's unordered read and reports nothing, so it's always correct.
  • Native (Rust): IcebergScanExec becomes multi-partition when an ordering is present (one sorted stream per file), and the planner wraps it in SortPreservingMergeExec. No changes to iceberg-rust — we just
    call its existing reader once per file instead of once for the whole batch.

The PR also adds two config flags for the Iceberg scan:

  • spark.comet.scan.icebergNative.sortMerge.enabled (default on) — report the sort order and do the per-partition merge. Only does anything when Iceberg's spark.sql.iceberg.planning.preserve-data-ordering is
    on (off by default).
  • spark.comet.scan.icebergNative.reportPartitioning.enabled (default off) — report key-grouped partitioning for storage-partitioned joins. Off by default while we build out coverage for the adaptive-execution
    partition-pushdown path.

Note: A global ORDER BY still keeps its final sort — a per-partition merge isn't a cluster-wide order — so that case is unchanged.

How are these changes tested?

  • Native unit tests in iceberg_scan.rs: multi-partition with a reported ordering, single-partition without one.
  • End-to-end suite CometIcebergSortMergeReadSuite over real Iceberg tables (local Hadoop catalog, sort order set via the Iceberg Java API, one file per insert).

@parthchandra

Copy link
Copy Markdown
Contributor Author

This PR has some followups -

  • Sort keys not in the projection. Today, if a sort column isn't selected, we skip the merge and read unordered. Todo: read the sort columns anyway (even when not selected), merge on them, then drop them before returning.
  • Transform sort orders (bucket/truncate/etc.). We only handle plain-column sort orders for now. todo: also handle transformed sort keys, matching Iceberg's per-type comparison exactly.
  • Read concurrency. The merge reads files more lazily than today's unordered read. As a followup benchmark, and if many-small-file reads regress, read files ahead into small buffers.
  • Single-file partitions. A partition with one file still gets wrapped in the merge (a no-op). Skip the wrapper in that case.
  • Storage-partitioned join confidence tests, then default reportPartitioning on. Add tests for: (a) a heavily skewed key so the replicate path runs, (b) different bucket counts per side (e.g. bucket(8) vs bucket(4)) so a partition reducer kicks in, (c) three-way and self joins on the same key. Once green, flip the flag on by default.

@parthchandra

Copy link
Copy Markdown
Contributor Author

@anuragmantri, @peter-toth , you might be interested in looking at this. k-way merge to maintain the ordered property of sorted Iceberg tables. Feedback, especially about test coverage, would be highly appreciated.

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.

Sort-aware Iceberg reads in Comet via a per-partition streaming merge

1 participant