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.