diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 1c11bd8de6d..473006aa378 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -529,6 +529,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 1cf2f06d213..d00c7f6d2e2 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -177,6 +177,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 8fcaf964ca4..2109292025c 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1241,6 +1241,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) .with_sql_text_pool(&exec_context.spark_plan) + .with_jvm_output_sorts(&exec_context.spark_plan) .with_task_context(exec_context.task_context.clone()) .with_class_loader(exec_context.class_loader.clone()) .with_shuffle_partition_pusher( @@ -3657,6 +3658,7 @@ mod native_sort_spill_tests { share: usize, executor_cores: usize, spill_before_output: Option, + jvm_consumer: bool, ) -> OutputTrace { let shape = source.shape; let rows = source.rows; @@ -3708,7 +3710,8 @@ mod native_sort_spill_tests { }, )]) .unwrap(); - let sort: Arc = Arc::new(SortExec::new(ordering, child)); + let sort: Arc = + Arc::new(SortExec::new(ordering, child).with_spill_before_output(jvm_consumer)); let mut stream = sort.execute(0, session.task_ctx()).unwrap(); let mut trace = OutputTrace { reserved: vec![], @@ -3744,8 +3747,14 @@ mod native_sort_spill_tests { let share = SPILL_BEFORE_OUTPUT_SHARE; for rows_per_batch in [8192, 3] { let case = format!("{shape:?} rows={rows} rows_per_batch={rows_per_batch}"); - let held = - sort_and_trace(KibRows::new(shape, rows, rows_per_batch), share, 8, None).await; + let held = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + None, + true, + ) + .await; eprintln!( "{case} off: spill_count={} peak={}MiB reserved at 0%={}MiB 50%={}MiB 90%={}MiB", held.spill_count(), @@ -3762,6 +3771,7 @@ mod native_sort_spill_tests { share, 8, Some(share / 4), + true, ) .await; eprintln!( @@ -3782,11 +3792,31 @@ mod native_sort_spill_tests { share, 8, Some(held.peak_reserved), + true, ) .await; assert_eq!(below.spill_count(), 0, "{case}"); assert_eq!(below.reserved.len(), held.reserved.len(), "{case}"); assert!(below.reserved_at(0.9) >= input_bytes(rows) / 2, "{case}"); + + let native_consumer = sort_and_trace( + KibRows::new(shape, rows, rows_per_batch), + share, + 8, + Some(share / 4), + false, + ) + .await; + assert_eq!(native_consumer.spill_count(), 0, "{case}"); + assert_eq!( + native_consumer.reserved.len(), + held.reserved.len(), + "{case}" + ); + assert!( + native_consumer.reserved_at(0.9) >= input_bytes(rows) / 2, + "{case}" + ); } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1febe577337..63a983baef9 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -314,6 +314,29 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + jvm_output_sorts: std::collections::HashSet, +} + +pub(crate) fn jvm_output_sorts(root: &Operator) -> std::collections::HashSet { + let mut sorts = std::collections::HashSet::new(); + let mut pending = vec![root]; + while let Some(op) = pending.pop() { + match op.op_struct.as_ref() { + Some(OpStruct::Sort(_)) => { + sorts.insert(op.plan_id); + } + Some( + OpStruct::Projection(_) + | OpStruct::Filter(_) + | OpStruct::Limit(_) + | OpStruct::Sample(_) + | OpStruct::Expand(_) + | OpStruct::Explode(_), + ) => pending.extend(op.children.iter()), + _ => {} + } + } + sorts } impl Default for PhysicalPlanner { @@ -333,12 +356,18 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + jvm_output_sorts: std::collections::HashSet::new(), } } /// Load the SQL text pool from the root operator of the plan about to be planned. Must be /// called with the *root* operator: the JVM only populates the pool there, and /// `QueryContext.sql_text_idx` values are indices into it. + pub fn with_jvm_output_sorts(mut self, root: &Operator) -> Self { + self.jvm_output_sorts = jvm_output_sorts(root); + self + } + pub fn with_sql_text_pool(mut self, root: &Operator) -> Self { self.sql_text_pool = root .sql_text_pool @@ -1783,7 +1812,8 @@ impl PhysicalPlanner { LexOrdering::new(exprs?).unwrap(), Arc::clone(&child.native_plan), ) - .with_fetch(fetch), + .with_fetch(fetch) + .with_spill_before_output(self.jvm_output_sorts.contains(&spark_plan.plan_id)), ); if let Some(skip) = sort.skip.filter(|&n| n > 0).map(|n| n as usize) { @@ -5391,7 +5421,7 @@ mod tests { use crate::execution::operators::{ExecutionError, PartitionedRankLimitExec, WindowFnKind}; use crate::execution::planner::{ - convert_spark_types_to_arrow_schema, literal_to_array_ref, + convert_spark_types_to_arrow_schema, jvm_output_sorts, literal_to_array_ref, parse_file_scan_tasks_from_common, }; use crate::execution::shuffle::CometPartitioning; @@ -6755,6 +6785,138 @@ mod tests { } } + fn op(plan_id: u32, op_struct: OpStruct, children: Vec) -> Operator { + Operator { + plan_id, + sql_text_pool: vec![], + children, + op_struct: Some(op_struct), + } + } + + fn sort_op(plan_id: u32, child: Operator) -> Operator { + let order = spark_expression::Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(spark_expression::Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(create_proto_datatype()), + })), + query_context: None, + expr_id: None, + })), + direction: spark_expression::SortDirection::Ascending as i32, + null_ordering: spark_expression::NullOrdering::NullsFirst as i32, + }, + ))), + query_context: None, + expr_id: None, + }; + op( + plan_id, + OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![order], + fetch: None, + skip: None, + }), + vec![child], + ) + } + + fn leaf(plan_id: u32) -> Operator { + let mut scan = create_scan(); + scan.plan_id = plan_id; + scan + } + + #[test] + fn jvm_output_sorts_are_sorts_reached_from_the_root_through_streaming_operators() { + use std::collections::HashSet; + let ids = |v: &[u32]| v.iter().copied().collect::>(); + + assert_eq!(jvm_output_sorts(&sort_op(1, leaf(0))), ids(&[1])); + + let projected = op( + 3, + OpStruct::Projection(Default::default()), + vec![op( + 2, + OpStruct::Filter(Default::default()), + vec![sort_op(1, leaf(0))], + )], + ); + assert_eq!(jvm_output_sorts(&projected), ids(&[1])); + + let nested = sort_op( + 2, + op( + 3, + OpStruct::Window(Default::default()), + vec![sort_op(1, leaf(0))], + ), + ); + assert_eq!(jvm_output_sorts(&nested), ids(&[2])); + + for consumer in [ + OpStruct::ShuffleWriter(Default::default()), + OpStruct::Window(Default::default()), + OpStruct::HashAgg(Default::default()), + OpStruct::ParquetWriter(Default::default()), + OpStruct::WindowGroupLimit(Default::default()), + ] { + let plan = op(2, consumer, vec![sort_op(1, leaf(0))]); + assert!(jvm_output_sorts(&plan).is_empty(), "{plan:?}"); + } + + let join = op( + 5, + OpStruct::SortMergeJoin(Default::default()), + vec![sort_op(1, leaf(0)), sort_op(2, leaf(3))], + ); + assert!(jvm_output_sorts(&join).is_empty()); + assert!( + jvm_output_sorts(&op(6, OpStruct::Projection(Default::default()), vec![join])) + .is_empty() + ); + } + + fn sort_execs(plan: &Arc, out: &mut Vec) { + if let Some(sort) = plan.downcast_ref::() { + out.push(sort.spill_before_output()); + } + for child in plan.children() { + sort_execs(child, out); + } + } + + fn planned_sort_flags(root: &Operator) -> Vec { + let planner = PhysicalPlanner::default().with_jvm_output_sorts(root); + let (_, _, plan) = planner.create_plan(root, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + flags + } + + #[test] + fn planner_lets_only_a_sort_read_by_the_jvm_spill_before_output() { + assert_eq!(planned_sort_flags(&sort_op(1, leaf(0))), vec![true]); + assert_eq!( + planned_sort_flags(&create_filter(sort_op(1, leaf(0)), 3)), + vec![true] + ); + assert_eq!( + planned_sort_flags(&sort_op(2, sort_op(1, leaf(0)))), + vec![true, false] + ); + let sort = sort_op(1, leaf(0)); + let planner = PhysicalPlanner::default(); + let (_, _, plan) = planner.create_plan(&sort, &mut vec![], 1).unwrap(); + let mut flags = vec![]; + sort_execs(&plan.native_plan, &mut flags); + assert_eq!(flags, vec![false]); + } + fn create_scan() -> Operator { Operator { plan_id: 0, diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index 51514d63b94..ce3173f5f73 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -1241,6 +1241,7 @@ pub struct SortExec { /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. /// If `fetch` is `None`, this will be `None`. filter: Option>>, + spill_before_output: bool, } impl SortExec { @@ -1260,9 +1261,19 @@ impl SortExec { common_sort_prefix: sort_prefix, cache: Arc::new(cache), filter: None, + spill_before_output: false, } } + pub fn with_spill_before_output(mut self, spill_before_output: bool) -> Self { + self.spill_before_output = spill_before_output; + self + } + + pub fn spill_before_output(&self) -> bool { + self.spill_before_output + } + /// Whether this `SortExec` preserves partitioning of the children pub fn preserve_partitioning(&self) -> bool { self.preserve_partitioning @@ -1336,6 +1347,7 @@ impl SortExec { fetch: self.fetch, cache: Arc::clone(&self.cache), filter: self.filter.clone(), + spill_before_output: self.spill_before_output, } } @@ -1724,10 +1736,14 @@ impl ExecutionPlan for SortExec { execution_options.sort_spill_reservation_bytes; let in_place_bytes = execution_options.sort_in_place_threshold_bytes; - let spill_before_output = context - .session_config() - .get_extension::() - .map_or(0, |threshold| threshold.0); + let spill_before_output = if self.spill_before_output { + context + .session_config() + .get_extension::() + .map_or(0, |threshold| threshold.0) + } else { + 0 + }; let compression = context.session_config().spill_compression(); let runtime = context.runtime_env(); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -1867,7 +1883,8 @@ impl ExecutionPlan for SortExec { Ok(Some(Arc::new( SortExec::new(updated_exprs, make_with_child(projection, self.input())?) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ))) } @@ -1947,7 +1964,8 @@ impl ExecutionPlan for SortExec { let new_sort = Arc::new( SortExec::new(self.expr.clone(), new_child) .with_fetch(self.fetch()) - .with_preserve_partitioning(self.preserve_partitioning()), + .with_preserve_partitioning(self.preserve_partitioning()) + .with_spill_before_output(self.spill_before_output), ) as Arc; Ok(FilterPushdownPropagation { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 03bd2701754..55b0d46131d 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -237,6 +237,13 @@ object CometConf extends ShimCometConf { createExecEnabledConfig("sortMergeJoin", defaultValue = true) val COMET_EXEC_AGGREGATE_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("aggregate", defaultValue = true) + val COMET_EXEC_SORT_AGGREGATE_ENABLED: ConfigEntry[Boolean] = + createExecEnabledConfig( + "sortAggregate", + defaultValue = true, + notes = Some( + "When enabled, a SortAggregate that Comet can run is converted to a native hash " + + "aggregate, with a native sort restoring its output ordering")) val COMET_EXEC_COLLECT_LIMIT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("collectLimit", defaultValue = true) val COMET_EXEC_COALESCE_ENABLED: ConfigEntry[Boolean] = @@ -486,6 +493,15 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.filter.shortCircuitJvmDispatch.enabled") + .category(CATEGORY_EXEC) + .doc("Whether a native filter evaluates a top-level conjunct that runs through the JVM " + + "codegen dispatcher only on rows where all preceding conjuncts are true, by rewriting " + + "`a AND b` into `CASE WHEN a THEN b END`. Applies to deterministic conditions only.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") 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 65b85540e41..15e0737df44 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -22,8 +22,8 @@ package org.apache.comet.rules import scala.collection.mutable.ListBuffer import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction, AggregateMode, Average, Count, Final, Max, MaxMinBy, Min, Partial, PartialMerge, Sum} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag @@ -35,7 +35,7 @@ import org.apache.spark.sql.comet.shims.ShimCometEmptyRelation import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} @@ -150,6 +150,17 @@ object CometExecRule { */ val ENGINE_CHOICE_SPARK_TAG: TreeNodeTag[Unit] = TreeNodeTag[Unit]("comet.engineChoiceSpark") + /** Tag set on a native hash aggregate converted from a `SortAggregateExec`. */ + val CONVERTED_SORT_AGGREGATE_TAG: TreeNodeTag[Unit] = + TreeNodeTag[Unit]("comet.convertedSortAggregate") + + /** + * Tag set on the sort inserted above a converted `SortAggregateExec` for a consumer that + * requires the ordering the sort aggregate provided. + */ + val SORT_AGGREGATE_ORDER_TAG: TreeNodeTag[Unit] = + TreeNodeTag[Unit]("comet.sortAggregateOrder") + /** * Serializes the native plan of each block of adjacent native operators into its topmost * operator. Blocks that already hold a serialized plan are left as they are, so this can run @@ -587,6 +598,9 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) convertToComet(s, CometShuffleExchangeExec) .getOrElse(preserveSparkAggregateBuffers(s)) + case agg: SortAggregateExec if CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.get(conf) => + convertSortAggregate(agg).getOrElse(agg) + case op => // if all children are native (or if this is a leaf node) then see if there is a // registered handler for creating a fully native plan @@ -644,7 +658,7 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } plan.transformUp { case op => - val converted = convertNode(op) + val converted = convertNode(restoreSortAggregateOrdering(op)) // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions // when the broadcast child has a Comet plan underneath. This enables exchange reuse // between the DPP subquery and the join's CometBroadcastExchangeExec because both @@ -928,6 +942,107 @@ case class CometExecRule(session: SparkSession, wholePlan: Boolean = false) } } + private def leadsToConvertedSortAggregate(plan: SparkPlan): Boolean = + plan.getTagValue(CometExecRule.CONVERTED_SORT_AGGREGATE_TAG).isDefined || + (plan.children.size == 1 && leadsToConvertedSortAggregate(plan.children.head)) + + /** + * Spark's EnsureRequirements left out the sort a consumer needs when a `SortAggregateExec` + * below it already provided the ordering. Once that aggregate is a native hash aggregate, the + * sort goes back on exactly the edges whose requirement is no longer met, before the consumer + * itself is converted. + */ + private def restoreSortAggregateOrdering(op: SparkPlan): SparkPlan = { + if (op.isInstanceOf[CometPlan]) return op + val required = op.requiredChildOrdering + if (required.length != op.children.length || required.forall(_.isEmpty)) return op + var changed = false + val children = op.children.zip(required).map { case (child, ordering) => + if (ordering.nonEmpty && leadsToConvertedSortAggregate(child) && + !SortOrder.orderingSatisfies(child.outputOrdering, ordering)) { + changed = true + val sort = SortExec(ordering, global = false, child) + child.logicalLink.foreach(sort.setLogicalLink) + val restored = if (child.isInstanceOf[CometNativeExec]) { + tryConvertToComet(sort, CometSortExec).getOrElse(sort) + } else { + sort + } + restored.setTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG, ()) + restored + } else { + child + } + } + if (changed) op.withNewChildren(children) else op + } + + private def orderInsensitive(fn: AggregateFunction): Boolean = fn match { + case _: Min | _: Max | _: Count | _: Sum | _: Average | _: MaxMinBy => true + case _ => false + } + + /** + * A `SortAggregateExec` as a native hash aggregate. Spark plans one when an aggregation buffer + * is not mutable, such as MIN or MAX over strings, and runs it over input sorted by the + * grouping keys. The hash aggregate starts from an `ObjectHashAggregateExec` with the same + * fields, which Spark runs correctly over any input order, so any operator reverted to Spark + * later stays correct. Only aggregates whose result does not depend on the input order are + * converted: Spark's sort is stable, so FIRST, LAST and the like see the rows of a group in + * their input order, which the native hash aggregate does not guarantee. MAX_BY and MIN_BY + * depend on it only among rows tied on the ordering, where they are non-deterministic in Spark + * too and already differ from it on the native hash aggregate path. The sort by the grouping + * keys directly below the aggregate is then dropped. A consumer that relied on the ordering of + * the sort aggregate gets a sort back, see [[restoreSortAggregateOrdering]]. + */ + private def convertSortAggregate(agg: SortAggregateExec): Option[SparkPlan] = { + val required = agg.requiredChildOrdering.headOption.getOrElse(Nil) + val sortedByKeys = agg.child match { + case sort: CometSortExec => + sort.sortOrder.length == required.length && + sort.sortOrder.zip(required).forall { case (a, b) => + a.semanticEquals(b) + } && + (sort.originalPlan match { + case s: SortExec => !s.global + case _ => false + }) + case _ => false + } + val input = if (sortedByKeys) agg.child.children.head else agg.child + if (!agg.aggregateExpressions.forall(e => orderInsensitive(e.aggregateFunction))) { + withFallbackReason( + agg, + "SortAggregate with an aggregate function that depends on the input order stays in Spark") + return None + } + if (!input.isInstanceOf[CometNativeExec]) { + return None + } + val hashAgg = ObjectHashAggregateExec( + agg.requiredChildDistributionExpressions, + agg.isStreaming, + agg.numShufflePartitions, + agg.groupingExpressions, + agg.aggregateExpressions, + agg.aggregateAttributes, + agg.initialInputBufferOffset, + agg.resultExpressions, + input) + hashAgg.copyTagsFrom(agg) + val converted = tryConvertToComet(hashAgg, CometObjectHashAggregateExec) + converted.foreach(_.setTagValue(CometExecRule.CONVERTED_SORT_AGGREGATE_TAG, ())) + if (converted.isEmpty) { + hashAgg + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .foreach(reasons => withFallbackReasons(agg, reasons)) + if (!hasFallbackReason(agg)) { + withFallbackReason(agg, "SortAggregate could not be converted to a native hash aggregate") + } + } + converted + } + /** Convert a Spark plan to a Comet plan using the specified serde handler */ private def convertToComet(op: SparkPlan, handler: CometOperatorSerde[_]): Option[SparkPlan] = { val converted = tryConvertToComet(op, handler) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index ad53cceb31b..4583c2bce00 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -129,27 +129,28 @@ abstract class CometMaxMinBy[T <: MaxMinBy] extends CometAggregateExpressionSerd " Results may differ from Spark in that case.") override def getUnsupportedReasons(): Seq[String] = Seq( - "The value and ordering must both be fixed-length types (boolean, integral, floating-point," + - " decimal, date, or timestamp). A variable-length or nested type such as string, binary, or" + - " struct falls back to Spark.") + "The ordering must be a fixed-length type (boolean, integral, floating-point, decimal, date," + + " or timestamp). The value may also be a string with the default UTF8_BINARY collation." + + " Other variable-length or nested types such as binary or struct fall back to Spark.") override def getSupportLevel(expr: T): SupportLevel = { - // Both the value and ordering must be fixed-length types. + // The ordering must be a fixed-length type. The value may also be a UTF8_BINARY string: the + // native side only stores it as Arrow row bytes and never compares it. // - // On its own a variable-length type never reaches here: Spark only uses HashAggregate (the - // aggregate operator Comet accelerates) when the aggregation buffer is mutable, and the buffer - // holds both the running value and the running ordering, so a StringType in either position - // forces SortAggregate, which Comet does not convert. + // A string in the buffer makes Spark plan a SortAggregate, which Comet converts to a native + // hash aggregate (see `CometExecRule.convertSortAggregate`), or an ObjectHashAggregate when a + // TypedImperativeAggregate sits in the same aggregate. // - // The check is still load-bearing, because a TypedImperativeAggregate elsewhere in the same - // aggregate switches Spark to ObjectHashAggregate, which Comet does convert. In that shape a - // string ordering would otherwise be compared by Arrow's row format as raw UTF-8 bytes, while + // A string ordering stays in Spark: Arrow's row format would compare raw UTF-8 bytes, while // Spark compares collation sort keys. See the fallback cases in max_by.sql. // // The native side compares the ordering column via Arrow's row format, which supports all of // the fixed-length orderable types allowed below. - if (!AggSerde.minMaxDataTypeSupported(expr.valueExpr.dataType)) { - Unsupported(Some(s"Unsupported value data type: ${expr.valueExpr.dataType}")) + val valueType = expr.valueExpr.dataType + val stringValue = + valueType.isInstanceOf[StringType] && !AggSerde.isStringCollationType(valueType) + if (!stringValue && !AggSerde.minMaxDataTypeSupported(valueType)) { + Unsupported(Some(s"Unsupported value data type: $valueType")) } else if (!AggSerde.minMaxDataTypeSupported(expr.orderingExpr.dataType)) { Unsupported(Some(s"Unsupported ordering data type: ${expr.orderingExpr.dataType}")) } else { @@ -1395,7 +1396,7 @@ object CometMode extends CometAggregateExpressionSerde[Mode] with CometTypeShim } } -object AggSerde { +object AggSerde extends CometTypeShim { import org.apache.spark.sql.types._ def minMaxDataTypeSupported(dt: DataType): Boolean = { @@ -1436,7 +1437,13 @@ object AggSerde { /** Shared support level for `Min` / `Max` based on the result data type. */ def minMaxSupportLevel(dt: DataType): SupportLevel = { - if (!minMaxDataTypeSupported(dt)) { + if (dt.isInstanceOf[StringType]) { + if (isStringCollationType(dt)) { + Unsupported(Some(s"Unsupported collated string type: $dt")) + } else { + Compatible() + } + } else if (!minMaxDataTypeSupported(dt)) { Unsupported(Some(s"Unsupported data type: $dt")) } else if ((dt == FloatType || dt == DoubleType) && COMET_EXEC_STRICT_FLOATING_POINT.get()) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala index 98c1105edc1..de3ecb31206 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala @@ -65,6 +65,15 @@ case class CometColumnarToRowExec(child: SparkPlan) override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override def supportCodegen: Boolean = !readsBroadcast(child) + + private def readsBroadcast(plan: SparkPlan): Boolean = plan match { + case _: CometBroadcastExchangeExec => true + case stage: BroadcastQueryStageExec => readsBroadcast(stage.plan) + case reused: ReusedExchangeExec => readsBroadcast(reused.child) + case _ => false + } + // `ColumnarToRowExec` processes the input RDD directly, which is kind of a leaf node in the // codegen stage and needs to do the limit check. protected override def canCheckLimitNotReached: Boolean = true 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 0718697a794..1333c8c567f 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 @@ -53,13 +53,14 @@ import org.apache.spark.util.SerializableConfiguration import org.apache.spark.util.io.ChunkedByteBuffer import com.google.common.base.Objects -import com.google.protobuf.CodedOutputStream +import com.google.protobuf.{CodedOutputStream, Message} import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, ConfigEntry, ContribServices} import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.rules.CometExecRule -import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.{CometOperatorSerde, Compatible, ExprOuterClass, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} @@ -1428,15 +1429,90 @@ object CometFilterExec extends CometOperatorSerde[FilterExec] { val cond = exprToProto(op.condition, op.child.output) if (cond.isDefined && childOp.nonEmpty) { + val predicate = + if (CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.get() && + op.condition.deterministic) { + shortCircuitJvmDispatch(cond.get) + } else { + cond.get + } val filterBuilder = OperatorOuterClass.Filter .newBuilder() - .setPredicate(cond.get) + .setPredicate(predicate) Some(builder.setFilter(filterBuilder).build()) } else { None } } + def shortCircuitJvmDispatch(predicate: Expr): Expr = { + val conjuncts = flattenAnd(predicate) + val expensive = conjuncts.map(containsJvmScalarUdf) + if (!expensive.drop(1).contains(true)) { + predicate + } else { + val segments = ArrayBuffer(ArrayBuffer(conjuncts.head)) + conjuncts.zip(expensive).tail.foreach { case (conjunct, isExpensive) => + if (isExpensive) segments += ArrayBuffer(conjunct) else segments.last += conjunct + } + segments.tail.foldLeft(balancedAnd(segments.head.toIndexedSeq)) { (guard, segment) => + Expr + .newBuilder() + .setCaseWhen( + ExprOuterClass.CaseWhen + .newBuilder() + .addWhen(guard) + .addThen(balancedAnd(segment.toIndexedSeq))) + .build() + } + } + } + + private def flattenAnd(predicate: Expr): Seq[Expr] = { + val leaves = ArrayBuffer.empty[Expr] + var stack: List[Expr] = predicate :: Nil + while (stack.nonEmpty) { + val current = stack.head + stack = stack.tail + if (current.hasAnd) { + stack = current.getAnd.getLeft :: current.getAnd.getRight :: stack + } else { + leaves += current + } + } + leaves.toSeq + } + + private def balancedAnd(conjuncts: IndexedSeq[Expr]): Expr = { + if (conjuncts.length == 1) { + conjuncts.head + } else { + val mid = conjuncts.length / 2 + Expr + .newBuilder() + .setAnd( + ExprOuterClass.BinaryExpr + .newBuilder() + .setLeft(balancedAnd(conjuncts.slice(0, mid))) + .setRight(balancedAnd(conjuncts.slice(mid, conjuncts.length)))) + .build() + } + } + + private def containsJvmScalarUdf(message: Message): Boolean = message match { + case e: Expr if e.hasJvmScalarUdf => true + case _ => + message.getAllFields.values.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case children: java.util.List[_] => + children.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case _ => false + } + case _ => false + } + } + override def createExec(nativeOp: Operator, op: FilterExec): CometNativeExec = { CometFilterExec(nativeOp, op, op.output, op.condition, op.child, SerializedPlan(None)) } diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql index 8aedc4f5f17..b67a4247eb8 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/first_last.sql @@ -183,7 +183,7 @@ SELECT first(val) IGNORE NULLS FROM test_single_null -- first IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +query expect_fallback(depends on the input order) SELECT grp, first(i_val) IGNORE NULLS, first(l_val) IGNORE NULLS, @@ -327,7 +327,7 @@ SELECT last(val) IGNORE NULLS FROM test_single_null -- last IGNORE NULLS: multiple data types -- ============================================================ -query expect_fallback(SortAggregate is not supported) +query expect_fallback(depends on the input order) SELECT grp, last(i_val) IGNORE NULLS, last(l_val) IGNORE NULLS, diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql index 9c8354f06ab..b8270db9e1a 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/max_by.sql @@ -17,9 +17,8 @@ -- max_by(x, y) returns the value of x associated with the maximum value of y. -- --- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate --- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length --- value types such as string force SortAggregate and fall back to Spark. +-- The value (x) must be a fixed-length type or a string. A string value makes Spark plan a +-- SortAggregate, which Comet converts to a native hash aggregate. -- -- Ordering values are kept unique within each group so results are deterministic (max_by is -- non-deterministic when several rows tie on the maximum ordering). @@ -260,14 +259,11 @@ query SELECT grp, max_by(v, ord) FROM mb_signed_zero GROUP BY grp ORDER BY grp -- ============================================================ --- Variable-length value or ordering falls back to Spark +-- Variable-length ordering falls back to Spark -- --- A plain max_by over a string is planned as SortAggregate, which Comet never converts, so the --- serde's type check is not what stops it. Pairing it with a TypedImperativeAggregate switches --- Spark to ObjectHashAggregate, which Comet does convert, and then getSupportLevel is the only --- thing that keeps the aggregate off the native path. That matters most for a string *ordering*: --- Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys, so this --- check is load-bearing for correctness and not just an optimisation. +-- A string value runs natively: as a plain aggregate through the converted SortAggregate, and +-- next to a TypedImperativeAggregate through ObjectHashAggregate. A string ordering stays in +-- Spark: Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys. -- ============================================================ statement @@ -277,8 +273,31 @@ statement INSERT INTO mb_varlen VALUES (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2') -query expect_fallback(Unsupported value data type) +query SELECT grp, max_by(s, v), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER BY grp query expect_fallback(Unsupported ordering data type) SELECT grp, max_by(v, s), percentile(v, 0.5) FROM mb_varlen GROUP BY grp ORDER BY grp + +-- ============================================================ +-- String value with an integer ordering +-- ============================================================ + +statement +CREATE TABLE mb_str(s string, ord int, grp string) USING parquet + +statement +INSERT INTO mb_str VALUES + ('b', 1, 'g1'), ('', 5, 'g1'), (NULL, 3, 'g1'), + ('é日本', 2, 'g2'), (NULL, 9, 'g2'), ('😀', 4, 'g2'), + ('z', NULL, 'g3'), ('y', NULL, 'g3'), + ('x', 6, 'g4') + +query +SELECT grp, max_by(s, ord) FROM mb_str GROUP BY grp ORDER BY grp + +query +SELECT max_by(s, ord) FROM mb_str + +query +SELECT grp, max_by(s, ord), max(s), count(*) FROM mb_str GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql index d015acc35d3..55828e5c65c 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_by.sql @@ -17,9 +17,8 @@ -- min_by(x, y) returns the value of x associated with the minimum value of y. -- --- The value (x) must be a fixed-length type: Spark only uses HashAggregate (the aggregate --- operator Comet accelerates) when the aggregation buffer is mutable, so variable-length --- value types such as string force SortAggregate and fall back to Spark. +-- The value (x) must be a fixed-length type or a string. A string value makes Spark plan a +-- SortAggregate, which Comet converts to a native hash aggregate. -- -- Ordering values are kept unique within each group so results are deterministic (min_by is -- non-deterministic when several rows tie on the minimum ordering). @@ -245,11 +244,11 @@ query SELECT grp, min_by(v, ord) FROM mnb_signed_zero GROUP BY grp ORDER BY grp -- ============================================================ --- Variable-length value or ordering falls back to Spark +-- Variable-length ordering falls back to Spark -- --- See the equivalent section in max_by.sql for why this needs a TypedImperativeAggregate --- alongside it: without one, Spark plans SortAggregate and Comet never sees the aggregate at all, --- so the serde's type check would go untested. +-- A string value runs natively: as a plain aggregate through the converted SortAggregate, and +-- next to a TypedImperativeAggregate through ObjectHashAggregate. A string ordering stays in +-- Spark: Arrow's row format compares raw UTF-8 bytes, while Spark compares collation sort keys. -- ============================================================ statement @@ -259,8 +258,31 @@ statement INSERT INTO mnb_varlen VALUES (1, 'a', 'g1'), (2, 'b', 'g1'), (3, 'c', 'g2') -query expect_fallback(Unsupported value data type) +query SELECT grp, min_by(s, v), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER BY grp query expect_fallback(Unsupported ordering data type) SELECT grp, min_by(v, s), percentile(v, 0.5) FROM mnb_varlen GROUP BY grp ORDER BY grp + +-- ============================================================ +-- String value with an integer ordering +-- ============================================================ + +statement +CREATE TABLE mnb_str(s string, ord int, grp string) USING parquet + +statement +INSERT INTO mnb_str VALUES + ('b', 1, 'g1'), ('', 5, 'g1'), (NULL, 3, 'g1'), + ('é日本', 2, 'g2'), (NULL, 9, 'g2'), ('😀', 4, 'g2'), + ('z', NULL, 'g3'), ('y', NULL, 'g3'), + ('x', 6, 'g4') + +query +SELECT grp, min_by(s, ord) FROM mnb_str GROUP BY grp ORDER BY grp + +query +SELECT min_by(s, ord) FROM mnb_str + +query +SELECT grp, min_by(s, ord), max(s), count(*) FROM mnb_str GROUP BY grp ORDER BY grp diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql index dbec88e2f66..9fcde06186e 100644 --- a/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql +++ b/spark/src/test/resources/sql-tests/expressions/aggregate/min_max.sql @@ -21,7 +21,7 @@ CREATE TABLE test_min_max(i int, d double, s string, grp string) USING parquet statement INSERT INTO test_min_max VALUES (1, 1.5, 'b', 'x'), (3, 3.5, 'a', 'x'), (2, 2.5, 'c', 'y'), (NULL, NULL, NULL, 'y'), (-1, -1.5, 'z', 'x') -query expect_fallback(SortAggregate is not supported) +query SELECT min(i), max(i), min(d), max(d), min(s), max(s) FROM test_min_max query diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index e82831ec426..04b3746a5d1 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,11 +31,11 @@ import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} -import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec} +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.functions.{avg, col, count_distinct, expr, sum} import org.apache.spark.sql.internal.SQLConf @@ -3448,4 +3448,198 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + private val stringRows: Seq[(Integer, String)] = Seq( + (1, "b"), + (1, ""), + (1, null), + (1, "a"), + (2, null), + (2, null), + (3, "\u00e9t\u00e9"), + (3, "ete"), + (3, "\u65e5\u672c"), + (4, "\uff61"), + (4, "\ud83d\ude00"), + (4, "z"), + (5, ""), + (null, "x"), + (null, "\u00ff")) + + private def withStrings(f: => Unit): Unit = + withParquetTable(stringRows, "str_tbl")(f) + + private def sortAggregates(plan: SparkPlan): Seq[SortAggregateExec] = + collect(plan) { case a: SortAggregateExec => a } + + private def nativeAggregates(plan: SparkPlan): Seq[CometHashAggregateExec] = + collect(plan) { case a: CometHashAggregateExec => a } + + for (aqe <- Seq("false", "true")) { + test(s"min and max over strings run natively, grouped and ungrouped (AQE=$aqe)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withStrings { + for (query <- Seq( + "SELECT min(_2), max(_2), count(_2) FROM str_tbl", + "SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1", + "SELECT _1, max(_2) FROM str_tbl WHERE _1 = 2 GROUP BY _1", + "SELECT max(_2) FROM str_tbl WHERE _1 > 100", + "SELECT _1, min(_2), sum(_1) FROM str_tbl GROUP BY _1")) { + val (sparkPlan, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(sparkPlan).nonEmpty, s"$query:\n$sparkPlan") + assert(sortAggregates(cometPlan).isEmpty, s"$query:\n$cometPlan") + assert(nativeAggregates(cometPlan).nonEmpty, s"$query:\n$cometPlan") + } + } + } + } + } + + test("min and max over strings compare UTF-8 bytes as Spark does") { + withStrings { + checkSparkAnswerAndOperator(sql("SELECT min(_2), max(_2) FROM str_tbl WHERE _1 = 4")) + checkSparkAnswerAndOperator(sql("SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1")) + } + } + + test("a converted sort aggregate keeps the ordering a sort-merge join relies on") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + withParquetTable((0 until 50).map(i => (i % 6, s"v$i")), "str_right") { + val query = + """SELECT a._1, a.m, r._2 FROM (SELECT _1, max(_2) AS m FROM str_tbl GROUP BY _1) a + |JOIN str_right r ON a._1 = r._1""".stripMargin + val (_, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + val join = collect(cometPlan) { case j: CometSortMergeJoinExec => j } + assert(join.size == 1, s"plan:\n$cometPlan") + val orderSorts = collect(cometPlan) { + case s: CometSortExec + if s.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => + s + } + assert(orderSorts.nonEmpty, s"plan:\n$cometPlan") + } + } + } + } + + test("a converted sort aggregate gets no sort without a consumer that needs its ordering") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + val (_, cometPlan) = checkSparkAnswer(sql("SELECT _1, max(_2) FROM str_tbl GROUP BY _1")) + val shuffles = collect(cometPlan) { case s: CometShuffleExchangeExec => s } + assert(shuffles.nonEmpty, s"plan:\n$cometPlan") + assert(shuffles.forall(s => !s.child.isInstanceOf[CometSortExec]), s"plan:\n$cometPlan") + assert(nativeAggregates(cometPlan).size == 2, s"plan:\n$cometPlan") + assert( + collect(cometPlan) { + case s: CometSortExec + if s.getTagValue(CometExecRule.SORT_AGGREGATE_ORDER_TAG).isDefined => + s + }.isEmpty, + s"plan:\n$cometPlan") + assert(collect(cometPlan) { case s: CometSortExec => s }.isEmpty, s"plan:\n$cometPlan") + } + } + } + + test("a sort aggregate Comet cannot run natively stays in Spark") { + withStrings { + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT _1, max_by(_2, _2) FROM str_tbl GROUP BY _1")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } + + test("an order-sensitive sort aggregate over input sorted beyond its keys stays in Spark") { + withStrings { + val (_, cometPlan) = checkSparkAnswer( + sql("SELECT _1, first(_2) FROM (SELECT * FROM str_tbl DISTRIBUTE BY _1 SORT BY _1, _2) " + + "GROUP BY _1")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } + + test("a converted sort aggregate reverted by the cost-based choice runs as an object hash") { + withSQLConf( + CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key -> "true", + CometConf.COMET_EXEC_COST_BASED_ENGINES_COST_TABLE.key -> "agg.flat.comet=1000000,0,0", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withStrings { + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT _1, min(_2), max(_2) FROM str_tbl GROUP BY _1")) + assert(sortAggregates(cometPlan).isEmpty, s"plan:\n$cometPlan") + assert( + collect(cometPlan) { case a: ObjectHashAggregateExec => a }.nonEmpty, + s"plan:\n$cometPlan") + } + } + } + + private val byRows: Seq[(Integer, String, Integer)] = Seq( + (1, "b", 1), + (1, "", 5), + (1, null, 3), + (2, "\u00e9\u65e5\u672c", 2), + (2, null, 9), + (2, "\ud83d\ude00", 4), + (3, "z", null), + (3, "y", null), + (4, "x", 6), + (null, "w", 7)) + + for (aqe <- Seq("false", "true"); fn <- Seq("max_by", "min_by")) { + test(s"$fn over a string value runs natively, grouped and ungrouped (AQE=$aqe)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withParquetTable(byRows, "by_tbl") { + for (query <- Seq( + s"SELECT _1, $fn(_2, _3) FROM by_tbl GROUP BY _1", + s"SELECT $fn(_2, _3) FROM by_tbl", + s"SELECT _1, $fn(_2, _3), max(_2), count(*) FROM by_tbl GROUP BY _1", + s"SELECT $fn(_2, _3) FROM by_tbl WHERE _1 = 3")) { + val (sparkPlan, cometPlan) = checkSparkAnswer(sql(query)) + assert(sortAggregates(sparkPlan).nonEmpty, s"$query:\n$sparkPlan") + assert(sortAggregates(cometPlan).isEmpty, s"$query:\n$cometPlan") + assert(nativeAggregates(cometPlan).nonEmpty, s"$query:\n$cometPlan") + } + } + } + } + } + + test("max_by and min_by over a string value pick one of the values tied on the ordering") { + val rows = (0 until 400).map(i => (i % 4, s"v${i % 7}", if (i % 3 == 0) 10 else i % 3)) + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3") { + withParquetTable(rows, "by_ties") { + for ((fn, tiedOrd) <- Seq("max_by" -> 10, "min_by" -> 1)) { + val df = sql(s"SELECT _1, $fn(_2, _3) FROM by_ties GROUP BY _1") + val result = df.collect().map(r => r.getInt(0) -> r.getString(1)).toMap + assert(nativeAggregates(df.queryExecution.executedPlan).nonEmpty) + val allowed = rows.filter(_._3 == tiedOrd).groupBy(_._1).map { case (k, v) => + k -> v.map(_._2).toSet + } + assert(result.keySet == allowed.keySet, s"$fn: $result") + result.foreach { case (k, v) => assert(allowed(k).contains(v), s"$fn group $k: $v") } + } + } + } + } + + test("max_by over a string value keeps a native partial with a Spark final") { + withSQLConf(CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false") { + withParquetTable(byRows, "by_tbl") { + checkSparkAnswer(sql("SELECT _1, max_by(_2, _3), min_by(_2, _3) FROM by_tbl GROUP BY _1")) + } + } + } + + test("an ungrouped order-sensitive sort aggregate over an ordered input stays in Spark") { + withParquetTable(byRows, "by_tbl") { + val (_, cometPlan) = checkSparkAnswer( + sql("SELECT max_by(_2, _1), first(_2) FROM (SELECT * FROM by_tbl ORDER BY _3, _2)")) + assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") + } + } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index b79b412cc03..36e2a8c25e8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -136,6 +136,41 @@ class CometExecSuite extends CometTestBase { } } + test("a native sort spills before output only when the JVM reads its output") { + def sortSpills(plan: SparkPlan): Seq[Long] = + collect(plan) { case s: CometSortExec => s.metrics("spill_count").value } + withSQLConf( + CometConf.COMET_EXEC_SORT_SPILL_BEFORE_OUTPUT_THRESHOLD.key -> "1", + CometConf.COMET_EXEC_WINDOW_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withParquetTable((0 until 20000).map(i => (i % 37, s"v${i * 7919 % 20000}", i)), "a") { + withParquetTable((0 until 5000).map(i => (i % 37, s"w$i")), "b") { + val (_, jvmRead) = checkSparkAnswer(sql("SELECT _1, _2, _3 FROM a ORDER BY _2, _3")) + assert(sortSpills(jvmRead).nonEmpty) + assert(sortSpills(jvmRead).forall(_ > 0), jvmRead) + + val (_, filtered) = + checkSparkAnswer(sql("SELECT _2, _3 + 1 FROM a WHERE _1 > 3 ORDER BY _2, _3")) + assert(sortSpills(filtered).nonEmpty) + assert(sortSpills(filtered).forall(_ > 0), filtered) + + val (_, windowRead) = checkSparkAnswer( + sql("SELECT _1, _2, row_number() OVER (PARTITION BY _1 ORDER BY _2, _3) FROM a")) + assert(collect(windowRead) { case w: CometWindowExec => w }.nonEmpty, windowRead) + assert(sortSpills(windowRead).nonEmpty) + assert(sortSpills(windowRead).forall(_ == 0), windowRead) + + val (_, joinRead) = checkSparkAnswer( + sql("SELECT a._1, a._2, b._2 FROM a JOIN b ON a._1 = b._1 AND a._3 < 300")) + assert(collect(joinRead) { case j: CometSortMergeJoinExec => j }.nonEmpty, joinRead) + assert(sortSpills(joinRead).nonEmpty) + assert(sortSpills(joinRead).forall(_ == 0), joinRead) + } + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala new file mode 100644 index 00000000000..e95e8455f72 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala @@ -0,0 +1,168 @@ +/* + * 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.exec + +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.comet.{CometFilterExec, CometNativeScanExec, CometScanExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + +import org.apache.comet.CometConf +import org.apache.comet.serde.ExprOuterClass.Expr + +class CometFilterShortCircuitSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val backreference = "^(\\\\d)\\\\1" + + override def beforeAll(): Unit = { + super.beforeAll() + spark.udf.register( + "strict_even", + (x: Int) => { + if (x % 3 != 0) throw new IllegalStateException(s"strict_even evaluated on $x") + x % 2 == 0 + }) + spark.udf.register("is_even", (x: Int) => x % 2 == 0) + } + + private def withData(f: => Unit): Unit = { + withTable("t") { + sql("""CREATE TABLE t USING parquet AS + |SELECT + | IF(id % 7 = 0, NULL, CAST(id AS INT)) AS x, + | IF(id % 5 = 0, NULL, id % 4 = 0) AS b, + | IF(id % 11 = 0, NULL, CAST(id * 37 % 1000 AS STRING)) AS s + |FROM range(4000)""".stripMargin) + f + } + } + + private def filterPredicates(df: DataFrame): Seq[Expr] = + collect(df.queryExecution.executedPlan) { case f: CometFilterExec => + f.nativeOp.getFilter.getPredicate + } + + private def caseWhenDepth(e: Expr): Int = + if (e.hasCaseWhen) 1 + caseWhenDepth(e.getCaseWhen.getWhen(0)) else 0 + + private def scanDataFilters(df: DataFrame): Seq[Expression] = + collect(df.queryExecution.executedPlan) { + case s: CometNativeScanExec => s.dataFilters + case s: CometScanExec => s.dataFilters + }.flatten + + private def checkRewritten(query: String, depth: Int = 1): Unit = { + val df = sql(query) + checkSparkAnswerAndOperator(df) + val predicates = filterPredicates(df) + assert(predicates.nonEmpty, df.queryExecution.executedPlan) + assert(predicates.map(caseWhenDepth).max == depth, predicates) + } + + private def checkNotRewritten(query: String): Unit = { + val df = sql(query) + checkSparkAnswer(df) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen), filterPredicates(df)) + } + + test("expensive conjunct is evaluated only on rows passing preceding conjuncts") { + withData { + checkRewritten("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + checkRewritten("SELECT x FROM t WHERE x > 10 AND x % 3 = 0 AND strict_even(x) AND x < 3000") + } + } + + test("disabled rewrite evaluates the expensive conjunct on every row") { + withData { + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + val e = intercept[Exception](df.collect()) + assert( + e.toString.contains("strict_even evaluated on") || + Option(e.getCause).exists(_.toString.contains("strict_even evaluated on"))) + } + } + } + + test("three-valued logic matches Spark") { + withData { + checkRewritten("SELECT x, b, s FROM t WHERE b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE NOT b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE (b OR x IS NULL) AND is_even(x)") + checkRewritten(s"SELECT x, b, s FROM t WHERE b AND s RLIKE '$backreference'") + checkRewritten(s"SELECT x, b, s FROM t WHERE x % 3 = 0 AND NOT (s RLIKE '$backreference')") + checkRewritten("SELECT x, b, s FROM t WHERE x % 2 = 0 AND (NOT is_even(x) OR b)") + } + } + + test("nested AND/OR and multiple expensive conjuncts") { + withData { + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND (b OR s RLIKE '$backreference') AND x > 100") + checkRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) AND (is_even(x) AND s RLIKE '$backreference')", + depth = 2) + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND strict_even(x) AND b AND s RLIKE '$backreference'", + depth = 2) + checkNotRewritten(s"SELECT x, s FROM t WHERE x % 3 = 0 OR s RLIKE '$backreference'") + checkNotRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) OR (is_even(x) AND s RLIKE '$backreference')") + } + } + + test("no rewrite without a preceding conjunct or JVM dispatch") { + withData { + checkNotRewritten("SELECT x FROM t WHERE is_even(x)") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND b AND s IS NOT NULL") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND s LIKE '1%'") + } + } + + test("nondeterministic condition is not rewritten") { + withData { + val df = sql( + "SELECT x FROM t WHERE monotonically_increasing_id() % 5 < 4 AND x % 3 = 0 AND is_even(x)") + assert(filterPredicates(df).nonEmpty, df.queryExecution.executedPlan) + assert( + filterPredicates(df).forall(p => !p.hasCaseWhen), + df.queryExecution.optimizedPlan.toString + df.queryExecution.executedPlan) + checkSparkAnswer(df) + } + } + + test("scan pushdown is unchanged") { + withData { + val query = "SELECT x FROM t WHERE x > 10 AND b AND is_even(x)" + val rewritten = sql(query) + var baselineFilters: Seq[Expression] = Nil + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql(query) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + baselineFilters = scanDataFilters(df) + } + assert(filterPredicates(rewritten).exists(_.hasCaseWhen)) + assert(scanDataFilters(rewritten).nonEmpty) + assert(scanDataFilters(rewritten) == baselineFilters) + checkSparkAnswerAndOperator(rewritten) + } + } +} 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 8bb8d5ddf70..392ad516a00 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,11 +31,12 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometColumnarToRowExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} +import org.apache.spark.sql.execution.{InputAdapter, LocalTableScanExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} @@ -1479,6 +1480,56 @@ class CometJoinSuite extends CometTestBase { } } + for (costBased <- Seq("true", "false"); maxFields <- Seq("3", "100")) { + test( + "a Spark broadcast hash join over a skewed join reads a Comet broadcast stage " + + s"(costBasedEngines=$costBased, codegen.maxFields=$maxFields)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.SKEW_JOIN_ENABLED.key -> "true", + SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR.key -> "2", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100", + SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> maxFields, + CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key -> costBased) { + withParquetTable( + (0 until 4000).map(i => (if (i % 10 < 8) 0 else i % 50, i % 10, i)), + "skew_f") { + withParquetTable((0 until 500).map(i => (i % 50, i)), "skew_g") { + withParquetTable((0 until 10).map(i => (i, s"name_$i")), "skew_d") { + val query = + """SELECT /*+ BROADCAST(d), MERGE(f, g) */ f._1, f._2, f._3, g._2, d._2 + |FROM skew_f f JOIN skew_g g ON f._1 = g._1 + |LEFT JOIN skew_d d ON f._2 = d._1 + |DISTRIBUTE BY f._3""".stripMargin + val (_, plan) = checkSparkAnswer(sql(query)) + val skewed = collect(plan) { case j: SortMergeJoinExec if j.isSkewJoin => j } + assert(skewed.nonEmpty, s"expected a skewed Spark sort-merge join:\n$plan") + val sparkBroadcastJoins = collect(plan) { case j: BroadcastHashJoinExec => j } + assert(sparkBroadcastJoins.size == 1, s"plan:\n$plan") + val build = sparkBroadcastJoins.head.right + val readsCometBroadcast = build.collectFirst { case c: CometColumnarToRowExec => + c.child match { + case InputAdapter(stage) => stage + case stage => stage + } + } + assert( + readsCometBroadcast.exists(_.isInstanceOf[BroadcastQueryStageExec]), + s"expected the join to read a Comet broadcast stage:\n$plan") + assert(!build.isInstanceOf[WholeStageCodegenExec], s"plan:\n$plan") + } + } + } + } + } + } + test("BroadcastNestedLoopJoin with unequal filter") { withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { // Include NULL keys: predicate `_1 > _1` returns NULL with a NULL operand, so 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 25ec4c2e0eb..2e2ff2673e6 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala @@ -19,6 +19,7 @@ package org.apache.comet.rules +import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame} import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.expressions.aggregate.Partial @@ -27,7 +28,7 @@ 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.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} -import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{col, sum} @@ -43,6 +44,9 @@ import org.apache.comet.rules.EngineCostTable.CostClass._ class CostBasedEngineChoiceSuite extends CometTestBase { + override protected def sparkConf: SparkConf = + super.sparkConf.set(CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key, "false") + private val flag = CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key private val costTable = CometConf.COMET_EXEC_COST_BASED_ENGINES_COST_TABLE.key private val expensiveNativeSort = costTable -> "sort.flat.comet=1000,0,0" @@ -1669,4 +1673,22 @@ class CostBasedEngineChoiceSuite extends CometTestBase { } } } + + test("a sort aggregate converted to a native hash aggregate is priced as an object hash") { + withTables { + withSQLConf( + CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + flag -> "true") { + val plan = run("SELECT k, max(s) FROM t GROUP BY k") + assert(count(plan) { case a: SortAggregateExec => a } == 0, s"plan:\n$plan") + val aggs = nodes(plan).collect { case a: CometHashAggregateExec => a } + assert(aggs.size == 2, s"plan:\n$plan") + aggs.foreach { agg => + assert(agg.originalPlan.isInstanceOf[ObjectHashAggregateExec], s"plan:\n$plan") + assert(model.costClasses(agg) == Seq(Agg, AggObjectHash), s"plan:\n$plan") + } + } + } + } } diff --git a/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala index 37253b2485c..5de75ae186b 100644 --- a/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/WideRowSortFallbackSuite.scala @@ -43,6 +43,7 @@ class WideRowSortFallbackSuite extends CometTestBase { override protected def sparkConf: SparkConf = super.sparkConf + .set(CometConf.COMET_EXEC_SORT_AGGREGATE_ENABLED.key, "false") .set(shuffleMinLeaves, "0") .set(CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key, "false")