diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 93d39973b0e..8666d8c80fb 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -53,7 +53,6 @@ The following operator families fall back to Spark and are not on the current ro omitted from the tables below and may be reconsidered based on demand: - **Structured Streaming operators** (`StateStoreSaveExec`, `StateStoreRestoreExec`, `StreamingSymmetricHashJoinExec`, and similar): Comet targets batch execution. -- **Cartesian / cross joins** (`CartesianProductExec`): rare and expensive, with little acceleration benefit. - **Pickled (non-Arrow) Python UDFs** (`BatchEvalPythonExec`): Comet accelerates Arrow-based Python UDFs only ([#4234](https://github.com/apache/datafusion-comet/pull/4234)). - **Typed Dataset operators** (`DeserializeToObjectExec`, `SerializeFromObjectExec`, `MapElementsExec`, `MapPartitionsExec`, `MapGroupsExec`, `CoGroupExec`, `AppendColumnsExec`, and similar): produced by `map`, `mapPartitions`, `groupByKey`, `cogroup`, and other typed `Dataset` transformations. They exist to run user JVM functions on JVM objects, which Comet cannot do natively. The operators above them can still run natively: set `spark.comet.convert.typedDataset.enabled=true` and Comet converts the output of a typed operation to Arrow. This is disabled by default because it can be slower than Spark when the operators above do little work, such as an aggregate over a few groups, which Spark compiles together with the typed operation into one loop. @@ -102,6 +101,7 @@ omitted from the tables below and may be reconsidered based on demand: | `ShuffledHashJoinExec` | ✅ | | | `SortMergeJoinExec` | ✅ | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark. | | `BroadcastNestedLoopJoinExec` | ✅ | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)). | +| `CartesianProductExec` | ✅ | Opt-in (`spark.comet.exec.cartesianProduct.enabled`). Only condition-free products of single-partition inputs, such as after `COALESCE(1)`; never a bare file scan. | `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 494e7acdff9..880d083ba27 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -497,6 +497,7 @@ fn op_name(op: &OpStruct) -> &'static str { OpStruct::CsvScan(_) => "CsvScan", OpStruct::ShuffleScan(_) => "ShuffleScan", OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin", + OpStruct::CrossJoin(_) => "CrossJoin", OpStruct::Sample(_) => "Sample", OpStruct::MergeRows(_) => "MergeRows", OpStruct::RangeScan(_) => "RangeScan", diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index 0b91f27187d..64687195a30 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! Builders for the native join operators (hash join, sort-merge join, and broadcast nested loop -//! join), and the join-key and join-filter planning they share. +//! Builders for the native join operators (hash join, sort-merge join, broadcast nested loop +//! join, and cross join), and the join-key and join-filter planning they share. use std::{collections::HashMap, sync::Arc}; @@ -37,7 +37,8 @@ use datafusion::{ physical_plan::{ filter::FilterExec, joins::{ - utils::JoinFilter, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + utils::JoinFilter, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, }, projection::ProjectionExec, ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions, @@ -344,6 +345,48 @@ impl OperatorBuilder for BroadcastNestedLoopJoinBuilder { } } +/// Builder for `CrossJoin` (Spark's `CartesianProductExec`). +pub struct CrossJoinBuilder; + +impl OperatorBuilder for CrossJoinBuilder { + fn build( + &self, + spark_plan: &Operator, + inputs: &mut Vec>>>, + partition_count: usize, + planner: &PhysicalPlanner, + ) -> PlanCreationResult { + let children = &spark_plan.children; + + // No keys or condition. The join type only satisfies the helper; CrossJoinExec emits left + // then right columns, as Spark does. + let (join_params, scans, shuffle_scans) = planner.parse_join_parameters( + inputs, + children, + &[], + &[], + JoinType::Inner as i32, + &None, + partition_count, + )?; + + let left = Arc::clone(&join_params.left.native_plan); + let right = Arc::clone(&join_params.right.native_plan); + + let cross_join = Arc::new(CrossJoinExec::new(left, right)); + + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new( + spark_plan.plan_id, + cross_join, + vec![join_params.left, join_params.right], + )), + )) + } +} + impl PhysicalPlanner { /// Keep Spark metric identities for the small set of nodes whose children /// runtime-filter placement can replace. Stop at native/Spark tree boundaries. diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs index c80a40a290c..ccc7f4afbf9 100644 --- a/native/core/src/execution/planner/operator_registry.rs +++ b/native/core/src/execution/planner/operator_registry.rs @@ -63,6 +63,7 @@ pub enum OperatorType { MergeRows, Explode, BroadcastNestedLoopJoin, + CrossJoin, WindowGroupLimit, } @@ -151,7 +152,9 @@ impl OperatorRegistry { /// Register join operators fn register_join_operators(&mut self) { - use super::join::{BroadcastNestedLoopJoinBuilder, HashJoinBuilder, SortMergeJoinBuilder}; + use super::join::{ + BroadcastNestedLoopJoinBuilder, CrossJoinBuilder, HashJoinBuilder, SortMergeJoinBuilder, + }; self.builders .insert(OperatorType::HashJoin, Box::new(HashJoinBuilder)); @@ -161,6 +164,8 @@ impl OperatorRegistry { OperatorType::BroadcastNestedLoopJoin, Box::new(BroadcastNestedLoopJoinBuilder), ); + self.builders + .insert(OperatorType::CrossJoin, Box::new(CrossJoinBuilder)); } /// Register window operators @@ -236,6 +241,7 @@ fn get_operator_type(spark_operator: &Operator) -> Option { OpStruct::CsvScan(_) => Some(OperatorType::CsvScan), OpStruct::ShuffleScan(_) => Some(OperatorType::ShuffleScan), OpStruct::BroadcastNestedLoopJoin(_) => Some(OperatorType::BroadcastNestedLoopJoin), + OpStruct::CrossJoin(_) => Some(OperatorType::CrossJoin), OpStruct::RangeScan(_) => None, OpStruct::Sample(_) => None, // Not yet in OperatorType enum OpStruct::MergeRows(_) => Some(OperatorType::MergeRows), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index c4c25d9f73d..73d95ba1b0b 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -71,6 +71,7 @@ message Operator { RangeScan range_scan = 121; MergeRows merge_rows = 122; ArrowPythonUdf arrow_python_udf = 123; + CrossJoin cross_join = 124; // Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete // scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched // by `type_url` on the native side. Using a single permanent field -- rather than a new oneof @@ -1156,6 +1157,10 @@ message BroadcastNestedLoopJoin { optional spark.spark_expression.Expr condition = 3; } +// Spark's CartesianProductExec without a condition (one with a condition falls back to Spark). +message CrossJoin { +} + enum JoinType { Inner = 0; LeftOuter = 1; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..009191aaf85 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -361,6 +361,14 @@ object CometConf extends ShimCometConf { createExecEnabledConfig("broadcastNestedLoopJoin", defaultValue = true) val COMET_EXEC_SORT_MERGE_JOIN_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("sortMergeJoin", defaultValue = true) + val COMET_EXEC_CARTESIAN_PRODUCT_ENABLED: ConfigEntry[Boolean] = + createExecEnabledConfig( + "cartesianProduct", + defaultValue = false, + notes = Some( + "Opt-in: DataFusion's CrossJoinExec buffers the left input and cannot spill. Only " + + "condition-free products of single-partition inputs run natively, such as after " + + "COALESCE(1); never a bare file scan")) val COMET_EXEC_EXISTENCE_JOIN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.existenceJoin.enabled") .category(CATEGORY_ENABLE_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d79927f0c04..543de3bce5f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -52,7 +52,7 @@ import org.apache.spark.sql.execution.datasources.v2.json.JsonScan import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.execution.datasources.v2.text.TextScan import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, Exchange, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -120,6 +120,7 @@ object CometExecRule { classOf[BroadcastNestedLoopJoinExec] -> CometBroadcastNestedLoopJoinExec, classOf[ShuffledHashJoinExec] -> CometHashJoinExec, classOf[SortMergeJoinExec] -> CometSortMergeJoinExec, + classOf[CartesianProductExec] -> CometCartesianProductExec, classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, classOf[RangeExec] -> CometRangeExec, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3f331fc5b9c..0461365a122 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -45,7 +45,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanExecBase import org.apache.spark.sql.execution.exchange.ReusedExchangeExec -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType} @@ -3273,6 +3273,75 @@ object CometBroadcastNestedLoopJoinExec extends CometOperatorSerde[BroadcastNest } } +case class CometCartesianProductExec( + override val nativeOp: Operator, + override val originalPlan: SparkPlan, + override val output: Seq[Attribute], + override val left: SparkPlan, + override val right: SparkPlan, + override val serializedPlanOpt: SerializedPlan) + extends CometBinaryExec { + + // Inherits Spark's UnknownPartitioning from originalPlan, so exchange planning is unchanged. + + override def withNewChildrenInternal(newLeft: SparkPlan, newRight: SparkPlan): SparkPlan = + this.copy(left = newLeft, right = newRight) + + override def stringArgs: Iterator[Any] = Iterator(left, right) + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometCartesianProductExec => + this.output == other.output && + this.left == other.left && + this.right == other.right && + this.serializedPlanOpt == other.serializedPlanOpt + case _ => + false + } + } + + override def hashCode(): Int = Objects.hashCode(output, left, right) + + override lazy val metrics: Map[String, SQLMetric] = + CometMetricNode.joinMetrics(sparkContext) +} + +object CometCartesianProductExec extends CometOperatorSerde[CartesianProductExec] { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_CARTESIAN_PRODUCT_ENABLED) + + override def getSupportLevel(op: CartesianProductExec): SupportLevel = { + // CartesianProductExec carries its own condition, and CrossJoinExec has no post-product filter. + if (op.condition.isDefined) { + Unsupported(Some("CartesianProduct with a join condition is not supported natively")) + } else if (op.left.outputPartitioning.numPartitions != 1 || + op.right.outputPartitioning.numPartitions != 1) { + // A native plan runs per partition and cannot add an exchange here, so a cross product is + // only correct when both inputs are a single partition. + Unsupported( + Some( + "Native CartesianProduct requires single-partition inputs (a multi-partition " + + "cross product cannot be expressed per-partition)")) + } else { + Compatible(None) + } + } + + override def convert( + op: CartesianProductExec, + builder: Operator.Builder, + childOp: Operator*): Option[Operator] = { + val joinBuilder = OperatorOuterClass.CrossJoin.newBuilder() + Some(builder.setCrossJoin(joinBuilder).build()) + } + + override def createExec(nativeOp: Operator, op: CartesianProductExec): CometNativeExec = { + CometCartesianProductExec(nativeOp, op, op.output, op.left, op.right, SerializedPlan(None)) + } +} + object CometBroadcastHashJoinExec extends CometOperatorSerde[HashJoin] with CometHashJoin { override def enabledConfig: Option[ConfigEntry[Boolean]] = diff --git a/spark/src/test/resources/sql-tests/join/cartesian_product.sql b/spark/src/test/resources/sql-tests/join/cartesian_product.sql new file mode 100644 index 00000000000..00ee9efe7f2 --- /dev/null +++ b/spark/src/test/resources/sql-tests/join/cartesian_product.sql @@ -0,0 +1,78 @@ +-- 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. + +-- Tests for CartesianProductExec, which Spark plans for a cross join with no equi key when +-- nothing is broadcast. Comet runs it natively only for a condition-free product of +-- single-partition inputs. A file scan reports UnknownPartitioning(0), so each side is +-- coalesced to one partition. AQE is off so it cannot coalesce the REPARTITION(2) case. + +-- Config: spark.comet.exec.cartesianProduct.enabled=true +-- Config: spark.sql.autoBroadcastJoinThreshold=-1 +-- Config: spark.sql.adaptive.autoBroadcastJoinThreshold=-1 +-- Config: spark.sql.adaptive.enabled=false + +statement +CREATE TABLE cp_left(k int, v int) USING parquet + +statement +INSERT INTO cp_left VALUES (0, 10), (NULL, 11), (2, NULL), (NULL, NULL) + +statement +CREATE TABLE cp_right(k int, v int) USING parquet + +statement +INSERT INTO cp_right VALUES (0, 100), (NULL, NULL) + +statement +CREATE TABLE cp_empty(k int, v int) USING parquet + +statement +CREATE TABLE cp_large(k int, v int) USING parquet + +-- More rows than the default batch size of 8192. +statement +INSERT INTO cp_large SELECT id, id FROM range(10000) + +-- Native, with NULLs on both sides +query +SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM cp_left) l +CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_right) r + +-- Native, empty left side (the side CrossJoinExec buffers) +query +SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM cp_empty) l +CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_right) r + +-- Native, empty right side +query +SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM cp_left) l +CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_empty) r + +-- Native, right side spanning several batches +query +SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM cp_left) l +CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_large) r + +-- A join condition falls back to Spark +query expect_fallback(CartesianProduct with a join condition is not supported natively) +SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM cp_left) l +JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_right) r ON l.v > r.v + +-- A multi-partition input falls back to Spark +query expect_fallback(Native CartesianProduct requires single-partition inputs) +SELECT * FROM (SELECT /*+ REPARTITION(2) */ * FROM cp_left) l +CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM cp_right) r diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index b39d16f1235..710ef5887d2 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeRefer import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometCartesianProductExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.SortAggregateExec @@ -154,6 +154,25 @@ class CometJoinSuite extends CometTestBase { } } + test("cross join stays on Spark by default (native path is opt-in)") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((0, 10), (1, 11), (2, 12)), "l") { + withParquetTable(Seq((0, 100), (1, 101)), "r") { + val df = sql( + "SELECT * FROM (SELECT /*+ COALESCE(1) */ * FROM l) l " + + "CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM r) r") + checkSparkAnswer(df) + assert( + !df.queryExecution.executedPlan.exists(_.isInstanceOf[CometCartesianProductExec]), + "cross join must stay on Spark when cartesianProduct.enabled is false") + } + } + } + } + test("join - self join") { val df1 = testData.select(testData("key")).as("df1") val df2 = testData.select(testData("key")).as("df2") diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCartesianProductBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCartesianProductBenchmark.scala new file mode 100644 index 00000000000..6d49037855c --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCartesianProductBenchmark.scala @@ -0,0 +1,193 @@ +/* + * 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.spark.sql.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, SparkSession} +import org.apache.spark.sql.comet.{CometCartesianProductExec, CometExec, CometPlan} +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.joins.CartesianProductExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Benchmark for `CartesianProductExec`. Comet runs it natively only for a condition-free product + * of single-partition inputs, so each side is one Parquet file coalesced to one partition. The + * cases cover a small left side, a large left side (`CrossJoinExec` buffers the left) and a + * square product, each run with Spark, Comet with the product in Spark, and Comet. To run it: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometCartesianProductBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometCartesianProductBenchmark-**results.txt". + */ +object CometCartesianProductBenchmark extends CometBenchmarkBase { + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometCartesianProductBenchmark") + .set("spark.master", "local[5]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + // Off-heap memory is disabled, so Comet needs on-heap mode enabled. + .set("spark.comet.exec.onHeap.enabled", "true") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + sparkSession.conf.set("spark.sql.shuffle.partitions", "2") + // No broadcast, so Spark plans a CartesianProductExec. + sparkSession.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") + sparkSession.conf.set(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key, "-1") + // Keep each one-file scan to one task. + sparkSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, (1L << 30).toString) + + sparkSession + } + + private case class CartesianCase(name: String, leftRows: Long, rightRows: Long) + + private val SparkCaseName = "Spark" + private val ProductInSparkCaseName = "Comet, product in Spark" + private val CometCaseName = "Comet" + + private val sparkConfigs = Seq(CometConf.COMET_ENABLED.key -> "false") + + private val productInSparkConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_CARTESIAN_PRODUCT_ENABLED.key -> "false") + + private val cometConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_CARTESIAN_PRODUCT_ENABLED.key -> "true") + + // A file scan reports UnknownPartitioning(0), which the native product rejects, so each side is + // coalesced to one partition. The sum reads both sides. + private val query = + "SELECT sum(l.v + r.w) FROM (SELECT /*+ COALESCE(1) */ * FROM lhs) l " + + "CROSS JOIN (SELECT /*+ COALESCE(1) */ * FROM rhs) r" + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val cases = Seq( + CartesianCase("small left, large right", leftRows = 64, rightRows = 256 * 1024), + CartesianCase("large left, small right", leftRows = 256 * 1024, rightRows = 64), + CartesianCase("square", leftRows = 4096, rightRows = 4096)) + + withTempPath { dir => + withTempTable("lhs", "rhs") { + def writeTable(view: String, path: String, df: DataFrame): Unit = { + df.coalesce(1).write.parquet(s"${dir.getAbsolutePath}/$path") + spark.read.parquet(s"${dir.getAbsolutePath}/$path").createOrReplaceTempView(view) + } + + cases.foreach { c => + writeTable("lhs", s"lhs_${c.leftRows}", spark.range(c.leftRows).selectExpr("id AS v")) + writeTable("rhs", s"rhs_${c.rightRows}", spark.range(c.rightRows).selectExpr("id AS w")) + + val productRows = c.leftRows * c.rightRows + val name = s"Cartesian product - ${c.name} (${c.leftRows} x ${c.rightRows})" + runBenchmark(name) { + verify(name, productRows) + val benchmark = new Benchmark(name, productRows, output = output) + benchmark.addCase(SparkCaseName)(_ => runQuery(sparkConfigs)) + benchmark.addCase(ProductInSparkCaseName)(_ => runQuery(productInSparkConfigs)) + benchmark.addCase(CometCaseName)(_ => runQuery(cometConfigs)) + benchmark.run() + } + } + } + } + } + + private def runQuery(configs: Seq[(String, String)]): Unit = + withSQLConf(configs: _*) { + spark.sql(query).noop() + } + + /** + * Fails unless each arm runs the product it claims, the Comet arm is fully native, and all arms + * return the same sum. + */ + private def verify(name: String, productRows: Long): Unit = { + val arms = Seq( + SparkCaseName -> sparkConfigs, + ProductInSparkCaseName -> productInSparkConfigs, + CometCaseName -> cometConfigs) + val sums = arms.map { case (arm, configs) => + // A local, because `withSQLConf` returns `Unit` on Spark 3.4 and 3.5. + var sum: Any = null + withSQLConf(configs: _*) { + val df = spark.sql(query) + sum = df.collect().head.get(0) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + def fail(problem: String): Nothing = + throw new AssertionError(s"$name, $arm: $problem\n${plan.treeString}") + + val products = collect(plan) { + case p: CartesianProductExec => p + case p: CometExec if p.originalPlan.isInstanceOf[CartesianProductExec] => p + } + val expected = + if (arm == CometCaseName) classOf[CometCartesianProductExec] + else classOf[CartesianProductExec] + products match { + case Seq(node) if expected.isInstance(node) => + if (arm == SparkCaseName && + node.metrics("numOutputRows").value != productRows) { + fail(s"the product produced ${node.metrics("numOutputRows").value} rows") + } + case _ => fail(s"expected one ${expected.getSimpleName}") + } + if (arm == CometCaseName) { + // Query stages are leaves to the checker, so check the plan inside each one as well. + val stagePlans = collect(plan) { case s: QueryStageExec => s.plan } + val wrappers = Seq(classOf[QueryStageExec], classOf[AQEShuffleReadExec]) + (plan +: stagePlans) + .flatMap(findFirstNonCometOperator(_, wrappers: _*)) + .headOption + .foreach(op => + fail(s"the plan is not fully native, first non-Comet operator: ${op.nodeName}")) + } + if (arm == ProductInSparkCaseName && collect(plan) { case p: CometPlan => p }.isEmpty) { + fail("Comet runs no operator") + } + } + arm -> sum + } + val (_, sparkSum) = sums.head + sums.tail.foreach { case (arm, sum) => + if (sum != sparkSum) { + throw new AssertionError(s"$name: Spark returned $sparkSum, $arm returned $sum") + } + } + } +}