From 207f06a45c24ed15be8a372a4d507f7f04a4787b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 9 Oct 2026 19:06:20 -0600 Subject: [PATCH 01/19] fix: keep a streamed side sorted outside the native plan in order through outer hash joins Spark's hash joins report the streamed side's order as their own, so Spark drops a sort above the join that the order satisfies. DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only when the probe input declares an order, and otherwise moves them after the matched rows of each batch. Comet builds every hash join on the build side, so Spark's left outer join built on the right and right outer join built on the left both run as that right outer join, and a streamed side that enters the native plan from the JVM, a relation sorted before it was cached for instance, arrives through a ScanExec that declares no order. The join then returned those rows out of order with no sort left in the plan to fix them. Send the streamed side's sort order with the join for those two shapes, and have the planner declare it on the probe input through a pass-through SortedInputExec when the input declares none. The declaration is placed on that one input rather than on every ScanExec because an order visible to the whole native plan would also change how DataFusion's aggregates group. A sort order that cannot be expressed natively sends the join back to Spark. With the in-memory cache on by default since #5634, the native cache scan reaches this with default settings. --- native/core/src/execution/operators/mod.rs | 2 + .../src/execution/operators/sorted_input.rs | 120 ++++++++++++++++++ native/core/src/execution/planner.rs | 40 +++++- native/proto/src/proto/operator.proto | 5 + .../apache/spark/sql/comet/operators.scala | 25 ++++ .../comet/exec/CometInMemoryCacheSuite.scala | 73 ++++++++++- .../apache/comet/exec/CometJoinSuite.scala | 59 ++++++++- 7 files changed, 319 insertions(+), 5 deletions(-) create mode 100644 native/core/src/execution/operators/sorted_input.rs diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index a15a528f90b..02cdea9326f 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -41,5 +41,7 @@ mod csv_scan; pub mod projection; mod scan; mod shuffle_scan; +mod sorted_input; pub use csv_scan::init_csv_datasource_exec; pub use shuffle_scan::ShuffleScanExec; +pub use sorted_input::SortedInputExec; diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs new file mode 100644 index 00000000000..ffdf7a5bc12 --- /dev/null +++ b/native/core/src/execution/operators/sorted_input.rs @@ -0,0 +1,120 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! An input that declares the order its rows arrive in. + +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::{ + apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; + +/// Passes `input` through unchanged and declares to DataFusion the order its rows arrive in. +/// +/// Spark's plan knows the order of an input that enters the native plan from the JVM, a relation +/// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none, +/// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only +/// when the probe input declares an order. Otherwise it moves them after the matched rows of each +/// batch, which loses an order Spark has already dropped a sort for. The planner puts this below +/// that join's probe side with the order Spark reports for it, and nowhere else: an order that +/// the whole native plan could see would also change how DataFusion's aggregates group. +#[derive(Debug)] +pub struct SortedInputExec { + input: Arc, + ordering: LexOrdering, + cache: Arc, +} + +impl SortedInputExec { + pub fn new(input: Arc, ordering: LexOrdering) -> Self { + let input_properties = input.properties(); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(input.schema(), [ordering.clone()]), + input_properties.output_partitioning().clone(), + input_properties.emission_type, + input_properties.boundedness, + )); + Self { + input, + ordering, + cache, + } + } +} + +impl DisplayAs for SortedInputExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "SortedInputExec: ordering=[{}]", self.ordering) + } +} + +impl ExecutionPlan for SortedInputExec { + fn name(&self) -> &str { + "SortedInputExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots(self.ordering.iter().map(|sort| Arc::clone(&sort.expr)), f) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("SortedInputExec requires one child"); + } + Ok(Arc::new(Self::new( + children.remove(0), + self.ordering.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index b88b477931d..68e7b4dfa12 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -37,6 +37,7 @@ mod lance_scan; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::DynamicFilterJoinExec; use crate::execution::operators::IcebergScanExec; +use crate::execution::operators::SortedInputExec; use crate::execution::operators::TopKReaderFilterExec; use crate::execution::{ operators::{ @@ -2629,8 +2630,21 @@ impl PhysicalPlanner { *probe = Self::prepare_probe_filter_for_runtime_reader(Arc::clone(probe)); } - let left = Arc::clone(&join_params.left.native_plan); - let right = Arc::clone(&join_params.right.native_plan); + let mut left = Arc::clone(&join_params.left.native_plan); + let mut right = Arc::clone(&join_params.right.native_plan); + + // Spark dropped any sort above this join that the streamed side's order + // satisfies, and the native join keeps that order only if the probe input + // declares it; see `SortedInputExec`. + if !join.streamed_sort_orders.is_empty() { + let probe = if join.build_side == BuildSide::BuildLeft as i32 { + &mut right + } else { + &mut left + }; + *probe = + self.with_declared_order(Arc::clone(probe), &join.streamed_sort_orders)?; + } // Null-aware anti-join must run in CollectLeft mode. In Partitioned mode // each partition only sees per-partition null/emptiness state, which can @@ -2921,6 +2935,27 @@ impl PhysicalPlanner { Arc::new(prepared) } + /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless it already + /// declares one. Spark's plan knows the order of an input that enters the native plan from + /// the JVM, but the `ScanExec` it arrives through declares none. + fn with_declared_order( + &self, + plan: Arc, + sort_orders: &[Expr], + ) -> Result, ExecutionError> { + if plan.output_ordering().is_some() { + return Ok(plan); + } + let exprs = sort_orders + .iter() + .map(|expr| self.create_sort_expr(expr, plan.schema())) + .collect::, _>>()?; + match LexOrdering::new(exprs) { + Some(ordering) => Ok(Arc::new(SortedInputExec::new(plan, ordering))), + None => Ok(plan), + } + } + /// Attach after choosing the final build side, including the projection emitted /// by `swap_inputs`, without running DataFusion's physical optimizer. pub(crate) fn apply_join_dynamic_filter( @@ -6006,6 +6041,7 @@ mod tests { build_side: 0, null_aware_anti_join: false, dynamic_filter_enabled: false, + streamed_sort_orders: vec![], })), }; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 97c65f0da36..0d087366ac4 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1026,6 +1026,11 @@ 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; + // The streamed side's sort order, as `SortOrder` expressions over its output, when Spark + // reports that order for the join and the native join keeps it only if the probe input + // declares it: a left outer join built on the right, or a right outer join built on the left. + // Empty otherwise. + repeated spark.spark_expression.Expr streamed_sort_orders = 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 bcae614556b..e2dc9dee896 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 @@ -2864,6 +2864,30 @@ trait CometHashJoin { val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output)) val rightKeys = join.rightKeys.map(exprToProto(_, join.right.output)) + // Spark reports the streamed side's order as the join's, and drops a sort above the join + // that it satisfies. DataFusion's hash join keeps the unmatched probe rows of a right outer + // join in place only when the probe input declares an order, and otherwise moves them after + // the matched rows of each batch. Comet builds on the build side, so a left outer join built + // on the right and a right outer join built on the left both run as that right outer join, + // and an input sorted outside the native plan, a relation sorted before it was cached for + // instance, declares no order to DataFusion. Send the order so the native join declares it + // on the probe input; see SortedInputExec. Every other join type keeps the probe order as + // it is. + val streamedSortOrders = (join.joinType, join.buildSide) match { + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty => + val streamed = if (join.buildSide == BuildRight) join.left else join.right + val orders = join.outputOrdering.map(exprToProto(_, streamed.output)) + if (orders.exists(_.isEmpty)) { + withFallbackReason( + join, + "the streamed side's sort order, which this outer join has to keep, cannot be " + + "expressed natively") + return None + } + orders.flatten + case _ => Nil + } + if (leftKeys.forall(_.isDefined) && rightKeys.forall(_.isDefined) && childOp.nonEmpty) { @@ -2876,6 +2900,7 @@ trait CometHashJoin { else OperatorOuterClass.BuildSide.BuildRight) .setNullAwareAntiJoin(isNullAwareAntiJoin) .setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf)) + .addAllStreamedSortOrders(streamedSortOrders.asJava) condition.foreach(joinBuilder.setCondition) Some(builder.setHashJoin(joinBuilder).build()) } else { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index aba97d80219..0c9b5c58bd6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -37,11 +37,11 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlanInfo} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -3133,4 +3133,73 @@ class CometInMemoryCacheSuite extends CometTestBase { }.nonEmpty) } } + + test("hash joins over a sorted cache keep the streamed side's order") { + // CometInMemoryTableScan reports the cached plan's sort order, so Spark drops the sort it + // would otherwise plan above a join that keeps that order. DataFusion's hash join keeps the + // unmatched probe rows of a left outer join built on the right in place only when the probe + // input declares an order, which the native scan over the cache does not, so the native join + // declares the cache's order on its probe input. An inner join never moves probe rows. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + def query(hint: String, joinType: String): DataFrame = + streamed + .join(build.hint(hint), streamed("_1") === build("_1"), joinType) + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + + def sortedPartitions(df: DataFrame): Array[Boolean] = + df.queryExecution.toRdd + .mapPartitions { it => + val keys = it.map(_.getInt(0)).toArray + Iterator(keys.sameElements(keys.sorted)) + } + .collect() + + def sorts(plan: SparkPlan): Seq[SparkPlan] = + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s } + + val outerRows = (0 until 10000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + val innerRows = outerRows.filter(_.get(2) != null) + + Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach { + case (hint, joinClass) => + val outer = query(hint, "left_outer") + val outerPlan = outer.queryExecution.executedPlan + assert(collect(outerPlan) { case j if j.getClass == joinClass => j }.size == 1) + assert(collect(outerPlan) { case s: CometInMemoryTableScanExec => s }.size == 1) + assert(sorts(outerPlan).isEmpty, outerPlan) + assert(outer.collect().toSet == outerRows.toSet) + assert(sortedPartitions(outer).forall(identity), outerPlan) + + val inner = query(hint, "inner") + val innerPlan = inner.queryExecution.executedPlan + assert(collect(innerPlan) { case j if j.getClass == joinClass => j }.size == 1) + assert(sorts(innerPlan).isEmpty, innerPlan) + assert(inner.collect().toSet == innerRows.toSet) + assert(sortedPartitions(inner).forall(identity), innerPlan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } } 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..ef621054127 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, Dynam 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.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, 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 @@ -2152,4 +2152,61 @@ class CometJoinSuite extends CometTestBase { } } } + + test("outer hash joins keep a streamed side sorted outside the native plan in order") { + // Spark drops the sort above the join because the join reports the streamed side's order, + // and the streamed side arrives sorted through a JVM input, which declares no order to + // DataFusion unless the planner declares it on the join's probe input; see + // CometHashJoin.doConvert. CometInMemoryCacheSuite covers the native cache scan. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + def sortedPartitions(df: DataFrame): Array[Boolean] = + df.queryExecution.toRdd + .mapPartitions { it => + val keys = it.map(_.getInt(0)).toArray + Iterator(keys.sameElements(keys.sorted)) + } + .collect() + + val expected = (0 until 10000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + + Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach { + case (hint, joinClass) => + val df = streamed + .join(build.hint(hint), streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j if j.getClass == joinClass => j }.size == 1, plan) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + assert(df.collect().toSet == expected.toSet) + assert(sortedPartitions(df).forall(identity), plan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } } From f74ae8f0cbb2cfcadbad662b0f503d5acc7890ca Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 15:58:01 -0600 Subject: [PATCH 02/19] test: cover the right outer shape, an aggregate above the join and the fallback From review of the hash join order declaration: - The join and cache tests also run a right outer join built on the left, which Comet does not swap, and compare rows as multisets, not sets. - A cache test groups by the key above the join, where the declared order puts both aggregates in sorted mode. - A join test covers the fallback when the streamed order cannot be serialized. - Planner tests cover declaring the order on the probe input for both shapes, and leaving a probe that already declares one alone. - SortedInputExec's doc says the declared order is visible above the join. --- .../src/execution/operators/sorted_input.rs | 9 +- native/core/src/execution/planner.rs | 94 ++++++++++++++++ .../comet/exec/CometInMemoryCacheSuite.scala | 104 +++++++++++++++--- .../apache/comet/exec/CometJoinSuite.scala | 87 ++++++++++++--- 4 files changed, 257 insertions(+), 37 deletions(-) diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index ffdf7a5bc12..e805c8c758b 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -36,9 +36,12 @@ use datafusion::physical_plan::{ /// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none, /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only /// when the probe input declares an order. Otherwise it moves them after the matched rows of each -/// batch, which loses an order Spark has already dropped a sort for. The planner puts this below -/// that join's probe side with the order Spark reports for it, and nowhere else: an order that -/// the whole native plan could see would also change how DataFusion's aggregates group. +/// batch, which loses an order Spark has already dropped a sort for. The planner puts this only +/// on that join's probe input, with the order Spark reports for it, but the join keeps the order +/// in its own output, so the operators above it in the native plan see it too: an aggregate +/// grouped by the key runs in sorted, streaming mode, and a sort aggregate takes the order as +/// already satisfied. The order therefore has to be one Spark guarantees, since a wrong one gives +/// wrong groups rather than only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 4106dbb9a8e..2b624f51e79 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6072,6 +6072,100 @@ mod tests { assert_eq!("ScanExec", hash_join_exec.children[1].native_plan.name()); } + /// The hash join in `plan`, under the projection that `swap_inputs` adds when Spark builds + /// on the right. + fn planned_hash_join( + plan: &Arc, + ) -> &datafusion::physical_plan::joins::HashJoinExec { + let join = if plan.is::() { + plan.children()[0] + } else { + plan + }; + join.downcast_ref::() + .expect("a hash join") + } + + fn sorted_outer_hash_join( + join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + children: Vec, + ) -> Operator { + Operator { + children, + op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + left_join_keys: vec![create_bound_reference(0)], + right_join_keys: vec![create_bound_reference(0)], + join_type: join_type as i32, + build_side: build_side as i32, + streamed_sort_orders: vec![create_sort_order(0, 3, false, true)], + ..Default::default() + })), + ..Default::default() + } + } + + #[test] + fn hash_join_declares_streamed_order_on_probe_input() { + // The two outer joins Comet sends the streamed side's order for both run as DataFusion's + // right join, whose probe input is its right one, whichever side Spark builds on. + for (join_type, build_side) in [ + ( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + ), + ( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + ), + ] { + let operator = + sorted_outer_hash_join(join_type, build_side, vec![create_scan(), create_scan()]); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let join = planned_hash_join(&planned.native_plan); + assert_eq!("ScanExec", join.left().name()); + let probe = join.right(); + let sorted = probe + .downcast_ref::() + .expect("the probe input declares the streamed side's order"); + assert_eq!("ScanExec", sorted.children()[0].name()); + assert_eq!( + "col_0@0 ASC", + probe.properties().output_ordering().unwrap().to_string() + ); + // The join keeps the order in its own output, so the operators above it see it too. + assert!(planned.native_plan.properties().output_ordering().is_some()); + } + } + + #[test] + fn hash_join_keeps_probe_input_that_declares_an_order() { + // A probe input that already declares an order, such as a sort in the native plan, is + // left as it is. + let sort = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![create_sort_order(0, 3, false, true)], + ..Default::default() + })), + ..Default::default() + }; + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + vec![create_scan(), sort], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + assert_eq!( + "SortExec", + planned_hash_join(&planned.native_plan).right().name() + ); + } + #[tokio::test] async fn spark_plan_dynamic_filter_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 5849c7eb7c6..569259a356c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -37,8 +37,9 @@ import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils @@ -3292,9 +3293,10 @@ class CometInMemoryCacheSuite extends CometTestBase { test("hash joins over a sorted cache keep the streamed side's order") { // CometInMemoryTableScan reports the cached plan's sort order, so Spark drops the sort it // would otherwise plan above a join that keeps that order. DataFusion's hash join keeps the - // unmatched probe rows of a left outer join built on the right in place only when the probe - // input declares an order, which the native scan over the cache does not, so the native join - // declares the cache's order on its probe input. An inner join never moves probe rows. + // unmatched probe rows of a left outer join built on the right, or of a right outer join built + // on the left, in place only when the probe input declares an order, which the native scan + // over the cache does not, so the native join declares the cache's order on its probe input. + // An inner join never moves probe rows. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -3308,11 +3310,21 @@ class CometInMemoryCacheSuite extends CometTestBase { streamed.count() val build = spark.table("small") - def query(hint: String, joinType: String): DataFrame = - streamed - .join(build.hint(hint), streamed("_1") === build("_1"), joinType) + def query(hint: String, joinType: String): DataFrame = { + val joined = if (joinType == "right_outer") { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } else { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } + joined .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) .sortWithinPartitions("k") + } + + def joins(plan: SparkPlan): Seq[(Class[_], BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j.getClass, j.buildSide) + case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide) + } def sortedPartitions(df: DataFrame): Array[Boolean] = df.queryExecution.toRdd @@ -3335,19 +3347,24 @@ class CometInMemoryCacheSuite extends CometTestBase { "shuffle_hash" -> classOf[CometHashJoinExec], "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach { case (hint, joinClass) => - val outer = query(hint, "left_outer") - val outerPlan = outer.queryExecution.executedPlan - assert(collect(outerPlan) { case j if j.getClass == joinClass => j }.size == 1) - assert(collect(outerPlan) { case s: CometInMemoryTableScanExec => s }.size == 1) - assert(sorts(outerPlan).isEmpty, outerPlan) - assert(outer.collect().toSet == outerRows.toSet) - assert(sortedPartitions(outer).forall(identity), outerPlan) + Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft).foreach { + case (joinType, buildSide) => + val outer = query(hint, joinType) + val outerPlan = outer.queryExecution.executedPlan + assert(joins(outerPlan) == Seq((joinClass, buildSide)), outerPlan) + assert(collect(outerPlan) { case s: CometInMemoryTableScanExec => + s + }.size == 1) + assert(sorts(outerPlan).isEmpty, outerPlan) + QueryTest.sameRows(outerRows, outer.collect().toSeq).foreach(fail(_)) + assert(sortedPartitions(outer).forall(identity), outerPlan) + } val inner = query(hint, "inner") val innerPlan = inner.queryExecution.executedPlan - assert(collect(innerPlan) { case j if j.getClass == joinClass => j }.size == 1) + assert(joins(innerPlan) == Seq((joinClass, BuildRight)), innerPlan) assert(sorts(innerPlan).isEmpty, innerPlan) - assert(inner.collect().toSet == innerRows.toSet) + QueryTest.sameRows(innerRows, inner.collect().toSeq).foreach(fail(_)) assert(sortedPartitions(inner).forall(identity), innerPlan) } } finally { @@ -3357,4 +3374,59 @@ class CometInMemoryCacheSuite extends CometTestBase { } } } + + test("an aggregate on the key above a hash join over a sorted cache groups each key once") { + // The order the join declares on its probe input reaches past the join, so the aggregate + // above it in the same native plan groups in sorted, streaming mode. That is right while + // Spark's reported order holds, and a wrong order would split a key across groups, which only + // a comparison of the rows as a multiset catches. Small batches make each group span several. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40 by 2).map(m => (m * 1000, m)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + // Fifty rows per key, and a match in the build side for every even key. + val expected = (0 until 40).map { m => + Row(m * 1000, 50L, 50L * m + 49000L, if (m % 2 == 0) 50L else 0L) + } + + for { + hint <- Seq("shuffle_hash", "broadcast") + joinType <- Seq("left_outer", "right_outer") + } { + val joined = if (joinType == "right_outer") { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } else { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } + val df = joined + .groupBy(streamed("_1")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + val plan = df.queryExecution.executedPlan + assert( + collect(plan) { case j @ (_: CometHashJoinExec | _: CometBroadcastHashJoinExec) => + j + }.size == 1, + plan) + assert(collect(plan) { case a: CometHashAggregateExec => a }.size == 2, plan) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } } 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 2a2d70ab0b3..1c31323e090 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -26,11 +26,11 @@ 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.{CometTestBase, DataFrame, QueryTest, 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.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, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} @@ -2287,23 +2287,74 @@ class CometJoinSuite extends CometTestBase { Row(k, i, if (k % 10 == 0) k / 10 else null) } - Seq( - "shuffle_hash" -> classOf[CometHashJoinExec], - "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach { - case (hint, joinClass) => - val df = streamed - .join(build.hint(hint), streamed("_1") === build("_1"), "left_outer") - .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) - .sortWithinPartitions("k") - val plan = df.queryExecution.executedPlan - assert(collect(plan) { case j if j.getClass == joinClass => j }.size == 1, plan) - assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) - assert( - collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, - plan) - assert(df.collect().toSet == expected.toSet) - assert(sortedPartitions(df).forall(identity), plan) + // Both shapes run as DataFusion's right join: Comet swaps a join built on the right + // and leaves one built on the left as it is. + def outerJoin(hint: String, joinType: String): DataFrame = { + val joined = if (joinType == "left_outer") { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } else { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } + joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") } + + for { + (hint, joinClass) <- Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]) + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + } { + val df = outerJoin(hint, joinType) + val plan = df.queryExecution.executedPlan + val joins: Seq[(Class[_], BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j.getClass, j.buildSide) + case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide) + } + assert(joins == Seq((joinClass, buildSide)), plan) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(sortedPartitions(df).forall(identity), plan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("an outer hash join stays in Spark when the streamed side's order cannot be expressed") { + // The native join keeps the streamed side's order only by declaring it on its probe input, so + // a join whose order Comet cannot serialize stays in Spark rather than lose the order. With + // Add disabled, Comet cannot serialize the order below. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Add") -> "false") { + withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions($"_1" + 1).cache() + streamed.count() + val build = spark.table("small") + val (_, plan) = checkSparkAnswerAndFallbackReason( + streamed.join( + build.hint("shuffle_hash"), + streamed("_1") === build("_1"), + "left_outer"), + "the streamed side's sort order, which this outer join has to keep, cannot be " + + "expressed natively") + assert(collect(plan) { case j: ShuffledHashJoinExec => j }.size == 1, plan) + assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, plan) } finally { spark.catalog.clearCache() } From 2f95eefb3606ce96a3bdc2771c478cc2b4d26ce6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 06:47:42 -0600 Subject: [PATCH 03/19] fix: declare the streamed order only over operators that keep a JVM input's order From review: - with_declared_order walks down from the probe input through operators that report keeping their input's order and declares only when it reaches a ScanExec or ShuffleScanExec. DataFusion reports a right hash join as keeping its probe order even when it moves unmatched rows, so the walk stops at any hash join that is not an inner join. The null-aware anti join stops it, so an order Spark reports above NOT IN no longer reaches the aggregates above the outer join. - doConvert sends the longest prefix of the streamed order that Comet can express, since declaring any order keeps the probe rows in place, and falls back only when it cannot express the first key. - SortedInputExec's doc names the top-k sort as another consumer, and the CometHashJoin comment names the null-aware anti join as the exception. - Tests: dwsmith1983's NOT IN repro, a residual join condition, a descending, nulls last streamed order with null keys, the prefix case, the aggregate test's join and exchange checks, and planner tests for the walk and the planned join type. --- .../src/execution/operators/sorted_input.rs | 11 +- native/core/src/execution/planner.rs | 99 ++++++++- native/proto/src/proto/operator.proto | 8 +- .../apache/spark/sql/comet/operators.scala | 18 +- .../comet/exec/CometInMemoryCacheSuite.scala | 22 +- .../apache/comet/exec/CometJoinSuite.scala | 197 +++++++++++++++--- 6 files changed, 303 insertions(+), 52 deletions(-) diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index e805c8c758b..d3650cab03f 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -37,11 +37,12 @@ use datafusion::physical_plan::{ /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only /// when the probe input declares an order. Otherwise it moves them after the matched rows of each /// batch, which loses an order Spark has already dropped a sort for. The planner puts this only -/// on that join's probe input, with the order Spark reports for it, but the join keeps the order -/// in its own output, so the operators above it in the native plan see it too: an aggregate -/// grouped by the key runs in sorted, streaming mode, and a sort aggregate takes the order as -/// already satisfied. The order therefore has to be one Spark guarantees, since a wrong one gives -/// wrong groups rather than only rows out of order. +/// on that join's probe input, with the order Spark reports for it, and only when the rows reach +/// that input from the JVM through operators that keep their order. The join keeps the order in +/// its own output, so the operators above it in the native plan see it too: an aggregate grouped +/// by the key runs in sorted, streaming mode, a sort aggregate takes the order as already +/// satisfied, and a top-k sort whose keys start with it stops reading early. A wrong order would +/// therefore give wrong groups or wrong rows, not only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 2b624f51e79..a6825685611 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2953,13 +2953,14 @@ impl PhysicalPlanner { /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless it already /// declares one. Spark's plan knows the order of an input that enters the native plan from - /// the JVM, but the `ScanExec` it arrives through declares none. + /// the JVM, but the `ScanExec` it arrives through declares none. Only that order is declared: + /// when an operator in between may have moved rows, `plan` is left as it is. fn with_declared_order( &self, plan: Arc, sort_orders: &[Expr], ) -> Result, ExecutionError> { - if plan.output_ordering().is_some() { + if plan.output_ordering().is_some() || !Self::keeps_jvm_input_order(&plan) { return Ok(plan); } let exprs = sort_orders @@ -2972,6 +2973,32 @@ impl PhysicalPlanner { } } + /// Whether `plan`'s rows come from an input outside the native plan through operators that + /// keep their input's order, so that the order Spark reports for them still holds. Any other + /// operator may have moved rows, such as the null-aware anti join, whose streamed side Comet + /// builds on, and declaring Spark's order over it would be wrong. + fn keeps_jvm_input_order(plan: &Arc) -> bool { + if plan.is::() || plan.is::() { + return true; + } + // DataFusion reports every right hash join as keeping its probe order, but keeps it only + // when the probe declares one, and no input on this path declares an order. + if let Some(join) = plan.downcast_ref::() { + if join.join_type() != &DFJoinType::Inner { + return false; + } + } + let mut kept = plan + .children() + .into_iter() + .zip(plan.maintains_input_order()) + .filter_map(|(child, keeps)| keeps.then_some(child)); + match (kept.next(), kept.next()) { + (Some(child), None) => Self::keeps_jvm_input_order(child), + _ => false, + } + } + /// Attach after choosing the final build side, including the projection emitted /// by `swap_inputs`, without running DataFusion's physical optimizer. pub(crate) fn apply_join_dynamic_filter( @@ -6125,6 +6152,7 @@ mod tests { .create_plan(&operator, &mut vec![], 1) .unwrap(); let join = planned_hash_join(&planned.native_plan); + assert_eq!(&datafusion::common::JoinType::Right, join.join_type()); assert_eq!("ScanExec", join.left().name()); let probe = join.right(); let sorted = probe @@ -6140,6 +6168,73 @@ mod tests { } } + #[test] + fn hash_join_declares_order_through_operators_that_keep_it() { + let projection = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Projection(spark_operator::Projection { + project_list: vec![create_bound_reference(0)], + })), + ..Default::default() + }; + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + vec![create_scan(), projection], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!("SortedInputExec", probe.name()); + assert_eq!("ProjectionExec", probe.children()[0].name()); + } + + #[test] + fn hash_join_declares_no_order_over_an_operator_that_may_move_rows() { + let hash_join = |join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + null_aware_anti_join: bool| Operator { + children: vec![create_scan(), create_scan()], + op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + left_join_keys: vec![create_bound_reference(0)], + right_join_keys: vec![create_bound_reference(0)], + join_type: join_type as i32, + build_side: build_side as i32, + null_aware_anti_join, + ..Default::default() + })), + ..Default::default() + }; + for probe_input in [ + // Comet builds a null-aware anti join on Spark's streamed side, so its output is not + // in the order Spark reports for it. + hash_join( + spark_operator::JoinType::LeftAnti, + spark_operator::BuildSide::BuildRight, + true, + ), + // A right join whose probe declares no order moves its unmatched rows, although + // DataFusion reports it as keeping its probe order. + hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + false, + ), + ] { + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + vec![create_scan(), probe_input], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!("HashJoinExec", probe.name()); + } + } + #[test] fn hash_join_keeps_probe_input_that_declares_an_order() { // A probe input that already declares an order, such as a sort in the native plan, is diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 0d087366ac4..392d159bf88 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1026,10 +1026,10 @@ 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; - // The streamed side's sort order, as `SortOrder` expressions over its output, when Spark - // reports that order for the join and the native join keeps it only if the probe input - // declares it: a left outer join built on the right, or a right outer join built on the left. - // Empty otherwise. + // The streamed side's sort order, or the longest prefix of it Comet can express, as + // `SortOrder` expressions over its output, when Spark reports that order for the join and the + // native join keeps it only if the probe input declares it: a left outer join built on the + // right, or a right outer join built on the left. Empty otherwise. repeated spark.spark_expression.Expr streamed_sort_orders = 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 e2dc9dee896..44936cdd11a 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 @@ -2872,19 +2872,29 @@ trait CometHashJoin { // and an input sorted outside the native plan, a relation sorted before it was cached for // instance, declares no order to DataFusion. Send the order so the native join declares it // on the probe input; see SortedInputExec. Every other join type keeps the probe order as - // it is. + // it is, except the null-aware anti join: Comet builds that one on its streamed side, so its + // output is not in the order Spark reports for it, and the planner declares no order above + // it. + // + // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix + // of the order that Comet can express is enough. Nothing here can tell whether an operator + // above relies on the order, so a join whose first key Comet cannot express stays in Spark. val streamedSortOrders = (join.joinType, join.buildSide) match { case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty => val streamed = if (join.buildSide == BuildRight) join.left else join.right - val orders = join.outputOrdering.map(exprToProto(_, streamed.output)) - if (orders.exists(_.isEmpty)) { + val orders = join.outputOrdering.iterator + .map(exprToProto(_, streamed.output)) + .takeWhile(_.isDefined) + .flatten + .toSeq + if (orders.isEmpty) { withFallbackReason( join, "the streamed side's sort order, which this outer join has to keep, cannot be " + "expressed natively") return None } - orders.flatten + orders case _ => Nil } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 569259a356c..1a4b35edf57 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3401,8 +3401,10 @@ class CometInMemoryCacheSuite extends CometTestBase { } for { - hint <- Seq("shuffle_hash", "broadcast") - joinType <- Seq("left_outer", "right_outer") + (hint, joinClass) <- Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]) + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) } { val joined = if (joinType == "right_outer") { build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) @@ -3413,12 +3415,20 @@ class CometInMemoryCacheSuite extends CometTestBase { .groupBy(streamed("_1")) .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) val plan = df.queryExecution.executedPlan + val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j, j.buildSide) + case j: CometBroadcastHashJoinExec => (j, j.buildSide) + } + assert( + joins.map { case (j, side) => (j.getClass, side) } == Seq((joinClass, buildSide)), + plan) + // A final aggregate after a shuffle would merge a key a partial aggregate split, so + // both aggregates have to sit above the join with no exchange in between. + def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined assert( - collect(plan) { case j @ (_: CometHashJoinExec | _: CometBroadcastHashJoinExec) => - j - }.size == 1, + collect(plan) { case a: CometHashAggregateExec => a }.count(aboveJoin) == 2, plan) - assert(collect(plan) { case a: CometHashAggregateExec => a }.size == 2, plan) + assert(!collect(plan) { case e: ShuffleExchangeLike => e }.exists(aboveJoin), plan) assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) } 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 1c31323e090..1b44b32ceb8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -26,10 +26,10 @@ 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, QueryTest, Row} +import org.apache.spark.sql.{Column, CometTestBase, DataFrame, QueryTest, 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.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 @@ -2255,6 +2255,18 @@ class CometJoinSuite extends CometTestBase { } } + // Whether every partition of `df` is in `order`, which is bound to its output. + private def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = { + val ordering = new InterpretedOrdering(order) + df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part => + part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) } + } + } + + // The order of the integer key in the first output column. + private def keyOrder(direction: SortDirection): SortOrder = + SortOrder(BoundReference(0, IntegerType, nullable = true), direction) + test("outer hash joins keep a streamed side sorted outside the native plan in order") { // Spark drops the sort above the join because the join reports the streamed side's order, // and the streamed side arrives sorted through a JVM input, which declares no order to @@ -2274,26 +2286,23 @@ class CometJoinSuite extends CometTestBase { streamed.count() val build = spark.table("small") - def sortedPartitions(df: DataFrame): Array[Boolean] = - df.queryExecution.toRdd - .mapPartitions { it => - val keys = it.map(_.getInt(0)).toArray - Iterator(keys.sameElements(keys.sorted)) - } - .collect() - - val expected = (0 until 10000).map { i => + // A residual condition, which keeps only the build rows past 3, sends more probe rows + // down the join's unmatched path. + def expected(withCondition: Boolean): Seq[Row] = (0 until 10000).map { i => val k = i % 100 - Row(k, i, if (k % 10 == 0) k / 10 else null) + val matched = k % 10 == 0 && (!withCondition || k / 10 > 3) + Row(k, i, if (matched) k / 10 else null) } // Both shapes run as DataFusion's right join: Comet swaps a join built on the right // and leaves one built on the left as it is. - def outerJoin(hint: String, joinType: String): DataFrame = { + def outerJoin(hint: String, joinType: String, withCondition: Boolean): DataFrame = { + val keys = streamed("_1") === build("_1") + val on = if (withCondition) keys && build("_2") > 3 else keys val joined = if (joinType == "left_outer") { - streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + streamed.join(build.hint(hint), on, joinType) } else { - build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + build.hint(hint).join(streamed, on, joinType) } joined .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) @@ -2305,8 +2314,9 @@ class CometJoinSuite extends CometTestBase { "shuffle_hash" -> classOf[CometHashJoinExec], "broadcast" -> classOf[CometBroadcastHashJoinExec]) (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + withCondition <- Seq(false, true) } { - val df = outerJoin(hint, joinType) + val df = outerJoin(hint, joinType, withCondition) val plan = df.queryExecution.executedPlan val joins: Seq[(Class[_], BuildSide)] = collect(plan) { case j: CometHashJoinExec => (j.getClass, j.buildSide) @@ -2317,8 +2327,8 @@ class CometJoinSuite extends CometTestBase { assert( collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) - assert(sortedPartitions(df).forall(identity), plan) + QueryTest.sameRows(expected(withCondition), df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } } finally { spark.catalog.clearCache() @@ -2328,10 +2338,11 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join stays in Spark when the streamed side's order cannot be expressed") { - // The native join keeps the streamed side's order only by declaring it on its probe input, so - // a join whose order Comet cannot serialize stays in Spark rather than lose the order. With - // Add disabled, Comet cannot serialize the order below. + test("an outer hash join declares the part of the streamed order Comet can express") { + // Declaring any order keeps the probe rows in place, so a prefix of the order is enough. + // Nothing tells whether an operator above relies on the order, so a join whose first key Comet + // cannot express stays in Spark rather than lose it. With Add disabled, Comet cannot express + // `_1 + 1`. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -2341,20 +2352,41 @@ class CometJoinSuite extends CometTestBase { CometConf.getExprEnabledConfigKey("Add") -> "false") { withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + def sortedCache(order: Column*): DataFrame = { + val cached = + spark.table("big").repartition(2, $"_1").sortWithinPartitions(order: _*).cache() + cached.count() + cached + } + val build = spark.table("small") try { - val streamed = - spark.table("big").repartition(2, $"_1").sortWithinPartitions($"_1" + 1).cache() - streamed.count() - val build = spark.table("small") - val (_, plan) = checkSparkAnswerAndFallbackReason( - streamed.join( + val unexpressed = sortedCache($"_1" + 1) + val (_, fallbackPlan) = checkSparkAnswerAndFallbackReason( + unexpressed.join( build.hint("shuffle_hash"), - streamed("_1") === build("_1"), + unexpressed("_1") === build("_1"), "left_outer"), "the streamed side's sort order, which this outer join has to keep, cannot be " + "expressed natively") - assert(collect(plan) { case j: ShuffledHashJoinExec => j }.size == 1, plan) - assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, plan) + assert( + collect(fallbackPlan) { case j: ShuffledHashJoinExec => j }.size == 1, + fallbackPlan) + assert(collect(fallbackPlan) { case j: CometHashJoinExec => j }.isEmpty, fallbackPlan) + + val prefixed = sortedCache($"_1", $"_1" + 1) + val df = prefixed + .join(build.hint("shuffle_hash"), prefixed("_1") === build("_1"), "left_outer") + .select(prefixed("_1").as("k"), prefixed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: CometHashJoinExec => j }.size == 1, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 1000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } finally { spark.catalog.clearCache() } @@ -2362,4 +2394,107 @@ class CometJoinSuite extends CometTestBase { } } } + + test("outer hash joins keep a descending, nulls last streamed order with null keys") { + // Null keys match nothing, so all of them take the unmatched path, at the end of each + // partition. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + 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") { + try { + val streamed = spark + .table("big") + .repartition(2, $"_1") + .sortWithinPartitions($"_1".desc_nulls_last) + .cache() + streamed.count() + val build = spark.table("small").hint("shuffle_hash") + val expected = data.map { case (k, i) => + Row(k.getOrElse(null), i, k.filter(_ % 10 == 0).map(_ / 10).getOrElse(null)) + } + Seq("left_outer", "right_outer").foreach { joinType => + val joined = if (joinType == "left_outer") { + streamed.join(build, streamed("_1") === build("_1"), joinType) + } else { + build.join(streamed, build("_1") === streamed("_1"), joinType) + } + val df = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions($"k".desc_nulls_last) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: CometHashJoinExec => j }.size == 1, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Descending)), plan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("NOT IN below an outer hash join over a sorted cache keeps each group once") { + // Spark reports the streamed side's order for a null-aware anti join, but Comet builds that + // join on its streamed side, so the order does not hold above it. Declaring it on the outer + // join's probe input would put the aggregates in sorted mode over unsorted rows and split + // groups. Sparse keys keep DataFusion off its perfect hash build path, which keeps the build + // batches in order and would hide that, and small batches spread each partition over many. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + |FROM (SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)) a + |LEFT JOIN small b ON a.k = b._1 + |GROUP BY a.k""".stripMargin + def rows(): Seq[(Int, Long, Long, Long)] = + sql(query) + .collect() + .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) + .toSeq + .sorted + val cometRows = rows() + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[(Int, Long, Long, Long)] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = rows() + } + assert( + cometRows.map(_._1).distinct.size == cometRows.size, + "a key is grouped twice") + assert(cometRows == sparkRows) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("cached") + } + } + } + } + } + } } From 2e0fb634695567ac1b238cea8a36584a48d32fbf Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 09:45:58 -0600 Subject: [PATCH 04/19] fix: send no streamed order when a Comet join that may move rows produces it A cache reports the order of the plan that wrote it, so the output of a native null-aware anti join, an outer hash join whose probe order the native planner did not declare, or a nested loop join can reach an outer hash join as a JVM input whose reported order does not hold. Declaring that order put the aggregates above the join in sorted mode over unsorted rows, which split groups. CometHashJoin.doConvert now walks down the streamed side, through CometScanWrapper and into the plan that wrote any cache it reads, and sends no order when such a join produces it with no sort in between. --- .../apache/spark/sql/comet/operators.scala | 46 ++- .../apache/comet/exec/CometJoinSuite.scala | 265 ++++++++++++++++-- 2 files changed, 292 insertions(+), 19 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 44936cdd11a..9efaeffcba5 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 @@ -42,6 +42,7 @@ import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -2876,12 +2877,22 @@ trait CometHashJoin { // output is not in the order Spark reports for it, and the planner declares no order above // it. // + // The order Spark reports can be false to begin with. A Comet join that may move its + // streamed rows reports it all the same: the null-aware anti join, a right outer join like + // this one whose probe input declared no order, and the nested loop join, which keeps the + // order of neither input. A cache reports the order of the plan that wrote it, so the rows + // of such a join can reach this one through a cache, as a JVM input that the native planner + // cannot look behind. Declaring the order there would let an operator above rely on it, an + // aggregate in sorted mode for instance, which then splits groups. So nothing is declared + // when such a join feeds the streamed side with no sort in between. + // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix // of the order that Comet can express is enough. Nothing here can tell whether an operator // above relies on the order, so a join whose first key Comet cannot express stays in Spark. + val streamed = if (join.buildSide == BuildRight) join.left else join.right val streamedSortOrders = (join.joinType, join.buildSide) match { - case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty => - val streamed = if (join.buildSide == BuildRight) join.left else join.right + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) + if join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed) => val orders = join.outputOrdering.iterator .map(exprToProto(_, streamed.output)) .takeWhile(_.isDefined) @@ -2917,6 +2928,37 @@ trait CometHashJoin { None } } + + // Whether a Comet join that may move its streamed rows, while reporting Spark's order for + // them, feeds the order that `plan` reports with no sort in between. Whether an outer join + // declared the order on its probe input is up to the native planner, so every one built that + // way counts. A cache reports the order of the plan that wrote it, so the walk goes on into + // that plan. While CometExecRule converts a plan, a native plan's JVM inputs sit behind a + // CometScanWrapper, a leaf, so the walk looks through that too. + private def mayReportFalseOrder(plan: SparkPlan): Boolean = plan match { + case _ if plan.outputOrdering.isEmpty => false + case _: SortExec | _: CometSortExec => false + case join: CometBroadcastHashJoinExec + if join.isNullAwareAntiJoin || runsAsRightOuterJoin(join.joinType, join.buildSide) => + true + case join: CometHashJoinExec if runsAsRightOuterJoin(join.joinType, join.buildSide) => true + case _: CometBroadcastNestedLoopJoinExec => true + case wrapper: CometScanWrapper => mayReportFalseOrder(wrapper.originalPlan) + case scan: CometInMemoryTableScanExec => + mayReportFalseOrder(scan.originalPlan.relation.cachedPlan) + case scan: InMemoryTableScanExec => mayReportFalseOrder(scan.relation.cachedPlan) + case stage: QueryStageExec => mayReportFalseOrder(stage.plan) + case _ => plan.children.exists(mayReportFalseOrder) + } + + // A left outer join built on the right and a right outer join built on the left both run as + // DataFusion's right outer join, which keeps its unmatched probe rows in place only when the + // probe input declares an order. + private def runsAsRightOuterJoin(joinType: JoinType, buildSide: BuildSide): Boolean = + (joinType, buildSide) match { + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => true + case _ => false + } } case class CometBroadcastNestedLoopJoinExec( 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 1b44b32ceb8..35dc61bee73 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,11 +31,12 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, 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.{Inner, LeftOuter} import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, 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.columnar.InMemoryTableScanExec 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} @@ -2472,24 +2473,118 @@ class CometJoinSuite extends CometTestBase { |FROM (SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)) a |LEFT JOIN small b ON a.k = b._1 |GROUP BY a.k""".stripMargin - def rows(): Seq[(Int, Long, Long, Long)] = - sql(query) - .collect() - .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) - .toSeq - .sorted - val cometRows = rows() - // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. - var sparkRows = Seq.empty[(Int, Long, Long, Long)] - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = rows() - } + checkEachKeyGroupedOnce(query) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + + test("an outer hash join over a cache of a NOT IN keeps each group once") { + // The same NOT IN as the test above, but cached, so its rows reach the outer join from the + // JVM. The cache reports the order Spark reports for the null-aware anti join, which its rows + // are not in, and the outer join's probe input is then a plain JVM input as far as the native + // plan can see. The order must not be declared there either. A cached plan that AQE runs + // reports no order, so AQE is off here, as Spark 3.4 has it for every cached plan by default. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") + .cache() + .createOrReplaceTempView("anti") + spark.table("anti").count() + val cachedPlan = cachedPlanOf(spark.table("anti")) + assert( + cachedPlan.exists { + case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin + case _ => false + }, + s"the cache was not written by a native NOT IN:\n$cachedPlan") + assert( + !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), + "the cached NOT IN rows are in order, so nothing here can split a group") + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + |FROM anti a LEFT JOIN small b ON a.k = b._1 + |GROUP BY a.k""".stripMargin + checkEachKeyGroupedOnce(query) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("anti") + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + + test("an outer hash join over a cache of an undeclared outer join keeps each group once") { + // The IN subquery's semi join sits between the sorted cache and the first outer join, and the + // native planner declares no order over a semi join, so that outer join moves its unmatched + // rows, and the cache it writes is out of the order it reports. The second outer join reads + // that cache as a JVM input and must not declare the order either. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM (SELECT * FROM cached WHERE k IN (SELECT _1 FROM incl)) a + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) assert( - cometRows.map(_._1).distinct.size == cometRows.size, - "a key is grouped twice") - assert(cometRows == sparkRows) + cachedPlan.exists { + case join: CometBroadcastHashJoinExec => join.joinType == LeftOuter + case _ => false + }, + s"the cache was not written by a native outer join:\n$cachedPlan") + assert( + !partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), + "the cached outer join rows are in order, so nothing here can split a group") + val query = + """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) + |FROM joined j LEFT JOIN small c ON j.k = c._1 + |GROUP BY j.k""".stripMargin + checkEachKeyGroupedOnce(query) } finally { spark.catalog.clearCache() + spark.catalog.dropTempView("joined") spark.catalog.dropTempView("cached") } } @@ -2497,4 +2592,140 @@ class CometJoinSuite extends CometTestBase { } } } + + test("an outer hash join over a cache sorted above an outer hash join declares the order") { + // The sort in the plan that wrote the cache puts its rows in the order the cache reports, + // whatever the outer join below the sort did, so the second outer join declares that order. + // The sort is on (k, v), so Spark does not drop it as one the first join already satisfies. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM cached a LEFT JOIN small b ON a.k = b._1""".stripMargin) + .sortWithinPartitions("k", "v") + .cache() + .createOrReplaceTempView("sorted_joined") + spark.table("sorted_joined").count() + val cachedPlan = cachedPlanOf(spark.table("sorted_joined")) + assert( + cachedPlan.exists { + case join: CometBroadcastHashJoinExec => join.joinType == LeftOuter + case _ => false + }, + s"the cache was not written by a native outer join:\n$cachedPlan") + val df = sql("""SELECT /*+ BROADCAST(c) */ j.k, j.v, c._2 + |FROM sorted_joined j LEFT JOIN small c ON j.k = c._1""".stripMargin) + .sortWithinPartitions("k") + val joins = df.queryExecution.executedPlan.collect { + case join: CometBroadcastHashJoinExec => join + } + assert(joins.size == 1, df.queryExecution.executedPlan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0) + assert(partitionsInOrder(df, keyOrder(Ascending))) + var sparkRows = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = df.collect().toSeq + } + QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_)) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("sorted_joined") + spark.catalog.dropTempView("cached") + } + } + } + } + } + + test("an outer hash join over a cache of a nested loop join keeps each group once") { + // DataFusion's nested loop join emits each probe batch's matches build row by build row, so + // it keeps the order of neither input, while Spark reports the streamed side's order for an + // inner one. The cache it writes reports that order, and the outer join reading the cache + // must not declare it. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable(Seq(0, 1).map(Tuple1(_)), "two") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("""SELECT /*+ BROADCAST(t) */ a.k, a.v, t._1 AS x + |FROM cached a JOIN two t ON a.v % 2 <= t._1""".stripMargin) + .cache() + .createOrReplaceTempView("nlj") + spark.table("nlj").count() + val cachedPlan = cachedPlanOf(spark.table("nlj")) + assert( + cachedPlan.exists(_.isInstanceOf[CometBroadcastNestedLoopJoinExec]), + s"the cache was not written by a native nested loop join:\n$cachedPlan") + assert( + !partitionsInOrder(spark.table("nlj"), keyOrder(Ascending)), + "the cached nested loop join rows are in order, so nothing here can split a group") + checkEachKeyGroupedOnce( + """SELECT /*+ BROADCAST(c) */ n.k, count(*), sum(n.v), count(c._2) + |FROM nlj n LEFT JOIN small c ON n.k = c._1 + |GROUP BY n.k""".stripMargin) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("nlj") + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + + // Runs `query`, which groups by the integer key in its first column and returns three longs, + // with Comet and with Spark, and checks that Comet returns each key once and Spark's rows. + private def checkEachKeyGroupedOnce(query: String): Unit = { + def rows(): Seq[(Int, Long, Long, Long)] = + sql(query) + .collect() + .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) + .toSeq + .sorted + val cometRows = rows() + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[(Int, Long, Long, Long)] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = rows() + } + assert(cometRows.map(_._1).distinct.size == cometRows.size, "a key is grouped twice") + assert(cometRows == sparkRows) + } + + // The physical plan that wrote the cache `df` reads. + private def cachedPlanOf(df: DataFrame): SparkPlan = + df.queryExecution.executedPlan + .collectFirst { + case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan + case scan: InMemoryTableScanExec => scan.relation.cachedPlan + } + .getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) } From 79b4d336d04050f81ba469da780756c076adbd6a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 14:10:21 -0600 Subject: [PATCH 05/19] fix: walk an adaptive cache's final plan and keep the streamed order through semi joins Review feedback on the streamed order that a native outer hash join declares on its probe input. The serde walk that decides whether a cache or checkpoint can be trusted now: - reads the final plan of an adaptive cached plan, as Spark's cache scan does, since the adaptive plan reports no order itself and the walk stopped there; - follows only the streamed side of a hash join, and counts an outer join as declared only when it sent its order and its probe input is a scan in its own native plan; - no longer trusts a leaf that reports an order unless it is a scan or a range, which covers an RDD checkpoint whose plan is gone. The planner now walks through right semi, anti and mark hash joins as well as inner ones, since DataFusion builds their probe indices in ascending order whether or not the probe declares an order. Only a right outer join depends on the declaration. Tests cover a cached NOT IN with AQE off and on, through Spark's cache scan and through the native one, a checkpointed NOT IN, an outer join over a cache written by a declared outer join and by one over a sample, an outer join over a semi join, a NOT IN below an inner join's build side, and a broadcast outer join over a sorted cache with default settings. The test that compared Comet with itself now compares against the rows. --- native/core/src/execution/planner.rs | 108 +++-- .../apache/spark/sql/comet/operators.scala | 80 +++- .../comet/exec/CometInMemoryCacheSuite.scala | 125 +++++- .../apache/comet/exec/CometJoinSuite.scala | 379 +++++++++++++----- 4 files changed, 538 insertions(+), 154 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a6825685611..d1b905bacbd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2981,10 +2981,19 @@ impl PhysicalPlanner { if plan.is::() || plan.is::() { return true; } - // DataFusion reports every right hash join as keeping its probe order, but keeps it only - // when the probe declares one, and no input on this path declares an order. + // DataFusion reports every right hash join as keeping its probe order, but a right outer + // join keeps it only when the probe declares one, and no input on this path declares an + // order. An inner join emits its matches probe row by probe row, and a right semi, anti or + // mark join builds its probe indices in ascending order, so none of them depends on a + // declaration. Any other join type keeps no probe order at all. if let Some(join) = plan.downcast_ref::() { - if join.join_type() != &DFJoinType::Inner { + if !matches!( + join.join_type(), + DFJoinType::Inner + | DFJoinType::RightSemi + | DFJoinType::RightAnti + | DFJoinType::RightMark + ) { return false; } } @@ -6190,11 +6199,13 @@ mod tests { assert_eq!("ProjectionExec", probe.children()[0].name()); } - #[test] - fn hash_join_declares_no_order_over_an_operator_that_may_move_rows() { - let hash_join = |join_type: spark_operator::JoinType, - build_side: spark_operator::BuildSide, - null_aware_anti_join: bool| Operator { + /// A hash join over two scans, for use as the probe input of the join under test. + fn hash_join_over_scans( + join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + null_aware_anti_join: bool, + ) -> Operator { + Operator { children: vec![create_scan(), create_scan()], op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { left_join_keys: vec![create_bound_reference(0)], @@ -6205,33 +6216,82 @@ mod tests { ..Default::default() })), ..Default::default() - }; - for probe_input in [ + } + } + + #[test] + fn hash_join_declares_order_through_joins_that_keep_probe_order() { + // An inner join emits its matches probe row by probe row, and a right semi, anti or mark + // join emits its probe indices in ascending order. None of them depends on the probe + // declaring an order, so the order Spark reports for the outer join's streamed side + // still holds above them. + use datafusion::common::JoinType as DF; + use spark_operator::BuildSide::{BuildLeft, BuildRight}; + use spark_operator::JoinType::{Existence, Inner, LeftAnti, LeftSemi}; + for (join_type, build_side, planned_type) in [ + (Inner, BuildLeft, DF::Inner), + (Inner, BuildRight, DF::Inner), + (LeftSemi, BuildRight, DF::RightSemi), + (LeftAnti, BuildRight, DF::RightAnti), + (Existence, BuildRight, DF::RightMark), + ] { + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + BuildLeft, + vec![ + create_scan(), + hash_join_over_scans(join_type, build_side, false), + ], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!( + "SortedInputExec", + probe.name(), + "{join_type:?} {build_side:?}" + ); + let declared = planned_hash_join(probe.children()[0]); + assert_eq!(&planned_type, declared.join_type()); + } + } + + #[test] + fn hash_join_declares_no_order_over_an_operator_that_may_move_rows() { + use spark_operator::BuildSide::{BuildLeft, BuildRight}; + use spark_operator::JoinType::{FullOuter, LeftAnti, LeftSemi, RightOuter}; + for (probe_input, reason) in [ // Comet builds a null-aware anti join on Spark's streamed side, so its output is not // in the order Spark reports for it. - hash_join( - spark_operator::JoinType::LeftAnti, - spark_operator::BuildSide::BuildRight, - true, + ( + hash_join_over_scans(LeftAnti, BuildRight, true), + "a null-aware anti join", ), // A right join whose probe declares no order moves its unmatched rows, although // DataFusion reports it as keeping its probe order. - hash_join( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - false, + ( + hash_join_over_scans(RightOuter, BuildLeft, false), + "a right outer join", + ), + // The remaining types keep no probe order at all: a left semi join built on the left + // emits build rows, and a full outer join emits its unmatched build rows last. + ( + hash_join_over_scans(LeftSemi, BuildLeft, false), + "a left semi join", + ), + ( + hash_join_over_scans(FullOuter, BuildLeft, false), + "a full outer join", ), ] { - let operator = sorted_outer_hash_join( - spark_operator::JoinType::RightOuter, - spark_operator::BuildSide::BuildLeft, - vec![create_scan(), probe_input], - ); + let operator = + sorted_outer_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); let (_, _, planned) = PhysicalPlanner::default() .create_plan(&operator, &mut vec![], 1) .unwrap(); let probe = planned_hash_join(&planned.native_plan).right(); - assert_eq!("HashJoinExec", probe.name()); + assert_eq!("HashJoinExec", probe.name(), "{reason}"); } } 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 9efaeffcba5..a045474c7f0 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 @@ -40,9 +40,10 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanExecBase import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -2884,7 +2885,8 @@ trait CometHashJoin { // of such a join can reach this one through a cache, as a JVM input that the native planner // cannot look behind. Declaring the order there would let an operator above rely on it, an // aggregate in sorted mode for instance, which then splits groups. So nothing is declared - // when such a join feeds the streamed side with no sort in between. + // when such a join feeds the streamed side with no sort in between, or when a leaf reports + // an order that cannot be checked; see mayReportFalseOrder. // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix // of the order that Comet can express is enough. Nothing here can tell whether an operator @@ -2929,28 +2931,82 @@ trait CometHashJoin { } } - // Whether a Comet join that may move its streamed rows, while reporting Spark's order for - // them, feeds the order that `plan` reports with no sort in between. Whether an outer join - // declared the order on its probe input is up to the native planner, so every one built that - // way counts. A cache reports the order of the plan that wrote it, so the walk goes on into - // that plan. While CometExecRule converts a plan, a native plan's JVM inputs sit behind a - // CometScanWrapper, a leaf, so the walk looks through that too. + // Whether the order that `plan` reports may be false, because something that feeds it with no + // sort in between cannot be relied on to keep it. Comet's hash joins report the order of their + // streamed side, which Spark's own joins keep, but three Comet joins do not always: the + // null-aware anti join, which Comet builds on its streamed side; an outer join that runs as + // DataFusion's right outer join, which keeps its probe rows in place only if the planner + // declared an order on its probe input; and the nested loop join, which DataFusion runs build + // row by build row. So does a leaf that reports an order from state this walk cannot check, an + // RDD checkpointed from a plan that is gone for one. + // + // The walk follows the order, so a hash join, which passes on only its streamed side's, is + // walked through that side alone. A cache reports the order of the plan that wrote it, and the + // walk goes on into that plan, picked as Spark's scan picks it: an adaptive plan reports no + // order, but the scan reads the final plan behind it once there is one. While CometExecRule + // converts a plan, a native plan's JVM inputs sit behind a CometScanWrapper, a leaf, so the + // walk looks through that too. private def mayReportFalseOrder(plan: SparkPlan): Boolean = plan match { + case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => + mayReportFalseOrder(adaptive.executedPlan) case _ if plan.outputOrdering.isEmpty => false case _: SortExec | _: CometSortExec => false - case join: CometBroadcastHashJoinExec - if join.isNullAwareAntiJoin || runsAsRightOuterJoin(join.joinType, join.buildSide) => - true - case join: CometHashJoinExec if runsAsRightOuterJoin(join.joinType, join.buildSide) => true case _: CometBroadcastNestedLoopJoinExec => true + case join: CometBroadcastHashJoinExec if join.isNullAwareAntiJoin => true + case join: CometBroadcastHashJoinExec => + hashJoinMayReportFalseOrder( + join.nativeOp, + join.joinType, + join.buildSide, + join.left, + join.right) + case join: CometHashJoinExec => + hashJoinMayReportFalseOrder( + join.nativeOp, + join.joinType, + join.buildSide, + join.left, + join.right) + case join: HashJoin => + mayReportFalseOrder(if (join.buildSide == BuildRight) join.left else join.right) case wrapper: CometScanWrapper => mayReportFalseOrder(wrapper.originalPlan) case scan: CometInMemoryTableScanExec => mayReportFalseOrder(scan.originalPlan.relation.cachedPlan) case scan: InMemoryTableScanExec => mayReportFalseOrder(scan.relation.cachedPlan) case stage: QueryStageExec => mayReportFalseOrder(stage.plan) + // A scan reports an order that Spark already relies on. + case _: DataSourceScanExec | _: DataSourceV2ScanExecBase | _: RangeExec | _: CometLeafExec => + false + case _ if plan.children.isEmpty => true case _ => plan.children.exists(mayReportFalseOrder) } + // Whether a Comet hash join that reports an order may not keep it. Only the streamed side's + // order reaches the join's output, so a false order on the build side is of no concern, and an + // outer join of the two shapes has to have declared the order on its probe input as well. + private def hashJoinMayReportFalseOrder( + nativeOp: Operator, + joinType: JoinType, + buildSide: BuildSide, + left: SparkPlan, + right: SparkPlan): Boolean = { + val (streamed, streamedIndex) = if (buildSide == BuildRight) (left, 0) else (right, 1) + mayReportFalseOrder(streamed) || + (runsAsRightOuterJoin(joinType, buildSide) && !declaresStreamedOrder(nativeOp, streamedIndex)) + } + + // Whether the native planner declares the streamed order on the probe input of the join that + // `nativeOp` describes. It does so when the serde sent the order and the probe input reads from + // the JVM, which is the case when the join's plan has a scan there. A native operator over that + // scan is declared as well, as far as the native walk goes, but that is for the planner to say, + // so it counts as not declared. + private def declaresStreamedOrder(nativeOp: Operator, streamedIndex: Int): Boolean = + nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 && + nativeOp.getChildrenCount > streamedIndex && { + val probe = nativeOp.getChildren(streamedIndex) + probe.hasScan || probe.hasShuffleScan + } + // A left outer join built on the right and a right outer join built on the left both run as // DataFusion's right outer join, which keeps its unmatched probe rows in place only when the // probe input declares an order. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 1a4b35edf57..8378af44aec 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -38,6 +38,7 @@ import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.plans.LeftOuter import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} @@ -3326,14 +3327,6 @@ class CometInMemoryCacheSuite extends CometTestBase { case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide) } - def sortedPartitions(df: DataFrame): Array[Boolean] = - df.queryExecution.toRdd - .mapPartitions { it => - val keys = it.map(_.getInt(0)).toArray - Iterator(keys.sameElements(keys.sorted)) - } - .collect() - def sorts(plan: SparkPlan): Seq[SparkPlan] = collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s } @@ -3375,6 +3368,122 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + // Whether each partition of `df` has the integer key in its first column in ascending order. + private def sortedPartitions(df: DataFrame): Array[Boolean] = + df.queryExecution.toRdd + .mapPartitions { it => + val keys = it.map(_.getInt(0)).toArray + Iterator(keys.sameElements(keys.sorted)) + } + .collect() + + test("a broadcast outer hash join over a sorted cache keeps its order with default settings") { + // A cache written with no exchange in its plan is not adaptive, so with AQE on, as it is by + // default, the scan reports the cached plan's order just as it does with AQE off. Only the + // native cache scan is switched on here, and the broadcast join is the one Spark plans for + // a small build side. + withSQLConf(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = spark.table("big").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + val df = streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val expected = (0 until 10000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + // The plan AQE chose is final once the query has run. + val plan = df.queryExecution.executedPlan + assert(plan.isInstanceOf[AdaptiveSparkPlanExec], plan) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1 && joins.head.joinType == LeftOuter, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) + assert(sortedPartitions(df).forall(identity), plan) + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("an outer hash join over a native cache scan of a NOT IN keeps each group once") { + // The native cache scan reports the order of the plan that wrote the cache. A NOT IN is not in + // the order Spark reports for it, since Comet builds it on its streamed side, so the outer + // join that reads the cache must not declare the order, or the aggregates above it run in + // sorted mode over unsorted rows and split groups. Sparse keys and small batches make each + // group span several. A cache with an exchange in its plan is written under AQE, and the scan + // reads the final plan behind the adaptive one, so this runs with AQE off and on. + Seq(false, true).foreach { adaptive => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") + .cache() + .createOrReplaceTempView("anti") + spark.table("anti").count() + assert( + !sortedPartitions(spark.table("anti")).forall(identity), + "the cached NOT IN rows are in order, so nothing here can split a group") + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + |FROM anti a LEFT JOIN small b ON a.k = b._1 + |GROUP BY a.k""".stripMargin + def rows(): Seq[(Int, Long, Long, Long)] = + sql(query) + .collect() + .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) + .toSeq + .sorted + val cometRows = rows() + val plan = sql(query).queryExecution.executedPlan + assert( + collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, + s"the outer join does not read the native cache scan (AQE $adaptive):\n$plan") + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[(Int, Long, Long, Long)] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = rows() + } + assert( + cometRows.map(_._1).distinct.size == cometRows.size, + "a key is grouped twice") + assert(cometRows == sparkRows) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("anti") + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + } + test("an aggregate on the key above a hash join over a sorted cache groups each key once") { // The order the join declares on its probe input reaches past the join, so the aggregate // above it in the same native plan groups in sorted, streaming mode. That is right while 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 35dc61bee73..3fc5c62d516 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,10 +31,10 @@ 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.plans.{Inner, LeftOuter} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, RDDScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException @@ -2488,14 +2488,71 @@ class CometJoinSuite extends CometTestBase { // The same NOT IN as the test above, but cached, so its rows reach the outer join from the // JVM. The cache reports the order Spark reports for the null-aware anti join, which its rows // are not in, and the outer join's probe input is then a plain JVM input as far as the native - // plan can see. The order must not be declared there either. A cached plan that AQE runs - // reports no order, so AQE is off here, as Spark 3.4 has it for every cached plan by default. + // plan can see. The order must not be declared there either. A cache with an exchange in its + // plan is written under AQE from Spark 3.5, and Spark's scan then reports the order of the + // final plan behind the adaptive one, which reports none itself, so this runs with AQE off + // and on. + Seq(false, true).foreach { adaptive => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") + .cache() + .createOrReplaceTempView("anti") + spark.table("anti").count() + val cachedPlan = cachedPlanOf(spark.table("anti")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin + case _ => false + }.isDefined, + s"the cache was not written by a native NOT IN (AQE $adaptive):\n$cachedPlan") + assert( + !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), + "the cached NOT IN rows are in order, so nothing here can split a group") + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + |FROM anti a LEFT JOIN small b ON a.k = b._1 + |GROUP BY a.k""".stripMargin + checkEachKeyGroupedOnce(query) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("anti") + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + } + + test("an outer hash join over a checkpoint of a NOT IN keeps each group once") { + // A checkpoint is read by an RDD scan that reports the order of the plan it was taken from, + // as a cache does, but that plan is gone, so nothing can say whether the order holds. Here it + // is the order Spark reports for the null-aware anti join, which the rows are not in. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", SQLConf.COLUMN_BATCH_SIZE.key -> "8", CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { @@ -2509,19 +2566,16 @@ class CometJoinSuite extends CometTestBase { .createOrReplaceTempView("cached") spark.table("cached").count() sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") - .cache() + .localCheckpoint() .createOrReplaceTempView("anti") - spark.table("anti").count() - val cachedPlan = cachedPlanOf(spark.table("anti")) + val plan = spark.table("anti").queryExecution.executedPlan + val scans = collect(plan) { case scan: RDDScanExec => scan } assert( - cachedPlan.exists { - case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin - case _ => false - }, - s"the cache was not written by a native NOT IN:\n$cachedPlan") + scans.size == 1 && scans.head.outputOrdering.nonEmpty, + s"the checkpoint does not report the NOT IN's order:\n$plan") assert( !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), - "the cached NOT IN rows are in order, so nothing here can split a group") + "the checkpointed NOT IN rows are in order, so nothing here can split a group") val query = """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) |FROM anti a LEFT JOIN small b ON a.k = b._1 @@ -2538,11 +2592,12 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a cache of an undeclared outer join keeps each group once") { - // The IN subquery's semi join sits between the sorted cache and the first outer join, and the - // native planner declares no order over a semi join, so that outer join moves its unmatched - // rows, and the cache it writes is out of the order it reports. The second outer join reads - // that cache as a JVM input and must not declare the order either. + test("an outer hash join over a cache of an outer hash join with no declared order") { + // The planner declares the order on an outer join's probe input only through operators that + // keep their input's order, and the native sample does not report that it does. An outer + // join over a sample therefore moves its unmatched rows, although the serde sent its order, + // and the cache it writes is out of the order it reports. The second outer join reads that + // cache as a JVM input and must not declare the order either. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", @@ -2550,53 +2605,56 @@ class CometJoinSuite extends CometTestBase { CometConf.COMET_BATCH_SIZE.key -> "8", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w - |FROM (SELECT * FROM cached WHERE k IN (SELECT _1 FROM incl)) a - |LEFT JOIN small b ON a.k = b._1""".stripMargin) - .cache() - .createOrReplaceTempView("joined") - spark.table("joined").count() - val cachedPlan = cachedPlanOf(spark.table("joined")) - assert( - cachedPlan.exists { - case join: CometBroadcastHashJoinExec => join.joinType == LeftOuter - case _ => false - }, - s"the cache was not written by a native outer join:\n$cachedPlan") - assert( - !partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), - "the cached outer join rows are in order, so nothing here can split a group") - val query = - """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) - |FROM joined j LEFT JOIN small c ON j.k = c._1 - |GROUP BY j.k""".stripMargin - checkEachKeyGroupedOnce(query) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("joined") - spark.catalog.dropTempView("cached") - } + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM (SELECT * FROM cached TABLESAMPLE (90 PERCENT) REPEATABLE (7)) a + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && + join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 + case _ => false + }.isDefined, + s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") + assert( + find(cachedPlan)(_.isInstanceOf[CometSampleExec]).isDefined, + s"the outer join does not read a native sample:\n$cachedPlan") + assert( + !partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), + "the cached outer join rows are in order, so nothing here can split a group") + val query = + """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) + |FROM joined j LEFT JOIN small c ON j.k = c._1 + |GROUP BY j.k""".stripMargin + checkEachKeyGroupedOnce(query) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("joined") + spark.catalog.dropTempView("cached") } } } } } - test("an outer hash join over a cache sorted above an outer hash join declares the order") { - // The sort in the plan that wrote the cache puts its rows in the order the cache reports, - // whatever the outer join below the sort did, so the second outer join declares that order. - // The sort is on (k, v), so Spark does not drop it as one the first join already satisfies. + test("an outer hash join over a cache of a declared outer hash join declares the order too") { + // The first outer join's probe input is the sorted cache, a JVM input that the planner + // declares the order on, so the join keeps its unmatched rows in place and the cache it + // writes is in the order it reports. The second outer join can rely on that. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", @@ -2616,34 +2674,38 @@ class CometJoinSuite extends CometTestBase { spark.table("cached").count() sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w |FROM cached a LEFT JOIN small b ON a.k = b._1""".stripMargin) - .sortWithinPartitions("k", "v") .cache() - .createOrReplaceTempView("sorted_joined") - spark.table("sorted_joined").count() - val cachedPlan = cachedPlanOf(spark.table("sorted_joined")) + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) assert( - cachedPlan.exists { - case join: CometBroadcastHashJoinExec => join.joinType == LeftOuter + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && + join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 case _ => false - }, - s"the cache was not written by a native outer join:\n$cachedPlan") + }.isDefined, + s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") + assert( + partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), + "the cached outer join rows are out of order, so the first join did not keep them") val df = sql("""SELECT /*+ BROADCAST(c) */ j.k, j.v, c._2 - |FROM sorted_joined j LEFT JOIN small c ON j.k = c._1""".stripMargin) + |FROM joined j LEFT JOIN small c ON j.k = c._1""".stripMargin) .sortWithinPartitions("k") - val joins = df.queryExecution.executedPlan.collect { - case join: CometBroadcastHashJoinExec => join - } - assert(joins.size == 1, df.queryExecution.executedPlan) - assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0) - assert(partitionsInOrder(df, keyOrder(Ascending))) - var sparkRows = Seq.empty[Row] - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = df.collect().toSeq + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 2000).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) } - QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_)) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } finally { spark.catalog.clearCache() - spark.catalog.dropTempView("sorted_joined") + spark.catalog.dropTempView("joined") spark.catalog.dropTempView("cached") } } @@ -2651,11 +2713,10 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a cache of a nested loop join keeps each group once") { - // DataFusion's nested loop join emits each probe batch's matches build row by build row, so - // it keeps the order of neither input, while Spark reports the streamed side's order for an - // inner one. The cache it writes reports that order, and the outer join reading the cache - // must not declare it. + test("an outer hash join declares the streamed order over a semi join") { + // `k IN (SELECT ...)` runs as a left semi join, which emits its probe rows in ascending + // order whether or not the probe input declares an order. The outer join above it therefore + // declares the cache's order on the semi join's output, and keeps it. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", @@ -2663,7 +2724,7 @@ class CometJoinSuite extends CometTestBase { CometConf.COMET_BATCH_SIZE.key -> "8", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable(Seq(0, 1).map(Tuple1(_)), "two") { + withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { try { spark @@ -2674,25 +2735,35 @@ class CometJoinSuite extends CometTestBase { .cache() .createOrReplaceTempView("cached") spark.table("cached").count() - sql("""SELECT /*+ BROADCAST(t) */ a.k, a.v, t._1 AS x - |FROM cached a JOIN two t ON a.v % 2 <= t._1""".stripMargin) - .cache() - .createOrReplaceTempView("nlj") - spark.table("nlj").count() - val cachedPlan = cachedPlanOf(spark.table("nlj")) + val df = sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM (SELECT * FROM cached WHERE k IN (SELECT _1 FROM incl)) a + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + // Spark also pushes the IN to the other input of the outer join, so the outer join + // is found by its type and has to read the semi join over the cache. + val outer = collect(plan) { + case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join + } + assert(outer.size == 1, plan) assert( - cachedPlan.exists(_.isInstanceOf[CometBroadcastNestedLoopJoinExec]), - s"the cache was not written by a native nested loop join:\n$cachedPlan") + outer.head.left match { + case semi: CometBroadcastHashJoinExec => semi.joinType == LeftSemi + case _ => false + }, + plan) + assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) assert( - !partitionsInOrder(spark.table("nlj"), keyOrder(Ascending)), - "the cached nested loop join rows are in order, so nothing here can split a group") - checkEachKeyGroupedOnce( - """SELECT /*+ BROADCAST(c) */ n.k, count(*), sum(n.v), count(c._2) - |FROM nlj n LEFT JOIN small c ON n.k = c._1 - |GROUP BY n.k""".stripMargin) + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + val expected = (0 until 2000).filter(i => (i % 40) % 7 != 3).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } finally { spark.catalog.clearCache() - spark.catalog.dropTempView("nlj") spark.catalog.dropTempView("cached") } } @@ -2701,6 +2772,90 @@ class CometJoinSuite extends CometTestBase { } } + test("a NOT IN below the build side of an inner hash join leaves the streamed order declared") { + // An inner join passes on only its streamed side's order, so a NOT IN below its build side + // cannot affect it, and the outer join above the inner join still declares it. A build side + // behind an exchange reports no order, so it needs a build side that is a cache already + // partitioned on the key, where the shuffled hash join gets no exchange, and that reports the + // order of the NOT IN that wrote it. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40).map(k => (k * 1000, k)), "dim") { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + Seq("big" -> "cached", "dim" -> "dim_cached").foreach { case (table, view) => + spark + .table(table) + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView(view) + spark.table(view).count() + } + sql("SELECT * FROM dim_cached WHERE k NOT IN (SELECT _1 FROM excl)") + .cache() + .createOrReplaceTempView("anti") + spark.table("anti").count() + val cachedPlan = cachedPlanOf(spark.table("anti")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin + case _ => false + }.isDefined, + s"the cache was not written by a native NOT IN:\n$cachedPlan") + val query = + """SELECT /*+ SHUFFLE_HASH(b), BROADCAST(c) */ a.k, a.v, b.v AS x, c._2 AS w + |FROM cached a + |JOIN anti b ON a.k = b.k + |LEFT JOIN small c ON a.k = c._1""".stripMargin + val df = sql(query).sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val inner = collect(plan) { case join: CometHashJoinExec => join } + assert(inner.size == 1 && inner.head.joinType == Inner, plan) + assert(inner.head.buildSide == BuildRight, plan) + assert( + find(inner.head.right) { + case _: CometInMemoryTableScanExec | _: InMemoryTableScanExec => true + case _ => false + }.isDefined, + plan) + assert(collect(plan) { case e: ShuffleExchangeLike => e }.isEmpty, plan) + val outer = collect(plan) { + case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join + } + assert(outer.size == 1, plan) + assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + val expected = (0 until 2000) + .filter(i => !Set(1, 10, 20).contains(i % 40)) + .map { i => + val m = i % 40 + Row(m * 1000, i, m, if (m % 2 == 0) m else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } finally { + spark.catalog.clearCache() + spark.catalog.dropTempView("anti") + spark.catalog.dropTempView("dim_cached") + spark.catalog.dropTempView("cached") + } + } + } + } + } + } + } + // Runs `query`, which groups by the integer key in its first column and returns three longs, // with Comet and with Spark, and checks that Comet returns each key once and Spark's rows. private def checkEachKeyGroupedOnce(query: String): Unit = { @@ -2720,12 +2875,16 @@ class CometJoinSuite extends CometTestBase { assert(cometRows == sparkRows) } - // The physical plan that wrote the cache `df` reads. - private def cachedPlanOf(df: DataFrame): SparkPlan = - df.queryExecution.executedPlan - .collectFirst { - case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan - case scan: InMemoryTableScanExec => scan.relation.cachedPlan - } - .getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) + // The physical plan that wrote the cache `df` reads, as Spark's scan reports its order: for an + // adaptive cached plan that is the final plan behind it. + private def cachedPlanOf(df: DataFrame): SparkPlan = { + val cachedPlan = collectFirst(df.queryExecution.executedPlan) { + case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan + case scan: InMemoryTableScanExec => scan.relation.cachedPlan + }.getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) + cachedPlan match { + case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => adaptive.executedPlan + case other => other + } + } } From ce6a82fe932e960f10dee0db2c24c8e719d95cc5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 12 Oct 2026 14:10:21 -0600 Subject: [PATCH 06/19] docs: describe how native hash joins keep the streamed order Add a Hash Join Output Order section to the operator compatibility page: native hash joins keep the streamed side's order Spark relies on, an outer join whose first streamed sort key Comet cannot express falls back to Spark, and where the order might not hold Comet declares nothing. --- .../latest/compatibility/operators.md | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index afab52bfce8..41f41058424 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -92,6 +92,26 @@ 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 Output Order + +Spark reports the sort order of a hash join's streamed side as the join's own, and drops a sort +above the join that this order satisfies. Comet's native hash joins keep that order, including a +left or right outer join over a streamed side that was sorted outside the native plan, such as a +sorted relation read from the [in-memory cache](#in-memory-cache). + +Comet passes the order to the native join as sort keys, and an outer join needs at least the first +key to be one that Comet can express natively. Declaring any leading part of the order keeps the +rows in place, so the join does not need the rest. An outer join whose first sort key Comet cannot +express falls back to Spark and records a fallback reason (see +[Understanding Comet Plans](../understanding-comet-plans.md)), because Comet cannot tell whether +an operator above the join relies on the order. + +Where the order Spark reports might not hold, Comet declares nothing and an outer join can return +its unmatched rows out of order. The output of a `NOT IN` subquery or of a nested loop join +([#6898](https://github.com/apache/datafusion-comet/issues/6898)) is one such case, and so is a +cache or checkpoint that one of them wrote. Comet also declares nothing when it cannot follow the +order back to its source, as with an RDD checkpoint. + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of From a2ad967e98370c5ce773ef550fbe08b201b7cabf Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 09:05:10 -0600 Subject: [PATCH 07/19] fix: keep the streamed order to expressible keys and stop CometSampleExec reporting one Spark's SampleExec reports no output ordering, so CometSampleExec must not report its child's: a cache written by a join over a sample carried an order that nothing keeps. The streamed order now stops at a key whose type the native sort does not accept, as for a collated string, SortedInputExec keeps its input's equivalence properties, and the contributor guide says that an operator reporting an ordering must keep it natively. --- .../contributor-guide/adding_a_new_operator.md | 17 +++++++++++++++++ .../src/execution/operators/sorted_input.rs | 12 ++++++++---- native/core/src/execution/planner.rs | 9 +++++++++ .../spark/sql/comet/CometSampleExec.scala | 4 +--- .../org/apache/spark/sql/comet/operators.scala | 10 +++++++--- 5 files changed, 42 insertions(+), 10 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 85946f44fd9..f0371030526 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -357,6 +357,23 @@ case class CometProjectExec( } ``` +#### Reporting an Output Ordering + +An operator that overrides `outputOrdering`, as `CometFilterExec` does above with +`child.outputOrdering`, makes a claim that Spark acts on: it drops a sort above the operator that +the order satisfies. Report what the Spark operator reports and nothing more. `SampleExec`, for +example, reports no order, so `CometSampleExec` must not report its child's. + +The claim has to hold at run time as well. A native hash join declares the order Spark reports for +its streamed side to DataFusion, and an aggregate above the join may then run in sorted mode, or a +top-k sort may stop reading early, on the strength of that order. If the order is false, the +result is wrong groups or wrong rows, not only rows out of order. So an operator that reports an +ordering must keep it natively, and the native operator must say so by returning `true` from +`maintains_input_order` for that child, as DataFusion's filter and projection do. The planner +declares the order through an operator only when the operator says it keeps it. An operator that +cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in +`operators.scala`, which makes the hash join serde send no order over it. + #### Plan Identity and Exchange Reuse Spark's `ReuseExchangeAndSubquery` identifies equivalent plans through canonicalization and may diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index d3650cab03f..6411d579dd4 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -23,11 +23,11 @@ use std::sync::Arc; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{internal_err, Result}; use datafusion::execution::TaskContext; -use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; +use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::{ - apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - SendableRecordBatchStream, + apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, SendableRecordBatchStream, }; /// Passes `input` through unchanged and declares to DataFusion the order its rows arrive in. @@ -53,8 +53,12 @@ pub struct SortedInputExec { impl SortedInputExec { pub fn new(input: Arc, ordering: LexOrdering) -> Self { let input_properties = input.properties(); + // Start from the input's own properties, so the constants and equivalence classes it + // knows of stay visible to the operators above. + let mut equivalence_properties = input.equivalence_properties().clone(); + equivalence_properties.add_ordering(ordering.iter().cloned()); let cache = Arc::new(PlanProperties::new( - EquivalenceProperties::new_with_orderings(input.schema(), [ordering.clone()]), + equivalence_properties, input_properties.output_partitioning().clone(), input_properties.emission_type, input_properties.boundedness, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d1b905bacbd..9595ffb6770 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2955,6 +2955,15 @@ impl PhysicalPlanner { /// declares one. Spark's plan knows the order of an input that enters the native plan from /// the JVM, but the `ScanExec` it arrives through declares none. Only that order is declared: /// when an operator in between may have moved rows, `plan` is left as it is. + /// + /// The serde sends an order only when the streamed side's Comet operators report it, and it + /// checks that nothing under them moves rows (`mayReportFalseOrder`), while this walk asks + /// the native operators the same thing through `maintains_input_order`. The two agree as + /// long as a Comet operator that reports an order keeps it natively and says so. One that + /// does not lands here with `plan` unchanged, and the join then moves its unmatched rows + /// after Spark has dropped the sort above it. Nothing here can fall back at that point, so + /// the invariant is kept at the operator: see "Adding a New Operator" in the contributor + /// guide. fn with_declared_order( &self, plan: Arc, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala index dcdae2c779d..54391580a44 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet -import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.{SampleExec, SparkPlan} @@ -97,8 +97,6 @@ case class CometSampleExec( override def outputPartitioning: Partitioning = child.outputPartitioning - override def outputOrdering: Seq[SortOrder] = child.outputOrdering - override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = this.copy(child = newChild) 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 a045474c7f0..32375ebfece 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 @@ -64,7 +64,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedScalarSortElementType, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -2889,13 +2889,17 @@ trait CometHashJoin { // an order that cannot be checked; see mayReportFalseOrder. // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix - // of the order that Comet can express is enough. Nothing here can tell whether an operator - // above relies on the order, so a join whose first key Comet cannot express stays in Spark. + // of the order that Comet can express is enough. A key counts as expressible when the native + // sort accepts its type, as `supportedSortType` decides for a sort: a collated string is + // ordered by its collation in Spark but would be declared in raw byte order. Nothing here can + // tell whether an operator above relies on the order, so a join whose first key Comet cannot + // express stays in Spark. val streamed = if (join.buildSide == BuildRight) join.left else join.right val streamedSortOrders = (join.joinType, join.buildSide) match { case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) if join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed) => val orders = join.outputOrdering.iterator + .takeWhile(order => supportedScalarSortElementType(order.dataType)) .map(exprToProto(_, streamed.output)) .takeWhile(_.isDefined) .flatten From ddce25a6fe648e749a36aae55caec4813f8af99a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 09:05:10 -0600 Subject: [PATCH 08/19] test: cover the streamed order for decimal, string and timestamp keys Also assert that the aggregate test's join sent the order and emits its groups sorted, and replace the sample test with one for a cache that reports no order. --- .../comet/exec/CometInMemoryCacheSuite.scala | 9 ++ .../apache/comet/exec/CometJoinSuite.scala | 109 ++++++++++++++++-- 2 files changed, 107 insertions(+), 11 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 8378af44aec..6f67f0a7233 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3531,6 +3531,15 @@ class CometInMemoryCacheSuite extends CometTestBase { assert( joins.map { case (j, side) => (j.getClass, side) } == Seq((joinClass, buildSide)), plan) + // The serde sent the order, which is what lets the aggregates above run in sorted + // mode, and a sorted aggregate emits its groups in key order. + assert( + collect(plan) { + case j: CometHashJoinExec => j.nativeOp + case j: CometBroadcastHashJoinExec => j.nativeOp + }.forall(_.getHashJoin.getStreamedSortOrdersCount > 0), + plan) + assert(sortedPartitions(df).forall(identity), plan) // A final aggregate after a shuffle would merge a key a partial aggregate split, so // both aggregates have to sit above the join with no exchange in between. def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined 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 3fc5c62d516..e1289b7684d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -40,6 +40,7 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec 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.functions.{count, lit, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} @@ -2444,6 +2445,92 @@ class CometJoinSuite extends CometTestBase { } } + test("outer hash joins keep the streamed order for decimal, string and timestamp keys") { + // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde + // and the native sort expression handle the key, and the sorted aggregate above the join finds + // its group boundaries by comparing it. The timestamps run in a session time zone that is not + // UTC and include values before 1970. Struct, array and map keys are out of scope, since + // exprToProto rejects them by default and such a join stays in Spark. + val keyTypes = Seq( + "decimal(20,2)" -> Seq( + "-999999999999999999.99", + "-0.01", + "0.00", + "0.01", + "999999999999999999.99"), + "string" -> Seq("apple", "Äpfel", "naïve", "日本語", "中文", "한국어", "😀"), + "timestamp" -> Seq( + "1900-01-01 00:00:00", + "1969-12-31 23:59:59.999999", + "1970-01-01 00:00:00", + "2026-11-01 01:30:00", + "2200-06-15 12:00:00")) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + keyTypes.foreach { case (typeName, keys) => + val data = (0 until 2000).map(i => (keys(i % keys.size), i)) + val buildData = keys.zipWithIndex.collect { case (k, i) if i % 2 == 0 => (k, i) } + withParquetTable(data, "big") { + withParquetTable(buildData, "small") { + try { + val streamed = spark + .table("big") + .select($"_1".cast(typeName).as("_1"), $"_2") + .repartition(2, $"_1") + .sortWithinPartitions("_1") + .cache() + streamed.count() + val build = spark.table("small").select($"_1".cast(typeName).as("_1"), $"_2") + val order = SortOrder( + BoundReference(0, streamed.schema("_1").dataType, nullable = true), + Ascending) + + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("shuffle_hash") + if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + } + val df = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.size == 1, s"$typeName $joinType\n$plan") + assert( + joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, + s"$typeName $joinType\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$typeName $joinType\n$plan") + checkSparkAnswer(df) + assert(partitionsInOrder(df, order), s"$typeName $joinType\n$plan") + + // The aggregate runs in sorted mode over the declared order, and groups a key once + // only if the join keeps it. + val agg = joined + .groupBy(streamed("_1").as("k")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + checkSparkAnswer(agg) + assert(agg.count() == keys.size, s"$typeName $joinType groups a key twice") + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + } + test("NOT IN below an outer hash join over a sorted cache keeps each group once") { // Spark reports the streamed side's order for a null-aware anti join, but Comet builds that // join on its streamed side, so the order does not hold above it. Declaring it on the outer @@ -2592,12 +2679,11 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a cache of an outer hash join with no declared order") { - // The planner declares the order on an outer join's probe input only through operators that - // keep their input's order, and the native sample does not report that it does. An outer - // join over a sample therefore moves its unmatched rows, although the serde sent its order, - // and the cache it writes is out of the order it reports. The second outer join reads that - // cache as a JVM input and must not declare the order either. + test("a cache of an outer hash join over a native sample reports no order, as Spark does") { + // Spark's SampleExec reports no output ordering, so the sort above a join over a sample stays, + // and CometSampleExec reports none either. The native sample does not say it keeps its input's + // order, so the planner could not declare an order through it, and a cache that reported the + // order would carry a false claim to whatever reads it next. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", @@ -2625,17 +2711,18 @@ class CometJoinSuite extends CometTestBase { assert( find(cachedPlan) { case join: CometBroadcastHashJoinExec => - join.joinType == LeftOuter && - join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 + join.joinType == LeftOuter && join.nativeOp.getHashJoin.getStreamedSortOrdersCount == 0 case _ => false }.isDefined, - s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") + s"the cache was not written by a native outer join that sent no order:\n$cachedPlan") assert( find(cachedPlan)(_.isInstanceOf[CometSampleExec]).isDefined, s"the outer join does not read a native sample:\n$cachedPlan") assert( - !partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), - "the cached outer join rows are in order, so nothing here can split a group") + cachedPlan.outputOrdering.isEmpty, + s"the cache reports an order that nothing keeps:\n$cachedPlan") + // The sort above the cache stays, so the partitions come back sorted. + assert(partitionsInOrder(spark.table("joined").sortWithinPartitions("k"), keyOrder(Ascending))) val query = """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) |FROM joined j LEFT JOIN small c ON j.k = c._1 From 5adbeb4db085618747c8fc0404b8761fc300847b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 09:38:49 -0600 Subject: [PATCH 09/19] test: compare the key-type rows with Spark's as a multiset --- .../apache/comet/exec/CometJoinSuite.scala | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 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 e8645df8b6a..a6c30cd427e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -2475,6 +2475,17 @@ class CometJoinSuite extends CometTestBase { } } + // Compares the rows of a fresh `df` as a multiset with the ones Spark returns for it. The rows + // of a key tie under the sort, so their order within it differs between the two. + private def sameRowsAsSpark(df: => DataFrame): Unit = { + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = df.collect().toSeq + } + QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_)) + } + test("outer hash joins keep the streamed order for decimal, string and timestamp keys") { // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde // and the native sort expression handle the key, and the sorted aggregate above the join finds @@ -2494,7 +2505,7 @@ class CometJoinSuite extends CometTestBase { "1969-12-31 23:59:59.999999", "1970-01-01 00:00:00", "2026-11-01 01:30:00", - "2200-06-15 12:00:00")) + "2030-06-15 12:00:00")) withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -2529,9 +2540,10 @@ class CometJoinSuite extends CometTestBase { hinted.join(streamed, build("_1") === streamed("_1"), joinType) } } - val df = joined + def sorted: DataFrame = joined .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) .sortWithinPartitions("k") + val df = sorted val plan = df.queryExecution.executedPlan val joins = collect(plan) { case j: CometHashJoinExec => j } assert(joins.size == 1, s"$typeName $joinType\n$plan") @@ -2541,15 +2553,15 @@ class CometJoinSuite extends CometTestBase { assert( collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, s"$typeName $joinType\n$plan") - checkSparkAnswer(df) + sameRowsAsSpark(sorted) assert(partitionsInOrder(df, order), s"$typeName $joinType\n$plan") // The aggregate runs in sorted mode over the declared order, and groups a key once // only if the join keeps it. - val agg = joined + def agg: DataFrame = joined .groupBy(streamed("_1").as("k")) .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) - checkSparkAnswer(agg) + sameRowsAsSpark(agg) assert(agg.count() == keys.size, s"$typeName $joinType groups a key twice") } } finally { @@ -2752,7 +2764,10 @@ class CometJoinSuite extends CometTestBase { cachedPlan.outputOrdering.isEmpty, s"the cache reports an order that nothing keeps:\n$cachedPlan") // The sort above the cache stays, so the partitions come back sorted. - assert(partitionsInOrder(spark.table("joined").sortWithinPartitions("k"), keyOrder(Ascending))) + assert( + partitionsInOrder( + spark.table("joined").sortWithinPartitions("k"), + keyOrder(Ascending))) val query = """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) |FROM joined j LEFT JOIN small c ON j.k = c._1 From 8e751305434ce26d410b3c318f85c6892719eec7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 09:50:58 -0600 Subject: [PATCH 10/19] test: make the aggregate test fail when the join's order is not declared With 40 keys of 50 rows each, a join that moved its unmatched rows still fed the aggregate the keys in first-seen order, so the test passed without the declaration. Add a layout of 1000 keys with two rows each, where a batch holds alternately matched and unmatched keys, and assert that the serde sent the order and that each partition comes back in key order. --- .../comet/exec/CometInMemoryCacheSuite.scala | 125 ++++++++++-------- 1 file changed, 71 insertions(+), 54 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 23601cd8333..33de232ce46 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3488,7 +3488,10 @@ class CometInMemoryCacheSuite extends CometTestBase { // The order the join declares on its probe input reaches past the join, so the aggregate // above it in the same native plan groups in sorted, streaming mode. That is right while // Spark's reported order holds, and a wrong order would split a key across groups, which only - // a comparison of the rows as a multiset catches. Small batches make each group span several. + // a comparison of the rows as a multiset catches. With 40 keys and 8-row batches each group + // spans several batches. With 1000 keys a batch holds four keys, alternately matched and + // unmatched, so a join that moved its unmatched rows would also change the order in which the + // aggregate first sees the keys, and its output would not come back sorted. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -3496,62 +3499,76 @@ class CometInMemoryCacheSuite extends CometTestBase { SQLConf.COLUMN_BATCH_SIZE.key -> "8", CometConf.COMET_BATCH_SIZE.key -> "8", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40 by 2).map(m => (m * 1000, m)), "small") { - try { - val streamed = - spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() - streamed.count() - val build = spark.table("small") - - // Fifty rows per key, and a match in the build side for every even key. - val expected = (0 until 40).map { m => - Row(m * 1000, 50L, 50L * m + 49000L, if (m % 2 == 0) 50L else 0L) - } - - for { - (hint, joinClass) <- Seq( - "shuffle_hash" -> classOf[CometHashJoinExec], - "broadcast" -> classOf[CometBroadcastHashJoinExec]) - (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) - } { - val joined = if (joinType == "right_outer") { - build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) - } else { - streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + Seq(40, 1000).foreach { keys => + val data = (0 until 2000).map(i => ((i % keys) * 1000, i)) + withParquetTable(data, "big") { + withParquetTable((0 until keys by 2).map(m => (m * 1000, m)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + // A match in the build side for every even key. + val expected = data.groupBy(_._1).toSeq.map { case (k, rows) => + Row( + k, + rows.size.toLong, + rows.map(_._2.toLong).sum, + if ((k / 1000) % 2 == 0) { + rows.size.toLong + } else 0L) } - val df = joined - .groupBy(streamed("_1")) - .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) - val plan = df.queryExecution.executedPlan - val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) { - case j: CometHashJoinExec => (j, j.buildSide) - case j: CometBroadcastHashJoinExec => (j, j.buildSide) + + for { + (hint, joinClass) <- Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]) + (joinType, buildSide) <- Seq( + "left_outer" -> BuildRight, + "right_outer" -> BuildLeft) + } { + val joined = if (joinType == "right_outer") { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } else { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } + val df = joined + .groupBy(streamed("_1")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + val plan = df.queryExecution.executedPlan + val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j, j.buildSide) + case j: CometBroadcastHashJoinExec => (j, j.buildSide) + } + assert( + joins.map { case (j, side) => (j.getClass, side) } == Seq( + (joinClass, buildSide)), + plan) + // The serde sent the order, which is what lets the aggregates above run in + // sorted mode, and a sorted aggregate emits its groups in key order. + assert( + collect(plan) { + case j: CometHashJoinExec => j.nativeOp + case j: CometBroadcastHashJoinExec => j.nativeOp + }.forall(_.getHashJoin.getStreamedSortOrdersCount > 0), + plan) + assert(sortedPartitions(df).forall(identity), plan) + // A final aggregate after a shuffle would merge a key a partial aggregate split, + // so both aggregates have to sit above the join with no exchange in between. + def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined + assert( + collect(plan) { case a: CometHashAggregateExec => a }.count(aboveJoin) == 2, + plan) + assert( + !collect(plan) { case e: ShuffleExchangeLike => e }.exists(aboveJoin), + plan) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) } - assert( - joins.map { case (j, side) => (j.getClass, side) } == Seq((joinClass, buildSide)), - plan) - // The serde sent the order, which is what lets the aggregates above run in sorted - // mode, and a sorted aggregate emits its groups in key order. - assert( - collect(plan) { - case j: CometHashJoinExec => j.nativeOp - case j: CometBroadcastHashJoinExec => j.nativeOp - }.forall(_.getHashJoin.getStreamedSortOrdersCount > 0), - plan) - assert(sortedPartitions(df).forall(identity), plan) - // A final aggregate after a shuffle would merge a key a partial aggregate split, so - // both aggregates have to sit above the join with no exchange in between. - def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined - assert( - collect(plan) { case a: CometHashAggregateExec => a }.count(aboveJoin) == 2, - plan) - assert(!collect(plan) { case e: ShuffleExchangeLike => e }.exists(aboveJoin), plan) - assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + } finally { + spark.catalog.clearCache() } - } finally { - spark.catalog.clearCache() } } } From 65f406fe0fee685ae0af721aa715418ed294f149 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 10:41:45 -0600 Subject: [PATCH 11/19] fix: declare the order of a native range for the joins above it Spark's RangeExec reports id ascending, or descending for a negative step, and Spark drops a sort above an outer hash join on that basis. The native range is a leaf with no JVM input, so the planner declared no order on the join's probe input and the join moved its unmatched rows. LazyMemoryExec::add_ordering leaves the cached output ordering of the plan properties stale in DataFusion 55, so the planner wraps the range in a SortedInputExec instead, which now reports its input's metrics. --- .../src/execution/operators/sorted_input.rs | 7 ++++ native/core/src/execution/planner.rs | 13 +++++-- native/operators/src/lib.rs | 2 +- native/operators/src/range.rs | 34 +++++++++++++++++++ 4 files changed, 52 insertions(+), 4 deletions(-) diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index 6411d579dd4..a31053bd6e8 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -25,6 +25,7 @@ use datafusion::common::{internal_err, Result}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::MetricsSet; use datafusion::physical_plan::{ apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, @@ -101,6 +102,12 @@ impl ExecutionPlan for SortedInputExec { vec![true] } + /// The rows pass through, so the input's metrics are this operator's, and a range whose root + /// this wraps still reports its output rows. + fn metrics(&self) -> Option { + self.input.metrics() + } + fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ea7ca4fd05c..acdcb8ad21a 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -88,7 +88,7 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_operators::{ - range_exec, CometFilterExec, CometProjectionExec, ExpandExec, ExplodeExec, + range_exec, range_sort_expr, CometFilterExec, CometProjectionExec, ExpandExec, ExplodeExec, PartitionedRankLimitExec, SampleExec, WindowFnKind, }; use datafusion_comet_spark_expr::{ @@ -2065,7 +2065,7 @@ impl PhysicalPlanner { } OpStruct::RangeScan(range) => { // A leaf with no JVM input: each task produces its own partition of the range. - let range: Arc = Arc::new(range_exec( + let source: Arc = Arc::new(range_exec( range.start, range.step, range.num_elements, @@ -2073,10 +2073,17 @@ impl PhysicalPlanner { self.partition, self.session_ctx.copied_config().batch_size(), )?); + // The range's order is real, and Spark's `RangeExec` reports it, so declare it + // for the operators above, a hash join in particular: see `SortedInputExec`. + let ordering = LexOrdering::new([range_sort_expr(range.step)]); + let source: Arc = match ordering { + Some(ordering) => Arc::new(SortedInputExec::new(source, ordering)), + None => source, + }; Ok(( vec![], vec![], - Arc::new(SparkPlan::new(spark_plan.plan_id, range, vec![])), + Arc::new(SparkPlan::new(spark_plan.plan_id, source, vec![])), )) } OpStruct::Sample(sample) => { diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs index b5f762025ad..97c80a97895 100644 --- a/native/operators/src/lib.rs +++ b/native/operators/src/lib.rs @@ -37,6 +37,6 @@ pub use expand::ExpandExec; pub use explode::ExplodeExec; pub use filter::CometFilterExec; pub use projection::CometProjectionExec; -pub use range::range_exec; +pub use range::{range_exec, range_sort_expr}; pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind}; pub use sample::SampleExec; diff --git a/native/operators/src/range.rs b/native/operators/src/range.rs index 3766c82b1a6..4ba4f6d6429 100644 --- a/native/operators/src/range.rs +++ b/native/operators/src/range.rs @@ -16,8 +16,11 @@ // under the License. use arrow::array::{Int64Array, RecordBatch}; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::common::Result; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec}; use datafusion::physical_plan::Partitioning; use parking_lot::RwLock; @@ -58,6 +61,23 @@ pub fn range_exec( Ok(exec) } +/// The order a range's values come out in. They are `partition_start + i * step` and never +/// overflow, so they ascend for a positive step and descend for a negative one, which is the order +/// Spark's `RangeExec` reports. +/// +/// `LazyMemoryExec::add_ordering` does not refresh the plan properties' cached output ordering in +/// DataFusion 55, so `output_ordering()` would still be `None` after it. The planner declares this +/// order by wrapping the range in a `SortedInputExec` instead. +pub fn range_sort_expr(step: i64) -> PhysicalSortExpr { + PhysicalSortExpr::new( + Arc::new(Column::new("id", 0)), + SortOptions { + descending: step < 0, + nulls_first: step > 0, + }, + ) +} + /// The first value of partition `partition` and the number of values in it, computed the way /// Spark's `RangeExec` does (`initRange` in its generated code): in `BigInteger` arithmetic, with /// the partition's start and end clamped to the `Long` range. @@ -247,6 +267,20 @@ mod tests { assert_eq!(collect(i64::MIN, step, n, 1, 1000).await, expected); } + #[tokio::test] + async fn test_sort_expr_matches_the_values() { + for step in [1i64, 3, -1, -7] { + let n = num_elements(0, 100 * step, step); + let values = collect(0, step, n, 2, 8).await; + let expr = range_sort_expr(step); + assert_eq!(expr.options.descending, step < 0, "step {step}"); + let mut sorted = values.clone(); + sorted.sort_by_key(|v| if expr.options.descending { -v } else { *v }); + // Both partitions are in order, and the partitions follow each other. + assert_eq!(values, sorted, "step {step}"); + } + } + #[test] fn test_partition_bounds() { // range(0, 10, 1, 3): Spark puts 3, 3 and 4 values in the three partitions. From 9c26cd8990fa227eedd950f1f20415e8681f54fb Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 10:41:45 -0600 Subject: [PATCH 12/19] refactor: name the streamed side of a hash join once --- .../apache/spark/sql/comet/operators.scala | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 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 39c9c490ed7..62d05b765bb 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 @@ -2915,10 +2915,10 @@ trait CometHashJoin { // ordered by its collation in Spark but would be declared in raw byte order. Nothing here can // tell whether an operator above relies on the order, so a join whose first key Comet cannot // express stays in Spark. - val streamed = if (join.buildSide == BuildRight) join.left else join.right - val streamedSortOrders = (join.joinType, join.buildSide) match { - case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) - if join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed) => + val (streamed, _) = streamedSide(join.buildSide, join.left, join.right) + val streamedSortOrders = + if (runsAsRightOuterJoin(join.joinType, join.buildSide) && + join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed)) { val orders = join.outputOrdering.iterator .takeWhile(order => supportedScalarSortElementType(order.dataType)) .map(exprToProto(_, streamed.output)) @@ -2933,8 +2933,9 @@ trait CometHashJoin { return None } orders - case _ => Nil - } + } else { + Nil + } if (leftKeys.forall(_.isDefined) && rightKeys.forall(_.isDefined) && @@ -2993,7 +2994,7 @@ trait CometHashJoin { join.left, join.right) case join: HashJoin => - mayReportFalseOrder(if (join.buildSide == BuildRight) join.left else join.right) + mayReportFalseOrder(streamedSide(join.buildSide, join.left, join.right)._1) case wrapper: CometScanWrapper => mayReportFalseOrder(wrapper.originalPlan) case scan: CometInMemoryTableScanExec => mayReportFalseOrder(scan.originalPlan.relation.cachedPlan) @@ -3015,7 +3016,7 @@ trait CometHashJoin { buildSide: BuildSide, left: SparkPlan, right: SparkPlan): Boolean = { - val (streamed, streamedIndex) = if (buildSide == BuildRight) (left, 0) else (right, 1) + val (streamed, streamedIndex) = streamedSide(buildSide, left, right) mayReportFalseOrder(streamed) || (runsAsRightOuterJoin(joinType, buildSide) && !declaresStreamedOrder(nativeOp, streamedIndex)) } @@ -3032,6 +3033,13 @@ trait CometHashJoin { probe.hasScan || probe.hasShuffleScan } + // The streamed side of a hash join, which is the side it is not built on, and its child index. + private def streamedSide( + buildSide: BuildSide, + left: SparkPlan, + right: SparkPlan): (SparkPlan, Int) = + if (buildSide == BuildRight) (left, 0) else (right, 1) + // A left outer join built on the right and a right outer join built on the left both run as // DataFusion's right outer join, which keeps its unmatched probe rows in place only when the // probe input declares an order. From a4eb244de0b175ad95d213cfe5f71da8a63d78e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 10:41:45 -0600 Subject: [PATCH 13/19] test: cover a native range and share the cache test helpers Move partitionsInOrder, keyOrder, checkEachKeyGroupedOnce and cachedPlanOf to CometTestBase so both suites use one copy, and set up the sorted cache the join tests share once. --- .../comet/exec/CometInMemoryCacheSuite.scala | 47 +-- .../apache/comet/exec/CometJoinSuite.scala | 385 ++++++++---------- .../org/apache/spark/sql/CometTestBase.scala | 54 ++- 3 files changed, 227 insertions(+), 259 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 33de232ce46..4b43eda5d7b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -36,7 +36,7 @@ import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} -import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.LeftOuter import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} @@ -3350,7 +3350,7 @@ class CometInMemoryCacheSuite extends CometTestBase { }.size == 1) assert(sorts(outerPlan).isEmpty, outerPlan) QueryTest.sameRows(outerRows, outer.collect().toSeq).foreach(fail(_)) - assert(sortedPartitions(outer).forall(identity), outerPlan) + assert(partitionsInOrder(outer, keyOrder(Ascending)), outerPlan) } val inner = query(hint, "inner") @@ -3358,7 +3358,7 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(joins(innerPlan) == Seq((joinClass, BuildRight)), innerPlan) assert(sorts(innerPlan).isEmpty, innerPlan) QueryTest.sameRows(innerRows, inner.collect().toSeq).foreach(fail(_)) - assert(sortedPartitions(inner).forall(identity), innerPlan) + assert(partitionsInOrder(inner, keyOrder(Ascending)), innerPlan) } } finally { spark.catalog.clearCache() @@ -3368,15 +3368,6 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - // Whether each partition of `df` has the integer key in its first column in ascending order. - private def sortedPartitions(df: DataFrame): Array[Boolean] = - df.queryExecution.toRdd - .mapPartitions { it => - val keys = it.map(_.getInt(0)).toArray - Iterator(keys.sameElements(keys.sorted)) - } - .collect() - test("a broadcast outer hash join over a sorted cache keeps its order with default settings") { // A cache written with no exchange in its plan is not adaptive, so with AQE on, as it is by // default, the scan reports the cached plan's order just as it does with AQE off. Only the @@ -3406,7 +3397,7 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) - assert(sortedPartitions(df).forall(identity), plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } finally { spark.catalog.clearCache() } @@ -3446,32 +3437,18 @@ class CometInMemoryCacheSuite extends CometTestBase { .createOrReplaceTempView("anti") spark.table("anti").count() assert( - !sortedPartitions(spark.table("anti")).forall(identity), + !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), "the cached NOT IN rows are in order, so nothing here can split a group") val query = """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) |FROM anti a LEFT JOIN small b ON a.k = b._1 |GROUP BY a.k""".stripMargin - def rows(): Seq[(Int, Long, Long, Long)] = - sql(query) - .collect() - .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) - .toSeq - .sorted - val cometRows = rows() - val plan = sql(query).queryExecution.executedPlan - assert( - collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, - s"the outer join does not read the native cache scan (AQE $adaptive):\n$plan") - // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. - var sparkRows = Seq.empty[(Int, Long, Long, Long)] - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = rows() - } - assert( - cometRows.map(_._1).distinct.size == cometRows.size, - "a key is grouped twice") - assert(cometRows == sparkRows) + checkEachKeyGroupedOnce( + query, + plan => + assert( + collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, + s"the outer join does not read the native cache scan (AQE $adaptive):\n$plan")) } finally { spark.catalog.clearCache() spark.catalog.dropTempView("anti") @@ -3553,7 +3530,7 @@ class CometInMemoryCacheSuite extends CometTestBase { case j: CometBroadcastHashJoinExec => j.nativeOp }.forall(_.getHashJoin.getStreamedSortOrdersCount > 0), plan) - assert(sortedPartitions(df).forall(identity), plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) // A final aggregate after a shuffle would merge a key a partial aggregate split, // so both aggregates have to sit above the join with no exchange in between. def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined 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 a6c30cd427e..54af1ef7764 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeRefer import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, RDDScanExec, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec @@ -44,7 +44,7 @@ import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExcha import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{count, lit, sum} 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, IntegerType, LongType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} @@ -2287,18 +2287,6 @@ class CometJoinSuite extends CometTestBase { } } - // Whether every partition of `df` is in `order`, which is bound to its output. - private def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = { - val ordering = new InterpretedOrdering(order) - df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part => - part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) } - } - } - - // The order of the integer key in the first output column. - private def keyOrder(direction: SortDirection): SortOrder = - SortOrder(BoundReference(0, IntegerType, nullable = true), direction) - test("outer hash joins keep a streamed side sorted outside the native plan in order") { // Spark drops the sort above the join because the join reports the streamed side's order, // and the streamed side arrives sorted through a JVM input, which declares no order to @@ -2486,6 +2474,53 @@ class CometJoinSuite extends CometTestBase { QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_)) } + test("outer hash joins keep the order of a native range") { + // Spark's RangeExec reports id ascending, or descending for a negative step, and the sort above + // the join is dropped on that basis. The range is a native leaf, not a JVM input, so the + // planner has to declare its order itself or the join moves its unmatched rows. A broadcast + // join, since a shuffled one would put an exchange between the range and the join. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + withParquetTable((0 until 40).map(k => (k.toLong, k)), "small") { + val build = spark.table("small") + for { + (start, end, step) <- Seq((0L, 2000L, 1L), (1999L, -1L, -1L)) + joinType <- Seq("left_outer", "right_outer") + } { + val streamed = spark.range(start, end, step, 1) + val joined = if (joinType == "left_outer") { + streamed.join(build.hint("broadcast"), streamed("id") === build("_1"), joinType) + } else { + build.hint("broadcast").join(streamed, streamed("id") === build("_1"), joinType) + } + val direction = if (step > 0) Ascending else Descending + val df = joined + .select(streamed("id").as("k"), build("_2").as("w")) + .sortWithinPartitions(if (step > 0) $"k" else $"k".desc) + val plan = df.queryExecution.executedPlan + val context = s"$joinType step $step\n$plan" + assert(collect(plan) { case r: CometRangeExec => r }.size == 1, context) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + val expected = + (start until end by step).map(id => Row(id, if (id < 40) id.toInt else null)) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert( + partitionsInOrder( + df, + SortOrder(BoundReference(0, LongType, nullable = true), direction)), + context) + } + } + } + } + test("outer hash joins keep the streamed order for decimal, string and timestamp keys") { // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde // and the native sort expression handle the key, and the sorted aggregate above the join finds @@ -2573,12 +2608,11 @@ class CometJoinSuite extends CometTestBase { } } - test("NOT IN below an outer hash join over a sorted cache keeps each group once") { - // Spark reports the streamed side's order for a null-aware anti join, but Comet builds that - // join on its streamed side, so the order does not hold above it. Declaring it on the outer - // join's probe input would put the aggregates in sorted mode over unsorted rows and split - // groups. Sparse keys keep DataFusion off its perfect hash build path, which keeps the build - // batches in order and would hide that, and small batches spread each partition over many. + // The setup several cache tests share: `big(_1, _2)` with 2000 rows over 40 sparse keys and + // `small` with a row for every even key, and `big` as `cached(k, v)`, partitioned on the key, + // sorted within each partition and cached, with 8-row batches and the native cache on. `views` + // names the temporary views that `f` creates, which are dropped afterwards. + private def withSortedCache(views: String*)(f: => Unit): Unit = { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", @@ -2586,29 +2620,40 @@ class CometJoinSuite extends CometTestBase { CometConf.COMET_BATCH_SIZE.key -> "8", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - val query = - """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + f + } finally { + spark.catalog.clearCache() + (views :+ "cached").foreach(spark.catalog.dropTempView) + } + } + } + } + } + + test("NOT IN below an outer hash join over a sorted cache keeps each group once") { + // Spark reports the streamed side's order for a null-aware anti join, but Comet builds that + // join on its streamed side, so the order does not hold above it. Declaring it on the outer + // join's probe input would put the aggregates in sorted mode over unsorted rows and split + // groups. Sparse keys keep DataFusion off its perfect hash build path, which keeps the build + // batches in order and would hide that, and small batches spread each partition over many. + withSortedCache() { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) |FROM (SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)) a |LEFT JOIN small b ON a.k = b._1 |GROUP BY a.k""".stripMargin - checkEachKeyGroupedOnce(query) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("cached") - } - } - } + checkEachKeyGroupedOnce(query) } } } @@ -2726,60 +2771,35 @@ class CometJoinSuite extends CometTestBase { // and CometSampleExec reports none either. The native sample does not say it keeps its input's // order, so the planner could not declare an order through it, and a cache that reported the // order would carry a false claim to whatever reads it next. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + withSortedCache("joined") { + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w |FROM (SELECT * FROM cached TABLESAMPLE (90 PERCENT) REPEATABLE (7)) a |LEFT JOIN small b ON a.k = b._1""".stripMargin) - .cache() - .createOrReplaceTempView("joined") - spark.table("joined").count() - val cachedPlan = cachedPlanOf(spark.table("joined")) - assert( - find(cachedPlan) { - case join: CometBroadcastHashJoinExec => - join.joinType == LeftOuter && join.nativeOp.getHashJoin.getStreamedSortOrdersCount == 0 - case _ => false - }.isDefined, - s"the cache was not written by a native outer join that sent no order:\n$cachedPlan") - assert( - find(cachedPlan)(_.isInstanceOf[CometSampleExec]).isDefined, - s"the outer join does not read a native sample:\n$cachedPlan") - assert( - cachedPlan.outputOrdering.isEmpty, - s"the cache reports an order that nothing keeps:\n$cachedPlan") - // The sort above the cache stays, so the partitions come back sorted. - assert( - partitionsInOrder( - spark.table("joined").sortWithinPartitions("k"), - keyOrder(Ascending))) - val query = - """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && join.nativeOp.getHashJoin.getStreamedSortOrdersCount == 0 + case _ => false + }.isDefined, + s"the cache was not written by a native outer join that sent no order:\n$cachedPlan") + assert( + find(cachedPlan)(_.isInstanceOf[CometSampleExec]).isDefined, + s"the outer join does not read a native sample:\n$cachedPlan") + assert( + cachedPlan.outputOrdering.isEmpty, + s"the cache reports an order that nothing keeps:\n$cachedPlan") + // The sort above the cache stays, so the partitions come back sorted. + assert( + partitionsInOrder(spark.table("joined").sortWithinPartitions("k"), keyOrder(Ascending))) + val query = + """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) |FROM joined j LEFT JOIN small c ON j.k = c._1 |GROUP BY j.k""".stripMargin - checkEachKeyGroupedOnce(query) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("joined") - spark.catalog.dropTempView("cached") - } - } - } + checkEachKeyGroupedOnce(query) } } @@ -2787,61 +2807,38 @@ class CometJoinSuite extends CometTestBase { // The first outer join's probe input is the sorted cache, a JVM input that the planner // declares the order on, so the join keeps its unmatched rows in place and the cache it // writes is in the order it reports. The second outer join can rely on that. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + withSortedCache("joined") { + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w |FROM cached a LEFT JOIN small b ON a.k = b._1""".stripMargin) - .cache() - .createOrReplaceTempView("joined") - spark.table("joined").count() - val cachedPlan = cachedPlanOf(spark.table("joined")) - assert( - find(cachedPlan) { - case join: CometBroadcastHashJoinExec => - join.joinType == LeftOuter && - join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 - case _ => false - }.isDefined, - s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") - assert( - partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), - "the cached outer join rows are out of order, so the first join did not keep them") - val df = sql("""SELECT /*+ BROADCAST(c) */ j.k, j.v, c._2 + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && + join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 + case _ => false + }.isDefined, + s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") + assert( + partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), + "the cached outer join rows are out of order, so the first join did not keep them") + val df = sql("""SELECT /*+ BROADCAST(c) */ j.k, j.v, c._2 |FROM joined j LEFT JOIN small c ON j.k = c._1""".stripMargin) - .sortWithinPartitions("k") - val plan = df.queryExecution.executedPlan - val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } - assert(joins.size == 1, plan) - assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) - val expected = (0 until 2000).map { i => - val m = i % 40 - Row(m * 1000, i, if (m % 2 == 0) m else null) - } - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) - assert(partitionsInOrder(df, keyOrder(Ascending)), plan) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("joined") - spark.catalog.dropTempView("cached") - } - } + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 2000).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } } @@ -2849,57 +2846,33 @@ class CometJoinSuite extends CometTestBase { // `k IN (SELECT ...)` runs as a left semi join, which emits its probe rows in ascending // order whether or not the probe input declares an order. The outer join above it therefore // declares the cache's order on the semi join's output, and keeps it. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - val df = sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + withSortedCache() { + withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { + val df = sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w |FROM (SELECT * FROM cached WHERE k IN (SELECT _1 FROM incl)) a |LEFT JOIN small b ON a.k = b._1""".stripMargin) - .sortWithinPartitions("k") - val plan = df.queryExecution.executedPlan - // Spark also pushes the IN to the other input of the outer join, so the outer join - // is found by its type and has to read the semi join over the cache. - val outer = collect(plan) { - case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join - } - assert(outer.size == 1, plan) - assert( - outer.head.left match { - case semi: CometBroadcastHashJoinExec => semi.joinType == LeftSemi - case _ => false - }, - plan) - assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) - assert( - collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, - plan) - val expected = (0 until 2000).filter(i => (i % 40) % 7 != 3).map { i => - val m = i % 40 - Row(m * 1000, i, if (m % 2 == 0) m else null) - } - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) - assert(partitionsInOrder(df, keyOrder(Ascending)), plan) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("cached") - } - } + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + // Spark also pushes the IN to the other input of the outer join, so the outer join + // is found by its type and has to read the semi join over the cache. + val outer = collect(plan) { + case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join + } + assert(outer.size == 1, plan) + assert( + outer.head.left match { + case semi: CometBroadcastHashJoinExec => semi.joinType == LeftSemi + case _ => false + }, + plan) + assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 2000).filter(i => (i % 40) % 7 != 3).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } } } @@ -2987,36 +2960,4 @@ class CometJoinSuite extends CometTestBase { } } } - - // Runs `query`, which groups by the integer key in its first column and returns three longs, - // with Comet and with Spark, and checks that Comet returns each key once and Spark's rows. - private def checkEachKeyGroupedOnce(query: String): Unit = { - def rows(): Seq[(Int, Long, Long, Long)] = - sql(query) - .collect() - .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) - .toSeq - .sorted - val cometRows = rows() - // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. - var sparkRows = Seq.empty[(Int, Long, Long, Long)] - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - sparkRows = rows() - } - assert(cometRows.map(_._1).distinct.size == cometRows.size, "a key is grouped twice") - assert(cometRows == sparkRows) - } - - // The physical plan that wrote the cache `df` reads, as Spark's scan reports its order: for an - // adaptive cached plan that is the final plan behind it. - private def cachedPlanOf(df: DataFrame): SparkPlan = { - val cachedPlan = collectFirst(df.queryExecution.executedPlan) { - case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan - case scan: InMemoryTableScanExec => scan.relation.cachedPlan - }.getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) - cachedPlan match { - case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => adaptive.executedPlan - case other => other - } - } } diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 7fc3e66c169..4c21ee394a5 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -38,17 +38,19 @@ import org.apache.parquet.schema.{MessageType, MessageTypeParser} import org.apache.spark._ import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER} import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.{BoundReference, InterpretedOrdering, SortDirection, SortOrder} import org.apache.spark.sql.catalyst.plans.logical import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.CometPlanChecker +import org.apache.spark.sql.comet.{CometInMemoryTableScanExec, CometPlanChecker} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.internal._ import org.apache.spark.sql.sources.{BaseRelation, TableScan} import org.apache.spark.sql.test._ -import org.apache.spark.sql.types.{DecimalType, StructType} +import org.apache.spark.sql.types.{DecimalType, IntegerType, StructType} import org.apache.comet._ import org.apache.comet.shims.ShimCometSparkSessionExtensions @@ -1634,4 +1636,52 @@ abstract class CometTestBase (expected.failed.toOption, actual.failed.toOption) } } + + // Whether every partition of `df` is in `order`, which is bound to its output. + protected def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = { + val ordering = new InterpretedOrdering(order) + df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part => + part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) } + } + } + + // The order of the integer key in the first output column. + protected def keyOrder(direction: SortDirection): SortOrder = + SortOrder(BoundReference(0, IntegerType, nullable = true), direction) + + // Runs `query`, which groups by the integer key in its first column and returns three longs, + // with Comet and with Spark, and checks that Comet returns each key once and Spark's rows. + // `checkPlan` sees the plan Comet ran. + protected def checkEachKeyGroupedOnce( + query: String, + checkPlan: SparkPlan => Unit = _ => ()): Unit = { + def rows(): Seq[(Int, Long, Long, Long)] = + sql(query) + .collect() + .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) + .toSeq + .sorted + val cometRows = rows() + checkPlan(sql(query).queryExecution.executedPlan) + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[(Int, Long, Long, Long)] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = rows() + } + assert(cometRows.map(_._1).distinct.size == cometRows.size, "a key is grouped twice") + assert(cometRows == sparkRows) + } + + // The physical plan that wrote the cache `df` reads, as Spark's scan reports its order: for an + // adaptive cached plan that is the final plan behind it. + protected def cachedPlanOf(df: DataFrame): SparkPlan = { + val cachedPlan = collectFirst(df.queryExecution.executedPlan) { + case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan + case scan: InMemoryTableScanExec => scan.relation.cachedPlan + }.getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) + cachedPlan match { + case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => adaptive.executedPlan + case other => other + } + } } From 3533ca6c47b7cd3f0a68c82ed22c50f372661bd7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 15:50:45 -0600 Subject: [PATCH 14/19] test: drop the imports the shared test helpers left unused Moving partitionsInOrder and keyOrder to CometTestBase left InterpretedOrdering and SortDirection unused in CometJoinSuite, which scalafix and the strict Scala warnings build reject. --- spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 54af1ef7764..a0580c11570 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.SparkException import org.apache.spark.sql.{Column, CometTestBase, DataFrame, QueryTest, 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.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, IsNotNull, SortOrder} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join From c52fbb863b87c6e684c21b0a49e9094d96681ea1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 15:50:45 -0600 Subject: [PATCH 15/19] test: cover a streamed double order with signed zeros and NaN An outer hash join over a streamed side sorted by a double key with -0.0, 0.0, NaN and the infinities keeps its order, groups each key once, and returns Spark's rows for a top-k above it. The planner declares the order on the key normalized as the native sort normalizes it, and the top-k runs in a native plan of its own over the join's output. --- .../apache/comet/exec/CometJoinSuite.scala | 84 +++++++++++++++++-- 1 file changed, 79 insertions(+), 5 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 a0580c11570..f00742330c5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -44,7 +44,7 @@ import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExcha import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{count, lit, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, IntegerType, LongType, MetadataBuilder, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} @@ -2463,15 +2463,16 @@ class CometJoinSuite extends CometTestBase { } } - // Compares the rows of a fresh `df` as a multiset with the ones Spark returns for it. The rows - // of a key tie under the sort, so their order within it differs between the two. - private def sameRowsAsSpark(df: => DataFrame): Unit = { + // Compares the rows of a fresh `df` with the ones Spark returns for it, as a multiset unless + // `isSorted`. The rows of a key tie under a sort, so their order within it can differ between + // the two. + private def sameRowsAsSpark(df: => DataFrame, isSorted: Boolean = false): Unit = { // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. var sparkRows = Seq.empty[Row] withSQLConf(CometConf.COMET_ENABLED.key -> "false") { sparkRows = df.collect().toSeq } - QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_)) + QueryTest.sameRows(sparkRows, df.collect().toSeq, isSorted).foreach(fail(_)) } test("outer hash joins keep the order of a native range") { @@ -2608,6 +2609,79 @@ class CometJoinSuite extends CometTestBase { } } + test("outer hash joins keep a streamed double order with signed zeros and NaN") { + // Spark sorts -0.0 and 0.0 as equal and NaN above every other value. The planner declares the + // order on the key normalized as the native sort normalizes it, so an operator above the join + // compares keys as Spark does. A top-k above the join runs in a native plan of its own, over + // the join's output, so it cannot stop early on the declared order, and one sorting by the key + // alone reads only the first rows of each partition, which holds only if the join keeps them in + // order. A broadcast join, since Spark joins on the normalized key and a shuffled join would + // shuffle the streamed side again. + val keys = + Seq(Double.NegativeInfinity, -1.5, -0.0, 0.0, 1.5, Double.PositiveInfinity, Double.NaN) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + // The build side holds -0.0 and not 0.0, so a streamed 0.0 matches only on the normalized + // key. + val buildData = Seq((Double.NegativeInfinity, 0), (-0.0, 1), (1.5, 2), (Double.NaN, 3)) + withParquetTable((0 until 2000).map(i => (keys(i % keys.size), i)), "big") { + withParquetTable(buildData, "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_2").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + val order = SortOrder(BoundReference(0, DoubleType, nullable = true), Ascending) + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("broadcast") + if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + } + def rows: DataFrame = + joined.select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + val df = rows.sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(rows.sortWithinPartitions("k")) + assert(partitionsInOrder(df, order), context) + + def agg: DataFrame = joined + .groupBy(streamed("_1").as("k")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + sameRowsAsSpark(agg) + val groupKeys = agg.collect().map(_.getDouble(0)).toSeq + assert(groupKeys.distinct.size == groupKeys.size, s"$joinType groups a key twice") + + // The first 600 rows by the key end among the zeros, which tie, so which of them + // are taken depends on v alone. The first 300 descending take every NaN, which tie + // too, and the rest from the infinities. A limit on the key alone ends where the + // zeros start, since which of them it took would be arbitrary. + sameRowsAsSpark(rows.orderBy("k", "v").limit(600), isSorted = true) + sameRowsAsSpark(rows.orderBy($"k".desc, $"v").limit(300), isSorted = true) + val belowZero = (0 until 2000).count(i => keys(i % keys.size) < -1.0) + sameRowsAsSpark(rows.select("k").orderBy("k").limit(belowZero), isSorted = true) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + // The setup several cache tests share: `big(_1, _2)` with 2000 rows over 40 sparse keys and // `small` with a row for every even key, and `big` as `cached(k, v)`, partitioned on the key, // sorted within each partition and cached, with 8-row batches and the native cache on. `views` From b50b2cdce66eefe7f80cf40e16f986eb3258bfa7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 15:50:45 -0600 Subject: [PATCH 16/19] fix: send a streamed order that starts with an array or a struct The check on the streamed order's keys took only the scalar types that the native sort accepts for a single column, so a join whose streamed order starts with an array or a struct went back to Spark. exprToProto takes those keys, and declines a nested key under a non-default null order itself. The check was there for the collated string, so it now rejects a key that holds one, at the top level or nested, as supportedSortType does for a sort. --- .../apache/spark/sql/comet/operators.scala | 14 ++-- .../apache/comet/exec/CometJoinSuite.scala | 69 ++++++++++++++++++- 2 files changed, 73 insertions(+), 10 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 62d05b765bb..63870f8c1f0 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 @@ -64,7 +64,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedScalarSortElementType, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -2910,17 +2910,17 @@ trait CometHashJoin { // an order that cannot be checked; see mayReportFalseOrder. // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix - // of the order that Comet can express is enough. A key counts as expressible when the native - // sort accepts its type, as `supportedSortType` decides for a sort: a collated string is - // ordered by its collation in Spark but would be declared in raw byte order. Nothing here can - // tell whether an operator above relies on the order, so a join whose first key Comet cannot - // express stays in Spark. + // of the order that Comet can express is enough. A key that exprToProto takes counts as + // expressible, unless it holds a collated string, at the top level or nested, which Spark + // orders by its collation but the native side would declare in raw byte order. That is the + // check `supportedSortType` makes for a sort. Nothing here can tell whether an operator above + // relies on the order, so a join whose first key Comet cannot express stays in Spark. val (streamed, _) = streamedSide(join.buildSide, join.left, join.right) val streamedSortOrders = if (runsAsRightOuterJoin(join.joinType, join.buildSide) && join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed)) { val orders = join.outputOrdering.iterator - .takeWhile(order => supportedScalarSortElementType(order.dataType)) + .takeWhile(order => !hasNonDefaultStringCollation(order.dataType)) .map(exprToProto(_, streamed.output)) .takeWhile(_.isDefined) .flatten 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 f00742330c5..2b730403d72 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -42,7 +42,7 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec 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.functions.{count, lit, sum} +import org.apache.spark.sql.functions.{array, count, lit, struct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType} @@ -2526,8 +2526,8 @@ class CometJoinSuite extends CometTestBase { // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde // and the native sort expression handle the key, and the sorted aggregate above the join finds // its group boundaries by comparing it. The timestamps run in a session time zone that is not - // UTC and include values before 1970. Struct, array and map keys are out of scope, since - // exprToProto rejects them by default and such a join stays in Spark. + // UTC and include values before 1970. Double, array and struct keys have tests of their own, + // and Spark cannot sort by a map. val keyTypes = Seq( "decimal(20,2)" -> Seq( "-999999999999999999.99", @@ -2682,6 +2682,69 @@ class CometJoinSuite extends CometTestBase { } } + test("outer hash joins keep a streamed order that starts with an array or a struct") { + // exprToProto takes an array or a struct sort key, so a join whose streamed order starts with + // one sends it, stays native, and the planner declares it on the JVM input. Spark drops a sort + // right below a join, so the sorted rows come from a cache. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => (i % 40, i)), "big") { + withParquetTable((0 until 40 by 2).map(k => (k, k / 2)), "small") { + val build = spark.table("small") + // An array of strings, the one array type that Comet reads from Spark's cache. + Seq( + "array" -> array(($"_1" % 7).cast("string"), $"_1".cast("string")), + "struct" -> struct($"_1" % 7, $"_1")).foreach { case (keyName, key) => + try { + val streamed = spark + .table("big") + .select(key.as("s"), $"_1", $"_2") + .repartition(2, $"_1") + .sortWithinPartitions("s") + .cache() + streamed.count() + val order = SortOrder( + BoundReference(0, streamed.schema("s").dataType, nullable = true), + Ascending) + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("broadcast") + if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + } + def df: DataFrame = joined + .select(streamed("s"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("s") + val plan = df.queryExecution.executedPlan + val context = s"$keyName $joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + assert(partitionsInOrder(df, order), context) + sameRowsAsSpark( + joined + .groupBy(streamed("s")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2")))) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + } + // The setup several cache tests share: `big(_1, _2)` with 2000 rows over 40 sparse keys and // `small` with a row for every even key, and `big` as `cached(k, v)`, partitioned on the key, // sorted within each partition and cached, with 8-row batches and the native cache on. `views` From 204440adeeda33d6bef030f6b2d9df867a5907ea Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 15:50:45 -0600 Subject: [PATCH 17/19] fix: declare the order of a sorted bucketed table read by the native scan With spark.sql.legacy.bucketedTableScan.outputOrdering on, Spark reports a sorted bucketed table's sort columns as the scan's order when every bucket has one file. The serde trusted that order and sent it, but the planner declared an order only over a JVM input, so an outer hash join over the native scan moved its unmatched rows after Spark had dropped the sort above it. A scan of files reads the files of its one file group one after the other, so the planner now counts it as a source whose order holds. --- native/core/src/execution/planner.rs | 61 +++++++++++++---- .../apache/comet/exec/CometJoinSuite.scala | 65 +++++++++++++++++++ 2 files changed, 112 insertions(+), 14 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index acdcb8ad21a..8d41cb6e5eb 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -97,6 +97,8 @@ use datafusion_comet_spark_expr::{ CsvWriteOptions, EvalMode, ListPositionsExpr, SparkArraysZipFunc, SparkBloomFilterVersion, SparkListAgg, SparkPercentile, Subquery, SumInteger, ToCsv, }; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::TableSchema; use iceberg::expr::Bind; @@ -3080,23 +3082,25 @@ impl PhysicalPlanner { /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless it already /// declares one. Spark's plan knows the order of an input that enters the native plan from - /// the JVM, but the `ScanExec` it arrives through declares none. Only that order is declared: - /// when an operator in between may have moved rows, `plan` is left as it is. + /// the JVM, and of a scan of sorted files, but the `ScanExec` or the file scan the rows + /// arrive through declares none. Only that order is declared: when an operator in between may + /// have moved rows, `plan` is left as it is. /// /// The serde sends an order only when the streamed side's Comet operators report it, and it /// checks that nothing under them moves rows (`mayReportFalseOrder`), while this walk asks /// the native operators the same thing through `maintains_input_order`. The two agree as - /// long as a Comet operator that reports an order keeps it natively and says so. One that - /// does not lands here with `plan` unchanged, and the join then moves its unmatched rows - /// after Spark has dropped the sort above it. Nothing here can fall back at that point, so - /// the invariant is kept at the operator: see "Adding a New Operator" in the contributor - /// guide. + /// long as a Comet operator that reports an order keeps it natively and says so, and a Comet + /// leaf that reports one is a source this walk knows or declares its order itself, as the + /// native range does. One that does not lands here with `plan` unchanged, and the join then + /// moves its unmatched rows after Spark has dropped the sort above it. Nothing here can fall + /// back at that point, so the invariant is kept at the operator: see "Adding a New Operator" + /// in the contributor guide. fn with_declared_order( &self, plan: Arc, sort_orders: &[Expr], ) -> Result, ExecutionError> { - if plan.output_ordering().is_some() || !Self::keeps_jvm_input_order(&plan) { + if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) { return Ok(plan); } let exprs = sort_orders @@ -3109,14 +3113,23 @@ impl PhysicalPlanner { } } - /// Whether `plan`'s rows come from an input outside the native plan through operators that - /// keep their input's order, so that the order Spark reports for them still holds. Any other - /// operator may have moved rows, such as the null-aware anti join, whose streamed side Comet - /// builds on, and declaring Spark's order over it would be wrong. - fn keeps_jvm_input_order(plan: &Arc) -> bool { + /// Whether `plan`'s rows come from a source whose order Spark knows through operators that + /// keep their input's order, so that the order Spark reports for them still holds. A source is + /// an input from the JVM, or a scan of files, which reads the files of its one file group one + /// after the other: Spark reports an order for a file scan only when every partition reads one + /// sorted file, as for a sorted bucketed table with one file per bucket. Any other operator may + /// have moved rows, such as the null-aware anti join, whose streamed side Comet builds on, and + /// declaring Spark's order over it would be wrong. + fn keeps_source_order(plan: &Arc) -> bool { if plan.is::() || plan.is::() { return true; } + if let Some(scan) = plan.downcast_ref::() { + return scan + .data_source() + .downcast_ref::() + .is_some_and(|config| config.file_groups.len() == 1); + } // DataFusion reports every right hash join as keeping its probe order, but a right outer // join keeps it only when the probe declares one, and no input on this path declares an // order. An inner join emits its matches probe row by probe row, and a right semi, anti or @@ -3139,7 +3152,7 @@ impl PhysicalPlanner { .zip(plan.maintains_input_order()) .filter_map(|(child, keeps)| keeps.then_some(child)); match (kept.next(), kept.next()) { - (Some(child), None) => Self::keeps_jvm_input_order(child), + (Some(child), None) => Self::keeps_source_order(child), _ => false, } } @@ -6564,6 +6577,26 @@ mod tests { ); } + #[test] + fn declared_order_reaches_a_scan_of_one_file_group() { + // A scan of files reads the files of its one file group one after the other, so the order + // Spark reports for it, a sorted bucketed table's, holds. With more than one file group, + // DataFusion may hand a file to another partition. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let scan = |groups: usize| -> Arc { + let source = Arc::new(ParquetSource::new(Arc::clone(&schema))) as Arc; + let file_groups = (0..groups) + .map(|i| FileGroup::new(vec![PartitionedFile::new(format!("f{i}.parquet"), 8)])) + .collect(); + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file_groups(file_groups) + .build(); + Arc::new(DataSourceExec::new(Arc::new(config))) + }; + assert!(PhysicalPlanner::keeps_source_order(&scan(1))); + assert!(!PhysicalPlanner::keeps_source_order(&scan(2))); + } + #[tokio::test] async fn spark_plan_dynamic_filter_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; 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 2b730403d72..f3b74a87233 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -2522,6 +2522,71 @@ class CometJoinSuite extends CometTestBase { } } + test("outer hash joins keep the order of a sorted bucketed table read by the native scan") { + // With spark.sql.legacy.bucketedTableScan.outputOrdering on, Spark reports a bucketed table's + // sort columns as the scan's order when every bucket has one file, and the sort above the + // join is dropped on that basis. The native scan reads each bucket's file in order, but it is + // a native leaf, not a JVM input, so the planner has to count it as a source whose order + // Spark reports, or the join moves its unmatched rows. A broadcast join, since a shuffled one + // would put an exchange between the scan and the join. + withTable("sorted_bucketed") { + // A single writer task writes one file per bucket. + (0 until 2000) + .map(i => (i % 40, i)) + .toDF("k", "v") + .coalesce(1) + .write + .format("parquet") + .bucketBy(2, "k") + .sortBy("k") + .saveAsTable("sorted_bucketed") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.BUCKETING_ENABLED.key -> "true", + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false", + SQLConf.LEGACY_BUCKETED_TABLE_SCAN_OUTPUT_ORDERING.key -> "true") { + withParquetTable((0 until 40 by 2).map(k => (k, k / 2)), "small") { + val streamed = spark.table("sorted_bucketed") + val build = spark.table("small") + val scanOrder = collectFirst(streamed.queryExecution.executedPlan) { + case scan: CometNativeScanExec if scan.bucketedScan => scan.outputOrdering + } + assert(scanOrder.exists(_.nonEmpty), streamed.queryExecution.executedPlan) + Seq("left_outer", "right_outer").foreach { joinType => + val joined = if (joinType == "left_outer") { + streamed.join(build.hint("broadcast"), streamed("k") === build("_1"), joinType) + } else { + build.hint("broadcast").join(streamed, streamed("k") === build("_1"), joinType) + } + val df = joined + .select(streamed("k"), streamed("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1, context) + assert( + collect(joins.head) { + case s: CometNativeScanExec if s.bucketedScan => s + }.size == 1, + context) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + val expected = (0 until 2000).map { i => + val k = i % 40 + Row(k, i, if (k % 2 == 0) k / 2 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), context) + } + } + } + } + } + test("outer hash joins keep the streamed order for decimal, string and timestamp keys") { // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde // and the native sort expression handle the key, and the sorted aggregate above the join finds From 939a05a9349d1e55ce8301d856b5608cc74d86fd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 15 Oct 2026 15:50:45 -0600 Subject: [PATCH 18/19] docs: qualify where native hash joins keep the streamed order The operator compatibility guide now says the order is kept where Comet can show that it holds, and points to #6922 for the caches it still leaves undeclared. The contributor guide says what a leaf that reports an ordering has to do natively. --- docs/source/contributor-guide/adding_a_new_operator.md | 6 ++++++ docs/source/user-guide/latest/compatibility/operators.md | 9 +++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index f0371030526..face1e4c2bd 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -374,6 +374,12 @@ declares the order through an operator only when the operator says it keeps it. cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in `operators.scala`, which makes the hash join serde send no order over it. +A leaf that reports an ordering is where the order starts, so the planner has to know that the +native operator produces its rows in that order. It knows that of an input from the JVM and of a +scan of files, which reads the files of each partition one after the other, and the native range +declares its order itself through `SortedInputExec`. A new leaf that reports an ordering must +declare it natively the same way. + #### Plan Identity and Exchange Reuse Spark's `ReuseExchangeAndSubquery` identifies equivalent plans through canonicalization and may diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 5c0eaf6fd22..a3ed25638c7 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -94,9 +94,9 @@ Both results are valid under Spark's semantics for these functions. ## Hash Join Output Order Spark reports the sort order of a hash join's streamed side as the join's own, and drops a sort -above the join that this order satisfies. Comet's native hash joins keep that order, including a -left or right outer join over a streamed side that was sorted outside the native plan, such as a -sorted relation read from the [in-memory cache](#in-memory-cache). +above the join that this order satisfies. Comet's native hash joins keep that order where Comet can +show that it holds. That includes a left or right outer join over a streamed side that was sorted +outside the native plan, such as a sorted relation read from the [in-memory cache](#in-memory-cache). Comet passes the order to the native join as sort keys, and an outer join needs at least the first key to be one that Comet can express natively. Declaring any leading part of the order keeps the @@ -109,7 +109,8 @@ Where the order Spark reports might not hold, Comet declares nothing and an oute its unmatched rows out of order. The output of a `NOT IN` subquery or of a nested loop join ([#6898](https://github.com/apache/datafusion-comet/issues/6898)) is one such case, and so is a cache or checkpoint that one of them wrote. Comet also declares nothing when it cannot follow the -order back to its source, as with an RDD checkpoint. +order back to its source, as with an RDD checkpoint, and over some caches written by an outer hash +join whose order does hold ([#6922](https://github.com/apache/datafusion-comet/issues/6922)). ## Window Functions From 1227aa6c8c26fa8808a2f45e1efac5cdcf13a7bd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 17 Oct 2026 07:33:06 -0600 Subject: [PATCH 19/19] fix: drop the native NOT IN order handling that the ordered-side fallback makes unreachable Comet now leaves a null-aware anti join over an ordered streamed side in Spark (#6810), so a native NOT IN never reports an order and its rows cannot reach an outer hash join with a false one. Remove the walk's null-aware case, the tests that needed a native NOT IN to write a cache or checkpoint, and the documentation that named NOT IN as a source of a false order. --- .../latest/compatibility/operators.md | 6 +- .../apache/spark/sql/comet/operators.scala | 27 +-- .../comet/exec/CometInMemoryCacheSuite.scala | 55 ----- .../apache/comet/exec/CometJoinSuite.scala | 208 +----------------- 4 files changed, 23 insertions(+), 273 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 855b77d9b7d..aaa95d859b5 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -117,7 +117,7 @@ for example the output of a window, a sort-merge join, or a cached table sorted 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. -## Hash Join Output Order +### Output Order Spark reports the sort order of a hash join's streamed side as the join's own, and drops a sort above the join that this order satisfies. Comet's native hash joins keep that order where Comet can @@ -132,9 +132,9 @@ express falls back to Spark and records a fallback reason (see an operator above the join relies on the order. Where the order Spark reports might not hold, Comet declares nothing and an outer join can return -its unmatched rows out of order. The output of a `NOT IN` subquery or of a nested loop join +its unmatched rows out of order. The output of a nested loop join ([#6898](https://github.com/apache/datafusion-comet/issues/6898)) is one such case, and so is a -cache or checkpoint that one of them wrote. Comet also declares nothing when it cannot follow the +cache or checkpoint that one wrote. Comet also declares nothing when it cannot follow the order back to its source, as with an RDD checkpoint, and over some caches written by an outer hash join whose order does hold ([#6922](https://github.com/apache/datafusion-comet/issues/6922)). 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 59b368c04aa..733f56aaded 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 @@ -2928,16 +2928,15 @@ trait CometHashJoin { // and an input sorted outside the native plan, a relation sorted before it was cached for // instance, declares no order to DataFusion. Send the order so the native join declares it // on the probe input; see SortedInputExec. Every other join type keeps the probe order as - // it is, except the null-aware anti join: Comet builds that one on its streamed side, so its - // output is not in the order Spark reports for it, and the planner declares no order above - // it. + // it is. The null-aware anti join is the exception, as Comet builds that one on its streamed + // side, but it runs natively only when that side reports no order, so it reports none. // // The order Spark reports can be false to begin with. A Comet join that may move its - // streamed rows reports it all the same: the null-aware anti join, a right outer join like - // this one whose probe input declared no order, and the nested loop join, which keeps the - // order of neither input. A cache reports the order of the plan that wrote it, so the rows - // of such a join can reach this one through a cache, as a JVM input that the native planner - // cannot look behind. Declaring the order there would let an operator above rely on it, an + // streamed rows reports it all the same: a right outer join like this one whose probe input + // declared no order, and the nested loop join, which keeps the order of neither input. A + // cache reports the order of the plan that wrote it, so the rows of such a join can reach + // this one through a cache, as a JVM input that the native planner cannot look behind. + // Declaring the order there would let an operator above rely on it, an // aggregate in sorted mode for instance, which then splits groups. So nothing is declared // when such a join feeds the streamed side with no sort in between, or when a leaf reports // an order that cannot be checked; see mayReportFalseOrder. @@ -2992,12 +2991,11 @@ trait CometHashJoin { // Whether the order that `plan` reports may be false, because something that feeds it with no // sort in between cannot be relied on to keep it. Comet's hash joins report the order of their - // streamed side, which Spark's own joins keep, but three Comet joins do not always: the - // null-aware anti join, which Comet builds on its streamed side; an outer join that runs as - // DataFusion's right outer join, which keeps its probe rows in place only if the planner - // declared an order on its probe input; and the nested loop join, which DataFusion runs build - // row by build row. So does a leaf that reports an order from state this walk cannot check, an - // RDD checkpointed from a plan that is gone for one. + // streamed side, which Spark's own joins keep, but two Comet joins do not always: an outer join + // that runs as DataFusion's right outer join, which keeps its probe rows in place only if the + // planner declared an order on its probe input; and the nested loop join, which DataFusion runs + // build row by build row. So does a leaf that reports an order from state this walk cannot + // check, an RDD checkpointed from a plan that is gone for one. // // The walk follows the order, so a hash join, which passes on only its streamed side's, is // walked through that side alone. A cache reports the order of the plan that wrote it, and the @@ -3011,7 +3009,6 @@ trait CometHashJoin { case _ if plan.outputOrdering.isEmpty => false case _: SortExec | _: CometSortExec => false case _: CometBroadcastNestedLoopJoinExec => true - case join: CometBroadcastHashJoinExec if join.isNullAwareAntiJoin => true case join: CometBroadcastHashJoinExec => hashJoinMayReportFalseOrder( join.nativeOp, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 4b43eda5d7b..c300365f536 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3406,61 +3406,6 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("an outer hash join over a native cache scan of a NOT IN keeps each group once") { - // The native cache scan reports the order of the plan that wrote the cache. A NOT IN is not in - // the order Spark reports for it, since Comet builds it on its streamed side, so the outer - // join that reads the cache must not declare the order, or the aggregates above it run in - // sorted mode over unsorted rows and split groups. Sparse keys and small batches make each - // group span several. A cache with an exchange in its plan is written under AQE, and the scan - // reads the final plan behind the adaptive one, so this runs with AQE off and on. - Seq(false, true).foreach { adaptive => - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") - .cache() - .createOrReplaceTempView("anti") - spark.table("anti").count() - assert( - !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), - "the cached NOT IN rows are in order, so nothing here can split a group") - val query = - """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) - |FROM anti a LEFT JOIN small b ON a.k = b._1 - |GROUP BY a.k""".stripMargin - checkEachKeyGroupedOnce( - query, - plan => - assert( - collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, - s"the outer join does not read the native cache scan (AQE $adaptive):\n$plan")) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("anti") - spark.catalog.dropTempView("cached") - } - } - } - } - } - } - } - test("an aggregate on the key above a hash join over a sorted cache groups each key once") { // The order the join declares on its probe input reaches past the join, so the aggregate // above it in the same native plan groups in sorted, streaming mode. That is right while 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 f5003451b68..6ae7b1081f9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -35,11 +35,10 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeRefer import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, RDDScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.SortAggregateExec -import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} @@ -3017,11 +3016,12 @@ class CometJoinSuite extends CometTestBase { } test("NOT IN below an outer hash join over a sorted cache keeps each group once") { - // Spark reports the streamed side's order for a null-aware anti join, but Comet builds that - // join on its streamed side, so the order does not hold above it. Declaring it on the outer - // join's probe input would put the aggregates in sorted mode over unsorted rows and split - // groups. Sparse keys keep DataFusion off its perfect hash build path, which keeps the build - // batches in order and would hide that, and small batches spread each partition over many. + // Spark reports the streamed side's order for a null-aware anti join, but a native one does + // not keep it, so Comet leaves a NOT IN over an ordered side in Spark. Its order then holds + // and the outer join above may declare it. Declaring it over a native NOT IN would put the + // aggregates in sorted mode over unsorted rows and split groups. Sparse keys keep DataFusion + // off its perfect hash build path, which keeps the build batches in order and would hide + // that, and small batches spread each partition over many. withSortedCache() { withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { val query = @@ -3034,114 +3034,6 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a cache of a NOT IN keeps each group once") { - // The same NOT IN as the test above, but cached, so its rows reach the outer join from the - // JVM. The cache reports the order Spark reports for the null-aware anti join, which its rows - // are not in, and the outer join's probe input is then a plain JVM input as far as the native - // plan can see. The order must not be declared there either. A cache with an exchange in its - // plan is written under AQE from Spark 3.5, and Spark's scan then reports the order of the - // final plan behind the adaptive one, which reports none itself, so this runs with AQE off - // and on. - Seq(false, true).foreach { adaptive => - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") - .cache() - .createOrReplaceTempView("anti") - spark.table("anti").count() - val cachedPlan = cachedPlanOf(spark.table("anti")) - assert( - find(cachedPlan) { - case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin - case _ => false - }.isDefined, - s"the cache was not written by a native NOT IN (AQE $adaptive):\n$cachedPlan") - assert( - !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), - "the cached NOT IN rows are in order, so nothing here can split a group") - val query = - """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) - |FROM anti a LEFT JOIN small b ON a.k = b._1 - |GROUP BY a.k""".stripMargin - checkEachKeyGroupedOnce(query) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("anti") - spark.catalog.dropTempView("cached") - } - } - } - } - } - } - } - - test("an outer hash join over a checkpoint of a NOT IN keeps each group once") { - // A checkpoint is read by an RDD scan that reports the order of the plan it was taken from, - // as a cache does, but that plan is gone, so nothing can say whether the order holds. Here it - // is the order Spark reports for the null-aware anti join, which the rows are not in. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", - CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)") - .localCheckpoint() - .createOrReplaceTempView("anti") - val plan = spark.table("anti").queryExecution.executedPlan - val scans = collect(plan) { case scan: RDDScanExec => scan } - assert( - scans.size == 1 && scans.head.outputOrdering.nonEmpty, - s"the checkpoint does not report the NOT IN's order:\n$plan") - assert( - !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)), - "the checkpointed NOT IN rows are in order, so nothing here can split a group") - val query = - """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) - |FROM anti a LEFT JOIN small b ON a.k = b._1 - |GROUP BY a.k""".stripMargin - checkEachKeyGroupedOnce(query) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("anti") - spark.catalog.dropTempView("cached") - } - } - } - } - } - } - test("a cache of an outer hash join over a native sample reports no order, as Spark does") { // Spark's SampleExec reports no output ordering, so the sort above a join over a sample stays, // and CometSampleExec reports none either. The native sample does not say it keeps its input's @@ -3252,88 +3144,4 @@ class CometJoinSuite extends CometTestBase { } } } - - test("a NOT IN below the build side of an inner hash join leaves the streamed order declared") { - // An inner join passes on only its streamed side's order, so a NOT IN below its build side - // cannot affect it, and the outer join above the inner join still declares it. A build side - // behind an exchange reports no order, so it needs a build side that is a cache already - // partitioned on the key, where the shuffled hash join gets no exchange, and that reports the - // order of the NOT IN that wrote it. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.COLUMN_BATCH_SIZE.key -> "8", - CometConf.COMET_BATCH_SIZE.key -> "8", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40).map(k => (k * 1000, k)), "dim") { - withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { - withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { - try { - Seq("big" -> "cached", "dim" -> "dim_cached").foreach { case (table, view) => - spark - .table(table) - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView(view) - spark.table(view).count() - } - sql("SELECT * FROM dim_cached WHERE k NOT IN (SELECT _1 FROM excl)") - .cache() - .createOrReplaceTempView("anti") - spark.table("anti").count() - val cachedPlan = cachedPlanOf(spark.table("anti")) - assert( - find(cachedPlan) { - case join: CometBroadcastHashJoinExec => join.isNullAwareAntiJoin - case _ => false - }.isDefined, - s"the cache was not written by a native NOT IN:\n$cachedPlan") - val query = - """SELECT /*+ SHUFFLE_HASH(b), BROADCAST(c) */ a.k, a.v, b.v AS x, c._2 AS w - |FROM cached a - |JOIN anti b ON a.k = b.k - |LEFT JOIN small c ON a.k = c._1""".stripMargin - val df = sql(query).sortWithinPartitions("k") - val plan = df.queryExecution.executedPlan - val inner = collect(plan) { case join: CometHashJoinExec => join } - assert(inner.size == 1 && inner.head.joinType == Inner, plan) - assert(inner.head.buildSide == BuildRight, plan) - assert( - find(inner.head.right) { - case _: CometInMemoryTableScanExec | _: InMemoryTableScanExec => true - case _ => false - }.isDefined, - plan) - assert(collect(plan) { case e: ShuffleExchangeLike => e }.isEmpty, plan) - val outer = collect(plan) { - case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join - } - assert(outer.size == 1, plan) - assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) - assert( - collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, - plan) - val expected = (0 until 2000) - .filter(i => !Set(1, 10, 20).contains(i % 40)) - .map { i => - val m = i % 40 - Row(m * 1000, i, m, if (m % 2 == 0) m else null) - } - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) - assert(partitionsInOrder(df, keyOrder(Ascending)), plan) - } finally { - spark.catalog.clearCache() - spark.catalog.dropTempView("anti") - spark.catalog.dropTempView("dim_cached") - spark.catalog.dropTempView("cached") - } - } - } - } - } - } - } }