Skip to content

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

Open
parthchandra wants to merge 2 commits into
apache:mainfrom
parthchandra:stream-merge
Open

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

Conversation

@parthchandra

@parthchandra parthchandra commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes ##5337

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.

@parthchandra

Copy link
Copy Markdown
Contributor Author

Added more followup issues in parent issue - #5323

@parthchandra
parthchandra marked this pull request as ready for review August 13, 2026 20:56

@anuragmantri anuragmantri left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pinging here @parthchandra. It is great to see there is interest in this optimization in the comet project.

I reviewed the general implementation and the test coverage in this PR. I don't have any major comments, most are clarifications and minor suggestions. That said, I'm new to Comet codebase and would like someone more familiar to also take a look.

Comment on lines +896 to +898
private def isReportable(order: SortOrder, output: Seq[Attribute]): Boolean =
isIdentityProjected(order, output) && exprToProto(order, output).isDefined

@anuragmantri anuragmantri Aug 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As identified in the Iceberg PR and the design doc #5323, UUID orders differently in Iceberg than in Spark's comparator, and the identity case for UUID needs to follow Iceberg's byte ordering specifically. This gate doesn't check the sort column's type at all. I believe we could rely on upstream apache/iceberg#16750 to not report ordering on UUID. Is my understanding correct?

Comment on lines +878 to +879
* We read scanExec.ordering (the raw reported order), not scanExec.outputOrdering. Spark blanks
* outputOrdering when a partition holds more than one file -- the case this merge handles.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe using scanExec.ordering is to bypass the check in Spark 3.4 and 3.5 which was fixed in Spark 4.2 by SPARK-55715 and this is needed to support all these Spark versions. @peter-toth, would you please also take a look at this to see if this is safe?

impl IcebergScanMetrics {
fn new(metrics: &ExecutionPlanMetricsSet) -> Self {
Self {
baseline: BaselineMetrics::new(metrics, 0),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude flagged this:

BaselineMetrics::new(metrics, 0) hardcodes partition 0, but this operator is now multi-partition in the ordered path. Should this thread through the actual partition index from execute()?

* cached per catalog name, so a shared name would bind every test to the first warehouse), and
* its tables are dropped in a `finally` so a failing test cannot leak a table into a later one.
*/
class CometIcebergSortMergeReadSuite

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I reviewed the tests, very nice coverage already. I would also a test that deletes some rows from one file in a multi-file sorted partition, to cover the merge alongside MOR deletes.

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.

2 participants