You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/16 00:44:09 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6835-a45445c6ef770b2d008b63ef42478bfdbc3f1e01 created (now b663de1e6c)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6835-a45445c6ef770b2d008b63ef42478bfdbc3f1e01
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at b663de1e6c fix: keep a streamed side sorted outside the native plan in order through outer hash joins (#6835)
This branch includes the following new commits:
new b663de1e6c fix: keep a streamed side sorted outside the native plan in order through outer hash joins (#6835)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: keep a streamed side sorted outside the native plan in order through outer hash joins (#6835)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6835-a45445c6ef770b2d008b63ef42478bfdbc3f1e01
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit b663de1e6ca00657f6bb6f76b831e50cb49218f7
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Oct 15 23:40:37 2026 +0000
fix: keep a streamed side sorted outside the native plan in order through outer hash joins (#6835)
* 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.
* 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.
* 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.
* 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.
* 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.
* 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.
* 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.
* 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.
* test: compare the key-type rows with Spark's as a multiset
* 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.
* 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.
* refactor: name the streamed side of a hash join once
* 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.
* 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.
* 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.
* 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.
* 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.
* 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.
---
.../contributor-guide/adding_a_new_operator.md | 23 +
.../user-guide/latest/compatibility/operators.md | 21 +
native/core/src/execution/operators/mod.rs | 2 +
.../core/src/execution/operators/sorted_input.rs | 135 ++++
native/core/src/execution/planner.rs | 344 +++++++-
native/operators/src/lib.rs | 2 +-
native/operators/src/range.rs | 34 +
native/proto/src/proto/operator.proto | 5 +
.../apache/spark/sql/comet/CometSampleExec.scala | 4 +-
.../org/apache/spark/sql/comet/operators.scala | 149 +++-
.../comet/exec/CometInMemoryCacheSuite.scala | 269 ++++++-
.../org/apache/comet/exec/CometJoinSuite.scala | 892 ++++++++++++++++++++-
.../scala/org/apache/spark/sql/CometTestBase.scala | 54 +-
13 files changed, 1911 insertions(+), 23 deletions(-)
diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md
index 85946f44fd..face1e4c2b 100644
--- a/docs/source/contributor-guide/adding_a_new_operator.md
+++ b/docs/source/contributor-guide/adding_a_new_operator.md
@@ -357,6 +357,29 @@ 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.
+
+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 75c2210ef8..5e42277d5f 100644
--- a/docs/source/user-guide/latest/compatibility/operators.md
+++ b/docs/source/user-guide/latest/compatibility/operators.md
@@ -109,6 +109,27 @@ 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 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
+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, 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
Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of
diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs
index a15a528f90..02cdea9326 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 0000000000..a31053bd6e
--- /dev/null
+++ b/native/core/src/execution/operators/sorted_input.rs
@@ -0,0 +1,135 @@
+// 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::{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,
+};
+
+/// 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 only
+/// 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<dyn ExecutionPlan>,
+ ordering: LexOrdering,
+ cache: Arc<PlanProperties>,
+}
+
+impl SortedInputExec {
+ pub fn new(input: Arc<dyn ExecutionPlan>, 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(
+ equivalence_properties,
+ 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<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ apply_expression_roots(self.ordering.iter().map(|sort| Arc::clone(&sort.expr)), f)
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ 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<MetricsSet> {
+ self.input.metrics()
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ 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<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ self.input.execute(partition, context)
+ }
+}
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index d871c941dd..9e9ca37e42 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -40,6 +40,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::{
@@ -87,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::{
@@ -96,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;
@@ -2089,7 +2092,7 @@ impl PhysicalPlanner {
}
OpStruct::RangeScan(range) => {
// A leaf with no JVM input: each task produces its own partition of the range.
- let range: Arc<dyn ExecutionPlan> = Arc::new(range_exec(
+ let source: Arc<dyn ExecutionPlan> = Arc::new(range_exec(
range.start,
range.step,
range.num_elements,
@@ -2097,10 +2100,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<dyn ExecutionPlan> = 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) => {
@@ -2781,8 +2791,21 @@ impl PhysicalPlanner {
*probe = Self::prepare_probe_for_runtime_filters(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
@@ -3091,6 +3114,83 @@ impl PhysicalPlanner {
Ok(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, 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, 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<dyn ExecutionPlan>,
+ sort_orders: &[Expr],
+ ) -> Result<Arc<dyn ExecutionPlan>, ExecutionError> {
+ if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) {
+ return Ok(plan);
+ }
+ let exprs = sort_orders
+ .iter()
+ .map(|expr| self.create_sort_expr(expr, plan.schema()))
+ .collect::<Result<Vec<_>, _>>()?;
+ match LexOrdering::new(exprs) {
+ Some(ordering) => Ok(Arc::new(SortedInputExec::new(plan, ordering))),
+ None => Ok(plan),
+ }
+ }
+
+ /// 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<dyn ExecutionPlan>) -> bool {
+ if plan.is::<ScanExec>() || plan.is::<ShuffleScanExec>() {
+ return true;
+ }
+ if let Some(scan) = plan.downcast_ref::<DataSourceExec>() {
+ return scan
+ .data_source()
+ .downcast_ref::<FileScanConfig>()
+ .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
+ // 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::<HashJoinExec>() {
+ if !matches!(
+ join.join_type(),
+ DFJoinType::Inner
+ | DFJoinType::RightSemi
+ | DFJoinType::RightAnti
+ | DFJoinType::RightMark
+ ) {
+ 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_source_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(
@@ -6289,6 +6389,7 @@ mod tests {
build_side: 0,
null_aware_anti_join: false,
dynamic_filter_enabled: false,
+ streamed_sort_orders: vec![],
})),
};
@@ -6303,6 +6404,239 @@ 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<dyn ExecutionPlan>,
+ ) -> &datafusion::physical_plan::joins::HashJoinExec {
+ let join = if plan.is::<datafusion::physical_plan::projection::ProjectionExec>() {
+ plan.children()[0]
+ } else {
+ plan
+ };
+ join.downcast_ref::<datafusion::physical_plan::joins::HashJoinExec>()
+ .expect("a hash join")
+ }
+
+ fn sorted_outer_hash_join(
+ join_type: spark_operator::JoinType,
+ build_side: spark_operator::BuildSide,
+ children: Vec<Operator>,
+ ) -> 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!(&datafusion::common::JoinType::Right, join.join_type());
+ assert_eq!("ScanExec", join.left().name());
+ let probe = join.right();
+ let sorted = probe
+ .downcast_ref::<crate::execution::operators::SortedInputExec>()
+ .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_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());
+ }
+
+ /// 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)],
+ 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()
+ }
+ }
+
+ #[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_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_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(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(), "{reason}");
+ }
+ }
+
+ #[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()
+ );
+ }
+
+ #[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<dyn ExecutionPlan> {
+ let source = Arc::new(ParquetSource::new(Arc::clone(&schema))) as Arc<dyn FileSource>;
+ 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/native/operators/src/lib.rs b/native/operators/src/lib.rs
index b5f762025a..97c80a9789 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 3766c82b1a..4ba4f6d642 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.
diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
index ac04e9e232..48e67ec7c1 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -1099,6 +1099,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, 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;
}
message SortMergeJoin {
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 dcdae2c779..54391580a4 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 3ee11e4978..a4fb261ed7 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,8 +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}
@@ -62,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, hasNonDefaultStringCollation, isStringCollationType, supportedSortType}
import org.apache.comet.serde.operator.CometSink
import org.apache.comet.shims.MergeRowsMetricsShim
@@ -2906,6 +2908,56 @@ 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, 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.
+ //
+ // 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, 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. 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 => !hasNonDefaultStringCollation(order.dataType))
+ .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
+ } else {
+ Nil
+ }
+
if (leftKeys.forall(_.isDefined) &&
rightKeys.forall(_.isDefined) &&
childOp.nonEmpty) {
@@ -2918,12 +2970,105 @@ 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 {
None
}
}
+
+ // 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 _: 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(streamedSide(join.buildSide, join.left, join.right)._1)
+ 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) = streamedSide(buildSide, left, right)
+ 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
+ }
+
+ // 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.
+ 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/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index 4dd68e6b5f..4b43eda5d7 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -36,13 +36,15 @@ 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}
-import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, 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
-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}
@@ -3288,4 +3290,265 @@ 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, 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",
+ 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 = {
+ 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 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) =>
+ 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(partitionsInOrder(outer, keyOrder(Ascending)), outerPlan)
+ }
+
+ val inner = query(hint, "inner")
+ val innerPlan = inner.queryExecution.executedPlan
+ assert(joins(innerPlan) == Seq((joinClass, BuildRight)), innerPlan)
+ assert(sorts(innerPlan).isEmpty, innerPlan)
+ QueryTest.sameRows(innerRows, inner.collect().toSeq).foreach(fail(_))
+ assert(partitionsInOrder(inner, keyOrder(Ascending)), innerPlan)
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ 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(partitionsInOrder(df, keyOrder(Ascending)), 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(
+ !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
+ // 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. 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",
+ 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") {
+ 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)
+ }
+
+ 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(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
+ 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()
+ }
+ }
+ }
+ }
+ }
+ }
}
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 d96d4f7b8e..f3b74a8723 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
@@ -28,21 +28,23 @@ import org.apache.hadoop.fs.Path
import org.apache.parquet.hadoop.ParquetFileReader
import org.apache.parquet.hadoop.util.HadoopInputFile
import org.apache.spark.SparkException
-import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.{Column, CometTestBase, DataFrame, 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.plans.Inner
+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
-import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec}
-import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec}
+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
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.{array, count, lit, struct, 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, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType}
import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus}
@@ -2284,4 +2286,880 @@ 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")
+
+ // 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
+ 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, 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), on, joinType)
+ } else {
+ build.hint(hint).join(streamed, on, 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)
+ withCondition <- Seq(false, true)
+ } {
+ 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)
+ 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(withCondition), df.collect().toSeq).foreach(fail(_))
+ assert(partitionsInOrder(df, keyOrder(Ascending)), plan)
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ 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",
+ 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") {
+ 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 unexpressed = sortedCache($"_1" + 1)
+ val (_, fallbackPlan) = checkSparkAnswerAndFallbackReason(
+ unexpressed.join(
+ build.hint("shuffle_hash"),
+ 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(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()
+ }
+ }
+ }
+ }
+ }
+
+ 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()
+ }
+ }
+ }
+ }
+ }
+
+ // 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, isSorted).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 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
+ // its group boundaries by comparing it. The timestamps run in a session time zone that is not
+ // 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",
+ "-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",
+ "2030-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)
+ }
+ }
+ 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")
+ 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")
+ 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.
+ def agg: DataFrame = joined
+ .groupBy(streamed("_1").as("k"))
+ .agg(count(lit(1)), sum(streamed("_2")), count(build("_2")))
+ sameRowsAsSpark(agg)
+ assert(agg.count() == keys.size, s"$typeName $joinType groups a key twice")
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+ }
+
+ 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()
+ }
+ }
+ }
+ }
+ }
+
+ 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`
+ // 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",
+ 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()
+ 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)
+ }
+ }
+ }
+
+ 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
+ // 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.
+ 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)
+ |FROM joined j LEFT JOIN small c ON j.k = c._1
+ |GROUP BY j.k""".stripMargin
+ checkEachKeyGroupedOnce(query)
+ }
+ }
+
+ 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.
+ 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
+ |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)
+ }
+ }
+
+ 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.
+ 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)
+ }
+ }
+ }
+
+ 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")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
}
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 7fc3e66c16..4c21ee394a 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
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org