diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 45fc26caee9..6bb3f0dcd59 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -22,7 +22,7 @@ package org.apache.comet.rules import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch} import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan} @@ -74,7 +74,7 @@ case class EliminateRedundantTransitions(session: SparkSession) } private def _apply(plan: SparkPlan): SparkPlan = { - val eliminatedPlan = plan transformUp { + val eliminatedPlan = stripIcebergWriteInputTransition(plan) transformUp { case ColumnarToRowExec(shuffleExchangeExec: CometShuffleExchangeExec) if plan.conf.adaptiveExecutionEnabled => shuffleExchangeExec @@ -169,6 +169,57 @@ case class EliminateRedundantTransitions(session: SparkSession) } } + /** + * `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg commit message) but + * consumes Arrow batches from its child over FFI, so Spark's + * `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *underneath* + * it. Strip that transition so `doExecuteColumnar` sees the columnar child directly; + * `CometIcebergNativeWrite.requiresNativeChildren` already guarantees the child was + * Comet-native when the write was converted. + * + * The write deliberately does not tag itself as a `ColumnarToRowTransition` to suppress the + * insertion: Spark leaves such a node untouched, so the whole subtree below the write is never + * visited and the transitions the rest of that subtree needs are never inserted + * (https://github.com/apache/datafusion-comet/issues/5689). + * + * This runs as a separate pass *before* the main `transformUp`, not as an arm inside it, + * because the write's input transition has to be removed before the generic transition + * cancellation can consume it. `transformUp` visits children first, so the + * `ColumnarToRowExec(CometSparkToColumnarExec)` arm would otherwise rewrite the write's child + * before the write itself is visited, and that arm is destructive at this boundary: + * - over a row source it drops the `CometSparkToColumnarExec` as well, leaving the write with + * a row child and no Arrow producer at all; + * - over a Spark-columnar source it keeps a `ColumnarToRowExec` but drops the Arrow bridge, + * so the write would be handed Spark `ColumnarVector`s where the FFI adapter requires + * `CometVector`s. That shape is reachable whenever `spark.comet.sparkToColumnar.enabled` + * admits the write's source: `CometSparkToColumnarExec.createExec` wraps it in a + * `CometScanWrapper` (a `CometNativeExec`, so `requiresNativeChildren` accepts it), and + * `CometExecRule` then unwraps the placeholder, leaving the bridge directly beneath the + * write. + */ + private def stripIcebergWriteInputTransition(plan: SparkPlan): SparkPlan = plan.transform { + case w: CometIcebergWriteExec => + stripColumnarToRow(w.child).map(child => w.withNewChildren(Seq(child))).getOrElse(w) + } + + /** + * Unwraps a columnar-to-row transition, returning the columnar child underneath it, or `None` + * when `plan` is not such a transition. + * + * The plain `ColumnarToRowExec` is what Spark's insertion pass produces and is the only form + * observed in the suites. The two Comet variants are handled as well because this rule is part + * of `postColumnarTransitions` and AQE applies those rules once per materialised stage plus + * once for the final plan, so it can be handed a plan whose transitions an earlier pass already + * rewrote. Missing a variant would not be caught at planning time -- the write would fail at + * runtime with `requires a columnar (Comet native) child`. + */ + private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match { + case CometNativeColumnarToRowExec(child) => Some(child) + case CometColumnarToRowExec(child) => Some(child) + case ColumnarToRowExec(child) => Some(child) + case _ => None + } + /** * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python * UDF operator can consume directly. By the time this rule runs the earlier diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 2cafe2cf49a..25e0c5873c0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.UnsafeProjection import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.types.BinaryType import org.apache.spark.sql.vectorized.ColumnarBatch @@ -68,13 +68,7 @@ case class CometIcebergWriteExec( @transient table: AnyRef, partitionSpecId: Int) extends CometNativeExec - with UnaryExecNode - // We consume Arrow batches (via FFI) and emit row-shaped commit messages, so we are a - // columnar-to-row transition. Without this trait Spark's - // `ApplyColumnarRulesAndInsertTransitions` wedges a `CometNativeColumnarToRowExec` between - // us and the Comet-native child, which would then fail `child.executeColumnar()` in - // `doExecuteColumnar`. - with ColumnarToRowTransition { + with UnaryExecNode { override def originalPlan: SparkPlan = child @@ -86,6 +80,16 @@ case class CometIcebergWriteExec( // Native exec emits a single Binary column; the surrounding command framework expects rows, so // the outer commit exec calls executeCollect on us. supportsColumnar = false keeps Spark from // inserting a ColumnarToRow that would clash with our (Nil-output-like) row contract. + // + // We do consume Arrow batches from `child` over FFI, so Spark's + // `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *below* us + // (our child is columnar-only, we are row-based). `EliminateRedundantTransitions` strips that + // transition back off so `doExecuteColumnar` sees the Comet-native child directly. + // + // Tagging this node as a `ColumnarToRowTransition` would also stop Spark inserting it, but at + // the cost of correctness: `ensureOutputsRowBased` returns such a node untouched, so nothing + // below the write is visited and the transitions the rest of that subtree needs are never + // inserted. See https://github.com/apache/datafusion-comet/issues/5689. override def supportsColumnar: Boolean = false override def executeCollect(): Array[InternalRow] = { diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 66ea7216786..8095e6feb05 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -32,8 +32,9 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.Row import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} @@ -696,6 +697,57 @@ class CometIcebergWriteActionSuite } } + // https://github.com/apache/datafusion-comet/issues/5689: the Iceberg CoW rewrite plan mixes a + // Spark-columnar `BatchScan` with row-based joins/filters underneath the write, so the write's + // subtree needs Spark to insert columnar-to-row transitions inside it. With AQE on those + // transitions are re-inserted when each stage materialises, which is why every other Iceberg + // suite in this file misses the problem; with AQE off the plan gets exactly one insertion pass + // and a missed transition fails the task with + // `ColumnarBatch cannot be cast to InternalRow`. + test("native acceleration: ReplaceData (CoW DELETE) with AQE disabled") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_delete_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.delete.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert( + "native_cow_delete_no_aqe", + Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4, "us-east", 40.0))) + } + // The IN-subquery is what puts a row-based join between the columnar CoW scan and the + // write, so the subtree genuinely needs a transition rather than being Comet-native + // end to end. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val snapshot = withNativeEnabled { + captureWrite("native_cow_delete_no_aqe") { + spark.sql("DELETE FROM cat.db.native_cow_delete_no_aqe WHERE id IN " + + "(SELECT col1 FROM VALUES (2) AS t(col1))") + } + } + assert( + snapshot.snapshotDelta == 1L, + s"expected 1 commit via native path, got ${snapshot.snapshotDelta}") + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert( + nativeExecs.nonEmpty, + "expected >= 1 CometIcebergWriteExec in captured plans, got 0. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + // The transition Spark inserts below the write must have been stripped again, otherwise + // `doExecuteColumnar` has no columnar child to pull Arrow batches from. + assert( + nativeExecs.forall(_.child.supportsColumnar), + "the native write must sit directly on its columnar child, got " + + nativeExecs.map(_.child.nodeName).mkString(", ")) + } + assertRows("native_cow_delete_no_aqe", Seq(1, 3, 4)) + } + } + test("native acceleration: ReplaceData (CoW UPDATE)") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => @@ -752,6 +804,73 @@ class CometIcebergWriteActionSuite } } + // The remaining two CoW shapes from issue #5689, with AQE off so the plan gets exactly one + // transition-insertion pass. `captureWrite` runs `assertColumnarContract` on every captured + // plan, so these pin the same invariant as the DELETE case over the UPDATE and MERGE rewrite + // shapes, which put different operators between the columnar CoW scan and the write. + test("native acceleration: ReplaceData (CoW UPDATE) with AQE disabled") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_update_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.update.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert( + "native_cow_update_no_aqe", + Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0))) + } + // As in the DELETE case, the IN-subquery is what puts a row-based join between the + // Spark-columnar CoW scan and the write. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assertNativeWriteEngages("native_cow_update_no_aqe", Seq(1, 2, 3)) { + spark.sql( + "UPDATE cat.db.native_cow_update_no_aqe SET amount = amount * 2 " + + "WHERE id IN (SELECT col1 FROM VALUES (2) AS t(col1))") + } + } + val r = spark + .sql("SELECT amount FROM cat.db.native_cow_update_no_aqe WHERE id = 2") + .collect() + assert(r.length == 1 && r(0).getDouble(0) == 40.0, s"got ${r.toSeq}") + } + } + + test("native acceleration: ReplaceData (CoW MERGE) with AQE disabled") { + // Unlike the unpartitioned MERGE above, this one *does* engage natively: partitioning the + // table puts a `CometColumnarExchange` (REBALANCE_PARTITIONS_BY_COL) between the JVM + // `MergeRowsExec` and the write, so the write's own child is Comet-native and + // `requiresNativeChildren` is satisfied even though `MergeRowsExec` itself is not. + // + // That makes this the strongest of the three CoW shapes for issue #5689: the subtree below + // the write mixes a Spark-columnar `BatchScan`, a row-based `MergeRowsExec` and Comet + // operators, so it needs transitions inserted in three different places. + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_merge_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.merge.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert("native_cow_merge_no_aqe", Seq((1, "us-east", 10.0), (2, "us-west", 20.0))) + } + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assertNativeWriteEngages("native_cow_merge_no_aqe", Seq(1, 2, 3)) { + spark.sql(""" + |MERGE INTO cat.db.native_cow_merge_no_aqe t + |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION ALL + | SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.amount = s.amount + |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id, s.region, s.amount) + |""".stripMargin) + } + } + } + } + test("native acceleration: complex types (struct, array, map) round-trip with field IDs") { assumeNativeAcceleration() withIcebergCatalog { _ => @@ -1633,6 +1752,10 @@ class CometIcebergWriteActionSuite private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) val plans = capturePlans(spark)(action) + // Every write in this suite gets the columnar contract checked, not just the ones written to + // exercise it: a Comet rewrite that hides part of a subtree from Spark's transition-insertion + // pass is a whole class of bug (issue #5689) and the check is free once the plans are here. + plans.foreach(assertColumnarContract) WriteSnapshot(countSnapshots(tableName) - before, plans) } @@ -1688,6 +1811,31 @@ class CometIcebergWriteActionSuite assume(icebergAvailable, "Iceberg not available in classpath") } + /** + * Every row-consuming operator must receive row-based input. Spark guarantees that by inserting + * `ColumnarToRow` transitions in `ApplyColumnarRulesAndInsertTransitions`; an operator that + * Comet rewrote in a way that skips the insertion pass shows up here as a row-based node with a + * columnar-only child, and would fail at runtime with a `ColumnarBatch cannot be cast to + * InternalRow` `ClassCastException` rather than at planning time. + * + * `CometIcebergWriteExec` is the one legitimate exception: it is row-based on the outside but + * pulls Arrow batches from its Comet-native child over FFI (see the class docstring), so the + * transition below it is deliberately stripped again by `EliminateRedundantTransitions`. + */ + private def assertColumnarContract(plan: SparkPlan): Unit = { + val violations = collectWithSubqueries(plan) { + case p + if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] && + !p.isInstanceOf[CometIcebergWriteExec] && + p.children.exists(c => c.supportsColumnar && !c.supportsRowBased) => + p + } + assert( + violations.isEmpty, + "row-based operators consuming columnar-only children (missing ColumnarToRow): " + + s"${violations.map(_.nodeName).mkString(", ")}. Plan:\n$plan") + } + /** * Flip [[CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED]] for the duration of `action`. * diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index fc530a2d2e1..7c5631640c8 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -29,12 +29,20 @@ import org.apache.iceberg.hadoop.{HadoopConfigurable, HadoopFileIO} import org.apache.iceberg.io.{FileIO, InputFile, OutputFile} import org.apache.iceberg.util.SerializableSupplier import org.apache.spark.SparkConf +import org.apache.spark.rdd.RDD import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.IcebergWriteExec +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSparkToColumnarExec, IcebergWriteExec} +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, ColumnarToRowExec, LeafExecNode, SparkPlan} +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.rules.EliminateRedundantTransitions import org.apache.comet.serde.{Compatible, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.CometIcebergNativeWrite class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTestBase { @@ -881,6 +889,71 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes fail(s"expected Unsupported for $tableName, got $other") } } + + /** + * Runs Spark's transition insertion followed by [[EliminateRedundantTransitions]] over a + * hand-built `CometIcebergWriteExec -> CometSparkToColumnarExec -> source` plan and returns the + * write's final child. + * + * Hand-built rather than driven through SQL because the shape depends on + * `spark.comet.sparkToColumnar.enabled` admitting the write's source operator, and the set of + * admitted operators is itself configurable. What matters is the rule's behaviour at that + * boundary, which this pins directly. + */ + private def writeChildAfterTransitionRules(source: SparkPlan): SparkPlan = { + val write = CometIcebergWriteExec( + Operator.newBuilder().build(), + CometSparkToColumnarExec(source), + batchWrite = null, + table = null, + partitionSpecId = 0) + val withTransitions = ApplyColumnarRulesAndInsertTransitions(Seq.empty, false).apply(write) + // Spark must insert a columnar-to-row transition below the row-based write; if it stops doing + // so the rest of the assertion is vacuous. + assert( + withTransitions.asInstanceOf[CometIcebergWriteExec].child.isInstanceOf[ColumnarToRowExec], + s"expected an inserted ColumnarToRowExec below the write, got:\n$withTransitions") + EliminateRedundantTransitions(spark) + .apply(withTransitions) + .asInstanceOf[CometIcebergWriteExec] + .child + } + + // https://github.com/apache/datafusion-comet/issues/5689: the write's input transition has to + // be stripped before the generic `ColumnarToRowExec(CometSparkToColumnarExec)` cancellation + // consumes it, otherwise that arm removes the Arrow bridge the write's FFI input depends on. + // Both source representations are covered because the cancellation treats them differently: + // over a row source it drops the bridge outright, over a Spark-columnar source it keeps a + // transition but leaves the write reading Spark `ColumnarVector`s instead of `CometVector`s. + test("row source keeps its Arrow bridge under the native Iceberg write") { + val source = TransitionProbeLeaf(columnar = false) + val child = writeChildAfterTransitionRules(source) + assert( + child == CometSparkToColumnarExec(source), + s"expected the write to sit directly on CometSparkToColumnarExec, got:\n$child") + } + + test("Spark-columnar source keeps its Arrow bridge under the native Iceberg write") { + val source = TransitionProbeLeaf(columnar = true) + val child = writeChildAfterTransitionRules(source) + assert( + child == CometSparkToColumnarExec(source), + s"expected the write to sit directly on CometSparkToColumnarExec, got:\n$child") + } +} + +/** + * Planning-only leaf used by the transition-boundary tests: `columnar` selects between the two + * source representations that can sit under a `CometSparkToColumnarExec`. Never executed. + */ +case class TransitionProbeLeaf(columnar: Boolean) extends LeafExecNode { + override def output: Seq[Attribute] = Seq( + AttributeReference("id", IntegerType, nullable = false)()) + override def supportsColumnar: Boolean = columnar + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException("planning-only node") + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = + throw new UnsupportedOperationException("planning-only node") } /**