From 044b4dc9e6263e1f91ab3b50499debe7e5099e9a Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 09:09:09 +0700 Subject: [PATCH 01/11] fix: sort a native hash join's output when the probe order it reports is not kept Spark reports a hash join's streamed-side ordering and may drop a sort above the join on that basis. DataFusion keeps unmatched probe rows of a right join in order only when it sees the probe input sorted, which it does not when that order comes from outside the native plan. The new HashJoin output_ordering field carries the ordering, and the native planner adds a sort above the join when the join does not already produce it. --- native/core/src/execution/planner.rs | 343 ++++++++++++++++++++++++-- native/proto/src/proto/operator.proto | 3 + 2 files changed, 325 insertions(+), 21 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 221d6b05360..06a0bf5897a 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1235,6 +1235,27 @@ impl PhysicalPlanner { } } + /// Returns a sort of `plan` by the SortOrder exprs in `ordering`, or `None` when `ordering` + /// is empty or the plan's equivalence properties already satisfy it. + fn sort_unless_ordered( + &self, + plan: Arc, + ordering: &[Expr], + ) -> Result>, ExecutionError> { + let schema = plan.schema(); + let sort_exprs = ordering + .iter() + .map(|expr| self.create_sort_expr(expr, Arc::clone(&schema))) + .collect::, _>>()?; + let Some(lex_ordering) = LexOrdering::new(sort_exprs.clone()) else { + return Ok(None); + }; + if plan.equivalence_properties().ordering_satisfy(sort_exprs)? { + return Ok(None); + } + Ok(Some(Arc::new(SortExec::new(lex_ordering, plan)))) + } + #[allow(clippy::too_many_arguments)] pub fn create_binary_expr( &self, @@ -2661,21 +2682,16 @@ impl PhysicalPlanner { // Exception: null-aware anti-join requires LeftAnti + build-right semantics // (which matches DataFusion's default), and swap_inputs would turn LeftAnti // into RightAnti, which DataFusion rejects with null_aware=true. - if join.build_side == BuildSide::BuildLeft as i32 || join.null_aware_anti_join { + let (join_root, mut additional_native_plans) = if join.build_side + == BuildSide::BuildLeft as i32 + || join.null_aware_anti_join + { let hash_join = Self::apply_join_dynamic_filter( hash_join, join.dynamic_filter_enabled && !join.null_aware_anti_join, self.session_ctx.copied_config().options(), )?; - Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new( - spark_plan.plan_id, - hash_join, - vec![join_params.left, join_params.right], - )), - )) + (hash_join, vec![]) } else { let swapped_hash_join = hash_join.as_ref().swap_inputs(PartitionMode::Partitioned)?; @@ -2690,18 +2706,32 @@ impl PhysicalPlanner { // a projection was added to the hash join additional_native_plans.push(Arc::clone(swapped_hash_join.children()[0])); } + (swapped_hash_join, additional_native_plans) + }; - Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new_with_additional( - spark_plan.plan_id, - swapped_hash_join, - vec![join_params.left, join_params.right], - additional_native_plans, - )), - )) - } + // Spark reports the streamed side's ordering and may have removed a sort above + // the join on that basis. DataFusion keeps unmatched probe rows in probe order + // only when it sees the probe input sorted, so sort the output otherwise. + let native_plan = match self + .sort_unless_ordered(Arc::clone(&join_root), &join.output_ordering)? + { + Some(sort) => { + additional_native_plans.push(join_root); + sort + } + None => join_root, + }; + + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new_with_additional( + spark_plan.plan_id, + native_plan, + vec![join_params.left, join_params.right], + additional_native_plans, + )), + )) } OpStruct::Window(wnd) => { let (scans, shuffle_scans, child) = @@ -6012,6 +6042,7 @@ mod tests { build_side: 0, null_aware_anti_join: false, dynamic_filter_enabled: false, + output_ordering: vec![], })), }; @@ -6306,6 +6337,276 @@ mod tests { assert_eq!("AggregateExec", planned.native_plan.name()); } + fn typed_bound_reference(index: i32, datatype: &spark_expression::DataType) -> Expr { + Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index, + datatype: Some(datatype.clone()), + })), + ..Default::default() + } + } + + fn ascending_nulls_first(child: Expr) -> Expr { + Expr { + expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder { + child: Some(Box::new(child)), + direction: spark_expression::SortDirection::Ascending as i32, + null_ordering: spark_expression::NullOrdering::NullsFirst as i32, + }))), + ..Default::default() + } + } + + fn two_column_scan(datatype: &spark_expression::DataType) -> Operator { + Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields: vec![datatype.clone(); 2], + source: String::new(), + })), + ..Default::default() + } + } + + /// A hash join on the first column of each side that reports `output_ordering`. + fn hash_join_on_first_columns( + join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + children: [Operator; 2], + key_type: &spark_expression::DataType, + output_ordering: Vec, + ) -> Operator { + Operator { + plan_id: 1, + children: children.into(), + op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + left_join_keys: vec![typed_bound_reference(0, key_type)], + right_join_keys: vec![typed_bound_reference(0, key_type)], + join_type: join_type as i32, + build_side: build_side as i32, + output_ordering, + ..Default::default() + })), + ..Default::default() + } + } + + /// Executes a planned join with one Int32 (key, value) batch in place of each child scan. + async fn execute_join( + planned: &crate::execution::spark_plan::SparkPlan, + task_ctx: Arc, + sides: [(Vec, Vec); 2], + ) -> RecordBatch { + use datafusion::common::tree_node::{Transformed, TreeNode}; + + let replacements: Vec<(Arc, Arc)> = planned + .children + .iter() + .zip(sides) + .map(|(child, (keys, values))| { + let batch = RecordBatch::try_new( + child.schema(), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap(); + let input: Arc = + MemorySourceConfig::try_new_exec(&[vec![batch]], child.schema(), None).unwrap(); + (Arc::clone(&child.native_plan), input) + }) + .collect(); + let native_plan = Arc::clone(&planned.native_plan) + .transform_up(|node| { + Ok( + match replacements + .iter() + .find(|(scan, _)| Arc::ptr_eq(scan, &node)) + { + Some((_, input)) => Transformed::yes(Arc::clone(input)), + None => Transformed::no(node), + }, + ) + }) + .unwrap() + .data; + let results = collect(native_plan.execute(0, task_ctx).unwrap()) + .await + .unwrap(); + arrow::compute::concat_batches(&planned.native_plan.schema(), &results).unwrap() + } + + fn additional_plan_names(planned: &crate::execution::spark_plan::SparkPlan) -> Vec<&str> { + planned + .additional_native_plans + .iter() + .map(|plan| plan.name()) + .collect() + } + + #[tokio::test] + async fn hash_join_sorts_output_when_probe_order_is_unknown() { + // Spark LeftOuter + BuildRight runs as a swapped DataFusion Right join. The probe scan + // declares no ordering, so DataFusion emits matched rows before unmatched ones. + let int_type = create_proto_datatype(); + let op = hash_join_on_first_columns( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + [two_column_scan(&int_type), two_column_scan(&int_type)], + &int_type, + vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], + ); + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + + let output = execute_join( + &planned, + task_ctx, + [(vec![1, 5, 2], vec![10, 50, 20]), (vec![5], vec![500])], + ) + .await; + assert_eq!( + &Int32Array::from(vec![1, 2, 5]) as &dyn Array, + output.column(0).as_ref() + ); + assert_eq!( + &Int32Array::from(vec![None, None, Some(5)]) as &dyn Array, + output.column(2).as_ref() + ); + assert_eq!("SortExec", planned.native_plan.name()); + assert_eq!( + vec!["HashJoinExec", "ProjectionExec"], + additional_plan_names(&planned) + ); + } + + #[tokio::test] + async fn hash_join_sorts_output_on_the_unswapped_build_left_path() { + // Spark RightOuter + BuildLeft runs unswapped. The probe key is the first column of the + // right side, at index 2 of the join output. + let int_type = create_proto_datatype(); + let op = hash_join_on_first_columns( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + [two_column_scan(&int_type), two_column_scan(&int_type)], + &int_type, + vec![ascending_nulls_first(typed_bound_reference(2, &int_type))], + ); + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + + let output = execute_join( + &planned, + task_ctx, + [(vec![5], vec![500]), (vec![1, 5, 2], vec![10, 50, 20])], + ) + .await; + assert_eq!( + &Int32Array::from(vec![1, 2, 5]) as &dyn Array, + output.column(2).as_ref() + ); + assert_eq!( + &Int32Array::from(vec![None, None, Some(5)]) as &dyn Array, + output.column(0).as_ref() + ); + assert_eq!("SortExec", planned.native_plan.name()); + assert_eq!(vec!["HashJoinExec"], additional_plan_names(&planned)); + } + + #[test] + fn hash_join_keeps_native_probe_order_without_a_sort() { + // A native sort below the probe side is visible to DataFusion, which then keeps unmatched + // probe rows in order. Float keys check that the ordering is normalized like the sort's. + let double_type = spark_expression::DataType { + type_id: 6, + type_info: None, + }; + for key_type in [create_proto_datatype(), double_type] { + for (join_type, build_side, probe_key_index, root) in [ + ( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + 0, + "ProjectionExec", + ), + ( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + 2, + "HashJoinExec", + ), + ] { + let sorted_probe = Operator { + children: vec![two_column_scan(&key_type)], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![ascending_nulls_first(typed_bound_reference( + 0, &key_type, + ))], + ..Default::default() + })), + ..Default::default() + }; + let children = if build_side == spark_operator::BuildSide::BuildRight { + [sorted_probe, two_column_scan(&key_type)] + } else { + [two_column_scan(&key_type), sorted_probe] + }; + let op = hash_join_on_first_columns( + join_type, + build_side, + children, + &key_type, + vec![ascending_nulls_first(typed_bound_reference( + probe_key_index, + &key_type, + ))], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap(); + assert_eq!( + root, + planned.native_plan.name(), + "{join_type:?} {key_type:?}" + ); + } + } + } + + #[test] + fn hash_join_without_output_ordering_adds_no_sort() { + let int_type = create_proto_datatype(); + for (join_type, build_side, root) in [ + ( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + "ProjectionExec", + ), + ( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + "HashJoinExec", + ), + ] { + let op = hash_join_on_first_columns( + join_type, + build_side, + [two_column_scan(&int_type), two_column_scan(&int_type)], + &int_type, + vec![], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap(); + assert_eq!(root, planned.native_plan.name(), "{join_type:?}"); + } + } + #[tokio::test] async fn projection_prunes_filter_output_and_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 97c65f0da36..43c04b9c106 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1026,6 +1026,9 @@ message HashJoin { // Opt-in completed-build filtering of probe batches and eligible native Parquet readers. // Older serialized plans leave it false. bool dynamic_filter_enabled = 10; + // SortOrder exprs bound to the join output that Spark reports as its output ordering. + // Empty when Spark reports no ordering or the native plan does not need to keep it. + repeated spark.spark_expression.Expr output_ordering = 11; } message SortMergeJoin { From 8dc4b90fd95e182df80b247b90aa78fa84417e8e Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 09:44:42 +0700 Subject: [PATCH 02/11] fix: send a hash join's reported ordering to the native planner For LeftOuter with BuildRight and RightOuter with BuildLeft, Spark reports the streamed side's ordering, and DataFusion runs the join as a right join that may emit unmatched probe rows after the matched ones. The serde now sends that ordering so the native planner can restore it, and falls back to Spark when the ordering cannot be sorted natively. --- .../apache/spark/sql/comet/operators.scala | 18 ++ .../apache/comet/exec/CometJoinSuite.scala | 173 +++++++++++++++++- 2 files changed, 186 insertions(+), 5 deletions(-) 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 bcae614556b..9ca5a0ca52f 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 @@ -2861,6 +2861,23 @@ trait CometHashJoin { } } + // DataFusion may move unmatched probe rows for these shapes, so the native planner needs the + // ordering Spark reports for the join to restore it. + val outputOrdering = (join.joinType, join.buildSide) match { + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty => + if (!supportedSortType(join, join.outputOrdering)) { + withFallbackReason(join, "Unsupported data type in hash join output ordering") + return None + } + val orders = join.outputOrdering.map(exprToProto(_, join.output)) + if (orders.exists(_.isEmpty)) { + withFallbackReason(join, "hash join output ordering not supported") + return None + } + orders.flatten + case _ => Nil + } + val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output)) val rightKeys = join.rightKeys.map(exprToProto(_, join.right.output)) @@ -2876,6 +2893,7 @@ trait CometHashJoin { else OperatorOuterClass.BuildSide.BuildRight) .setNullAwareAntiJoin(isNullAwareAntiJoin) .setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf)) + .addAllOutputOrdering(outputOrdering.asJava) condition.foreach(joinBuilder.setCondition) Some(builder.setHashJoin(joinBuilder).build()) } else { 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 369c2517636..49ff79811e7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -26,21 +26,21 @@ import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.SparkException -import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation -import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} +import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, InterpretedOrdering, IsNotNull, SortDirection, SortOrder} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBinaryExec, CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType, IntegerType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} @@ -1219,6 +1219,169 @@ class CometJoinSuite extends CometTestBase { } } + // A hash join reports its streamed side's ordering, so Spark drops a sort above it when the + // streamed side already arrives partitioned and sorted. Here that side is a cached table that + // reaches the native join through Comet's JVM to native conversion. + private def withSortedCache(data: DataFrame, order: Column*)(f: DataFrame => Unit): Unit = { + withSQLConf( + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val cached = data.repartition(2, data("_1")).sortWithinPartitions(order: _*).cache() + try { + cached.count() + f(cached) + } finally { + cached.unpersist() + } + } + } + + /** + * Checks the query's rows against Spark and that every partition is sorted by `order`, which is + * bound to the query's output. Returns the plan's single native hash join and its build side. + */ + private def checkJoinKeepsOrder( + query: => DataFrame, + order: SortOrder): (CometBinaryExec, BuildSide) = { + // withSQLConf returns Unit on Spark 3.x, so the expected rows are assigned inside it. + var expected = Seq.empty[String] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = query.collect().map(_.toString).sorted.toSeq + } + val df = query + val rows = df.collect() + assert(rows.map(_.toString).sorted.toSeq == expected) + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { + case j: CometHashJoinExec => (j, j.buildSide) + case j: CometBroadcastHashJoinExec => (j, j.buildSide) + } + assert(joins.length == 1, plan) + // Read before the plan runs again below, so the count covers the one run above. + assert(joins.head._1.metrics("output_rows").value == rows.length, plan) + val ordering = new InterpretedOrdering(Seq(order)) + df.queryExecution.toRdd.map(_.copy()).glom().collect().zipWithIndex.foreach { + case (part, i) => + val outOfOrder = part.zip(part.drop(1)).indexWhere { case (a, b) => ordering.gt(a, b) } + assert(outOfOrder < 0, s"partition $i is out of order at row $outOfOrder\n$plan") + } + joins.head + } + + private def keyOrder(dataType: DataType, direction: SortDirection): SortOrder = + SortOrder(BoundReference(0, dataType, nullable = true), direction) + + private val streamedOrderCases = + (for (hint <- Seq("shuffle_hash", "broadcast"); joinType <- Seq("left_outer", "right_outer")) + yield (hint, joinType, false, false)) ++ Seq( + ("shuffle_hash", "left_outer", true, false), + ("shuffle_hash", "right_outer", false, true)) + + for ((hint, joinType, withCondition, adaptive) <- streamedOrderCases) { + test( + s"$hint $joinType join keeps a cached streamed side's order, " + + s"condition=$withCondition, AQE=$adaptive") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1") { streamed => + val build = spark.table("small").hint(hint) + val keys = streamed("_1") === build("_1") + val on = if (withCondition) keys && build("_2") > 3 else keys + def query: DataFrame = { + val joined = + if (joinType == "left_outer") streamed.join(build, on, joinType) + else build.join(streamed, on, joinType) + joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + } + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val (_, buildSide) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) + assert(buildSide == (if (joinType == "left_outer") BuildRight else BuildLeft)) + } + } + } + } + } + } + + test("shuffle_hash join keeps a descending, nulls last streamed order with null keys") { + val data = (0 until 10000).map(i => (if (i % 7 == 0) None else Some(i % 100), i)) + withParquetTable(data, "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1".desc_nulls_last) { streamed => + val build = spark.table("small").hint("shuffle_hash") + checkJoinKeepsOrder( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions($"k".desc_nulls_last), + keyOrder(IntegerType, Descending)) + } + } + } + } + + test("shuffle_hash join keeps a descending streamed order on a double with NaN and -0.0") { + def d(i: Int): Option[Double] = i % 9 match { + case 0 => None + case 1 => Some(Double.NaN) + case 2 => Some(-0.0) + case 3 => Some(0.0) + case _ => Some(i % 50 - 25.5) + } + withParquetTable((0 until 10000).map(i => (i % 100, i, d(i))), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_3".desc_nulls_last) { streamed => + val build = spark.table("small").hint("shuffle_hash") + checkJoinKeepsOrder( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_3").as("d"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions($"d".desc_nulls_last), + keyOrder(DoubleType, Descending)) + } + } + } + } + + test("only outer hash joins that can move unmatched streamed rows send their ordering") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1") { streamed => + val build = spark.table("small").hint("shuffle_hash") + def query(joinType: String): DataFrame = streamed + .join(build, streamed("_1") === build("_1"), joinType) + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val (inner, _) = checkJoinKeepsOrder(query("inner"), keyOrder(IntegerType, Ascending)) + assert(inner.nativeOp.getHashJoin.getOutputOrderingCount == 0) + val (leftOuter, buildSide) = + checkJoinKeepsOrder(query("left_outer"), keyOrder(IntegerType, Ascending)) + assert(buildSide == BuildRight) + assert(leftOuter.nativeOp.getHashJoin.getOutputOrderingCount > 0) + } + } + } + } + + test("hash join falls back when its streamed ordering has an unsupported sort type") { + withParquetTable((0 until 1000).map(i => (i % 100, i, (i % 13, i.toString))), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_3") { streamed => + val build = spark.table("small").hint("shuffle_hash") + checkSparkAnswerAndFallbackReason( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_3").as("s"), streamed("_2").as("v"), build("_2").as("w")), + "Unsupported data type in hash join output ordering") + } + } + } + } + test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") { withSQLConf( SQLConf.PREFER_SORTMERGEJOIN.key -> "false", From 7f67a6d37cee32180742237e892f6afda5b0ee19 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 09:51:37 +0700 Subject: [PATCH 03/11] test: a native sort-merge join below a hash join adds no extra sort The ordering the merge join produces reaches the hash join through the swapped right join and its projection, so the planner keeps it without sorting again. --- native/core/src/execution/planner.rs | 50 ++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 06a0bf5897a..ea90e6a8e05 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6578,6 +6578,56 @@ mod tests { } } + #[test] + fn hash_join_over_a_native_sort_merge_join_adds_no_sort() { + // A sort merge join declares its key ordering, so a hash join streaming it keeps that + // ordering without a sort. Spark reports the merge join's left key as the ordering. + let int_type = create_proto_datatype(); + let key = || ascending_nulls_first(typed_bound_reference(0, &int_type)); + let sorted_scan = || Operator { + children: vec![two_column_scan(&int_type)], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![key()], + ..Default::default() + })), + ..Default::default() + }; + for smj_join_type in [ + spark_operator::JoinType::Inner, + spark_operator::JoinType::LeftOuter, + ] { + let sort_merge_join = Operator { + children: vec![sorted_scan(), sorted_scan()], + op_struct: Some(OpStruct::SortMergeJoin(spark_operator::SortMergeJoin { + left_join_keys: vec![typed_bound_reference(0, &int_type)], + right_join_keys: vec![typed_bound_reference(0, &int_type)], + join_type: smj_join_type as i32, + sort_options: vec![key()], + condition: None, + })), + ..Default::default() + }; + // The merge join's left key is column 0 of the hash join output. + let op = hash_join_on_first_columns( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + [sort_merge_join, two_column_scan(&int_type)], + &int_type, + vec![key()], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap(); + assert_ne!( + "SortExec", + planned.native_plan.name(), + "{smj_join_type:?}\n{}\n{:?}", + datafusion::physical_plan::displayable(planned.native_plan.as_ref()).indent(true), + planned.native_plan.properties().eq_properties + ); + } + } + #[test] fn hash_join_without_output_ordering_adds_no_sort() { let int_type = create_proto_datatype(); From 6ed1dca3725c5f769ca1defc9609d48330f646fe Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 11:32:20 +0700 Subject: [PATCH 04/11] fix: keep the streamed order of a null-aware anti join and widen the ordering tests A null-aware anti join runs unswapped, so DataFusion builds on Spark's streamed side and concatenates its batches in reverse; it now sends its ordering too and gets the same native sort. The tests cover multi-column and multi-batch orderings, the expression fallback, a natively sorted probe, join metrics with a sort on top, and the join class per hint. The compatibility guide describes how hash joins keep the order. --- .../latest/compatibility/operators.md | 6 + docs/source/user-guide/latest/operators.md | 4 +- native/core/src/execution/planner.rs | 7 +- native/proto/src/proto/operator.proto | 2 +- .../apache/spark/sql/comet/operators.scala | 36 ++-- .../apache/comet/exec/CometJoinSuite.scala | 186 +++++++++++++++--- 6 files changed, 188 insertions(+), 53 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 40732ced79f..a6519cf7663 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -91,6 +91,12 @@ rows by the grouping keys only, and Spark and Comet can leave rows with equal ke orders, so a group with more than one candidate value can return a different value than Spark. Both results are valid under Spark's semantics for these functions. +## Hash Join + +When Spark reports a hash join's output as ordered by its streamed side, Comet keeps that order. +When the order comes from outside the native plan, such as a cached table, Comet sorts the native +join output to restore it. + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index acefd4f0890..c950121e119 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -98,8 +98,8 @@ omitted from the tables below and may be reconsidered based on demand: | Operator | Status | Notes | | ----------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `BroadcastHashJoinExec` | ✅ | | -| `ShuffledHashJoinExec` | ✅ | | +| `BroadcastHashJoinExec` | ✅ | Falls back to Spark when an outer or `NOT IN` join reports an output ordering whose type or expression Comet cannot sort natively. | +| `ShuffledHashJoinExec` | ✅ | Falls back to Spark when an outer join reports an output ordering whose type or expression Comet cannot sort natively. | | `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)). | diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ea90e6a8e05..9ead1e5e264 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2711,7 +2711,8 @@ impl PhysicalPlanner { // Spark reports the streamed side's ordering and may have removed a sort above // the join on that basis. DataFusion keeps unmatched probe rows in probe order - // only when it sees the probe input sorted, so sort the output otherwise. + // only when it sees the probe input sorted, and does not keep a null-aware anti + // join's build order, so sort the output otherwise. let native_plan = match self .sort_unless_ordered(Arc::clone(&join_root), &join.output_ordering)? { @@ -6618,8 +6619,8 @@ mod tests { let (_, _, planned) = PhysicalPlanner::default() .create_plan(&op, &mut vec![], 1) .unwrap(); - assert_ne!( - "SortExec", + assert_eq!( + "ProjectionExec", planned.native_plan.name(), "{smj_join_type:?}\n{}\n{:?}", datafusion::physical_plan::displayable(planned.native_plan.as_ref()).indent(true), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 43c04b9c106..be46c7b13e2 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1027,7 +1027,7 @@ message HashJoin { // Older serialized plans leave it false. bool dynamic_filter_enabled = 10; // SortOrder exprs bound to the join output that Spark reports as its output ordering. - // Empty when Spark reports no ordering or the native plan does not need to keep it. + // Empty means the native plan adds no sort. repeated spark.spark_expression.Expr output_ordering = 11; } 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 9ca5a0ca52f..02bb7978218 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 @@ -2861,21 +2861,27 @@ trait CometHashJoin { } } - // DataFusion may move unmatched probe rows for these shapes, so the native planner needs the - // ordering Spark reports for the join to restore it. - val outputOrdering = (join.joinType, join.buildSide) match { - case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty => - if (!supportedSortType(join, join.outputOrdering)) { - withFallbackReason(join, "Unsupported data type in hash join output ordering") - return None - } - val orders = join.outputOrdering.map(exprToProto(_, join.output)) - if (orders.exists(_.isEmpty)) { - withFallbackReason(join, "hash join output ordering not supported") - return None - } - orders.flatten - case _ => Nil + // DataFusion may move unmatched probe rows for the outer shapes, and runs a null-aware anti + // join with Spark's streamed side as its build side, so the native planner needs the ordering + // Spark reports for the join to restore it. + val mayReorderStreamed = (join.joinType, join.buildSide) match { + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => true + case (LeftAnti, BuildRight) => isNullAwareAntiJoin + case _ => false + } + val outputOrdering = if (mayReorderStreamed && join.outputOrdering.nonEmpty) { + if (!supportedSortType(join, join.outputOrdering)) { + withFallbackReason(join, "Unsupported data type in hash join output ordering") + return None + } + val orders = join.outputOrdering.map(exprToProto(_, join.output)) + if (orders.exists(_.isEmpty)) { + withFallbackReason(join, "Unsupported expression in hash join output ordering") + return None + } + orders.flatten + } else { + Nil } val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output)) 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 49ff79811e7..767181251ae 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -30,11 +30,11 @@ import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, InterpretedOrdering, IsNotNull, SortDirection, SortOrder} -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, EliminateSorts} import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.comet.{CometBinaryExec, CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +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.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} @@ -1244,7 +1244,7 @@ class CometJoinSuite extends CometTestBase { */ private def checkJoinKeepsOrder( query: => DataFrame, - order: SortOrder): (CometBinaryExec, BuildSide) = { + order: SortOrder*): (CometBinaryExec, BuildSide) = { // withSQLConf returns Unit on Spark 3.x, so the expected rows are assigned inside it. var expected = Seq.empty[String] withSQLConf(CometConf.COMET_ENABLED.key -> "false") { @@ -1261,45 +1261,80 @@ class CometJoinSuite extends CometTestBase { assert(joins.length == 1, plan) // Read before the plan runs again below, so the count covers the one run above. assert(joins.head._1.metrics("output_rows").value == rows.length, plan) - val ordering = new InterpretedOrdering(Seq(order)) - df.queryExecution.toRdd.map(_.copy()).glom().collect().zipWithIndex.foreach { - case (part, i) => - val outOfOrder = part.zip(part.drop(1)).indexWhere { case (a, b) => ordering.gt(a, b) } - assert(outOfOrder < 0, s"partition $i is out of order at row $outOfOrder\n$plan") + assert(joins.head._1.metrics("build_input_rows").value > 0, plan) + // Spark dropped the sort above the join, so the order comes from the join alone. + val sortsAboveJoin = collect(plan) { + case s: SortExec => s + case s: CometSortExec => s + }.filter(s => find(s)(_ eq joins.head._1).isDefined) + assert(sortsAboveJoin.isEmpty, plan) + val ordering = new InterpretedOrdering(order) + // Propagates the session's confs, such as the batch size, to the tasks as collect() does. + val parts = SQLExecution.withSQLConfPropagated(spark) { + df.queryExecution.toRdd.map(_.copy()).glom().collect() + } + parts.zipWithIndex.foreach { case (part, i) => + val outOfOrder = part.zip(part.drop(1)).indexWhere { case (a, b) => ordering.gt(a, b) } + assert(outOfOrder < 0, s"partition $i is out of order at row $outOfOrder\n$plan") } joins.head } - private def keyOrder(dataType: DataType, direction: SortDirection): SortOrder = - SortOrder(BoundReference(0, dataType, nullable = true), direction) + private def keyOrder( + dataType: DataType, + direction: SortDirection, + ordinal: Int = 0): SortOrder = + SortOrder(BoundReference(ordinal, dataType, nullable = true), direction) + // (hint, join type, extra condition, AQE, native batch size) private val streamedOrderCases = (for (hint <- Seq("shuffle_hash", "broadcast"); joinType <- Seq("left_outer", "right_outer")) - yield (hint, joinType, false, false)) ++ Seq( - ("shuffle_hash", "left_outer", true, false), - ("shuffle_hash", "right_outer", false, true)) + yield (hint, joinType, false, false, 8192)) ++ Seq( + ("shuffle_hash", "left_outer", true, false, 8192), + ("shuffle_hash", "right_outer", false, true, 8192), + ("shuffle_hash", "left_outer", false, false, 100)) - for ((hint, joinType, withCondition, adaptive) <- streamedOrderCases) { + for ((hint, joinType, withCondition, adaptive, batchSize) <- streamedOrderCases) { test( s"$hint $joinType join keeps a cached streamed side's order, " + - s"condition=$withCondition, AQE=$adaptive") { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1") { streamed => - val build = spark.table("small").hint(hint) - val keys = streamed("_1") === build("_1") - val on = if (withCondition) keys && build("_2") > 3 else keys - def query: DataFrame = { - val joined = - if (joinType == "left_outer") streamed.join(build, on, joinType) - else build.join(streamed, on, joinType) - joined - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions("k") - } - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { - val (_, buildSide) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) - assert(buildSide == (if (joinType == "left_outer") BuildRight else BuildLeft)) + s"condition=$withCondition, AQE=$adaptive, batchSize=$batchSize") { + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1") { streamed => + val build = spark.table("small").hint(hint) + val keys = streamed("_1") === build("_1") + val on = if (withCondition) keys && build("_2") > 3 else keys + def query: DataFrame = { + val joined = + if (joinType == "left_outer") streamed.join(build, on, joinType) + else build.join(streamed, on, joinType) + joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + } + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val (join, buildSide) = + checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) + assert(buildSide == (if (joinType == "left_outer") BuildRight else BuildLeft)) + hint match { + case "broadcast" => assert(join.isInstanceOf[CometBroadcastHashJoinExec]) + case _ => assert(join.isInstanceOf[CometHashJoinExec]) + } + if (batchSize < 1000) { + // Several batches per partition reach the native join and its sort, and + // several come back to the JVM. + val partitions = 2 + assert(join.metrics("output_batches").value > partitions) + val plan = query.queryExecution.executedPlan + val batches = collectFirst(plan) { case c: ColumnarToRowTransition => + SQLExecution.withSQLConfPropagated(spark) { + c.child.executeColumnar().count() + } + } + assert(batches.exists(_ > partitions), plan) + } + } } } } @@ -1347,6 +1382,76 @@ class CometJoinSuite extends CometTestBase { } } + test("shuffle_hash join keeps a two-key streamed order") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1", $"_2".desc) { streamed => + val build = spark.table("small").hint("shuffle_hash") + checkJoinKeepsOrder( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions($"k", $"v".desc), + keyOrder(IntegerType, Ascending), + keyOrder(IntegerType, Descending, ordinal = 1)) + } + } + } + } + + // The streamed side is sorted inside the native plan, so DataFusion sees its ordering and keeps + // unmatched probe rows in place without an extra sort. EliminateSorts would drop that sort as + // one below a join. + test("shuffle_hash join keeps a natively sorted streamed side's order") { + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") + val build = spark.table("small").hint("shuffle_hash") + val (join, buildSide) = checkJoinKeepsOrder( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k"), + keyOrder(IntegerType, Ascending)) + assert(join.isInstanceOf[CometHashJoinExec]) + assert(buildSide == BuildRight) + } + } + } + } + + // Spark plans NOT IN as a null-aware anti join that streams the left side. Comet runs it with + // the left side as DataFusion's build side, which gives up that side's batch order. Sparse keys + // keep DataFusion on its hash map build, which stores the batches last to first. + test("null-aware anti join keeps a cached streamed side's order across batches") { + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "100") { + withParquetTable((0 until 10000).map(i => (i % 100 * 1000, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10000, i)), "small") { + withSortedCache(spark.table("big"), $"_1") { streamed => + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { + withTempView("streamed") { + streamed.createOrReplaceTempView("streamed") + val (join, _) = checkJoinKeepsOrder( + sql("SELECT _1 AS k, _2 AS v FROM streamed WHERE _1 NOT IN " + + "(SELECT _1 FROM small)").sortWithinPartitions("k"), + keyOrder(IntegerType, Ascending)) + join match { + case j: CometBroadcastHashJoinExec => assert(j.isNullAwareAntiJoin, j) + case j => fail(s"expected a null-aware broadcast hash join: $j") + } + } + } + } + } + } + } + } + test("only outer hash joins that can move unmatched streamed rows send their ordering") { withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { @@ -1382,6 +1487,23 @@ class CometJoinSuite extends CometTestBase { } } + test("hash join falls back when its streamed ordering has an unsupported expression") { + withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withSortedCache(spark.table("big"), $"_1" + 1) { streamed => + val build = spark.table("small").hint("shuffle_hash") + withSQLConf(CometConf.getExprEnabledConfigKey("Add") -> "false") { + checkSparkAnswerAndFallbackReason( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")), + "Unsupported expression in hash join output ordering") + } + } + } + } + } + test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") { withSQLConf( SQLConf.PREFER_SORTMERGEJOIN.key -> "false", From ef6623d4c4a34ab6a2324b5e3d5c4f9fdb43134e Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 12:16:12 +0700 Subject: [PATCH 05/11] test: tighten the hash join ordering tests and correct the NOT IN docs The null-aware anti join test now checks that its build spans several batches, a metric check that held for any batch size is gone, and a test comment claims only what the test checks. The guides say a NOT IN join is sorted too, and that only a NOT IN under OR, planned as an existence nested loop join, falls back to Spark. --- .../latest/compatibility/operators.md | 4 +-- docs/source/user-guide/latest/operators.md | 3 ++- .../apache/comet/exec/CometJoinSuite.scala | 25 ++++++++++++------- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index a6519cf7663..504733222fc 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -94,8 +94,8 @@ Both results are valid under Spark's semantics for these functions. ## Hash Join When Spark reports a hash join's output as ordered by its streamed side, Comet keeps that order. -When the order comes from outside the native plan, such as a cached table, Comet sorts the native -join output to restore it. +When DataFusion does not keep it, for example when it comes from a cached table or the join is a +`NOT IN`, Comet sorts the native join output to restore it. ## Window Functions diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index c950121e119..ef7a26acbb5 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -105,7 +105,8 @@ omitted from the tables below and may be reconsidered based on demand: `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual -(non-equi) join conditions, computed (non-column) join keys, and `NOT IN` fall back to Spark. +(non-equi) join conditions, and computed (non-column) join keys fall back to Spark. So does `NOT IN` +combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinExec`. ## Exchanges 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 767181251ae..74fe5c845e0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1322,10 +1322,8 @@ class CometJoinSuite extends CometTestBase { case _ => assert(join.isInstanceOf[CometHashJoinExec]) } if (batchSize < 1000) { - // Several batches per partition reach the native join and its sort, and - // several come back to the JVM. + // Several batches per partition come back to the JVM from the native join. val partitions = 2 - assert(join.metrics("output_batches").value > partitions) val plan = query.queryExecution.executedPlan val batches = collectFirst(plan) { case c: ColumnarToRowTransition => SQLExecution.withSQLConfPropagated(spark) { @@ -1399,9 +1397,8 @@ class CometJoinSuite extends CometTestBase { } } - // The streamed side is sorted inside the native plan, so DataFusion sees its ordering and keeps - // unmatched probe rows in place without an extra sort. EliminateSorts would drop that sort as - // one below a join. + // The streamed side is sorted inside the native plan, and the join's output keeps that order in + // every partition. EliminateSorts would drop that sort as one below a join. test("shuffle_hash join keeps a natively sorted streamed side's order") { withSQLConf( SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName, @@ -1436,14 +1433,24 @@ class CometJoinSuite extends CometTestBase { withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { withTempView("streamed") { streamed.createOrReplaceTempView("streamed") - val (join, _) = checkJoinKeepsOrder( + def query: DataFrame = sql("SELECT _1 AS k, _2 AS v FROM streamed WHERE _1 NOT IN " + - "(SELECT _1 FROM small)").sortWithinPartitions("k"), - keyOrder(IntegerType, Ascending)) + "(SELECT _1 FROM small)").sortWithinPartitions("k") + val (join, _) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) join match { case j: CometBroadcastHashJoinExec => assert(j.isNullAwareAntiJoin, j) case j => fail(s"expected a null-aware broadcast hash join: $j") } + // The build side spans several batches per partition. Read on a fresh plan that + // runs once, so the count covers that run alone. + val partitions = 2 + val df = query + assert(df.collect().nonEmpty) + val plan = df.queryExecution.executedPlan + val buildBatches = collectFirst(plan) { case j: CometBroadcastHashJoinExec => + j.metrics("build_input_batches").value + } + assert(buildBatches.exists(_ > partitions), plan) } } } From d1420ddcf898bfd35164573274e97df82ed1b9fe Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 14:26:53 +0700 Subject: [PATCH 06/11] test: a null-aware anti join's native output is sorted across build batches DataFusion builds a NOT IN join on Spark's streamed side and, on its hash path, concatenates those batches in reverse, so the planner sorts the output. The test feeds two sparse-key batches and checks the order. --- native/core/src/execution/planner.rs | 86 ++++++++++++++++++++++++---- 1 file changed, 76 insertions(+), 10 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 79c0aa4af03..56426db75e2 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6391,6 +6391,15 @@ mod tests { planned: &crate::execution::spark_plan::SparkPlan, task_ctx: Arc, sides: [(Vec, Vec); 2], + ) -> RecordBatch { + execute_join_batches(planned, task_ctx, sides.map(|side| vec![side])).await + } + + /// Executes a planned join with Int32 (key, value) batches in place of each child scan. + async fn execute_join_batches( + planned: &crate::execution::spark_plan::SparkPlan, + task_ctx: Arc, + sides: [Vec<(Vec, Vec)>; 2], ) -> RecordBatch { use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -6398,17 +6407,22 @@ mod tests { .children .iter() .zip(sides) - .map(|(child, (keys, values))| { - let batch = RecordBatch::try_new( - child.schema(), - vec![ - Arc::new(Int32Array::from(keys)), - Arc::new(Int32Array::from(values)), - ], - ) - .unwrap(); + .map(|(child, batches)| { + let batches: Vec = batches + .into_iter() + .map(|(keys, values)| { + RecordBatch::try_new( + child.schema(), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap() + }) + .collect(); let input: Arc = - MemorySourceConfig::try_new_exec(&[vec![batch]], child.schema(), None).unwrap(); + MemorySourceConfig::try_new_exec(&[batches], child.schema(), None).unwrap(); (Arc::clone(&child.native_plan), input) }) .collect(); @@ -6513,6 +6527,58 @@ mod tests { assert_eq!(vec!["HashJoinExec"], additional_plan_names(&planned)); } + #[tokio::test] + async fn hash_join_sorts_null_aware_anti_join_output() { + // A null-aware anti join runs unswapped in CollectLeft mode, so DataFusion builds on the + // streamed left side. Keys spanning 1024 or more values skip its perfect hash path, and + // the hash path emits the unmatched rows in reverse batch order. + let int_type = create_proto_datatype(); + let mut op = hash_join_on_first_columns( + spark_operator::JoinType::LeftAnti, + spark_operator::BuildSide::BuildRight, + [two_column_scan(&int_type), two_column_scan(&int_type)], + &int_type, + vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], + ); + let Some(OpStruct::HashJoin(join)) = op.op_struct.as_mut() else { + unreachable!() + }; + join.null_aware_anti_join = true; + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + + assert_eq!("SortExec", planned.native_plan.name()); + assert_eq!(vec!["HashJoinExec"], additional_plan_names(&planned)); + let hash_join = planned.native_plan.children()[0] + .downcast_ref::() + .unwrap(); + assert_eq!( + &datafusion::common::JoinType::LeftAnti, + hash_join.join_type() + ); + assert!(hash_join.null_aware); + + let output = execute_join_batches( + &planned, + task_ctx, + [ + vec![(vec![10, 20], vec![1, 2]), (vec![3000, 4000], vec![3, 4])], + vec![(vec![5], vec![500])], + ], + ) + .await; + assert_eq!( + &Int32Array::from(vec![10, 20, 3000, 4000]) as &dyn Array, + output.column(0).as_ref() + ); + assert_eq!( + &Int32Array::from(vec![1, 2, 3, 4]) as &dyn Array, + output.column(1).as_ref() + ); + } + #[test] fn hash_join_keeps_native_probe_order_without_a_sort() { // A native sort below the probe side is visible to DataFusion, which then keeps unmatched From 7336baccf95d2463a52d0fe0ee55aa2373ece3e2 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 9 Oct 2026 14:26:56 +0700 Subject: [PATCH 07/11] docs: say only a NOT IN combined with OR skips the native existence join A plain NOT IN runs as a native null-aware anti join; only one combined with another predicate through OR becomes an existence join that falls back to Spark. --- spark/src/main/scala/org/apache/comet/CometConf.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a187f8fa5f8..40a392c2bf4 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -325,7 +325,8 @@ object CometConf extends ShimCometConf { .doc( "Whether to enable native ExistenceJoin, produced when EXISTS / NOT EXISTS / IN is " + "combined with another predicate via OR. Sort-merge joins, residual (non-equi) join " + - "conditions, computed (non-column) join keys, and NOT IN fall back to Spark.") + "conditions, computed (non-column) join keys, and NOT IN combined with OR fall back " + + "to Spark.") .booleanConf .createWithDefault(true) val COMET_EXEC_AGGREGATE_ENABLED: ConfigEntry[Boolean] = From 93cf7605630101af240d40f5bb0e6e52e3457a2e Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sat, 10 Oct 2026 01:18:11 +0700 Subject: [PATCH 08/11] refactor: share the planner's sort-unless-ordered helper with the sort aggregate sort_unless_ordered now takes physical sort expressions. The hash join builds them at its call site, and the sort aggregate arm calls the same helper instead of its own ordering_satisfy check and SortExec wrap. --- native/core/src/execution/planner.rs | 48 ++++++++++++---------------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 56426db75e2..1b779aa1d29 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1235,22 +1235,17 @@ impl PhysicalPlanner { } } - /// Returns a sort of `plan` by the SortOrder exprs in `ordering`, or `None` when `ordering` - /// is empty or the plan's equivalence properties already satisfy it. + /// Returns a sort of `plan` by `ordering`, or `None` when `ordering` is empty or the plan's + /// equivalence properties already satisfy it. Operators whose Spark counterpart reports an + /// output ordering use this to sort their output when the native plan does not keep it. fn sort_unless_ordered( - &self, plan: Arc, - ordering: &[Expr], + ordering: Vec, ) -> Result>, ExecutionError> { - let schema = plan.schema(); - let sort_exprs = ordering - .iter() - .map(|expr| self.create_sort_expr(expr, Arc::clone(&schema))) - .collect::, _>>()?; - let Some(lex_ordering) = LexOrdering::new(sort_exprs.clone()) else { + let Some(lex_ordering) = LexOrdering::new(ordering.clone()) else { return Ok(None); }; - if plan.equivalence_properties().ordering_satisfy(sort_exprs)? { + if plan.equivalence_properties().ordering_satisfy(ordering)? { return Ok(None); } Ok(Some(Arc::new(SortExec::new(lex_ordering, plan)))) @@ -1722,14 +1717,8 @@ impl PhysicalPlanner { ) }) .collect(); - if !aggregate - .equivalence_properties() - .ordering_satisfy(ordering.clone())? + if let Some(sort) = Self::sort_unless_ordered(Arc::clone(&aggregate), ordering)? { - let sort: Arc = Arc::new(SortExec::new( - LexOrdering::new(ordering).unwrap(), - Arc::clone(&aggregate), - )); return Ok(( scans, shuffle_scans, @@ -2713,15 +2702,20 @@ impl PhysicalPlanner { // the join on that basis. DataFusion keeps unmatched probe rows in probe order // only when it sees the probe input sorted, and does not keep a null-aware anti // join's build order, so sort the output otherwise. - let native_plan = match self - .sort_unless_ordered(Arc::clone(&join_root), &join.output_ordering)? - { - Some(sort) => { - additional_native_plans.push(join_root); - sort - } - None => join_root, - }; + let join_schema = join_root.schema(); + let output_ordering = join + .output_ordering + .iter() + .map(|expr| self.create_sort_expr(expr, Arc::clone(&join_schema))) + .collect::, _>>()?; + let native_plan = + match Self::sort_unless_ordered(Arc::clone(&join_root), output_ordering)? { + Some(sort) => { + additional_native_plans.push(join_root); + sort + } + None => join_root, + }; Ok(( scans, From 1c9b7e9f1d3a0d75f122677996b2e743892cb0de Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sat, 10 Oct 2026 01:25:51 +0700 Subject: [PATCH 09/11] test: check the hash join ordering field for every shape that sends one The streamed order matrix and the null-aware anti join test assert that the native op carries an output ordering. The separate test keeps only the inner join check, since its outer half repeated a matrix case. --- .../org/apache/comet/exec/CometJoinSuite.scala | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) 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 74fe5c845e0..400cf4a8325 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1321,6 +1321,7 @@ class CometJoinSuite extends CometTestBase { case "broadcast" => assert(join.isInstanceOf[CometBroadcastHashJoinExec]) case _ => assert(join.isInstanceOf[CometHashJoinExec]) } + assert(join.nativeOp.getHashJoin.getOutputOrderingCount > 0) if (batchSize < 1000) { // Several batches per partition come back to the JVM from the native join. val partitions = 2 @@ -1438,7 +1439,9 @@ class CometJoinSuite extends CometTestBase { "(SELECT _1 FROM small)").sortWithinPartitions("k") val (join, _) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) join match { - case j: CometBroadcastHashJoinExec => assert(j.isNullAwareAntiJoin, j) + case j: CometBroadcastHashJoinExec => + assert(j.isNullAwareAntiJoin, j) + assert(j.nativeOp.getHashJoin.getOutputOrderingCount > 0) case j => fail(s"expected a null-aware broadcast hash join: $j") } // The build side spans several batches per partition. Read on a fresh plan that @@ -1459,21 +1462,17 @@ class CometJoinSuite extends CometTestBase { } } - test("only outer hash joins that can move unmatched streamed rows send their ordering") { + test("an inner hash join sends no streamed ordering") { withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { withSortedCache(spark.table("big"), $"_1") { streamed => val build = spark.table("small").hint("shuffle_hash") - def query(joinType: String): DataFrame = streamed - .join(build, streamed("_1") === build("_1"), joinType) + def query: DataFrame = streamed + .join(build, streamed("_1") === build("_1"), "inner") .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) .sortWithinPartitions("k") - val (inner, _) = checkJoinKeepsOrder(query("inner"), keyOrder(IntegerType, Ascending)) + val (inner, _) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) assert(inner.nativeOp.getHashJoin.getOutputOrderingCount == 0) - val (leftOuter, buildSide) = - checkJoinKeepsOrder(query("left_outer"), keyOrder(IntegerType, Ascending)) - assert(buildSide == BuildRight) - assert(leftOuter.nativeOp.getHashJoin.getOutputOrderingCount > 0) } } } From 7088411093ec9bc3c3136864bf7d1eddd7abbd29 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sun, 11 Oct 2026 21:09:57 +0700 Subject: [PATCH 10/11] test: keep the hash join probe order through runtime filter wrappers With runtime filters on, a probe-side filter or projection becomes CometFilterExec or CometProjectionExec before the join is built. Check that both keep the sort's ordering, so an outer hash join that reports its streamed ordering still adds no sort. --- native/core/src/execution/planner.rs | 107 +++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 34183d99181..4f684a1fb5c 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6649,6 +6649,113 @@ mod tests { } } + #[test] + fn hash_join_keeps_probe_order_through_a_runtime_filter_probe() { + // With runtime filters on, a probe-side filter or projection becomes its Comet wrapper. + // The wrapper must keep the sort's ordering so the join still needs no sort. + use datafusion::physical_plan::joins::HashJoinExec; + use datafusion_comet_operators::{CometFilterExec, CometProjectionExec}; + + let int_type = create_proto_datatype(); + let sorted_scan = || Operator { + children: vec![two_column_scan(&int_type)], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], + ..Default::default() + })), + ..Default::default() + }; + // Filters on the value column, so the key ordering can only come from the sort. + let filtered = || Operator { + children: vec![sorted_scan()], + op_struct: Some(OpStruct::Filter(spark_operator::Filter { + predicate: Some(Expr { + expr_struct: Some(Gt(Box::new(spark_expression::BinaryExpr { + left: Some(Box::new(typed_bound_reference(1, &int_type))), + right: Some(Box::new(Expr { + expr_struct: Some(Literal(spark_expression::Literal { + value: Some(literal::Value::IntVal(0)), + datatype: Some(int_type.clone()), + is_null: false, + })), + ..Default::default() + })), + }))), + ..Default::default() + }), + })), + ..Default::default() + }; + let projected = || Operator { + children: vec![sorted_scan()], + op_struct: Some(OpStruct::Projection(spark_operator::Projection { + project_list: vec![ + typed_bound_reference(0, &int_type), + typed_bound_reference(1, &int_type), + ], + })), + ..Default::default() + }; + for (join_type, build_side, probe_key_index, root) in [ + ( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + 0, + "ProjectionExec", + ), + ( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + 2, + "HashJoinExec", + ), + ] { + for (probe, probe_name) in [(filtered(), "FilterExec"), (projected(), "ProjectionExec")] + { + let children = if build_side == spark_operator::BuildSide::BuildRight { + [probe, two_column_scan(&int_type)] + } else { + [two_column_scan(&int_type), probe] + }; + let mut op = hash_join_on_first_columns( + join_type, + build_side, + children, + &int_type, + vec![ascending_nulls_first(typed_bound_reference( + probe_key_index, + &int_type, + ))], + ); + let Some(OpStruct::HashJoin(join)) = op.op_struct.as_mut() else { + unreachable!() + }; + join.dynamic_filter_enabled = true; + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap(); + let context = format!("{join_type:?} {probe_name}"); + assert_eq!(root, planned.native_plan.name(), "{context}"); + + // The probe is the right input of the native join in both shapes. + let join_node = if root == "HashJoinExec" { + &planned.native_plan + } else { + planned.native_plan.children()[0] + }; + let probe = join_node + .downcast_ref::() + .unwrap_or_else(|| panic!("{context}: {}", join_node.name())) + .right(); + assert_eq!(probe_name, probe.name(), "{context}"); + assert!( + probe.is::() || probe.is::(), + "{context}: probe was not rewritten for runtime filters" + ); + } + } + } + #[test] fn hash_join_over_a_native_sort_merge_join_adds_no_sort() { // A sort merge join declares its key ordering, so a hash join streaming it keeps that From ff7e1e310778c202a1e098994ce65eae4345901a Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 12 Oct 2026 22:45:45 +0700 Subject: [PATCH 11/11] fix: fall back to Spark for NOT IN over an ordered streamed side Comet runs a null-aware anti join unswapped in CollectLeft mode, so DataFusion builds on Spark's streamed side. On its hash map path the build batches are concatenated in reverse, and the unmatched rows come out in reverse batch order. Spark has already dropped any sort above the join because the join reports the streamed side's ordering, so sortWithinPartitions returns unsorted partitions and a sort aggregate above the join splits groups. Sorting the native output instead needed about four times the memory pool under fair_unified, so fall back to Spark when a null-aware anti join's streamed side reports an ordering, until the DataFusion fix (apache/datafusion#26226) reaches Comet. This drops the HashJoin output_ordering field and the planner sort from this branch. --- .../latest/compatibility/operators.md | 8 +- docs/source/user-guide/latest/operators.md | 4 +- native/core/src/execution/planner.rs | 575 +----------------- native/proto/src/proto/operator.proto | 3 - .../apache/spark/sql/comet/operators.scala | 36 +- .../apache/comet/exec/CometJoinSuite.scala | 383 ++++-------- 6 files changed, 180 insertions(+), 829 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 17f18e85d6c..af377967444 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -94,9 +94,11 @@ Both results are valid under Spark's semantics for these functions. ## Hash Join -When Spark reports a hash join's output as ordered by its streamed side, Comet keeps that order. -When DataFusion does not keep it, for example when it comes from a cached table or the join is a -`NOT IN`, Comet sorts the native join output to restore it. +Spark plans `NOT IN` as a null-aware anti join, and the native hash join builds it on Spark's +streamed side without keeping that side's order. When Spark reports the streamed side as ordered, +for example the output of a window, a sort-merge join, or a cached table sorted within its +partitions, Comet falls back to Spark for the join, so that a sort Spark removed above it on the +strength of that order is not lost. ## Window Functions diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 9e4caf27b62..5f50535842a 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -98,8 +98,8 @@ omitted from the tables below and may be reconsidered based on demand: | Operator | Status | Notes | | ----------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `BroadcastHashJoinExec` | ✅ | Falls back to Spark when an outer or `NOT IN` join reports an output ordering whose type or expression Comet cannot sort natively. | -| `ShuffledHashJoinExec` | ✅ | Falls back to Spark when an outer join reports an output ordering whose type or expression Comet cannot sort natively. | +| `BroadcastHashJoinExec` | ✅ | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering. | +| `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)). | diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5745b863d60..ae5916092b1 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1233,22 +1233,6 @@ impl PhysicalPlanner { } } - /// Returns a sort of `plan` by `ordering`, or `None` when `ordering` is empty or the plan's - /// equivalence properties already satisfy it. Operators whose Spark counterpart reports an - /// output ordering use this to sort their output when the native plan does not keep it. - fn sort_unless_ordered( - plan: Arc, - ordering: Vec, - ) -> Result>, ExecutionError> { - let Some(lex_ordering) = LexOrdering::new(ordering.clone()) else { - return Ok(None); - }; - if plan.equivalence_properties().ordering_satisfy(ordering)? { - return Ok(None); - } - Ok(Some(Arc::new(SortExec::new(lex_ordering, plan)))) - } - #[allow(clippy::too_many_arguments)] pub fn create_binary_expr( &self, @@ -1715,8 +1699,14 @@ impl PhysicalPlanner { ) }) .collect(); - if let Some(sort) = Self::sort_unless_ordered(Arc::clone(&aggregate), ordering)? + if !aggregate + .equivalence_properties() + .ordering_satisfy(ordering.clone())? { + let sort: Arc = Arc::new(SortExec::new( + LexOrdering::new(ordering).unwrap(), + Arc::clone(&aggregate), + )); return Ok(( scans, shuffle_scans, @@ -2669,16 +2659,21 @@ impl PhysicalPlanner { // Exception: null-aware anti-join requires LeftAnti + build-right semantics // (which matches DataFusion's default), and swap_inputs would turn LeftAnti // into RightAnti, which DataFusion rejects with null_aware=true. - let (join_root, mut additional_native_plans) = if join.build_side - == BuildSide::BuildLeft as i32 - || join.null_aware_anti_join - { + if join.build_side == BuildSide::BuildLeft as i32 || join.null_aware_anti_join { let hash_join = Self::apply_join_dynamic_filter( hash_join, join.dynamic_filter_enabled && !join.null_aware_anti_join, self.session_ctx.copied_config().options(), )?; - (hash_join, vec![]) + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new( + spark_plan.plan_id, + hash_join, + vec![join_params.left, join_params.right], + )), + )) } else { let swapped_hash_join = hash_join.as_ref().swap_inputs(PartitionMode::Partitioned)?; @@ -2693,38 +2688,18 @@ impl PhysicalPlanner { // a projection was added to the hash join additional_native_plans.push(Arc::clone(swapped_hash_join.children()[0])); } - (swapped_hash_join, additional_native_plans) - }; - // Spark reports the streamed side's ordering and may have removed a sort above - // the join on that basis. DataFusion keeps unmatched probe rows in probe order - // only when it sees the probe input sorted, and does not keep a null-aware anti - // join's build order, so sort the output otherwise. - let join_schema = join_root.schema(); - let output_ordering = join - .output_ordering - .iter() - .map(|expr| self.create_sort_expr(expr, Arc::clone(&join_schema))) - .collect::, _>>()?; - let native_plan = - match Self::sort_unless_ordered(Arc::clone(&join_root), output_ordering)? { - Some(sort) => { - additional_native_plans.push(join_root); - sort - } - None => join_root, - }; - - Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new_with_additional( - spark_plan.plan_id, - native_plan, - vec![join_params.left, join_params.right], - additional_native_plans, - )), - )) + Ok(( + scans, + shuffle_scans, + Arc::new(SparkPlan::new_with_additional( + spark_plan.plan_id, + swapped_hash_join, + vec![join_params.left, join_params.right], + additional_native_plans, + )), + )) + } } OpStruct::Window(wnd) => { let (scans, shuffle_scans, child) = @@ -6047,7 +6022,6 @@ mod tests { build_side: 0, null_aware_anti_join: false, dynamic_filter_enabled: false, - output_ordering: vec![], })), }; @@ -6342,499 +6316,6 @@ mod tests { assert_eq!("AggregateExec", planned.native_plan.name()); } - fn typed_bound_reference(index: i32, datatype: &spark_expression::DataType) -> Expr { - Expr { - expr_struct: Some(Bound(spark_expression::BoundReference { - index, - datatype: Some(datatype.clone()), - })), - ..Default::default() - } - } - - fn ascending_nulls_first(child: Expr) -> Expr { - Expr { - expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder { - child: Some(Box::new(child)), - direction: spark_expression::SortDirection::Ascending as i32, - null_ordering: spark_expression::NullOrdering::NullsFirst as i32, - }))), - ..Default::default() - } - } - - fn two_column_scan(datatype: &spark_expression::DataType) -> Operator { - Operator { - op_struct: Some(OpStruct::Scan(spark_operator::Scan { - fields: vec![datatype.clone(); 2], - source: String::new(), - })), - ..Default::default() - } - } - - /// A hash join on the first column of each side that reports `output_ordering`. - fn hash_join_on_first_columns( - join_type: spark_operator::JoinType, - build_side: spark_operator::BuildSide, - children: [Operator; 2], - key_type: &spark_expression::DataType, - output_ordering: Vec, - ) -> Operator { - Operator { - plan_id: 1, - children: children.into(), - op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { - left_join_keys: vec![typed_bound_reference(0, key_type)], - right_join_keys: vec![typed_bound_reference(0, key_type)], - join_type: join_type as i32, - build_side: build_side as i32, - output_ordering, - ..Default::default() - })), - ..Default::default() - } - } - - /// Executes a planned join with one Int32 (key, value) batch in place of each child scan. - async fn execute_join( - planned: &crate::execution::spark_plan::SparkPlan, - task_ctx: Arc, - sides: [(Vec, Vec); 2], - ) -> RecordBatch { - execute_join_batches(planned, task_ctx, sides.map(|side| vec![side])).await - } - - /// Executes a planned join with Int32 (key, value) batches in place of each child scan. - async fn execute_join_batches( - planned: &crate::execution::spark_plan::SparkPlan, - task_ctx: Arc, - sides: [Vec<(Vec, Vec)>; 2], - ) -> RecordBatch { - use datafusion::common::tree_node::{Transformed, TreeNode}; - - let replacements: Vec<(Arc, Arc)> = planned - .children - .iter() - .zip(sides) - .map(|(child, batches)| { - let batches: Vec = batches - .into_iter() - .map(|(keys, values)| { - RecordBatch::try_new( - child.schema(), - vec![ - Arc::new(Int32Array::from(keys)), - Arc::new(Int32Array::from(values)), - ], - ) - .unwrap() - }) - .collect(); - let input: Arc = - MemorySourceConfig::try_new_exec(&[batches], child.schema(), None).unwrap(); - (Arc::clone(&child.native_plan), input) - }) - .collect(); - let native_plan = Arc::clone(&planned.native_plan) - .transform_up(|node| { - Ok( - match replacements - .iter() - .find(|(scan, _)| Arc::ptr_eq(scan, &node)) - { - Some((_, input)) => Transformed::yes(Arc::clone(input)), - None => Transformed::no(node), - }, - ) - }) - .unwrap() - .data; - let results = collect(native_plan.execute(0, task_ctx).unwrap()) - .await - .unwrap(); - arrow::compute::concat_batches(&planned.native_plan.schema(), &results).unwrap() - } - - fn additional_plan_names(planned: &crate::execution::spark_plan::SparkPlan) -> Vec<&str> { - planned - .additional_native_plans - .iter() - .map(|plan| plan.name()) - .collect() - } - - #[tokio::test] - async fn hash_join_sorts_output_when_probe_order_is_unknown() { - // Spark LeftOuter + BuildRight runs as a swapped DataFusion Right join. The probe scan - // declares no ordering, so DataFusion emits matched rows before unmatched ones. - let int_type = create_proto_datatype(); - let op = hash_join_on_first_columns( - spark_operator::JoinType::LeftOuter, - spark_operator::BuildSide::BuildRight, - [two_column_scan(&int_type), two_column_scan(&int_type)], - &int_type, - vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], - ); - let ctx = SessionContext::new(); - let task_ctx = ctx.task_ctx(); - let planner = PhysicalPlanner::new(Arc::new(ctx), 0); - let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); - - let output = execute_join( - &planned, - task_ctx, - [(vec![1, 5, 2], vec![10, 50, 20]), (vec![5], vec![500])], - ) - .await; - assert_eq!( - &Int32Array::from(vec![1, 2, 5]) as &dyn Array, - output.column(0).as_ref() - ); - assert_eq!( - &Int32Array::from(vec![None, None, Some(5)]) as &dyn Array, - output.column(2).as_ref() - ); - assert_eq!("SortExec", planned.native_plan.name()); - assert_eq!( - vec!["HashJoinExec", "ProjectionExec"], - additional_plan_names(&planned) - ); - } - - #[tokio::test] - async fn hash_join_sorts_output_on_the_unswapped_build_left_path() { - // Spark RightOuter + BuildLeft runs unswapped. The probe key is the first column of the - // right side, at index 2 of the join output. - let int_type = create_proto_datatype(); - let op = hash_join_on_first_columns( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - [two_column_scan(&int_type), two_column_scan(&int_type)], - &int_type, - vec![ascending_nulls_first(typed_bound_reference(2, &int_type))], - ); - let ctx = SessionContext::new(); - let task_ctx = ctx.task_ctx(); - let planner = PhysicalPlanner::new(Arc::new(ctx), 0); - let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); - - let output = execute_join( - &planned, - task_ctx, - [(vec![5], vec![500]), (vec![1, 5, 2], vec![10, 50, 20])], - ) - .await; - assert_eq!( - &Int32Array::from(vec![1, 2, 5]) as &dyn Array, - output.column(2).as_ref() - ); - assert_eq!( - &Int32Array::from(vec![None, None, Some(5)]) as &dyn Array, - output.column(0).as_ref() - ); - assert_eq!("SortExec", planned.native_plan.name()); - assert_eq!(vec!["HashJoinExec"], additional_plan_names(&planned)); - } - - #[tokio::test] - async fn hash_join_sorts_null_aware_anti_join_output() { - // A null-aware anti join runs unswapped in CollectLeft mode, so DataFusion builds on the - // streamed left side. Keys spanning 1024 or more values skip its perfect hash path, and - // the hash path emits the unmatched rows in reverse batch order. - let int_type = create_proto_datatype(); - let mut op = hash_join_on_first_columns( - spark_operator::JoinType::LeftAnti, - spark_operator::BuildSide::BuildRight, - [two_column_scan(&int_type), two_column_scan(&int_type)], - &int_type, - vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], - ); - let Some(OpStruct::HashJoin(join)) = op.op_struct.as_mut() else { - unreachable!() - }; - join.null_aware_anti_join = true; - let ctx = SessionContext::new(); - let task_ctx = ctx.task_ctx(); - let planner = PhysicalPlanner::new(Arc::new(ctx), 0); - let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); - - assert_eq!("SortExec", planned.native_plan.name()); - assert_eq!(vec!["HashJoinExec"], additional_plan_names(&planned)); - let hash_join = planned.native_plan.children()[0] - .downcast_ref::() - .unwrap(); - assert_eq!( - &datafusion::common::JoinType::LeftAnti, - hash_join.join_type() - ); - assert!(hash_join.null_aware); - - let output = execute_join_batches( - &planned, - task_ctx, - [ - vec![(vec![10, 20], vec![1, 2]), (vec![3000, 4000], vec![3, 4])], - vec![(vec![5], vec![500])], - ], - ) - .await; - assert_eq!( - &Int32Array::from(vec![10, 20, 3000, 4000]) as &dyn Array, - output.column(0).as_ref() - ); - assert_eq!( - &Int32Array::from(vec![1, 2, 3, 4]) as &dyn Array, - output.column(1).as_ref() - ); - } - - #[test] - fn hash_join_keeps_native_probe_order_without_a_sort() { - // A native sort below the probe side is visible to DataFusion, which then keeps unmatched - // probe rows in order. Float keys check that the ordering is normalized like the sort's. - let double_type = spark_expression::DataType { - type_id: 6, - type_info: None, - }; - for key_type in [create_proto_datatype(), double_type] { - for (join_type, build_side, probe_key_index, root) in [ - ( - spark_operator::JoinType::LeftOuter, - spark_operator::BuildSide::BuildRight, - 0, - "ProjectionExec", - ), - ( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - 2, - "HashJoinExec", - ), - ] { - let sorted_probe = Operator { - children: vec![two_column_scan(&key_type)], - op_struct: Some(OpStruct::Sort(spark_operator::Sort { - sort_orders: vec![ascending_nulls_first(typed_bound_reference( - 0, &key_type, - ))], - ..Default::default() - })), - ..Default::default() - }; - let children = if build_side == spark_operator::BuildSide::BuildRight { - [sorted_probe, two_column_scan(&key_type)] - } else { - [two_column_scan(&key_type), sorted_probe] - }; - let op = hash_join_on_first_columns( - join_type, - build_side, - children, - &key_type, - vec![ascending_nulls_first(typed_bound_reference( - probe_key_index, - &key_type, - ))], - ); - let (_, _, planned) = PhysicalPlanner::default() - .create_plan(&op, &mut vec![], 1) - .unwrap(); - assert_eq!( - root, - planned.native_plan.name(), - "{join_type:?} {key_type:?}" - ); - } - } - } - - #[test] - fn hash_join_keeps_probe_order_through_a_runtime_filter_probe() { - // With runtime filters on, a probe-side filter or projection becomes its Comet wrapper. - // The wrapper must keep the sort's ordering so the join still needs no sort. - use datafusion::physical_plan::joins::HashJoinExec; - use datafusion_comet_operators::{CometFilterExec, CometProjectionExec}; - - let int_type = create_proto_datatype(); - let sorted_scan = || Operator { - children: vec![two_column_scan(&int_type)], - op_struct: Some(OpStruct::Sort(spark_operator::Sort { - sort_orders: vec![ascending_nulls_first(typed_bound_reference(0, &int_type))], - ..Default::default() - })), - ..Default::default() - }; - // Filters on the value column, so the key ordering can only come from the sort. - let filtered = || Operator { - children: vec![sorted_scan()], - op_struct: Some(OpStruct::Filter(spark_operator::Filter { - predicate: Some(Expr { - expr_struct: Some(Gt(Box::new(spark_expression::BinaryExpr { - left: Some(Box::new(typed_bound_reference(1, &int_type))), - right: Some(Box::new(Expr { - expr_struct: Some(Literal(spark_expression::Literal { - value: Some(literal::Value::IntVal(0)), - datatype: Some(int_type.clone()), - is_null: false, - })), - ..Default::default() - })), - }))), - ..Default::default() - }), - })), - ..Default::default() - }; - let projected = || Operator { - children: vec![sorted_scan()], - op_struct: Some(OpStruct::Projection(spark_operator::Projection { - project_list: vec![ - typed_bound_reference(0, &int_type), - typed_bound_reference(1, &int_type), - ], - })), - ..Default::default() - }; - for (join_type, build_side, probe_key_index, root) in [ - ( - spark_operator::JoinType::LeftOuter, - spark_operator::BuildSide::BuildRight, - 0, - "ProjectionExec", - ), - ( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - 2, - "HashJoinExec", - ), - ] { - for (probe, probe_name) in [(filtered(), "FilterExec"), (projected(), "ProjectionExec")] - { - let children = if build_side == spark_operator::BuildSide::BuildRight { - [probe, two_column_scan(&int_type)] - } else { - [two_column_scan(&int_type), probe] - }; - let mut op = hash_join_on_first_columns( - join_type, - build_side, - children, - &int_type, - vec![ascending_nulls_first(typed_bound_reference( - probe_key_index, - &int_type, - ))], - ); - let Some(OpStruct::HashJoin(join)) = op.op_struct.as_mut() else { - unreachable!() - }; - join.dynamic_filter_enabled = true; - let (_, _, planned) = PhysicalPlanner::default() - .create_plan(&op, &mut vec![], 1) - .unwrap(); - let context = format!("{join_type:?} {probe_name}"); - assert_eq!(root, planned.native_plan.name(), "{context}"); - - // The probe is the right input of the native join in both shapes. - let join_node = if root == "HashJoinExec" { - &planned.native_plan - } else { - planned.native_plan.children()[0] - }; - let probe = join_node - .downcast_ref::() - .unwrap_or_else(|| panic!("{context}: {}", join_node.name())) - .right(); - assert_eq!(probe_name, probe.name(), "{context}"); - assert!( - probe.is::() || probe.is::(), - "{context}: probe was not rewritten for runtime filters" - ); - } - } - } - - #[test] - fn hash_join_over_a_native_sort_merge_join_adds_no_sort() { - // A sort merge join declares its key ordering, so a hash join streaming it keeps that - // ordering without a sort. Spark reports the merge join's left key as the ordering. - let int_type = create_proto_datatype(); - let key = || ascending_nulls_first(typed_bound_reference(0, &int_type)); - let sorted_scan = || Operator { - children: vec![two_column_scan(&int_type)], - op_struct: Some(OpStruct::Sort(spark_operator::Sort { - sort_orders: vec![key()], - ..Default::default() - })), - ..Default::default() - }; - for smj_join_type in [ - spark_operator::JoinType::Inner, - spark_operator::JoinType::LeftOuter, - ] { - let sort_merge_join = Operator { - children: vec![sorted_scan(), sorted_scan()], - op_struct: Some(OpStruct::SortMergeJoin(spark_operator::SortMergeJoin { - left_join_keys: vec![typed_bound_reference(0, &int_type)], - right_join_keys: vec![typed_bound_reference(0, &int_type)], - join_type: smj_join_type as i32, - sort_options: vec![key()], - condition: None, - })), - ..Default::default() - }; - // The merge join's left key is column 0 of the hash join output. - let op = hash_join_on_first_columns( - spark_operator::JoinType::LeftOuter, - spark_operator::BuildSide::BuildRight, - [sort_merge_join, two_column_scan(&int_type)], - &int_type, - vec![key()], - ); - let (_, _, planned) = PhysicalPlanner::default() - .create_plan(&op, &mut vec![], 1) - .unwrap(); - assert_eq!( - "ProjectionExec", - planned.native_plan.name(), - "{smj_join_type:?}\n{}\n{:?}", - datafusion::physical_plan::displayable(planned.native_plan.as_ref()).indent(true), - planned.native_plan.properties().eq_properties - ); - } - } - - #[test] - fn hash_join_without_output_ordering_adds_no_sort() { - let int_type = create_proto_datatype(); - for (join_type, build_side, root) in [ - ( - spark_operator::JoinType::LeftOuter, - spark_operator::BuildSide::BuildRight, - "ProjectionExec", - ), - ( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - "HashJoinExec", - ), - ] { - let op = hash_join_on_first_columns( - join_type, - build_side, - [two_column_scan(&int_type), two_column_scan(&int_type)], - &int_type, - vec![], - ); - let (_, _, planned) = PhysicalPlanner::default() - .create_plan(&op, &mut vec![], 1) - .unwrap(); - assert_eq!(root, planned.native_plan.name(), "{join_type:?}"); - } - } - #[tokio::test] async fn projection_prunes_filter_output_and_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index be46c7b13e2..97c65f0da36 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1026,9 +1026,6 @@ message HashJoin { // Opt-in completed-build filtering of probe batches and eligible native Parquet readers. // Older serialized plans leave it false. bool dynamic_filter_enabled = 10; - // SortOrder exprs bound to the join output that Spark reports as its output ordering. - // Empty means the native plan adds no sort. - repeated spark.spark_expression.Expr output_ordering = 11; } message SortMergeJoin { 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 02bb7978218..b160456fded 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 @@ -2811,6 +2811,18 @@ trait CometHashJoin { return None } + // DataFusion builds a null-aware anti join on Spark's streamed side, and on its hash map path + // it emits the build batches last to first, so the output loses the ordering Spark reports + // for the join. Remove this fallback once Comet uses a DataFusion release that fixes + // https://github.com/apache/datafusion/issues/26226. + if (isNullAwareAntiJoin && join.outputOrdering.nonEmpty) { + withFallbackReason( + join, + "null-aware anti-join over an ordered streamed side is not supported natively, " + + "because the native join does not keep that order") + return None + } + val condition = join.condition.map { cond => val condProto = exprToProto(cond, join.left.output ++ join.right.output) if (condProto.isEmpty) { @@ -2861,29 +2873,6 @@ trait CometHashJoin { } } - // DataFusion may move unmatched probe rows for the outer shapes, and runs a null-aware anti - // join with Spark's streamed side as its build side, so the native planner needs the ordering - // Spark reports for the join to restore it. - val mayReorderStreamed = (join.joinType, join.buildSide) match { - case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => true - case (LeftAnti, BuildRight) => isNullAwareAntiJoin - case _ => false - } - val outputOrdering = if (mayReorderStreamed && join.outputOrdering.nonEmpty) { - if (!supportedSortType(join, join.outputOrdering)) { - withFallbackReason(join, "Unsupported data type in hash join output ordering") - return None - } - val orders = join.outputOrdering.map(exprToProto(_, join.output)) - if (orders.exists(_.isEmpty)) { - withFallbackReason(join, "Unsupported expression in hash join output ordering") - return None - } - orders.flatten - } else { - Nil - } - val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output)) val rightKeys = join.rightKeys.map(exprToProto(_, join.right.output)) @@ -2899,7 +2888,6 @@ trait CometHashJoin { else OperatorOuterClass.BuildSide.BuildRight) .setNullAwareAntiJoin(isNullAwareAntiJoin) .setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf)) - .addAllOutputOrdering(outputOrdering.asJava) condition.foreach(joinBuilder.setCondition) Some(builder.setHashJoin(joinBuilder).build()) } else { 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 2d0252de231..31974296b6f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -26,21 +26,22 @@ import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.SparkException -import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation -import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, InterpretedOrdering, IsNotNull, SortDirection, SortOrder} -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, EliminateSorts} +import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBinaryExec, CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, 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 import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType, IntegerType, MetadataBuilder, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} @@ -1321,293 +1322,175 @@ class CometJoinSuite extends CometTestBase { } } - // A hash join reports its streamed side's ordering, so Spark drops a sort above it when the - // streamed side already arrives partitioned and sorted. Here that side is a cached table that - // reaches the native join through Comet's JVM to native conversion. - private def withSortedCache(data: DataFrame, order: Column*)(f: DataFrame => Unit): Unit = { + // Spark plans NOT IN as a null-aware anti join that streams the left side and reports that + // side's ordering, so Spark drops a sort above the join when the side arrives sorted. Here that + // side is a cached table, sorted within its partitions, that reaches the native plan through + // Comet's JVM to native conversion. Comet runs the join with that side as DataFusion's build + // side. Sparse keys make DataFusion build a hash map, which takes the batches last to first; + // dense integer keys would take its array map path instead, which keeps them in order. + private def withSortedCache(data: DataFrame)(f: => Unit): Unit = { withSQLConf( CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "64", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760", SQLConf.SHUFFLE_PARTITIONS.key -> "2") { - val cached = data.repartition(2, data("_1")).sortWithinPartitions(order: _*).cache() + val cached = data.repartition(2, data("_1")).sortWithinPartitions("_1").cache() try { cached.count() - f(cached) + withTempView("streamed") { + cached.createOrReplaceTempView("streamed") + f + } } finally { cached.unpersist() } } } + private def withNotInTables(f: => Unit): Unit = { + withParquetTable((0 until 10000).map(i => (i % 100 * 1000, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10000, i)), "small") { + f + } + } + } + /** - * Checks the query's rows against Spark and that every partition is sorted by `order`, which is - * bound to the query's output. Returns the plan's single native hash join and its build side. + * Checks the query's rows against Spark, that no sort runs above a broadcast hash join, and + * that every partition is sorted by the query's first column, an integer. Returns the Comet + * plan. */ - private def checkJoinKeepsOrder( - query: => DataFrame, - order: SortOrder*): (CometBinaryExec, BuildSide) = { - // withSQLConf returns Unit on Spark 3.x, so the expected rows are assigned inside it. - var expected = Seq.empty[String] - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - expected = query.collect().map(_.toString).sorted.toSeq - } - val df = query - val rows = df.collect() - assert(rows.map(_.toString).sorted.toSeq == expected) - val plan = df.queryExecution.executedPlan - val joins = collect(plan) { - case j: CometHashJoinExec => (j, j.buildSide) - case j: CometBroadcastHashJoinExec => (j, j.buildSide) - } - assert(joins.length == 1, plan) - // Read before the plan runs again below, so the count covers the one run above. - assert(joins.head._1.metrics("output_rows").value == rows.length, plan) - assert(joins.head._1.metrics("build_input_rows").value > 0, plan) - // Spark dropped the sort above the join, so the order comes from the join alone. - val sortsAboveJoin = collect(plan) { - case s: SortExec => s - case s: CometSortExec => s - }.filter(s => find(s)(_ eq joins.head._1).isDefined) - assert(sortsAboveJoin.isEmpty, plan) - val ordering = new InterpretedOrdering(order) + private def checkSortedPartitions(query: => DataFrame): SparkPlan = { + val (_, plan) = checkSparkAnswer(query) + // Spark dropped the sort above the join, so the order comes from the plan below it. + val sorts = collect(plan) { case s @ (_: SortExec | _: CometSortExec) => s } + assert( + sorts.forall(_.find { + case _: BroadcastHashJoinExec | _: CometBroadcastHashJoinExec => true + case _ => false + }.isEmpty), + plan) // Propagates the session's confs, such as the batch size, to the tasks as collect() does. val parts = SQLExecution.withSQLConfPropagated(spark) { - df.queryExecution.toRdd.map(_.copy()).glom().collect() + query.queryExecution.toRdd.map(_.getInt(0)).glom().collect() } - parts.zipWithIndex.foreach { case (part, i) => - val outOfOrder = part.zip(part.drop(1)).indexWhere { case (a, b) => ordering.gt(a, b) } - assert(outOfOrder < 0, s"partition $i is out of order at row $outOfOrder\n$plan") + parts.zipWithIndex.foreach { case (keys, i) => + assert(keys.sameElements(keys.sorted), s"partition $i is out of order\n$plan") } - joins.head + plan } - private def keyOrder( - dataType: DataType, - direction: SortDirection, - ordinal: Int = 0): SortOrder = - SortOrder(BoundReference(ordinal, dataType, nullable = true), direction) - - // (hint, join type, extra condition, AQE, native batch size) - private val streamedOrderCases = - (for (hint <- Seq("shuffle_hash", "broadcast"); joinType <- Seq("left_outer", "right_outer")) - yield (hint, joinType, false, false, 8192)) ++ Seq( - ("shuffle_hash", "left_outer", true, false, 8192), - ("shuffle_hash", "right_outer", false, true, 8192), - ("shuffle_hash", "left_outer", false, false, 100)) + private val orderedNotInFallback = "null-aware anti-join over an ordered streamed side" - for ((hint, joinType, withCondition, adaptive, batchSize) <- streamedOrderCases) { - test( - s"$hint $joinType join keeps a cached streamed side's order, " + - s"condition=$withCondition, AQE=$adaptive, batchSize=$batchSize") { - withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1") { streamed => - val build = spark.table("small").hint(hint) - val keys = streamed("_1") === build("_1") - val on = if (withCondition) keys && build("_2") > 3 else keys - def query: DataFrame = { - val joined = - if (joinType == "left_outer") streamed.join(build, on, joinType) - else build.join(streamed, on, joinType) - joined - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions("k") - } - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { - val (join, buildSide) = - checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) - assert(buildSide == (if (joinType == "left_outer") BuildRight else BuildLeft)) - hint match { - case "broadcast" => assert(join.isInstanceOf[CometBroadcastHashJoinExec]) - case _ => assert(join.isInstanceOf[CometHashJoinExec]) - } - assert(join.nativeOp.getHashJoin.getOutputOrderingCount > 0) - if (batchSize < 1000) { - // Several batches per partition come back to the JVM from the native join. - val partitions = 2 - val plan = query.queryExecution.executedPlan - val batches = collectFirst(plan) { case c: ColumnarToRowTransition => - SQLExecution.withSQLConfPropagated(spark) { - c.child.executeColumnar().count() - } - } - assert(batches.exists(_ > partitions), plan) - } - } - } + for (adaptive <- Seq(false, true)) { + test(s"NOT IN over a sorted cached side falls back and keeps its order, AQE=$adaptive") { + withNotInTables { + withSortedCache(spark.table("big")) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val plan = checkSortedPartitions( + sql("SELECT _1 AS k, _2 AS v FROM streamed WHERE _1 NOT IN (SELECT _1 FROM small)") + .sortWithinPartitions("k")) + val joins = collect(plan) { case j: BroadcastHashJoinExec => j } + assert(joins.length == 1 && joins.head.isNullAwareAntiJoin, plan) + assert(collect(plan) { case j: CometBroadcastHashJoinExec => j }.isEmpty, plan) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains(orderedNotInFallback)), + plan) } } } } } - test("shuffle_hash join keeps a descending, nulls last streamed order with null keys") { - val data = (0 until 10000).map(i => (if (i % 7 == 0) None else Some(i % 100), i)) - withParquetTable(data, "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1".desc_nulls_last) { streamed => - val build = spark.table("small").hint("shuffle_hash") - checkJoinKeepsOrder( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions($"k".desc_nulls_last), - keyOrder(IntegerType, Descending)) + // Spark's sort aggregate trusts the same ordering, through an outer join that streams the NOT IN + // output, so a NOT IN that loses it splits groups. Comet aggregates are off so that Spark's sort + // aggregate runs over the joins, as it does for an aggregate Comet cannot convert. + test("NOT IN below an outer hash join keeps its groups and order under a GROUP BY") { + withNotInTables { + withParquetTable((0 until 50).map(i => (i * 2000, i)), "dim") { + withSortedCache(spark.table("big")) { + withSQLConf( + SQLConf.USE_OBJECT_HASH_AGG.key -> "false", + CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") { + val plan = checkSortedPartitions( + sql("""SELECT s.k, count(*) AS n, sort_array(collect_list(d._2)) AS w + |FROM (SELECT _1 AS k FROM streamed WHERE _1 NOT IN (SELECT _1 FROM small)) s + |LEFT JOIN dim d ON s.k = d._1 + |GROUP BY s.k""".stripMargin).sortWithinPartitions("k")) + assert(collect(plan) { case a: SortAggregateExec => a }.nonEmpty, plan) + } } } } } - test("shuffle_hash join keeps a descending streamed order on a double with NaN and -0.0") { - def d(i: Int): Option[Double] = i % 9 match { - case 0 => None - case 1 => Some(Double.NaN) - case 2 => Some(-0.0) - case 3 => Some(0.0) - case _ => Some(i % 50 - 25.5) - } - withParquetTable((0 until 10000).map(i => (i % 100, i, d(i))), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_3".desc_nulls_last) { streamed => - val build = spark.table("small").hint("shuffle_hash") - checkJoinKeepsOrder( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_3").as("d"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions($"d".desc_nulls_last), - keyOrder(DoubleType, Descending)) - } + test("NOT IN over an unordered streamed side runs natively") { + withNotInTables { + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { + val (_, plan) = checkSparkAnswerAndOperator( + sql("SELECT _1, _2 FROM big WHERE _1 NOT IN (SELECT _1 FROM small)")) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.length == 1 && joins.head.isNullAwareAntiJoin, plan) + assert(joins.head.outputOrdering.isEmpty, plan) } } } - test("shuffle_hash join keeps a two-key streamed order") { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1", $"_2".desc) { streamed => - val build = spark.table("small").hint("shuffle_hash") - checkJoinKeepsOrder( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions($"k", $"v".desc), - keyOrder(IntegerType, Ascending), - keyOrder(IntegerType, Descending, ordinal = 1)) - } - } - } - } + // With default settings the streamed side can be native from parquet. A window partitioned by + // `_1` sorts its input natively, and `k` carries `_1` through the window, so the NOT IN key is a + // window output that the optimizer cannot push below the sort. The ROWS frame keeps the window + // on DataFusion's streaming window operator, so the join's build side spans many batches. + private val windowedBig = + "SELECT _1, _2, min(_1) OVER (PARTITION BY _1 ORDER BY _2 ROWS UNBOUNDED PRECEDING) AS k " + + "FROM big" - // The streamed side is sorted inside the native plan, and the join's output keeps that order in - // every partition. EliminateSorts would drop that sort as one below a join. - test("shuffle_hash join keeps a natively sorted streamed side's order") { - withSQLConf( - SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName, - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.SHUFFLE_PARTITIONS.key -> "2") { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") - val build = spark.table("small").hint("shuffle_hash") - val (join, buildSide) = checkJoinKeepsOrder( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions("k"), - keyOrder(IntegerType, Ascending)) - assert(join.isInstanceOf[CometHashJoinExec]) - assert(buildSide == BuildRight) - } - } - } - } - - // Spark plans NOT IN as a null-aware anti join that streams the left side. Comet runs it with - // the left side as DataFusion's build side, which gives up that side's batch order. Sparse keys - // keep DataFusion on its hash map build, which stores the batches last to first. - test("null-aware anti join keeps a cached streamed side's order across batches") { - withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "100") { - withParquetTable((0 until 10000).map(i => (i % 100 * 1000, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10000, i)), "small") { - withSortedCache(spark.table("big"), $"_1") { streamed => - withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - withTempView("streamed") { - streamed.createOrReplaceTempView("streamed") - def query: DataFrame = - sql("SELECT _1 AS k, _2 AS v FROM streamed WHERE _1 NOT IN " + - "(SELECT _1 FROM small)").sortWithinPartitions("k") - val (join, _) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) - join match { - case j: CometBroadcastHashJoinExec => - assert(j.isNullAwareAntiJoin, j) - assert(j.nativeOp.getHashJoin.getOutputOrderingCount > 0) - case j => fail(s"expected a null-aware broadcast hash join: $j") - } - // The build side spans several batches per partition. Read on a fresh plan that - // runs once, so the count covers that run alone. - val partitions = 2 - val df = query - assert(df.collect().nonEmpty) - val plan = df.queryExecution.executedPlan - val buildBatches = collectFirst(plan) { case j: CometBroadcastHashJoinExec => - j.metrics("build_input_batches").value - } - assert(buildBatches.exists(_ > partitions), plan) - } - } + for (adaptive <- Seq(false, true)) { + test(s"NOT IN over a native window side falls back and keeps its order, AQE=$adaptive") { + withNotInTables { + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "64", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val plan = checkSortedPartitions( + sql(s"SELECT _1, _2 FROM ($windowedBig) t WHERE k NOT IN (SELECT _1 FROM small)") + .sortWithinPartitions("_1")) + assert(collect(plan) { case w: CometWindowExec => w }.nonEmpty, plan) + val joins = collect(plan) { case j: BroadcastHashJoinExec => j } + assert(joins.length == 1 && joins.head.isNullAwareAntiJoin, plan) + assert(collect(plan) { case j: CometBroadcastHashJoinExec => j }.isEmpty, plan) + // AQE plans the join again once the shuffle below the window completes. By then its + // build side is Spark's broadcast stage, so Comet skips the new join without a reason, + // and only the initial plan carries the reason the guard recorded. + if (!adaptive) { + assert( + new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains(orderedNotInFallback)), + plan) } } } } } - test("an inner hash join sends no streamed ordering") { - withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1") { streamed => - val build = spark.table("small").hint("shuffle_hash") - def query: DataFrame = streamed - .join(build, streamed("_1") === build("_1"), "inner") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions("k") - val (inner, _) = checkJoinKeepsOrder(query, keyOrder(IntegerType, Ascending)) - assert(inner.nativeOp.getHashJoin.getOutputOrderingCount == 0) - } - } - } - } - - test("hash join falls back when its streamed ordering has an unsupported sort type") { - withParquetTable((0 until 1000).map(i => (i % 100, i, (i % 13, i.toString))), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_3") { streamed => - val build = spark.table("small").hint("shuffle_hash") - checkSparkAnswerAndFallbackReason( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_3").as("s"), streamed("_2").as("v"), build("_2").as("w")), - "Unsupported data type in hash join output ordering") - } - } - } - } - - test("hash join falls back when its streamed ordering has an unsupported expression") { - withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { - withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { - withSortedCache(spark.table("big"), $"_1" + 1) { streamed => - val build = spark.table("small").hint("shuffle_hash") - withSQLConf(CometConf.getExprEnabledConfigKey("Add") -> "false") { - checkSparkAnswerAndFallbackReason( - streamed - .join(build, streamed("_1") === build("_1"), "left_outer") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")), - "Unsupported expression in hash join output ordering") - } - } + // The fallback only covers the null-aware join: a semi join over the same ordered side still + // runs natively, since DataFusion probes with that side and keeps its order. + test("LEFT SEMI join over the native window side used for NOT IN stays native") { + withNotInTables { + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "64") { + val plan = checkSortedPartitions( + sql(s"SELECT /*+ BROADCAST(small) */ t._1, t._2 FROM ($windowedBig) t " + + "LEFT SEMI JOIN small ON t.k = small._1").sortWithinPartitions("_1")) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.length == 1 && !joins.head.isNullAwareAntiJoin, plan) + assert(collect(plan) { case j: BroadcastHashJoinExec => j }.isEmpty, plan) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains(orderedNotInFallback)), + plan) } } }