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

[parthchandra]

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]

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]

@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]

Added more followup issues in parent issue - #5323

[anuragmantri]

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.

[anuragmantri]

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?

[parthchandra]

I was planning to address this in a follow up. But now I've added a type check rather than relying only on upstream. IcebergReflection.orderingUnsafeColumns returns UUID column names, and reportableOrdering now refuses any sort key in that set

[anuragmantri]

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?

[peter-toth]

Seems good to me, scanExec.ordering is the order reported by the source and that's what we need to keep with k-way merge if possible, and return as scanExec.outputOrdering so as Spark can use it to elimiate sorts.

[peter-toth]

So as we discussed offline CometIcebergNativeScanExec.scala.outputOrdering and this whole logic might not be needed if BatchScanExec -> CometIcebergNativeScanExec translation happend after EnsureRequirements runs (i.e. we eliminated the SortExecs based on BatchScanExec.outputOrdering).

And same for CometIcebergNativeScanExec.outputPartitioning...

[parthchandra]

@anuragmantri Correct. DataSourceV2ScanExecBase.outputOrdering blanks the ordering when a partition has more than one file on Spark 3.4–4.1 (relaxed in 4.2 by SPARK-55715). We need to keep this as long as we support versions before 4.2.

[parthchandra]

So as we discussed offline CometIcebergNativeScanExec.scala.outputOrdering and this whole logic might not be needed if BatchScanExec -> CometIcebergNativeScanExec translation happend after EnsureRequirements runs (i.e. we eliminated the SortExecs based on BatchScanExec.outputOrdering).

And same for CometIcebergNativeScanExec.outputPartitioning...

@peter-toth I dug in deeper and you are correct. EnsureRequirements does run before Comet converts the scan (confirmed in AdaptiveSparkPlanExec.queryStagePreparationRules and the non-AQE order), so it eliminates the shuffle on the BatchScanExec. An experiment confirmed it: with partitioning reporting off, the SMJ still has no Exchange. So I removed the outputPartitioning override and the reportPartitioning flag entirely.
outputOrdering is still needed because pre 4.2 Spark still has the check that blanks out BatchScanExec.outputOrdering for multi-file partitions. So we need to include the check to eliminate sort.

[anuragmantri]

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()?

[parthchandra]

Good catch Claude! Updated this.

[anuragmantri]

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.

[parthchandra]

added new test for this case

[andygrove]

One thing I'd like to dig into more before this lands: the reader fan-out in the ordered path.

SortPreservingMergeExec calls execute() on all of its input partitions up front, and IcebergScanExec::execute eagerly builds the ArrowReader and kicks off the read, so we end up with a live Parquet reader per data file in the Spark partition plus a buffered batch per stream inside the merge. data_file_concurrency_limit doesn't bound any of that anymore.

The comment in iceberg_scan.rs describes the fan-out as intrinsic to a k-way merge. That's true of the number of streams, but not of how many readers we have decoding concurrently or how much we buffer — a cascaded merge, or just falling back to the unordered read when a partition has too many files, would bound both.

The tables in the tests are fine, but the workload this targets is a sorted table, and a sorted table that has accumulated a lot of small commits is exactly where you get hundreds of files in one task. Since the merge reserves against Comet's memory pool, hitting the limit there is a query failure rather than a slowdown, which is a worse failure mode than the extra sort we're trying to avoid.

Could we add a config for the maximum files per partition we're willing to merge, and fall back to the unordered read above it? That keeps the default safe and lets people opt into deeper merges once we have a better sense of the memory profile.

[andygrove]

Following up on test coverage, because I want to make sure I'm reading this right.

As far as I can tell SupportsReportOrdering is what populates BatchScanExec.ordering, and that field is the only trigger for everything in this PR. Iceberg's implementation of it is apache/iceberg#16750, which is still open and targeted at Spark 4.1, so on the published Iceberg versions in our poms (1.5.2, 1.8.1, 1.10.0, 1.11.0) ordering is always None. That means reportableOrdering returns Nil, table_sort_orders goes over empty, native takes ordering = None, and we get the old single-partition unordered scan with no SortPreservingMergeExec.

If that's right, the effect in CI is bigger than the canceled plan assertions. The k-way merge itself never runs — every checkSparkAnswer in the new suite is exercising the read path we already had, so the only CI coverage of the new native code is the two Rust tests asserting partition_count().

There's a related wrinkle even on a reporting build: a global ORDER BY keeps its final sort, which would repair any mis-ordering the merge produced, so those tests catch dropped or duplicated rows but not wrong order. The order-sensitive assertions are all in the assume-gated SMJ/aggregate/window tests.

Am I understanding the situation correctly, or is there something in the CI setup I'm missing that does exercise the merge?

[andygrove]

One more, on the gate in serializePartitions:

val reportable = reportableOrdering(scanExec.ordering, output)
if (reportable.nonEmpty) {
  val protoOrders = reportable.map(exprToProto(_, output))
  if (protoOrders.forall(_.isDefined)) {
    commonBuilder.addAllTableSortOrders(protoOrders.map(_.get).asJava)
  }
}

The comment above it argues the forall always holds, and I agree that's true today. My worry is what happens if it ever stops being true, because the guard makes the unsafe direction the silent one. If we merge without telling Spark, we've wasted some work and nothing else. If we tell Spark and then don't merge, Spark drops the Sort and we return wrong results with no error.

That's not purely hypothetical, because the gate gets evaluated twice against two different points in time. outputOrdering is a lazy val forced during planning, this runs later off serializedPartitionData, and reportableOrdering reads COMET_ICEBERG_SORT_MERGE_ENABLED out of the active SQLConf each time. Two dynamic conf reads on separate threads at separate stages is the kind of thing that holds until it doesn't.

serializePartitions only has the one caller, so could we just pass outputOrdering in from the exec and have a single evaluation? Failing that I'd rather the mismatch threw than silently produced a proto that disagrees with what we told Spark.

[parthchandra]

That's not purely hypothetical, because the gate gets evaluated twice against two different points in time. outputOrdering is a lazy val forced during planning, this runs later off serializedPartitionData, and reportableOrdering reads COMET_ICEBERG_SORT_MERGE_ENABLED out of the active SQLConf each time. Two dynamic conf reads on separate threads at separate stages is the kind of thing that holds until it doesn't.

Fixed. This is evaluated once now.

serializePartitions only has the one caller, so could we just pass outputOrdering in from the exec and have a single evaluation? Failing that I'd rather the mismatch threw than silently produced a proto that disagrees with what we told Spark.

Fixed. This throws if a reported order can't be serialized, instead of the old silent if (forall(_.isDefined)) drop

[parthchandra]

Following up on test coverage, because I want to make sure I'm reading this right.

As far as I can tell SupportsReportOrdering is what populates BatchScanExec.ordering, and that field is the only trigger for everything in this PR. Iceberg's implementation of it is apache/iceberg#16750, which is still open and targeted at Spark 4.1, so on the published Iceberg versions in our poms (1.5.2, 1.8.1, 1.10.0, 1.11.0) ordering is always None. That means reportableOrdering returns Nil, table_sort_orders goes over empty, native takes ordering = None, and we get the old single-partition unordered scan with no SortPreservingMergeExec.

If that's right, the effect in CI is bigger than the canceled plan assertions. The k-way merge itself never runs — every checkSparkAnswer in the new suite is exercising the read path we already had, so the only CI coverage of the new native code is the two Rust tests asserting partition_count().

There's a related wrinkle even on a reporting build: a global ORDER BY keeps its final sort, which would repair any mis-ordering the merge produced, so those tests catch dropped or duplicated rows but not wrong order. The order-sensitive assertions are all in the assume-gated SMJ/aggregate/window tests.

Am I understanding the situation correctly, or is there something in the CI setup I'm missing that does exercise the merge?

You're reading this right. I cannot think of a way to mock Iceberg without the actual implementation (which is still in review). The tests cover the fallback, and are somewhat forward looking. In a way this feature is itself forward looking - once this feature is released in Iceberg and Spark starts to eliminate sorts based on what iceberg-java reports, Comet will produce absolutely garbage results unless we have this implementation in place.

[parthchandra]

One thing I'd like to dig into more before this lands: the reader fan-out in the ordered path.

SortPreservingMergeExec calls execute() on all of its input partitions up front, and IcebergScanExec::execute eagerly builds the ArrowReader and kicks off the read, so we end up with a live Parquet reader per data file in the Spark partition plus a buffered batch per stream inside the merge. data_file_concurrency_limit doesn't bound any of that anymore.

The comment in iceberg_scan.rs describes the fan-out as intrinsic to a k-way merge. That's true of the number of streams, but not of how many readers we have decoding concurrently or how much we buffer — a cascaded merge, or just falling back to the unordered read when a partition has too many files, would bound both.

The tables in the tests are fine, but the workload this targets is a sorted table, and a sorted table that has accumulated a lot of small commits is exactly where you get hundreds of files in one task. Since the merge reserves against Comet's memory pool, hitting the limit there is a query failure rather than a slowdown, which is a worse failure mode than the extra sort we're trying to avoid.

Could we add a config for the maximum files per partition we're willing to merge, and fall back to the unordered read above it? That keeps the default safe and lets people opt into deeper merges once we have a better sense of the memory profile.

I was planning to address this properly in a followup (#5343). But I can implement that in this PR if you think I should address it right away.

[mbutrovich]

First pass, looks to be in pretty good shape.

[mbutrovich]

Every test in this suite inserts 2 to 4 files per partition. Given the memory concern in iceberg_scan.rs is specifically about a partition with a large number of files (a live Parquet reader plus a buffered batch per file, all opened at once), would it be worth adding one test that inserts something like 50 to 100 single-row files into one partition? It would give the merge path a correctness check under the actual shape of the workload this feature targets, not just its happy-path shape, and would be a natural place to assert on memory pool usage or peak concurrent readers once #5343 lands.

[mbutrovich]

Following up on the fan-out discussion: I don't think this needs a wholly new algorithm, but it does need more than IcebergScanExec alone, since the operator that decides when to poll each partition is SortPreservingMergeExec, not this one. That operator has no notion of a bound and primes every input stream up front because polling is the only signal it has for "what's in this stream." Making IcebergScanExec::execute lazy internally doesn't avoid that, since SortPreservingMergeExec still calls poll_next on every partition almost immediately to seed its loser tree.

What would actually bound this is a custom merge operator that takes each file's min/max bound on the sort key alongside its FileScanTask, keeps a small active set merged the way SortPreservingMergeExec does today, and holds a priority queue of unopened files ordered by min bound, only calling execute/poll_next on the next pending file once its min bound could produce the next output value. That bounds concurrently-open readers by the overlap width of the file ranges at a given point in the merge rather than by total file count, which is the number that actually blows up on a sorted table with a lot of small commits.

The blocker today is data, not algorithm: FileScanTask in iceberg-rust doesn't carry column-level bounds. They exist upstream on the manifest-entry DataFile (lower_bounds/upper_bounds, already used for predicate pushdown), just not threaded down into the task struct the native scan gets. So this needs plumbing in iceberg-rust or an extra field fetched at scan-planning time on the JVM side, either way before a bound-driven admission scheme can work.

Given the risk, could we land @andygrove's simpler mitigation in this PR, a config for the max files to merge per partition that falls back to the unordered read above it, and keep the bound-driven design above as the concrete plan for #5343? Sort-merge ships on by default, and the workload it targets (a sorted table with accumulated small commits) is exactly where a partition ends up with hundreds of files, so I'd rather the default be safe now and the deeper optimization follow once the stats plumbing exists.