andygrove commented on code in PR #5331:
URL: https://github.com/apache/datafusion-comet/pull/5331#discussion_r4275347432
##########
spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala:
##########
@@ -1158,26 +1160,56 @@ object IcebergReflection extends Logging {
* types with either layout (e.g. geometry).
*/
def pageIndexUnsupportedColumns(schema: Any): Set[String] = {
+ topLevelColumnTypes(schema, "page-index-unsupported columns")
+ .getOrElse(Seq.empty)
+ .collect {
+ case (name, typeStr)
+ if typeStr.startsWith("decimal(") || typeStr == "uuid" ||
+ typeStr.startsWith("fixed[") || typeStr == "binary" =>
+ name
+ }
+ .toSet
+ }
+
+ /**
+ * Names of top-level columns whose Iceberg sort order can differ from
Spark's comparison of the
+ * Spark type they map to, so a native k-way merge keyed on the Spark value
could mis-order
+ * rows. Today that is UUID: Iceberg maps UUID to Spark StringType (see
TypeToSparkType) but
+ * sorts by its own UUID comparator, not by the canonical string, so the
file order and a string
+ * comparison can disagree. reportableOrdering refuses a sort key in this
set. v1 only reports
+ * identity, top-level sort keys, so top-level columns are enough.
+ *
+ * Returns None if the schema could not be read. That is the fail-closed
answer: we cannot rule
+ * out a UUID (or any other unsafe) sort key, so the caller must refuse the
ordering rather than
+ * assume it is safe -- reporting an ordering the native merge cannot
reproduce returns silently
+ * mis-ordered rows, because Spark has already dropped the Sort by then.
+ */
+ def orderingUnsafeColumns(schema: Any): Option[Set[String]] = {
+ topLevelColumnTypes(schema, "ordering-unsafe columns").map { columns =>
+ columns.collect { case (name, typeStr) if typeStr == "uuid" => name
}.toSet
+ }
+ }
+
+ /**
+ * Read the top-level columns of an Iceberg schema by reflection, yielding
(name, typeStr) per
+ * column. Returns None if the schema could not be read, letting each caller
pick its own
+ * fail-open vs fail-closed answer.
+ */
+ private def topLevelColumnTypes(schema: Any, context: String):
Option[Seq[(String, String)]] = {
import scala.jdk.CollectionConverters._
try {
val columns = getMethod(schema.getClass, "columns")
.invoke(schema)
.asInstanceOf[java.util.List[_]]
- columns.asScala.flatMap { column =>
+ Some(columns.asScala.iterator.map { column =>
Review Comment:
On Scala 2.12, `Iterator.toSeq` is `toStream`, so this returns a lazy
`Stream` and only the first column's reflection runs inside the `try`. The rest
run when `orderingUnsafeColumns` or `pageIndexUnsupportedColumns` call
`collect` on it, outside the `try`, so a failure there escapes `CometScanRule`
instead of becoming the `None` that refuses the ordering. I don't expect a
later column to fail when the first one didn't, but this is the method whose
job is to fail closed, and before the refactor `pageIndexUnsupportedColumns`
did the whole walk inside the `try`. Could this be `.toList`?
##########
spark/src/test/scala/org/apache/comet/CometIcebergSortMergeReadSuite.scala:
##########
@@ -0,0 +1,960 @@
+/*
+ * 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)()
+ private val gateC = AttributeReference("c", 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: reports the longest identity-projected prefix") {
+ // a is projected and reportable; the transform on b stops the prefix, so
only [a] is reported.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(Add(gateB,
Literal(1)), Ascending))
+ assert(
+ CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA,
gateB)) ===
+ Seq(SortOrder(gateA, Ascending)))
+ }
+
+ test("gate: a leading column outside the projection reports nothing") {
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateB)).isEmpty)
+ }
+
+ test("gate: a leading transform (non-AttributeReference) sort child reports
nothing") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateA)).isEmpty)
+ }
+
+ test("gate: absent or empty ordering reports nothing") {
+ 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) reports
nothing") {
+ // "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)
+ }
+
+ // orderingDecision decides fallback-vs-stay-native, which
reportableOrdering (the advertised
+ // prefix) alone does not capture: an empty prefix can mean "stay native,
read unordered" (safe)
+ // or "fall back to Spark" (a projected UUID/transform Spark may have
dropped a Sort on).
+
+ test("decision: a leading unprojected sort key stays native with no reported
ordering") {
+ // b (the sort key) is not selected. Nothing above the scan can require an
ordering that reaches
+ // it, so Spark dropped no Sort -- Comet stays native and simply reports
nothing.
+ val ordering = Seq(SortOrder(gateB, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && d.stayNative)
+ }
+
+ test("decision: a projected prefix then an unprojected key stays native with
the prefix") {
+ // a is projected + reportable; b is not selected. Report [a] and stay
native.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a two-field projected prefix then an unprojected key stays
native") {
+ // a and b are projected + reportable; c is not selected. Report [a, b]
and stay native.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending)))
+ assert(d.stayNative)
+ }
+
+ test("decision: an unprojected blocker before a projected UUID still stays
native") {
+ // First blocker wins: b (unprojected) stops the prefix at [a], so we stay
native even though a
+ // later field (c) is a projected unsafe/UUID column -- nothing above can
require past b.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateC), Set("c"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a projected prefix then a projected UUID falls back to
Spark") {
+ // a is projected + reportable; b is a projected UUID (unsafe). The prefix
is [a], but the
+ // blocker b is IN the output, so a required ordering could reach it and
Spark may already have
+ // dropped a Sort on it -- Comet must fall back even though a non-empty
prefix is reportable.
+ // This pins the `stayNative && ...` gate against a regression to the old
`reported.nonEmpty`.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("b"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && !d.stayNative)
+ }
+
+ test("decision: a leading UUID sort key falls back to Spark") {
+ // "a" is a projected UUID: a required ordering can reach it and Spark may
already have dropped
+ // a Sort on it, so Comet must fall back rather than read unordered.
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("a"))
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: a leading transform sort key falls back to Spark") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: no reported ordering stays native") {
+ assert(CometIcebergNativeScan.orderingDecision(None, Seq(gateA),
Set.empty).stayNative)
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // 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
Review Comment:
On the current apache/iceberg#16750 this table shape doesn't get an ordering
reported. Iceberg builds the grouping key from the projected columns only.
`Partitioning.groupingKeyType(projection(), specs())` drops a partition field
whose source column the query doesn't read. None of the queries on the
`PARTITIONED BY (p)` tables selects `p`, so they get ungrouped planning and no
ordering, and nine of the tests below cancel at `assumeOrderingReported`. It
also makes "results are identical with the sort-merge feature on and off"
compare the plain unordered read with itself. When I added `p` to those select
lists and ran against a build of that PR, all of them ran the merge or the sort
path and passed, apart from the two 70-file tests (next comment). Could we add
`p` to those queries and say why here?
##########
native/core/src/execution/operators/iceberg_scan.rs:
##########
@@ -292,46 +372,72 @@ impl IcebergScanExec {
return Ok(());
}
+ // Coalesce stats across partitions: each delete file gets one shared
cell, and concurrent
+ // partitions that need the same file await its single in-flight stat
instead of each issuing
+ // a HEAD (#6524). Create/fetch the per-path cells under a brief lock
(no await held), then
+ // drive the stats through the cells.
+ let cells: Vec<(String, Arc<tokio::sync::OnceCell<u64>>)> = {
Review Comment:
Nothing pins the sharing across partitions here.
`shared_delete_file_read_once_in_ordered_path` compares `bytes_scanned`, which
doesn't count the stat, so going back to a per-call map would leave every test
green. A cheap guard would be two `fill_delete_file_sizes` calls that share one
`size_cache`, deleting the delete file from disk in between. The second call
should still fill in the real size because it never stats, and with a per-call
cache it fails with `Failed to stat delete file`. Could we add that?
##########
spark/src/test/scala/org/apache/comet/CometIcebergSortMergeReadSuite.scala:
##########
@@ -0,0 +1,960 @@
+/*
+ * 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)()
+ private val gateC = AttributeReference("c", 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: reports the longest identity-projected prefix") {
+ // a is projected and reportable; the transform on b stops the prefix, so
only [a] is reported.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(Add(gateB,
Literal(1)), Ascending))
+ assert(
+ CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA,
gateB)) ===
+ Seq(SortOrder(gateA, Ascending)))
+ }
+
+ test("gate: a leading column outside the projection reports nothing") {
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateB)).isEmpty)
+ }
+
+ test("gate: a leading transform (non-AttributeReference) sort child reports
nothing") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateA)).isEmpty)
+ }
+
+ test("gate: absent or empty ordering reports nothing") {
+ 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) reports
nothing") {
+ // "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)
+ }
+
+ // orderingDecision decides fallback-vs-stay-native, which
reportableOrdering (the advertised
+ // prefix) alone does not capture: an empty prefix can mean "stay native,
read unordered" (safe)
+ // or "fall back to Spark" (a projected UUID/transform Spark may have
dropped a Sort on).
+
+ test("decision: a leading unprojected sort key stays native with no reported
ordering") {
+ // b (the sort key) is not selected. Nothing above the scan can require an
ordering that reaches
+ // it, so Spark dropped no Sort -- Comet stays native and simply reports
nothing.
+ val ordering = Seq(SortOrder(gateB, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && d.stayNative)
+ }
+
+ test("decision: a projected prefix then an unprojected key stays native with
the prefix") {
+ // a is projected + reportable; b is not selected. Report [a] and stay
native.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a two-field projected prefix then an unprojected key stays
native") {
+ // a and b are projected + reportable; c is not selected. Report [a, b]
and stay native.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending)))
+ assert(d.stayNative)
+ }
+
+ test("decision: an unprojected blocker before a projected UUID still stays
native") {
+ // First blocker wins: b (unprojected) stops the prefix at [a], so we stay
native even though a
+ // later field (c) is a projected unsafe/UUID column -- nothing above can
require past b.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateC), Set("c"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a projected prefix then a projected UUID falls back to
Spark") {
+ // a is projected + reportable; b is a projected UUID (unsafe). The prefix
is [a], but the
+ // blocker b is IN the output, so a required ordering could reach it and
Spark may already have
+ // dropped a Sort on it -- Comet must fall back even though a non-empty
prefix is reportable.
+ // This pins the `stayNative && ...` gate against a regression to the old
`reported.nonEmpty`.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("b"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && !d.stayNative)
+ }
+
+ test("decision: a leading UUID sort key falls back to Spark") {
+ // "a" is a projected UUID: a required ordering can reach it and Spark may
already have dropped
+ // a Sort on it, so Comet must fall back rather than read unordered.
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("a"))
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: a leading transform sort key falls back to Spark") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: no reported ordering stays native") {
+ assert(CometIcebergNativeScan.orderingDecision(None, Seq(gateA),
Set.empty).stayNative)
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // 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)") {
Review Comment:
Nothing in the suite would notice if the sort fallback stopped sorting. I
removed the `SortExec` from that branch of the planner so it returned the
unordered scan, and every test still passed. The tests on that path either end
in a global `ORDER BY`, which Spark keeps and which re-sorts whatever the scan
returns, or join and group data where order doesn't change the answer. Could we
add an order-sensitive one? Running "merges multiple sorted files per partition
and preserves order" with `maxFilesPerPartition` set to 1 does it. It takes the
sort path, its `ROW_NUMBER` depends on the order, and it failed with the
`SortExec` removed.
##########
spark/src/test/scala/org/apache/comet/CometIcebergSortMergeReadSuite.scala:
##########
@@ -0,0 +1,960 @@
+/*
+ * 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)()
+ private val gateC = AttributeReference("c", 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: reports the longest identity-projected prefix") {
+ // a is projected and reportable; the transform on b stops the prefix, so
only [a] is reported.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(Add(gateB,
Literal(1)), Ascending))
+ assert(
+ CometIcebergNativeScan.reportableOrdering(Some(ordering), Seq(gateA,
gateB)) ===
+ Seq(SortOrder(gateA, Ascending)))
+ }
+
+ test("gate: a leading column outside the projection reports nothing") {
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateB)).isEmpty)
+ }
+
+ test("gate: a leading transform (non-AttributeReference) sort child reports
nothing") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ assert(CometIcebergNativeScan.reportableOrdering(Some(ordering),
Seq(gateA)).isEmpty)
+ }
+
+ test("gate: absent or empty ordering reports nothing") {
+ 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) reports
nothing") {
+ // "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)
+ }
+
+ // orderingDecision decides fallback-vs-stay-native, which
reportableOrdering (the advertised
+ // prefix) alone does not capture: an empty prefix can mean "stay native,
read unordered" (safe)
+ // or "fall back to Spark" (a projected UUID/transform Spark may have
dropped a Sort on).
+
+ test("decision: a leading unprojected sort key stays native with no reported
ordering") {
+ // b (the sort key) is not selected. Nothing above the scan can require an
ordering that reaches
+ // it, so Spark dropped no Sort -- Comet stays native and simply reports
nothing.
+ val ordering = Seq(SortOrder(gateB, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && d.stayNative)
+ }
+
+ test("decision: a projected prefix then an unprojected key stays native with
the prefix") {
+ // a is projected + reportable; b is not selected. Report [a] and stay
native.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a two-field projected prefix then an unprojected key stays
native") {
+ // a and b are projected + reportable; c is not selected. Report [a, b]
and stay native.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set.empty)
+ assert(d.reported === Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending)))
+ assert(d.stayNative)
+ }
+
+ test("decision: an unprojected blocker before a projected UUID still stays
native") {
+ // First blocker wins: b (unprojected) stops the prefix at [a], so we stay
native even though a
+ // later field (c) is a projected unsafe/UUID column -- nothing above can
require past b.
+ val ordering =
+ Seq(SortOrder(gateA, Ascending), SortOrder(gateB, Ascending),
SortOrder(gateC, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateC), Set("c"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && d.stayNative)
+ }
+
+ test("decision: a projected prefix then a projected UUID falls back to
Spark") {
+ // a is projected + reportable; b is a projected UUID (unsafe). The prefix
is [a], but the
+ // blocker b is IN the output, so a required ordering could reach it and
Spark may already have
+ // dropped a Sort on it -- Comet must fall back even though a non-empty
prefix is reportable.
+ // This pins the `stayNative && ...` gate against a regression to the old
`reported.nonEmpty`.
+ val ordering = Seq(SortOrder(gateA, Ascending), SortOrder(gateB,
Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("b"))
+ assert(d.reported === Seq(SortOrder(gateA, Ascending)) && !d.stayNative)
+ }
+
+ test("decision: a leading UUID sort key falls back to Spark") {
+ // "a" is a projected UUID: a required ordering can reach it and Spark may
already have dropped
+ // a Sort on it, so Comet must fall back rather than read unordered.
+ val ordering = Seq(SortOrder(gateA, Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering), Seq(gateA,
gateB), Set("a"))
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: a leading transform sort key falls back to Spark") {
+ val ordering = Seq(SortOrder(Add(gateA, Literal(1)), Ascending))
+ val d = CometIcebergNativeScan.orderingDecision(Some(ordering),
Seq(gateA), Set.empty)
+ assert(d.reported.isEmpty && !d.stayNative)
+ }
+
+ test("decision: no reported ordering stays native") {
+ assert(CometIcebergNativeScan.orderingDecision(None, Seq(gateA),
Set.empty).stayNative)
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // 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") {
Review Comment:
Even with `p` selected, this test and the sort-fallback one below cancel on
apache/iceberg#16750. Iceberg caps the files it will merge in one task with
`spark.sql.iceberg.planning.preserve-data-ordering.max-files-per-task`, default
64, and reports no ordering at all for a task group over that. Copies of the
two tests with it set to 1000 ran a 70-way merge and a 70-file sort and passed.
Could both tests set it? Because both defaults are 64, the sort-fallback test
can't reach its path any other way.
##########
docs/source/user-guide/latest/iceberg.md:
##########
@@ -58,6 +58,26 @@ config
`spark.comet.scan.icebergNative.dataFileConcurrencyLimit`. This value def
maintain test behavior on Iceberg Java tests without `ORDER BY` clauses, but
we suggest increasing it to
values between 2 and 8 based on your workload.
+When a table is sorted and Iceberg reports its sort order (which requires
Iceberg's
+`spark.sql.iceberg.planning.preserve-data-ordering`), Comet reads each
partition's already-sorted
Review Comment:
No Iceberg release reports a sort order yet. 1.12.0 has neither
`SupportsReportOrdering` nor `preserve-data-ordering`, and both come from
apache/iceberg#16750, which is still open. Someone who sets the property on a
released Iceberg gets nothing, and this paragraph reads as if that's all it
takes. Could it say which Iceberg version is needed, or link the PR until there
is one? The `sortMerge.enabled` doc in `CometConf` needs the same note. It
would also help to list what Iceberg itself requires, since that decides
whether anything happens: `preserve-data-grouping` on, a partitioned table, a
query that reads the partition columns, and no more files in a task than
Iceberg's `max-files-per-task` (64 by default). With both defaults at 64, our
`maxFilesPerPartition` fallback only comes into play once someone raises
Iceberg's limit.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
