From 34a84119103e13f5be3e3a771995b4d0e8560d83 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 12 Oct 2026 14:23:23 +0100 Subject: [PATCH 1/2] fix: keep a Spark shuffle between two Spark operators under cost-based pricing With spark.comet.exec.costBasedEngines.enabled, BoundaryFormats prices every shuffle format with EngineCostModel. Between a Spark producer and a Spark consumer the model can price Comet's columnar shuffle (rows to Arrow, Comet write and read, Arrow to rows) below a Spark shuffle once rows are wide enough, so a Comet columnar shuffle was planned where no Comet operator reads or writes it. In spark_etl load_sat_variant_warehouse_destinations_info (union and sort in Spark under a REPARTITION before a parquet write) that shuffle took 2960 s of encoding and writing against 197 s for the Spark shuffle of the same stage. A shuffle between two Spark operators is no longer offered the columnar format; the Spark shuffle, with the same hash function, is always a candidate there, so no co-partitioning mode loses its feasible choices. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apache/comet/rules/BoundaryFormats.scala | 3 +- .../rules/CostBasedEngineChoiceSuite.scala | 62 ++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/BoundaryFormats.scala b/spark/src/main/scala/org/apache/comet/rules/BoundaryFormats.scala index f2725906ddf..bfb65739d18 100644 --- a/spark/src/main/scala/org/apache/comet/rules/BoundaryFormats.scala +++ b/spark/src/main/scala/org/apache/comet/rules/BoundaryFormats.scala @@ -365,7 +365,8 @@ object BoundaryFormats extends Logging with CometTypeShim { conversionsInto(input.consumer, NativeShuffle), CometHash)) } - if (current != SparkShuffle && columnarAvailable(input)) { + val betweenSparkOperators = !producerIsComet && input.consumer.contains(Engine.Spark) + if (current != SparkShuffle && !betweenSparkOperators && columnarAvailable(input)) { val write = if (producerIsComet) 2 else 1 candidates += (( ColumnarShuffle, diff --git a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala index 2e2ff2673e6..87b38712881 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.Partial import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.{ExpandExec, ProjectExec, SortExec, SparkPlan} +import org.apache.spark.sql.execution.{CommandResultExec, ExpandExec, ProjectExec, SortExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec @@ -153,6 +153,66 @@ class CostBasedEngineChoiceSuite extends CometTestBase { private def count[T](plan: SparkPlan)(pf: PartialFunction[SparkPlan, T]): Int = nodes(plan).collect(pf).size + /** The executed plan of `INSERT INTO target `, checking the rows it wrote. */ + private def insertPlan(schema: String, query: String): SparkPlan = { + val expected = sql(query).collect().toSeq + var plan: SparkPlan = null + withTable("target") { + sql(s"CREATE TABLE target ($schema) USING parquet") + val df = sql(s"INSERT INTO target $query") + checkAnswer(spark.table("target"), expected) + plan = df.queryExecution.executedPlan match { + case c: CommandResultExec => c.commandPhysicalPlan + case other => other + } + } + plan + } + + private def columnarBetweenSpark(plan: SparkPlan): Seq[Edge] = + edges(plan).filter(e => e.format == "columnar" && !e.consumerIsComet && !e.producerIsComet) + + for (aqe <- Seq("false", "true")) { + test(s"a shuffle between Spark operators stays a Spark shuffle at any price (AQE=$aqe)") { + withWide("w", 20) { + withAqe( + aqe, + flag -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false", + CometConf.COMET_EXEC_UNION_ENABLED.key -> "false", + CometConf.COMET_EXEC_SORT_ENABLED.key -> "false", + costTable -> "shuffleWrite.flat.spark=100000,0") { + val payload = (1 to 20).map(i => s"c$i") + val schema = ("k INT" +: payload.map(c => s"$c INT")).mkString(", ") + val query = + s"SELECT /*+ REPARTITION(37, k) */ * FROM (" + + s"SELECT k, ${payload.map(c => s"$c + 1 AS $c").mkString(", ")} FROM w " + + s"UNION ALL SELECT k, ${payload.mkString(", ")} FROM w) SORT BY k" + val plan = insertPlan(schema, query) + assert(columnarBetweenSpark(plan).isEmpty, s"columnar shuffle between Spark:\n$plan") + assert( + edges(plan).exists(e => e.format == "spark" && !e.producerIsComet), + s"expected a Spark shuffle under the write:\n$plan") + } + } + } + + test(s"a native producer still writes a Comet shuffle for a Spark write (AQE=$aqe)") { + withWide("w", 20) { + withAqe(aqe, flag -> "true", costTable -> "shuffleWrite.flat.spark=100000,0") { + val payload = (1 to 20).map(i => s"c$i") + val schema = ("k INT" +: payload.map(c => s"$c INT")).mkString(", ") + val plan = insertPlan( + schema, + s"SELECT /*+ REPARTITION(37, k) */ k, ${payload.mkString(", ")} FROM w") + assert( + edges(plan).exists(e => e.format != "spark" && e.producerIsComet), + s"expected a Comet shuffle written by a native operator:\n$plan") + } + } + } + } + for (aqe <- Seq("false", "true")) { test(s"native sorts read by Spark sort aggregates move only by their price (AQE=$aqe)") { withTables { From 3fa650292141ddc979250c02a3809795cf6bb175 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 12 Oct 2026 14:46:00 +0100 Subject: [PATCH 2/2] test: drop a needless string interpolator Co-Authored-By: Claude Opus 5.5 (1M context) --- .../org/apache/comet/rules/CostBasedEngineChoiceSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala index 87b38712881..3f463370824 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala @@ -185,7 +185,7 @@ class CostBasedEngineChoiceSuite extends CometTestBase { val payload = (1 to 20).map(i => s"c$i") val schema = ("k INT" +: payload.map(c => s"$c INT")).mkString(", ") val query = - s"SELECT /*+ REPARTITION(37, k) */ * FROM (" + + "SELECT /*+ REPARTITION(37, k) */ * FROM (" + s"SELECT k, ${payload.map(c => s"$c + 1 AS $c").mkString(", ")} FROM w " + s"UNION ALL SELECT k, ${payload.mkString(", ")} FROM w) SORT BY k" val plan = insertPlan(schema, query)