You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "parthchandra (via GitHub)" <gi...@apache.org> on 2026/10/01 22:18:47 UTC
Re: [PR] feat: Sort-aware Iceberg reads in Comet via a per-partition streaming merge [datafusion-comet]
parthchandra commented on code in PR #5331:
URL: https://github.com/apache/datafusion-comet/pull/5331#discussion_r4160366138
##########
spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala:
##########
@@ -1022,16 +1023,66 @@ case class CometScanRule(session: SparkSession)
}
}
+ // If Iceberg reports an ordering, EnsureRequirements may have already dropped the Sort
+ // above this scan (it decides that on the vanilla BatchScanExec, before Comet converts the
+ // scan). If the native scan cannot guarantee that ordering, reading unordered here would
+ // silently return wrong results, so stay on Spark -- its Iceberg reader produces the sorted
+ // output it promised. Evaluate the gate exactly once here and stash the result on the
+ // metadata; CometIcebergNativeScanExec.outputOrdering and the proto serde both read that
+ // stashed value, so the reported order cannot diverge from what native advertises.
+ val icebergReportsOrdering: Boolean = scanExec.ordering.exists(_.nonEmpty)
Review Comment:
Good catch, fixed. When the sort key isn't one of the selected columns, we no longer send the whole scan back to Spark. We keep the native scan and simply report no ordering (nothing above the scan can ask for an order on a column that isn't selected, so Spark hasn't dropped a sort). We still report the longest run of sort columns we can honour, and we only fall back to Spark when a selected column is unsafe — a UUID or a transform — or when we can't read the table schema. See `CometScanRule.scala:1000-1021` and the new orderingDecision in `CometIcebergNativeScan.scala:1000`.
##########
native/core/src/execution/planner.rs:
##########
@@ -1844,26 +1845,76 @@ impl PhysicalPlanner {
let metadata_location = common.metadata_location.clone();
let catalog_name = common.catalog_name.clone();
let tasks = parse_file_scan_tasks_from_common(common, &scan.file_scan_tasks)?;
+ let tasks_len = tasks.len();
let data_file_concurrency_limit = common.data_file_concurrency_limit as usize;
+ let max_files_per_partition = common.max_files_per_partition as usize;
- let iceberg_scan = IcebergScanExec::new(
+ // Table sort order Iceberg reported. Empty unless sortMerge is on and the order
+ // passed the identity gate in CometIcebergNativeScan. The SortOrder children are
+ // bound references into required_schema, so build the LexOrdering against it.
+ let ordering: Option<LexOrdering> = if common.table_sort_orders.is_empty() {
+ None
+ } else {
+ let exprs = common
+ .table_sort_orders
+ .iter()
+ .map(|expr| self.create_sort_expr(expr, Arc::clone(&required_schema)))
+ .collect::<Result<Vec<PhysicalSortExpr>, ExecutionError>>()?;
+ LexOrdering::new(exprs)
+ };
+
+ // A per-file-stream k-way merge opens one reader per file at once. Above the
+ // configured limit we instead read the partition unordered (bounded by
+ // data_file_concurrency_limit) and sort with a spillable SortExec, which bounds both
+ // open readers and memory. Both paths still produce sorted output, so the ordering
+ // Spark eliminated its Sort on is honoured either way. A limit of 0 (sortMerge
+ // disabled) always takes the sort path.
+ let use_merge = ordering.is_some() && tasks_len <= max_files_per_partition;
Review Comment:
Fixed. When a float or double column is a sort key but is not the last one, we no longer do the streaming merge — we read the files unordered and run the spillable sort instead, which sorts the way Spark does (treating -0.0 and +0.0 as equal). That removes the case where the merge trusted Iceberg's file order (which keeps -0.0 before +0.0) and produced a wrong result for things like ORDER BY ... LIMIT. A float/double that is the only key, or the last key, is still safe, so those still take the merge. See the decision at `planner.rs:1898-1908`, and three new plan tests: leading float → sort, trailing float → merge, single float → merge.
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala:
##########
@@ -893,6 +894,82 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit
Some(builder.setIcebergScan(icebergScanBuilder).build())
}
+ /**
+ * The part of an Iceberg-reported sort order that the native per-partition merge can honour, or
+ * Nil when the merge must stay off. Two callers use this one gate: the proto serialization
+ * (which turns on the native SortPreservingMergeExec) and
+ * CometIcebergNativeScanExec.outputOrdering (which tells Spark the scan is sorted). Sharing the
+ * gate means the two always agree.
+ *
+ * v1 accepts only identity sort fields on top-level columns that are in the projection. Each
+ * SortOrder child must be an AttributeReference in `output`, and must serialize to proto.
+ * Transform sort fields (bucket/truncate/...) are not AttributeReferences, so they fall through
+ * to Nil and we read unordered. Checking exprToProto here, not just in the proto path, keeps
+ * the two callers in step: outputOrdering never advertises an order the proto path would drop.
+ *
+ * We trust Iceberg on file-level sortedness. If it reports an ordering, SortOrderAnalyzer has
+ * already checked each file's sort_order_id matches the table order, so every file is sorted.
+ *
+ * We read scanExec.ordering (the raw reported order), not scanExec.outputOrdering. Spark blanks
Review Comment:
Fixed the comment.
##########
native/core/src/execution/operators/iceberg_scan.rs:
##########
@@ -707,6 +768,141 @@ mod tests {
.unwrap();
}
+ fn int_schema() -> arrow::datatypes::SchemaRef {
+ use arrow::datatypes::{DataType, Field, Schema};
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
+ }
+
+ fn single_col_ordering() -> Option<datafusion::physical_expr::LexOrdering> {
+ use arrow::compute::SortOptions;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+ LexOrdering::new(vec![PhysicalSortExpr {
+ expr: Arc::new(Column::new("a", 0)),
+ options: SortOptions::default(),
+ }])
+ }
+
+ // Builds a scan over three empty-delete tasks with the given reported ordering.
+ fn exec_with_ordering(
+ ordering: Option<datafusion::physical_expr::LexOrdering>,
+ ) -> IcebergScanExec {
+ use std::collections::HashMap;
+ let tasks = vec![
+ task_with_deletes(vec![]),
+ task_with_deletes(vec![]),
+ task_with_deletes(vec![]),
+ ];
+ IcebergScanExec::new(
+ "metadata.json".to_string(),
+ int_schema(),
+ HashMap::new(),
+ "cat".to_string(),
+ tasks,
+ 1,
+ ordering,
+ )
+ .unwrap()
+ }
+
+ // A reported ordering turns the scan into a multi-partition operator (one partition per task)
+ // so a SortPreservingMergeExec above can k-way merge the per-file sorted streams.
+ #[test]
+ fn reported_ordering_makes_scan_multi_partition() {
+ let exec = exec_with_ordering(single_col_ordering());
+ assert_eq!(exec.properties().partitioning.partition_count(), 3);
+ }
+
+ // Without a reported ordering the scan stays single-partition (Comet drives only execute(0),
+ // which must read every task), preserving the legacy unordered behaviour.
+ #[test]
+ fn no_ordering_keeps_single_partition() {
+ let exec = exec_with_ordering(None);
+ assert_eq!(exec.properties().partitioning.partition_count(), 1);
+ }
+
+ // The ordered scan reads each file as its own sorted partition and relies on
+ // SortPreservingMergeExec to k-way merge them into one globally sorted stream. This feeds known
+ // sorted partitions (with duplicate keys across partitions, and both asc and desc) into that
+ // merge with the same kind of LexOrdering the planner builds, and checks the output is globally
+ // sorted and complete. It is deterministic coverage of the merge that does not depend on an
+ // ordering-reporting Iceberg build (which is why the end-to-end suite's merge assertions cancel
+ // on the published Iceberg used in CI).
+ async fn merge_ints(input: Vec<Vec<i32>>, descending: bool) -> Vec<i32> {
Review Comment:
Yes please — go ahead and push that test to the branch. I'd rather take your version that you've already measured than reproduce it. Thank you!
##########
spark/src/test/scala/org/apache/comet/CometIcebergSortMergeReadSuite.scala:
##########
@@ -0,0 +1,885 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import java.util.concurrent.atomic.AtomicInteger
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.catalyst.expressions.{Add, Ascending, AttributeReference, Literal, SortOrder}
+import org.apache.spark.sql.comet.{CometIcebergNativeScanExec, CometSortExec}
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
+import org.apache.spark.sql.execution.{SortExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.types.IntegerType
+
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
+import org.apache.comet.serde.operator.CometIcebergNativeScan
+
+/**
+ * Tests for the sort-aware native Iceberg scan (branch `stream-merge`): the scan reports the
+ * Iceberg table sort order to Spark and does a per-partition streaming k-way merge of the
+ * already-sorted files, so Catalyst can drop the Sort and (with storage-partitioned join) the
+ * Exchange that a sort-merge join / grouped aggregate / window would otherwise need.
+ *
+ * The suite has two parts:
+ * - Unit tests of the [[CometIcebergNativeScan.reportableOrdering]] gate (no SparkSession
+ * required) -- the single decision shared by the proto serialization (which turns on the
+ * native SortPreservingMergeExec) and CometIcebergNativeScanExec.outputOrdering (which tells
+ * Spark the scan is sorted).
+ * - End-to-end tests over real Iceberg tables that exercise every Spark 4.0 mechanism which
+ * exploits already-sorted input to avoid a Sort or a shuffle (all keyed off
+ * `SortOrder.orderingSatisfies`, prefix semantics): `SupportsReportOrdering` ->
+ * `BatchScanExec.outputOrdering` and `EnsureRequirements` eliding the required-child Sort;
+ * `EliminateSorts` / `RemoveRedundantSorts`; `SortMergeJoinExec.requiredChildOrdering` +
+ * storage-partitioned join (`KeyGroupedPartitioning`,
+ * `spark.sql.sources.v2.bucketing.enabled`); `ReplaceHashWithSortAgg` / `SortAggregateExec`;
+ * `WindowExec`; `TakeOrderedAndProjectExec`.
+ *
+ * Two invariants determine the end-to-end assertions:
+ * 1. Correctness is checked unconditionally via `checkSparkAnswer` (Comet vs vanilla Spark).
+ * This is the primary guarantee: any k-way-merge defect (dropped, duplicated or mis-ordered
+ * rows, or an outputOrdering/outputPartitioning that does not match the rows the native
+ * operator actually produces) shows up as a result mismatch. It holds on any Iceberg build.
+ * 2. The strict "no Sort / no Exchange" plan assertions are the *target* of this feature.
+ * They only hold where the Iceberg build actually reports the ordering
+ * (`SupportsReportOrdering`, today an Iceberg fork feature -- the published/upstream Iceberg
+ * used in CI does not report it) and where the native scan reports `KeyGroupedPartitioning`.
+ * Each such test therefore runs the correctness check first, then `assume`s the reporting is
+ * active before asserting the plan shape, so it enforces the contract on a reporting build
+ * and is skipped (not failed) elsewhere. `sort = 0` is asserted only for *operator-required*
+ * orderings (SMJ / aggregate / window), never for a global `ORDER BY`, which Spark keeps
+ * regardless (the per-partition merge is not a global order).
+ *
+ * These cannot be SQL-file fixtures: setting an Iceberg sort order needs the Iceberg Java API
+ * (the Comet test session registers no Iceberg SQL extensions, so `WRITE ORDERED BY` will not
+ * parse), and the plan-shape assertions need access to the executed SparkPlan.
+ *
+ * Each test gets its own catalog name and temp warehouse (the Hadoop `SparkCatalog` instance is
+ * 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
+ extends CometTestBase
+ with CometIcebergTestBase
+ with AdaptiveSparkPlanHelper {
+
+ // ---------------------------------------------------------------------------------------------
+ // Unit tests of the reportableOrdering gate (merged from the former CometIcebergSortMergeSuite).
+ // v1 identity-scope gate as a pure function, independent of a live SparkSession or an Iceberg
+ // build that reports ordering. The flag defaults to enabled, so no SQLConf override is needed.
+ // ---------------------------------------------------------------------------------------------
+
+ private val gateA = AttributeReference("a", IntegerType)()
+ private val gateB = AttributeReference("b", IntegerType)()
+
+ test("gate: identity ordering on projected columns is reportable") {
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(
+ CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA, gateB)) === ordering)
+ }
+
+ test("gate: ordering on a column outside the projection falls back") {
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateB)).isEmpty)
+ }
+
+ test("gate: a transform (non-AttributeReference) sort child falls back") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA)).isEmpty)
+ }
+
+ test("gate: if any sort field is unreportable, the whole ordering falls back") {
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(Add(gateB, Literal(1)), Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA, gateB)).isEmpty)
+ }
+
+ test("gate: absent or empty ordering falls back") {
+ assert(CometIcebergNativeScan.reportableOrdering(None, Seq(gateA)).isEmpty)
+ assert(CometIcebergNativeScan.reportableOrdering(Some(Seq.empty), Seq(gateA)).isEmpty)
+ }
+
+ test("gate: a sort key on an ordering-unsafe column (e.g. UUID) falls back") {
+ // "a" stands in for a UUID column: Iceberg maps it to StringType but sorts by its own order,
+ // so it is unsafe to honour even though it looks like a plain string at the Spark level.
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(
+ CometIcebergNativeScan
+ .reportableOrdering(Some(ordering), Seq(gateA, gateB), Set("a"))
+ .isEmpty)
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // End-to-end fixtures.
+ // ---------------------------------------------------------------------------------------------
+
+ private val catalogCounter = new AtomicInteger(0)
+
+ // Storage-partitioned join config. preserve-data-grouping + v2 bucketing let Iceberg report
+ // KeyGroupedPartitioning on the BatchScanExec, and the join knobs force a sort-merge join over
+ // co-partitioned inputs so the Exchange can be eliminated. Comet does not report partitioning
+ // itself; Spark's EnsureRequirements eliminates the shuffle on the BatchScanExec before Comet
+ // converts the scan. Adaptive off keeps the executed plan stable for the counts below.
+ private val spjConf: Seq[(String, String)] = Seq(
+ "spark.sql.iceberg.planning.preserve-data-ordering" -> "true",
+ "spark.sql.iceberg.planning.preserve-data-grouping" -> "true",
+ "spark.sql.sources.v2.bucketing.enabled" -> "true",
+ "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true",
+ "spark.sql.requireAllClusterKeysForCoPartition" -> "false",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.join.preferSortMergeJoin" -> "true",
+ "spark.sql.adaptive.enabled" -> "false")
+
+ /**
+ * Runs `f` against a fresh, uniquely-named Hadoop catalog backed by a fresh temp warehouse,
+ * then drops the named tables (IF EXISTS) in a `finally` -- so tables are cleaned up even when
+ * the test body fails, and no two tests can collide on a table name. `f` receives the catalog
+ * name; tables live under the `db` namespace, e.g. `$cat.db.$table`.
+ */
+ private def withSortedTables(extraConf: Seq[(String, String)])(tables: String*)(
+ f: String => Unit): Unit = {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withTempIcebergDir { warehouseDir =>
+ val cat = s"sort_cat_${catalogCounter.incrementAndGet()}"
+ val cometConf = Seq(
+ s"spark.sql.catalog.$cat" -> "org.apache.iceberg.spark.SparkCatalog",
+ s"spark.sql.catalog.$cat.type" -> "hadoop",
+ s"spark.sql.catalog.$cat.warehouse" -> warehouseDir.getAbsolutePath,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true")
+ withSQLConf((cometConf ++ extraConf): _*) {
+ try f(cat)
+ finally tables.foreach(t => spark.sql(s"DROP TABLE IF EXISTS $cat.db.$t"))
+ }
+ }
+ }
+
+ /**
+ * Sets the table sort order via the Iceberg Java API. `cols` is (column, ascending); ascending
+ * uses Iceberg's default NULLS FIRST and descending its default NULLS LAST, matching the ORDER
+ * BY null-ordering the tests use.
+ */
+ private def replaceSortOrder(
+ cat: String,
+ namespace: String,
+ table: String,
+ cols: (String, Boolean)*): Unit = {
+ val catalog = spark.sessionState.catalogManager
+ .catalog(cat)
+ .asInstanceOf[org.apache.iceberg.spark.SparkCatalog]
+ val ident =
+ org.apache.spark.sql.connector.catalog.Identifier.of(Array(namespace), table)
+ val icebergTable = catalog
+ .loadTable(ident)
+ .asInstanceOf[org.apache.iceberg.spark.source.SparkTable]
+ .table()
+ var sortOrder = icebergTable.replaceSortOrder()
+ cols.foreach { case (c, asc) =>
+ sortOrder = if (asc) sortOrder.asc(c) else sortOrder.desc(c)
+ }
+ sortOrder.commit()
+ }
+
+ /**
+ * Sets a transform (bucket) sort order via the Iceberg Java API. Iceberg orders files by the
+ * bucket hash, so the reported sort field is a transform expression, not a plain column. Spark
+ * 4.0+ can convert a bucket transform ordering (V2ScanPartitioningAndOrdering threads the
+ * function catalog and V2ExpressionUtils special-cases BucketTransform); Spark 3.4 cannot, so
+ * the caller gates the test on isSpark40Plus. Either way Comet's reportableOrdering rejects the
+ * non-AttributeReference sort child (v1 is identity-scope only, #5339), so Comet must fall
+ * back.
+ */
+ private def replaceSortOrderBucket(
+ cat: String,
+ namespace: String,
+ table: String,
+ col: String,
+ numBuckets: Int): Unit = {
+ val catalog = spark.sessionState.catalogManager
+ .catalog(cat)
+ .asInstanceOf[org.apache.iceberg.spark.SparkCatalog]
+ val ident =
+ org.apache.spark.sql.connector.catalog.Identifier.of(Array(namespace), table)
+ val icebergTable = catalog
+ .loadTable(ident)
+ .asInstanceOf[org.apache.iceberg.spark.source.SparkTable]
+ .table()
+ icebergTable
+ .replaceSortOrder()
+ .asc(org.apache.iceberg.expressions.Expressions.bucket(col, numBuckets))
+ .commit()
+ }
+
+ /**
+ * Each string becomes a separate INSERT, hence a separate data file, so merging is required.
+ */
+ private def insertBatches(cat: String, table: String, batches: String*): Unit =
+ batches.foreach(values => spark.sql(s"INSERT INTO $cat.db.$table VALUES $values"))
+
+ private def nativeScans(plan: SparkPlan): Seq[CometIcebergNativeScanExec] =
+ collect(stripAQEPlan(plan)) { case s: CometIcebergNativeScanExec => s }
+
+ private def countSorts(plan: SparkPlan): Int =
+ collect(stripAQEPlan(plan)) {
+ case s: SortExec => s
+ case s: CometSortExec => s
+ }.size
+
+ private def countShuffles(plan: SparkPlan): Int =
+ collect(stripAQEPlan(plan)) {
+ case e: ShuffleExchangeExec => e
+ case e: CometShuffleExchangeExec => e
+ }.size
+
+ /** True once every native scan in the plan advertises the reported ordering. */
+ private def orderingReported(plan: SparkPlan): Boolean = {
+ val scans = nativeScans(plan)
+ scans.nonEmpty && scans.forall(_.outputOrdering.nonEmpty)
+ }
+
+ // NOTE ON CANCELED TESTS: the helper below CANCELS the test (ScalaTest `assume`, reported as
+ // "!!! CANCELED !!!", not a failure) when the Iceberg build on the classpath does not implement
+ // the DSv2 `SupportsReportOrdering` API. Published/upstream Iceberg (the runtime used in CI and
+ // the default mvn profiles) does not report a sort order, so `outputOrdering` comes back empty
+ // and there is no eliminated Sort to assert on. These tests therefore show as canceled there --
+ // that is expected, NOT a regression. The preceding `checkSparkAnswer` has already validated
+ // correctness; only the sort/shuffle-elimination plan assertion is skipped. Run against an
+ // ordering-reporting (fork) Iceberg build to exercise those assertions.
+
+ /** Cancels the test (see note above) unless the scan reported an ordering. */
+ private def assumeOrderingReported(plan: SparkPlan): Unit =
+ assume(
+ orderingReported(plan),
+ "current Iceberg build does not implement SupportsReportOrdering (no ordering reported); " +
+ "sort-elimination assertion skipped")
+
+ /**
+ * True if the Iceberg build on the classpath reports a sort order for `query`. Determined
+ * independently of Comet: the query is planned with the native Iceberg scan disabled, and we
+ * check whether Spark's own BatchScanExec advertises an outputOrdering. This is the signal that
+ * makes the fallback checks below meaningful -- when Iceberg does not report (the published
+ * Iceberg in CI), there is no reported ordering for Comet to decline, so those checks are
+ * skipped rather than firing on a scan that was never eligible to fall back.
+ */
+ private def icebergReportsOrdering(query: String): Boolean = {
+ // withSQLConf returns the block value on Spark 4.x but Unit on 3.4/3.5, so capture in a var
+ // rather than relying on its return value.
+ var reported = false
+ withSQLConf(CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "false") {
+ val plan = spark.sql(query).queryExecution.executedPlan
+ reported = collect(stripAQEPlan(plan)) { case b: BatchScanExec => b }
+ .exists(_.outputOrdering.nonEmpty)
+ }
+ reported
+ }
+
+ /**
+ * The Spark-fallback contract for a reported-but-unhonorable ordering: when the native scan
+ * cannot guarantee an ordering Iceberg reported, Comet must not convert the scan at all.
+ * Reading unordered natively would return silently-wrong results, because EnsureRequirements
+ * has already dropped the Sort above the scan on the strength of Iceberg's report. So the plan
+ * must contain no CometIcebergNativeScanExec -- the scan stays on Spark's Iceberg reader. Gated
+ * on a reporting Iceberg build (see icebergReportsOrdering); correctness is asserted separately
+ * by the caller.
+ */
+ private def assertFellBackToSpark(query: String, plan: SparkPlan): Unit = {
+ assume(
+ icebergReportsOrdering(query),
+ "current Iceberg build does not report an ordering; the Spark-fallback path is not exercised")
+ assert(
+ nativeScans(plan).isEmpty,
+ s"Comet reported an ordering it cannot honour instead of falling back to Spark:\n$plan")
+ }
+
+ // The reporting mechanism: SupportsReportOrdering -> CometIcebergNativeScanExec.outputOrdering
+ //
+ // NOTE ON TABLE SHAPE: Iceberg only reports a sort order for a table with a non-empty grouping
+ // key -- isOrderingEnabled in apache/iceberg#16750 is `!groupingKeyType().fields().isEmpty() &&
+ // canReportOrdering(...)`, and Iceberg's own testNoMergeReaderForUnpartitionedSortedTable asserts
+ // an unpartitioned sorted table reports nothing. So a merge test MUST use a partitioned table, or
+ // Iceberg reports no ordering, Comet takes the plain unordered read, and the merge never runs (the
+ // checkSparkAnswer then passes by comparing the unordered path against itself). Every merge test
+ // below partitions by a single-value column `p` so the ordering is reported while all files stay
+ // in one partition -- the shape that actually exercises the multi-file k-way merge -- and asserts
+ // assumeOrderingReported after the correctness check so a reporting build proves the merge ran.
+
+ test("native scan reports the table sort order for a multi-file sorted table") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ insertBatches(cat, "t", "(1,'a','P1'),(3,'c','P1')", "(2,'b','P1'),(4,'d','P1')")
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assume(nativeScans(plan).nonEmpty, "query did not use the native Iceberg scan")
+ assumeOrderingReported(plan)
+ assert(
+ nativeScans(plan).head.outputOrdering.head.child.references.exists(_.name == "id"),
+ s"expected the scan to report an ordering on id:\n$plan")
+ }
+ }
+
+ test("sort-merge disabled keeps the scan native and still reports the ordering (via sort)") {
+ withSortedTables(spjConf ++ Seq(CometConf.COMET_ICEBERG_SORT_MERGE_ENABLED.key -> "false"))(
+ "t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ insertBatches(cat, "t", "(1,'a','P1'),(3,'c','P1')", "(2,'b','P1'),(4,'d','P1')")
+
+ // Disabling only turns off the k-way merge, not the whole native scan: Comet still reads
+ // the table and still honours the reported order (via a spillable sort). Correctness must
+ // hold, and on an ordering-reporting Iceberg build the scan still advertises the order.
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assert(
+ nativeScans(plan).nonEmpty,
+ s"sort-merge disabled must not force the scan back to Spark:\n$plan")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ // K-way merge correctness. The tables are partitioned so Iceberg reports the ordering and the
+ // merge actually runs (see the NOTE ON TABLE SHAPE above); assumeOrderingReported proves that on a
+ // reporting build. Order sensitivity is verified with a window over the merged input rather than a
+ // global ORDER BY, because Spark keeps its final Sort for a global ORDER BY (a per-partition order
+ // does not satisfy a global one) and that re-sort would repair -- and hide -- a mis-ordered merge.
+
+ test("merges multiple sorted files per partition and preserves order") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 INT, data STRING) USING iceberg " +
+ "PARTITIONED BY (bucket(4, c1))")
+ replaceSortOrder(cat, "db", "t", "c1" -> true, "c2" -> true)
+ // Multiple files, each holding rows for every c1 bucket, so the k-way merge runs within a
+ // partition; c2 interleaves across the files so a mis-ordered merge changes the row numbering.
+ insertBatches(cat, "t", "(1,1,'a'),(2,1,'b')", "(1,3,'c'),(2,3,'d')", "(1,2,'e'),(2,2,'f')")
+
+ // ROW_NUMBER over the reported (c1, c2) order: Spark keeps that order (the window's required
+ // ordering is satisfied, no re-sort), so a mis-ordered merge yields wrong row numbers.
+ val query =
+ s"SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2) AS rn FROM $cat.db.t"
+ val (_, plan) = checkSparkAnswer(query)
+ assume(nativeScans(plan).nonEmpty, "query did not use the native Iceberg scan")
+ assumeOrderingReported(plan)
+ assert(
+ countSorts(plan) == 0,
+ s"the merge must satisfy the window ordering without a re-sort:\n$plan")
+ }
+ }
+
+ test("merge interleaves duplicate sort-key values across files") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ // The same id appears in several files; the merge must keep every row, not drop or mis-order.
+ insertBatches(
+ cat,
+ "t",
+ "(1,'a','P1'),(2,'b','P1')",
+ "(1,'c','P1'),(2,'d','P1')",
+ "(1,'e','P1'),(3,'f','P1')")
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id, data")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("merge applies merge-on-read deletes across a multi-file sorted partition") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg " +
+ "PARTITIONED BY (p) " +
+ "TBLPROPERTIES ('format-version'='2', 'write.delete.mode'='merge-on-read')")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ // Several files so a merge is required; then delete rows from some of them. On a v2
+ // merge-on-read table DELETE writes delete files rather than rewriting the data files, so
+ // the scan must apply the deletes while merging the still-sorted files.
+ insertBatches(
+ cat,
+ "t",
+ "(1,'a','P1'),(4,'d','P1')",
+ "(2,'b','P1'),(5,'e','P1')",
+ "(3,'c','P1'),(6,'f','P1')")
+ spark.sql(s"DELETE FROM $cat.db.t WHERE id IN (2, 5)")
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("merge honours NULLS FIRST on an ascending sort key") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 STRING, c3 STRING) USING iceberg " +
+ "PARTITIONED BY (c3)")
+ replaceSortOrder(cat, "db", "t", "c1" -> true) // ASC -> Iceberg default NULLS FIRST
+ insertBatches(
+ cat,
+ "t",
+ "(null,'x','P1'),(3,'c','P1')",
+ "(null,'y','P1'),(1,'a','P1'),(2,'b','P1')")
+
+ checkSparkAnswer(
+ s"SELECT c1, c2 FROM $cat.db.t WHERE c3 = 'P1' ORDER BY c1 ASC NULLS FIRST, c2")
+ }
+ }
+
+ test("merge honours NULLS LAST on a descending sort key") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 STRING, c3 STRING) USING iceberg " +
+ "PARTITIONED BY (c3)")
+ replaceSortOrder(cat, "db", "t", "c1" -> false) // DESC -> Iceberg default NULLS LAST
+ insertBatches(
+ cat,
+ "t",
+ "(null,'x','P1'),(1,'a','P1')",
+ "(null,'y','P1'),(3,'c','P1'),(2,'b','P1')")
+
+ checkSparkAnswer(
+ s"SELECT c1, c2 FROM $cat.db.t WHERE c3 = 'P1' ORDER BY c1 DESC NULLS LAST, c2")
+ }
+ }
+
+ test("merge on a descending sort order preserves order") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 INT, data STRING) USING iceberg " +
+ "PARTITIONED BY (bucket(4, c1))")
+ replaceSortOrder(cat, "db", "t", "c1" -> true, "c2" -> false) // c2 DESC, Iceberg NULLS LAST
+ insertBatches(cat, "t", "(1,3,'a'),(2,3,'b')", "(1,1,'c'),(2,1,'d')", "(1,2,'e'),(2,2,'f')")
+
+ // Window ordered DESC over the merged (c1, c2 DESC) order; a mis-ordered descending merge
+ // changes the row numbers (a global ORDER BY DESC would be re-sorted by Spark and hide it).
+ val query =
+ s"SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) AS rn FROM $cat.db.t"
+ val (_, plan) = checkSparkAnswer(query)
+ assume(nativeScans(plan).nonEmpty, "query did not use the native Iceberg scan")
+ assumeOrderingReported(plan)
+ assert(
+ countSorts(plan) == 0,
+ s"the descending merge must satisfy the window ordering without a re-sort:\n$plan")
+ }
+ }
+
+ test("merge on a multi-column sort order") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 STRING, c3 STRING, p STRING) USING iceberg " +
+ "PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "c3" -> true, "c1" -> true)
+ insertBatches(
+ cat,
+ "t",
+ "(1,'a','A','P1'),(3,'c','A','P1')",
+ "(2,'b','A','P1'),(1,'a','B','P1')",
+ "(2,'b','B','P1'),(3,'c','B','P1')")
+
+ val (_, plan) = checkSparkAnswer(s"SELECT c3, c1, c2 FROM $cat.db.t ORDER BY c3, c1")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("single file needs no merge") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ insertBatches(cat, "t", "(1,'a','P1'),(2,'b','P1')")
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("many small files in one partition merge correctly") {
+ // Raise the per-partition file limit so the k-way merge (not the sort fallback) runs with many
+ // files -- the shape this feature targets, where the merge opens one reader per file at once.
+ // Partitioned by a single value so Iceberg reports the ordering and the ~70-way merge actually
+ // runs; assumeOrderingReported proves it on a reporting build. checkSparkAnswer guards
+ // correctness; asserting on memory-pool usage / peak concurrent readers is a TODO for #5343.
+ val conf =
+ spjConf :+ (CometConf.COMET_ICEBERG_SORT_MERGE_MAX_FILES_PER_PARTITION.key -> "1000")
+ withSortedTables(conf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ val rows = (1 to 70).map(i => s"($i,'v$i','P1')")
+ insertBatches(cat, "t", rows: _*)
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("many small files in one partition stay correct (exercises the sort fallback)") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (id INT, data STRING, p STRING) USING iceberg PARTITIONED BY (p)")
+ replaceSortOrder(cat, "db", "t", "id" -> true)
+ // 70 single-row inserts -> 70 files in one partition, above the default maxFilesPerPartition
+ // (64). Partitioned so Iceberg reports the ordering; above the cap the native scan takes the
+ // fallback -- a single unordered read plus a spillable SortExec, not a 70-way merge -- which
+ // is the shape this feature targets (a sorted table with many small commits). assumeOrdering
+ // Reported proves the ordering was reported (so the sort-fallback path really ran) on a
+ // reporting build. checkSparkAnswer guards correctness; asserting on memory-pool usage / peak
+ // concurrent readers is a TODO for #5343.
+ val rows = (1 to 70).map(i => s"($i,'v$i','P1')")
+ insertBatches(cat, "t", rows: _*)
+
+ val (_, plan) = checkSparkAnswer(s"SELECT id, data FROM $cat.db.t ORDER BY id")
+ assumeOrderingReported(plan)
+ }
+ }
+
+ test("partitioned table with several files per partition") {
+ withSortedTables(spjConf)("t") { cat =>
+ spark.sql(
+ s"CREATE TABLE $cat.db.t (c1 INT, c2 STRING, c3 STRING) USING iceberg " +
+ "PARTITIONED BY (c3)")
+ replaceSortOrder(cat, "db", "t", "c1" -> true)
+ insertBatches(
+ cat,
+ "t",
+ "(1,'a','P1'),(3,'c','P1')",
+ "(2,'b','P1'),(4,'d','P1')",
+ "(5,'e','P2'),(7,'g','P2')",
+ "(6,'f','P2'),(8,'h','P2')")
+
+ checkSparkAnswer(s"SELECT c1, c2 FROM $cat.db.t WHERE c3 = 'P1' ORDER BY c1")
+ checkSparkAnswer(s"SELECT c1, c2 FROM $cat.db.t WHERE c3 = 'P2' ORDER BY c1")
+ }
+ }
+
+ test("sort key absent from the projection: falls back to an unordered read but stays correct") {
Review Comment:
You're right, the old test asserted nothing. Rewritten to check the real outcome: it now asserts the scan stays native and reports no ordering (CometIcebergSortMergeReadSuite.scala:643). I also added unit tests for the decision itself — stays-native for an unselected key, falls back for a selected UUID/transform, and the "unselected key before a UUID still stays native" case — starting at line 138.
##########
native/core/src/execution/operators/iceberg_scan.rs:
##########
@@ -175,18 +239,14 @@ impl IcebergScanExec {
fn execute_with_tasks(
&self,
tasks: Vec<FileScanTask>,
+ partition: usize,
context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
let output_schema = Arc::clone(&self.output_schema);
- let file_io = load_file_io(
- &self.catalog_properties,
- &self.metadata_location,
- &self.catalog_name,
- AccessMode::Read,
- )?;
+ let file_io = self.file_io.clone();
Review Comment:
Fair point. Logged a new issue for this - https://github.com/apache/datafusion-comet/issues/6524
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org