From 5116d05a689275285c46c4e4ed1823fbec0584c4 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Tue, 18 Aug 2026 18:24:34 +0530 Subject: [PATCH] feat(pwmj): support RightSemi/RightAnti existence joins --- datafusion/core/src/physical_planner.rs | 13 +- datafusion/core/tests/fuzz_cases/join_fuzz.rs | 78 +- .../src/joins/nested_loop_join.rs | 18 +- .../src/joins/piecewise_merge_join/exec.rs | 455 ++++++++-- .../piecewise_merge_join/existence_join.rs | 23 +- .../src/joins/piecewise_merge_join/mod.rs | 1 + .../right_existence_join.rs | 846 ++++++++++++++++++ .../src/joins/piecewise_merge_join/utils.rs | 42 +- datafusion/physical-plan/src/joins/utils.rs | 14 + datafusion/sqllogictest/test_files/pwmj.slt | 437 +++++++++ 10 files changed, 1751 insertions(+), 176 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/piecewise_merge_join/right_existence_join.rs diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 7002a4b04d957..853581493d3d8 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1604,16 +1604,13 @@ impl DefaultPhysicalPlanner { Arc::new(CrossJoinExec::new(physical_left, physical_right)) } else if num_range_filters == 1 && total_filters == 1 - // PWMJ supports classic joins and Left Semi/Anti existence joins. - // Right Semi/Anti and Mark joins are not implemented yet (they - // would require swapping the inputs so the marked side is buffered), - // so exclude them here and let them fall back to NestedLoopJoin. + // PWMJ supports classic joins and Semi/Anti existence joins. Mark + // joins are not implemented yet (they need an extra boolean column + // rather than a subset of one side's rows), so exclude them here + // and let them fall back to NestedLoopJoin. && !matches!( join_type, - JoinType::RightSemi - | JoinType::RightAnti - | JoinType::LeftMark - | JoinType::RightMark + JoinType::LeftMark | JoinType::RightMark ) && session_state .config_options() diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index fce6999fd0d77..586d975f050d4 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -1366,7 +1366,12 @@ fn make_staggered_batches_binary( // streamed side below is spread round-robin over several partitions as one-row batches, so // batches arrive in an order no static test pins down, and the counter that gates the final // pass is seeded from a partition count that deliberately disagrees with the `num_partitions` -// argument. +// argument. `RightSemi`/`RightAnti` take the mirror path -- every partition emits its own +// rows, decided against a single buffered key -- and are covered here too, over the same +// inputs, so the two halves are held to the same oracle. Their buffered side is fanned out as +// well, since it carries no single-partition requirement: each partition is folded to its own +// extreme on a separate task and those are then combined, and only randomization varies which +// partition holds the deciding key, or holds none at all. fn pwmj_kv_schema() -> Arc { Arc::new(Schema::new(vec![ @@ -1396,7 +1401,10 @@ fn pwmj_single_exec(ids: &[i32], keys: &[Option]) -> Arc .unwrap() } -/// Streamed side spread round-robin across `nparts` partitions, one row per batch. +/// Rows spread round-robin across `nparts` partitions, one row per batch, with any partition +/// that draws no row left holding a single empty batch. Used for the streamed side throughout, +/// and for the buffered side of the right existence joins, which place no single-partition +/// requirement on it. fn pwmj_parts_exec( ids: &[i32], keys: &[Option], @@ -1422,25 +1430,31 @@ fn pwmj_existence_plan( join_type: JoinType, ) -> Arc { // Matches `PiecewiseMergeJoinExec::required_input_ordering`: descending for `<`/`<=`, - // ascending for `>`/`>=`, NULLs first either way. - let sort_options = match op { - Operator::Lt | Operator::LtEq => SortOptions::new(true, true), - Operator::Gt | Operator::GtEq => SortOptions::new(false, true), - other => panic!("not a range operator: {other:?}"), + // ascending for `>`/`>=`, NULLs first either way. Right existence joins require no + // ordering at all -- they only read the buffered side's min/max -- so they are fed the + // left side unsorted, which is the input shape they will see in a real plan. + let buffered = match join_type { + JoinType::RightSemi | JoinType::RightAnti => left, + _ => { + let sort_options = match op { + Operator::Lt | Operator::LtEq => SortOptions::new(true, true), + Operator::Gt | Operator::GtEq => SortOptions::new(false, true), + other => panic!("not a range operator: {other:?}"), + }; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("k", 1)), + sort_options, + )]) + .unwrap(); + Arc::new(SortExec::new(ordering, left)) + } }; - let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( - Arc::new(Column::new("k", 1)), - sort_options, - )]) - .unwrap(); - let sorted_left = Arc::new(SortExec::new(ordering, left)); let on: (PhysicalExprRef, PhysicalExprRef) = (Arc::new(Column::new("k", 1)), Arc::new(Column::new("k", 1))); // `num_partitions` is 1 while the streamed side has up to 3: the final-pass counter must // come from the streamed side's partition count, not from this argument. Arc::new( - PiecewiseMergeJoinExec::try_new(sorted_left, right, on, op, join_type, 1) - .unwrap(), + PiecewiseMergeJoinExec::try_new(buffered, right, on, op, join_type, 1).unwrap(), ) } @@ -1475,10 +1489,13 @@ fn pwmj_nlj_oracle_plan( ) } -/// Executes every output partition concurrently and returns the surviving left `id`s, sorted. +/// Executes every output partition concurrently and returns the surviving `id`s, sorted. +/// Both halves of an existence join output an `id` as their first column: the left side's for +/// `LeftSemi`/`LeftAnti`, the right side's for `RightSemi`/`RightAnti`. /// -/// Concurrent rather than one partition at a time: the partitions share the watermark and race -/// to be the one that runs the final pass, which is the part a sequential drain cannot reach. +/// Concurrent rather than one partition at a time: for the left joins the partitions share the +/// watermark and race to be the one that runs the final pass, which is the part a sequential +/// drain cannot reach. async fn pwmj_collect_ids( plan: Arc, task_ctx: Arc, @@ -1524,6 +1541,10 @@ async fn fuzz_pwmj_existence_matches_nested_loop() { // A narrow key range forces duplicates and equal-boundary cases. let key_range = rng.random_range(1..6i32); let nparts = rng.random_range(1..4usize); + // Independent of the streamed side's, so the two counts disagree across seeds and the + // 1-partition buffered case is still drawn. Only the right existence joins can use it: + // the others require the buffered side coalesced and globally sorted. + let buffered_nparts = rng.random_range(1..4usize); let gen_keys = |n: usize, rng: &mut StdRng| -> Vec> { (0..n) @@ -1537,10 +1558,24 @@ async fn fuzz_pwmj_existence_matches_nested_loop() { let right_keys = gen_keys(right_len, &mut rng); for op in ops { - for join_type in [JoinType::LeftSemi, JoinType::LeftAnti] { + for join_type in [ + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + ] { + // Fanned out only for the right existence joins, whose buffered side is + // folded one partition per task and then combined -- a reduction the + // single-partition shape below never reaches. + let buffered = match join_type { + JoinType::RightSemi | JoinType::RightAnti => { + pwmj_parts_exec(&left_ids, &left_keys, buffered_nparts) + } + _ => pwmj_single_exec(&left_ids, &left_keys), + }; let got = pwmj_collect_ids( pwmj_existence_plan( - pwmj_single_exec(&left_ids, &left_keys), + buffered, pwmj_parts_exec(&right_ids, &right_keys, nparts), op, join_type, @@ -1562,7 +1597,8 @@ async fn fuzz_pwmj_existence_matches_nested_loop() { assert_eq!( got, want, "mismatch seed={seed} op={op:?} join_type={join_type:?} \ - nparts={nparts} left_keys={left_keys:?} right_keys={right_keys:?}" + nparts={nparts} buffered_nparts={buffered_nparts} \ + left_keys={left_keys:?} right_keys={right_keys:?}" ); } } diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index eb1df638c7dc5..d328c71d0bff2 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -32,8 +32,8 @@ use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::joins::SharedBitmapBuilder; use crate::joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceAsync, OnceFut, - build_join_schema, check_join_is_valid, estimate_join_statistics, - need_produce_right_in_final, + boolean_mask_from_filter, build_join_schema, check_join_is_valid, + estimate_join_statistics, need_produce_right_in_final, }; use crate::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, MetricType, MetricsSet, RatioMetrics, @@ -2793,20 +2793,6 @@ fn apply_filter_to_row_join_batch( Ok(bitmap_combined) } -/// Convert a boolean filter array into a unified mask bitmap. -/// -/// Caution: The filter result is NOT a bitmap; it contains true/false/null values. -/// For example, `1 < NULL` evaluates to NULL. Therefore, we must combine (AND) -/// the boolean array with its null bitmap to construct a unified bitmap. -#[inline] -fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray { - let (values, nulls) = filter_arr.clone().into_parts(); - match nulls { - Some(nulls) => BooleanArray::new(nulls.inner() & &values, None), - None => BooleanArray::new(values, None), - } -} - /// This function performs the following steps: /// 1. Apply filter to probe-side batch /// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 5183d3aa0feb7..25b781545ba62 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -18,16 +18,17 @@ use arrow::array::Array; use arrow::{ array::{ArrayRef, BooleanBufferBuilder, RecordBatch}, - compute::concat_batches, + compute::{concat, concat_batches}, util::bit_util, }; use arrow_schema::{SchemaRef, SortOptions}; use datafusion_common::not_impl_err; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{JoinSide, Result, internal_err}; +use datafusion_common::{JoinSide, Result, internal_datafusion_err, internal_err}; +use datafusion_common_runtime::SpawnedTask; use datafusion_execution::{ SendableRecordBatchStream, - memory_pool::{MemoryConsumer, MemoryReservation}, + memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}, }; use datafusion_expr::{JoinType, Operator}; use datafusion_physical_expr::equivalence::join_equivalence_properties; @@ -36,7 +37,7 @@ use datafusion_physical_expr::{ PhysicalSortExpr, }; use datafusion_physical_expr_common::physical_expr::fmt_sql; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use parking_lot::Mutex; use std::fmt::Formatter; use std::sync::Arc; @@ -47,10 +48,13 @@ use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::joins::piecewise_merge_join::classic_join::{ ClassicPWMJStream, PiecewiseMergeJoinStreamState, }; -use crate::joins::piecewise_merge_join::existence_join::ExistencePWMJStream; +use crate::joins::piecewise_merge_join::existence_join::{ + ExistencePWMJStream, extreme_key, +}; +use crate::joins::piecewise_merge_join::right_existence_join::RightExistencePWMJStream; use crate::joins::piecewise_merge_join::utils::{ - build_visited_indices_map, is_existence_join, is_right_existence_join, - is_supported_existence_join, + build_visited_indices_map, is_existence_join, is_supported_existence_join, + is_supported_right_existence_join, }; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; @@ -166,12 +170,11 @@ use crate::{ /// ``` /// /// ## Existence Joins (Semi, Anti, Mark) -/// Currently only `LeftSemi` and `LeftAnti` are supported. For these the marked side is -/// already the left (buffered) side, so no input swap is needed. The rest are rejected in -/// [`Self::try_new`]: `RightSemi`/`RightAnti`/`RightMark` mark the right side and need an -/// input swap, and `LeftMark` needs an extra boolean column rather than a filtered slice. +/// Every Semi/Anti join is supported; the Mark joins are rejected in [`Self::try_new`], as they +/// need an extra boolean column rather than a subset of one side's rows. The two sides are +/// served by different streams, because a single range predicate makes them different problems. /// -/// `LeftSemi`/`LeftAnti` are served by a dedicated stream, `ExistencePWMJStream` (see +/// `LeftSemi`/`LeftAnti` mark the buffered (left) side, which is `ExistencePWMJStream` (see /// `existence_join.rs`). Instead of materializing row pairs it records the matched set as a /// single index -- the start of the matched suffix of the buffered side -- and slices the /// buffered batch at that index once every streamed partition has been consumed. @@ -216,8 +219,26 @@ use crate::{ /// min value: 200 /// ``` /// -/// For both types of joins, the buffered side must be sorted ascending for `Operator::Lt` (<) or -/// `Operator::LtEq` (<=) and descending for `Operator::Gt` (>) or `Operator::GtEq` (>=). +/// `RightSemi`/`RightAnti` mark the streamed (right) side, which is +/// `RightExistencePWMJStream` (see `right_existence_join.rs`). Asking whether any buffered +/// row matches a given streamed row is, for a single range predicate, decided by one buffered +/// key -- the minimum for `<`/`<=`, the maximum for `>`/`>=`. So the buffered side is folded +/// down to that key as it arrives and never materialized, and each streamed batch is compared +/// against it, filtered, and emitted straight away rather than at the end. These are the only +/// join types here that require no ordered input and hold `O(1)` state. +/// +/// ```text +/// // Using the example of a less than `<` operation +/// let min = min_batch(buffered) // folded per batch, nothing retained +/// +/// for stream_row in stream_batch: +/// if min < stream_row: // some buffered row matches +/// output stream_row +/// ``` +/// +/// Except for `RightSemi`/`RightAnti`, the buffered side must be sorted ascending for +/// `Operator::Lt` (<) or `Operator::LtEq` (<=) and descending for `Operator::Gt` (>) or +/// `Operator::GtEq` (>=). /// /// # Partitioning Logic /// Piecewise Merge Join requires one buffered side partition + round robin partitioned stream side. A counter @@ -225,6 +246,11 @@ use crate::{ /// for processing the rest of the unmatched rows for Left and Full joins. The last partition that finishes /// execution will be responsible for outputting the unmatched rows. /// +/// `RightSemi`/`RightAnti` need no such coordination: each streamed row is decided on its own, +/// so every partition emits its own output and none has a final pass to run. They also place no +/// single-partition requirement on the buffered side -- a min/max combines across partitions, so +/// each is folded on its own task rather than funnelled through a `CoalescePartitionsExec`. +/// /// # Performance Explanation (cost) /// Piecewise Merge Join is used over Nested Loop Join due to its superior performance. Here is the breakdown: /// @@ -266,8 +292,12 @@ pub struct PiecewiseMergeJoinExec { pub join_type: JoinType, /// The schema once the join is applied schema: SchemaRef, - /// Buffered data + /// Buffered data, collected once and shared by every streamed partition. Unused by right + /// existence joins, which take `buffered_extreme_fut` instead. buffered_fut: OnceAsync, + /// Right existence joins only: the buffered side folded down to one key, so that side is + /// never materialized. Unused by every other join type. + buffered_extreme_fut: OnceAsync, /// Execution metrics metrics: ExecutionPlanMetricsSet, @@ -296,9 +326,8 @@ impl PiecewiseMergeJoinExec { join_type: JoinType, num_partitions: usize, ) -> Result { - // Left Semi/Anti are handled by `ExistencePWMJStream` (the marked side is - // already the buffered side, so no input swap is needed). Right existence joins - // and Mark joins are not yet supported. + // Semi/Anti joins are handled by the existence streams; Mark joins are not + // supported yet. if is_existence_join(join_type) && !is_supported_existence_join(join_type) { return not_impl_err!( "Existence join {join_type} is currently not supported for PiecewiseMergeJoin" @@ -308,22 +337,8 @@ impl PiecewiseMergeJoinExec { // Take the operator and enforce a sort order on the streamed + buffered side based on // the operator type. let sort_options = match operator { - Operator::Lt | Operator::LtEq => { - // For left existence joins the inputs will be swapped so the sort - // options are switched - if is_right_existence_join(join_type) { - SortOptions::new(false, true) - } else { - SortOptions::new(true, true) - } - } - Operator::Gt | Operator::GtEq => { - if is_right_existence_join(join_type) { - SortOptions::new(true, true) - } else { - SortOptions::new(false, true) - } - } + Operator::Lt | Operator::LtEq => SortOptions::new(true, true), + Operator::Gt | Operator::GtEq => SortOptions::new(false, true), _ => { return internal_err!( "Cannot contain non-range operator in PiecewiseMergeJoinExec" @@ -373,6 +388,7 @@ impl PiecewiseMergeJoinExec { join_type, schema, buffered_fut: Default::default(), + buffered_extreme_fut: Default::default(), metrics: ExecutionPlanMetricsSet::new(), left_child_plan_required_order, right_batch_required_orders, @@ -456,13 +472,17 @@ impl PiecewiseMergeJoinExec { // more testing. fn maintains_input_order(join_type: JoinType) -> Vec { match join_type { + // One output batch per streamed batch, in arrival order, with rows only ever + // *removed* by a filter, so each output partition keeps its streamed partition's + // order. The buffered side contributes no output column, hence `false` there. + // `RightMark` is excluded: it adds a `mark` column, so its orderings would not + // map across unchanged. + JoinType::RightSemi | JoinType::RightAnti => vec![false, true], // The existence side is expected to come in sorted JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { vec![false, false] } - JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { - vec![false, false] - } + JoinType::RightMark => vec![false, false], // Left, Right, Full, Inner Join is not guaranteed to maintain // input order as the streamed side will be sorted during // execution for `PiecewiseMergeJoin` @@ -474,6 +494,39 @@ impl PiecewiseMergeJoinExec { pub fn swap_inputs(&self) -> Result> { todo!() } + + /// Sets up the buffered-side collection used by the classic and left existence streams: + /// buffered partition 0 is executed here, and folded into one sorted batch when the + /// returned future is first polled. + /// + /// Right existence joins use `buffered_extreme_fut` instead, consuming every buffered + /// partition themselves. + fn buffered_side( + &self, + context: &Arc, + on_buffered: &PhysicalExprRef, + metrics: &BuildProbeJoinMetrics, + streamed_partitions: usize, + ) -> Result { + let buffered_fut = self.buffered_fut.try_once(|| { + let reservation = MemoryConsumer::new("PiecewiseMergeJoinInput") + .register(context.memory_pool()); + + let buffered_stream = self.buffered.execute(0, Arc::clone(context))?; + Ok(build_buffered_data( + buffered_stream, + Arc::clone(on_buffered), + metrics.clone(), + reservation, + build_visited_indices_map(self.join_type), + streamed_partitions, + )) + })?; + + Ok(BufferedSide::Initial(BufferedSideInitialState { + buffered_fut, + })) + } } impl ExecutionPlan for PiecewiseMergeJoinExec { @@ -502,23 +555,52 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + // Right existence joins reduce the buffered side to a min/max, which combines across + // partitions, so that side keeps whatever parallelism the plan gave it -- no + // `CoalescePartitionsExec` funnelling every buffered row through one thread. Every + // other join type walks the buffered side as a single sorted run and does need it. + let buffered = if is_supported_right_existence_join(self.join_type) { + Distribution::UnspecifiedDistribution + } else { + Distribution::SinglePartition + }; crate::InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, + buffered, Distribution::UnspecifiedDistribution, ]) } + fn benefits_from_input_partitioning(&self) -> Vec { + // Derived exactly as the default does, from this operator's own distribution + // requirements, so the two cannot drift apart as those change. + let mut benefits: Vec = self + .input_distribution_requirements() + .per_child_distributions() + .map(|dist| !matches!(dist, Distribution::SinglePartition)) + .collect(); + + // One deviation: right existence joins ask for `UnspecifiedDistribution` on the buffered + // side, which that rule reads as "worth fanning out". Folding a batch is one linear scan, + // which does not pay for a channel hop, so decline the round-robin `RepartitionExec`. + if is_supported_right_existence_join(self.join_type) + && let Some(buffered) = benefits.first_mut() + { + *buffered = false; + } + + benefits + } + fn required_input_ordering(&self) -> Vec> { - // Existence joins don't need to be sorted on one side. - if is_right_existence_join(self.join_type) { - // Unreachable: `try_new` rejects right existence joins, and this signature - // cannot return `Result`. They swap the inputs, so whoever implements them - // must require the order on the streamed side instead. - unimplemented!( - "required_input_ordering for right existence joins; guarded by try_new" - ) + // Right existence joins read nothing but a single min/max off the buffered side, which + // is `O(B)` from any order, so they require none. Every other join type walks the + // buffered side in order and does. + // + // The streamed side never carries a requirement: the classic and left existence + // streams sort each batch in memory, and the right existence stream needs no order. + if is_supported_right_existence_join(self.join_type) { + vec![None, None] } else { - // Sort the right side in memory, so we do not need to enforce any sorting vec![ Some(OrderingRequirements::from( self.left_child_plan_required_order.clone(), @@ -556,6 +638,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { // Re-set state. metrics: ExecutionPlanMetricsSet::new(), buffered_fut: Default::default(), + buffered_extreme_fut: Default::default(), })) } ChildrenPropertiesMode::Recompute => match &children[..] { @@ -621,59 +704,98 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { // has a single partition), otherwise the counter never reaches 1 and the final // pass is skipped. let streamed_partitions = self.streamed.output_partitioning().partition_count(); - let buffered_fut = self.buffered_fut.try_once(|| { - let reservation = MemoryConsumer::new("PiecewiseMergeJoinInput") - .register(context.memory_pool()); - - let buffered_stream = self.buffered.execute(0, Arc::clone(&context))?; - Ok(build_buffered_data( - buffered_stream, - Arc::clone(&on_buffered), - metrics.clone(), - reservation, - build_visited_indices_map(self.join_type), - streamed_partitions, - )) - })?; - - let streamed = self.streamed.execute(partition, Arc::clone(&context))?; let batch_size = context.session_config().batch_size(); - - let buffered_side = - BufferedSide::Initial(BufferedSideInitialState { buffered_fut }); - - if is_supported_existence_join(self.join_type) { - Ok(Box::pin(ExistencePWMJStream::try_new( - Arc::clone(&self.schema), - on_streamed, - self.join_type, - self.operator, - streamed, - buffered_side, - self.sort_options, - metrics, - batch_size, - ))) - } else if is_existence_join(self.join_type) { - // Right existence joins and Mark joins are rejected in `try_new`. - internal_err!( + match self.join_type { + // Right existence joins never read a buffered *row*, only a single min/max over + // the whole side, so they fold the buffered input away as it arrives instead of + // collecting it. + JoinType::RightSemi | JoinType::RightAnti => { + // `∃b. b < s` is decided by the smallest buffered key, `∃b. b > s` by the + // largest, so the operator alone picks the extreme. + let descending = matches!(self.operator, Operator::Gt | Operator::GtEq); + let extreme_fut = self.buffered_extreme_fut.try_once(|| { + let reservation = + MemoryConsumer::new("PiecewiseMergeJoinBufferedExtreme") + .register(context.memory_pool()); + + // Every buffered partition, not just partition 0: this join type does not + // require the buffered side coalesced, so it must consume all of it. + let buffered_partitions = + self.buffered.output_partitioning().partition_count(); + let buffered_streams = (0..buffered_partitions) + .map(|p| self.buffered.execute(p, Arc::clone(&context))) + .collect::>>()?; + Ok(build_buffered_extreme( + buffered_streams, + self.buffered.schema(), + Arc::clone(&on_buffered), + metrics.clone(), + reservation, + Arc::clone(context.memory_pool()), + descending, + )) + })?; + + let streamed = self.streamed.execute(partition, Arc::clone(&context))?; + + Ok(Box::pin(RightExistencePWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + extreme_fut, + metrics, + ))) + } + JoinType::LeftSemi | JoinType::LeftAnti => { + let buffered_side = self.buffered_side( + &context, + &on_buffered, + &metrics, + streamed_partitions, + )?; + let streamed = self.streamed.execute(partition, Arc::clone(&context))?; + + Ok(Box::pin(ExistencePWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + buffered_side, + self.sort_options, + metrics, + batch_size, + ))) + } + JoinType::LeftMark | JoinType::RightMark => internal_err!( "PiecewiseMergeJoin does not support existence join {} (should have been rejected in try_new)", self.join_type - ) - } else { - Ok(Box::pin(ClassicPWMJStream::try_new( - Arc::clone(&self.schema), - on_streamed, - self.join_type, - self.operator, - streamed, - buffered_side, - PiecewiseMergeJoinStreamState::WaitBufferedSide, - self.sort_options, - metrics, - batch_size, - ))) + ), + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + let buffered_side = self.buffered_side( + &context, + &on_buffered, + &metrics, + streamed_partitions, + )?; + let streamed = self.streamed.execute(partition, Arc::clone(&context))?; + + Ok(Box::pin(ClassicPWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + buffered_side, + PiecewiseMergeJoinStreamState::WaitBufferedSide, + self.sort_options, + metrics, + batch_size, + ))) + } } } @@ -816,6 +938,147 @@ impl BufferedSideData { } } +/// The entire buffered side of a right existence join, reduced to the one key that decides +/// every streamed row -- the minimum for `<`/`<=`, the maximum for `>`/`>=`. +/// +/// Right existence joins never look at a buffered *row*, so unlike [`BufferedSideData`] this +/// holds no batch: the buffered input is folded away as it streams in and the state is `O(1)` +/// however large that side is. See `right_existence_join.rs`. +pub(super) struct BufferedExtreme { + /// One-row array. The row is null exactly when no buffered key is non-null (an empty or + /// all-NULL buffered side), in which case nothing can ever match. + extreme: ArrayRef, + _reservation: MemoryReservation, +} + +impl BufferedExtreme { + pub(super) fn extreme(&self) -> &ArrayRef { + &self.extreme + } +} + +/// Reduces one buffered partition to a single extreme key, dropping each batch as it goes. +/// `None` when the partition produced no batch at all. +async fn partition_extreme( + mut buffered: SendableRecordBatchStream, + on_buffered: PhysicalExprRef, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + descending: bool, +) -> Result> { + let mut extreme: Option = None; + + while let Some(batch) = buffered.next().await.transpose()? { + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + + let keys = on_buffered.evaluate(&batch)?.into_array(batch.num_rows())?; + + // Reduced and dropped within this iteration, but as wide as the batch, and every + // partition folds one of these at once. Resized rather than grown, so what the pool + // sees is the widest key array in flight here and not their sum. + reservation.try_resize(keys.get_array_memory_size())?; + + let batch_extreme = extreme_key(&keys, descending)?; + + // Re-reduced as a pair rather than compared, so the running value is ordered exactly + // as each single reduction was. `extreme_key` ignores nulls, so a null from an + // all-NULL batch never displaces a real key. + extreme = Some(match extreme { + Some(running) => { + let pair = concat(&[running.as_ref(), batch_extreme.as_ref()])?; + extreme_key(&pair, descending)? + } + None => batch_extreme, + }); + } + + Ok(extreme) +} + +/// Folds every buffered partition down to one extreme key. `O(B)` time, and the only state it +/// retains is that one key -- no buffered batch is concatenated or held. The transient cost is +/// one key array per folding partition, which each task accounts against the pool for as long +/// as it holds it. +/// +/// A min/max combines across partitions, so this side needs no single-partition funnel -- +/// `input_distribution_requirements` asks for `UnspecifiedDistribution` here and the partitions +/// are folded independently, each in its own spawned task. +async fn build_buffered_extreme( + buffered_streams: Vec, + buffered_schema: SchemaRef, + on_buffered: PhysicalExprRef, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + memory_pool: Arc, + descending: bool, +) -> Result { + let tasks: Vec<_> = buffered_streams + .into_iter() + .enumerate() + .map(|(partition, stream)| { + let on_buffered = Arc::clone(&on_buffered); + let metrics = metrics.clone(); + // One reservation per task rather than a shared one: `MemoryReservation` is not + // shareable, and each task's transient is freed as soon as it finishes. + let reservation = MemoryConsumer::new(format!( + "PiecewiseMergeJoinBufferedFold[{partition}]" + )) + .register(&memory_pool); + SpawnedTask::spawn(partition_extreme( + stream, + on_buffered, + metrics, + reservation, + descending, + )) + }) + .collect(); + + // The tasks run concurrently; this only collects them. Awaiting in order is fine, and a + // failure propagates after the rest have been joined, since `SpawnedTask` aborts on drop. + let mut extreme: Option = None; + for task in tasks { + let partition_extreme = task.join_unwind().await.map_err(|e| { + internal_datafusion_err!("buffered extreme task failed: {e}") + })??; + if let Some(partition_extreme) = partition_extreme { + // Re-reduced as a pair, exactly as the batches within a partition were, so a value + // accumulated across partitions is ordered by the same rule. + extreme = Some(match extreme { + Some(running) => { + let pair = concat(&[running.as_ref(), partition_extreme.as_ref()])?; + extreme_key(&pair, descending)? + } + None => partition_extreme, + }); + } + } + + // No partition produced a batch, so there was nothing to reduce. Evaluating the key + // expression over an empty batch gives a correctly typed empty array, whose reduction is + // the null that represents "no buffered key". + let extreme = match extreme { + Some(extreme) => extreme, + None => { + let empty = RecordBatch::new_empty(buffered_schema); + let keys = on_buffered.evaluate(&empty)?.into_array(0)?; + extreme_key(&keys, descending)? + } + }; + + // The one-row extreme is all this join type ever holds -- a few bytes, and unrelated to the + // size of the buffered input. + let size = extreme.get_array_memory_size(); + reservation.try_grow(size)?; + metrics.build_mem_used.add(size); + + Ok(BufferedExtreme { + extreme, + _reservation: reservation, + }) +} + pub(super) enum BufferedSide { /// Indicates that build-side not collected yet Initial(BufferedSideInitialState), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs index b2c5212999f3b..9678499d914e9 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs @@ -18,9 +18,10 @@ //! PiecewiseMergeJoin stream specialized for existence joins. //! //! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `LeftSemi` or `LeftAnti`. -//! The other existence joins are rejected in `PiecewiseMergeJoinExec::try_new`: -//! `RightSemi`/`RightAnti`/`RightMark` mark the right side and so need an input swap, while -//! `LeftMark` marks the left side but needs an extra boolean column rather than a slice. +//! `RightSemi`/`RightAnti` mark the streamed side instead and are served by +//! `RightExistencePWMJStream` (see `right_existence_join.rs`); the Mark joins are rejected in +//! `PiecewiseMergeJoinExec::try_new`, as they need an extra boolean column rather than a +//! subset of one side's rows. //! //! # Motivation //! @@ -450,7 +451,7 @@ impl ExistencePWMJStream { /// Numeric, temporal, string, binary and boolean keys get a typed arrow kernel -- a linear /// scan that allocates nothing. Dictionary and nested keys fall to `min_max_batch_generic`, a /// `ScalarValue`-per-row comparator loop; specializing those is left to a follow-up. -fn extreme_key(values: &ArrayRef, descending: bool) -> Result { +pub(super) fn extreme_key(values: &ArrayRef, descending: bool) -> Result { let extreme = if descending { max_batch(values)? } else { @@ -771,11 +772,8 @@ mod tests { Ok(()) } - /// The unsupported existence joins must be rejected at construction, not deeper in. - /// `required_input_ordering` still has an `unimplemented!()` for right existence joins - /// and cannot return an error, so this test is what keeps that panic unreachable: if - /// someone opens the gate for RightSemi/RightAnti without also supplying an ordering - /// requirement, this fails instead of panicking the optimizer at runtime. + /// The Mark joins must be rejected at construction, not deeper in: `execute` has no way + /// to fall back, and the planner has already committed to this operator by then. #[test] fn try_new_rejects_unsupported_existence_joins() -> Result<()> { let left = build_table( @@ -793,12 +791,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, ); - for join_type in [ - JoinType::RightSemi, - JoinType::RightAnti, - JoinType::LeftMark, - JoinType::RightMark, - ] { + for join_type in [JoinType::LeftMark, JoinType::RightMark] { let err = PiecewiseMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs index 8c6815ad6c631..23fa3261a918e 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs @@ -22,4 +22,5 @@ pub use exec::PiecewiseMergeJoinExec; mod classic_join; mod exec; mod existence_join; +mod right_existence_join; mod utils; diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/right_existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/right_existence_join.rs new file mode 100644 index 0000000000000..80c91181f64ab --- /dev/null +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/right_existence_join.rs @@ -0,0 +1,846 @@ +// 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. + +//! PiecewiseMergeJoin stream specialized for right existence joins. +//! +//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `RightSemi` or +//! `RightAnti`. `LeftSemi`/`LeftAnti` are served by `ExistencePWMJStream` (see +//! `existence_join.rs`); the Mark joins are still rejected in +//! `PiecewiseMergeJoinExec::try_new`. +//! +//! # Algorithm +//! +//! Left and right existence joins mark opposite sides, and for a single range predicate +//! that difference is not symmetric — it collapses the work. +//! +//! `LeftSemi`/`LeftAnti` ask, for each *buffered* row, whether any streamed row matches, so +//! the answer depends on the whole streamed side and can only be emitted once it has all +//! been read. `RightSemi`/`RightAnti` ask the mirror question — for each *streamed* row, +//! does any buffered row match? — and with only `buffered_key OP streamed_key` to satisfy, +//! that is decided by a single buffered key: +//! +//! ```text +//! ∃b. b < s ⟺ min(b) < s ∃b. b > s ⟺ max(b) > s +//! ∃b. b <= s ⟺ min(b) <= s ∃b. b >= s ⟺ max(b) >= s +//! ``` +//! +//! The buffered side is therefore reduced to that one key -- `min_batch`/`max_batch`, folded +//! batch by batch as it streams in -- and each batch is dropped once folded. It is never +//! concatenated or retained, so the state this join holds is a single row however large that +//! side is. The reduction is shared, so every streamed partition reads the same key rather +//! than repeating the pass: +//! +//! ```text +//! operator `<`, buffered keys [NULL, 5, 9, 7] -> min = 5 +//! `b < s` holds for some b iff `5 < s` +//! ``` +//! +//! A min/max is `O(B)` from any order, so `required_input_ordering` returns nothing for either +//! child and no `SortExec` is planned. +//! +//! Every streamed row is then decided by comparing it against that one key, which is a +//! vectorized `cmp` kernel per batch and a filter. Nothing about a streamed row depends on +//! any other, so: +//! +//! * output is produced per batch as it arrives — no watermark, no final pass, and no +//! election among the streamed partitions, +//! * all N streamed partitions produce output, rather than one non-empty partition, +//! * the streamed side is never sorted, not even per batch. +//! +//! Rows whose join key is NULL never satisfy a comparison predicate. `min_batch`/`max_batch` +//! ignore NULLs, so the reduced key is null only when *every* buffered key is (or the buffered +//! side is empty) — and then no streamed row can match at all, which makes `RightSemi` empty +//! without reading the streamed side and `RightAnti` a passthrough of it. A NULL streamed key +//! makes its comparison NULL rather than false, which is "no match" — dropped by `RightSemi`, +//! kept by `RightAnti`. +//! +//! Picking the extreme and comparing against it must agree on ordering, and they do: arrow's +//! `cmp` kernels and its `min`/`max` accumulators both order floats by `total_cmp` +//! (`arrow-arith`'s `MinAccumulator` compares with `ArrowNativeTypeOp::is_lt`, seeded from +//! `MAX_TOTAL_ORDER`), so `-0.0` and NaN are treated identically on both sides. +//! +//! # Cost +//! +//! Let `B` be the buffered rows and `S` the streamed rows: `O(B)` to reduce the buffered side +//! plus `O(S)` to filter, and no sort on either side. The state retained across batches is one +//! row; while folding, each buffered partition also holds the key array of the batch it is +//! reducing, which it accounts against the memory pool for that long. +//! +//! This is why the shared state is [`BufferedExtreme`] and not `BufferedSideData`: every field +//! of the latter -- the concatenated batch, the key array, the visited-indices bitmap, the +//! final-pass counter -- would be dead here. +//! +//! [`PiecewiseMergeJoinExec`]: super::PiecewiseMergeJoinExec + +use std::sync::Arc; +use std::task::{Poll, ready}; + +use arrow::array::{Array, ArrayRef, RecordBatch, Scalar}; +use arrow::compute::filter_record_batch; +use arrow::compute::kernels::boolean::not; +use arrow::compute::kernels::cmp::{gt, gt_eq, lt, lt_eq}; +use arrow_schema::SchemaRef; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::PhysicalExprRef; +use futures::{Stream, StreamExt}; + +use crate::handle_state; +use crate::joins::piecewise_merge_join::exec::BufferedExtreme; +use crate::joins::utils::{ + BuildProbeJoinMetrics, OnceFut, StatefulStreamResult, boolean_mask_from_filter, +}; +use crate::stream::EmptyRecordBatchStream; + +pub(super) enum RightExistencePWMJStreamState { + /// Await the buffered side's reduction to a single key. + WaitBufferedExtreme, + /// Fetch streamed batches and emit the rows that do (`RightSemi`) or do not + /// (`RightAnti`) have a buffered match. + ScanStreamBatches, + Completed, +} + +pub(super) struct RightExistencePWMJStream { + /// Output schema, which for `RightSemi`/`RightAnti` is the streamed side's schema + schema: SchemaRef, + /// Physical expression evaluated on the streamed side. The buffered side's + /// equivalent is already evaluated when the buffered side is collected. + on_streamed: PhysicalExprRef, + /// `RightSemi` or `RightAnti` + join_type: JoinType, + /// Comparison operator + operator: Operator, + streamed: SendableRecordBatchStream, + /// Resolves to the whole buffered side reduced to one key. Shared with the other streamed + /// partitions, so that reduction happens exactly once. + buffered_extreme_fut: OnceFut, + state: RightExistencePWMJStreamState, + /// That key, held as a one-element [`Scalar`] so the comparison against a streamed batch is + /// one kernel call. `None` when the buffered side has no non-null key, i.e. nothing can ever + /// match. Only populated once `buffered_extreme_fut` has resolved. + buffered_extreme: Option>, + join_metrics: BuildProbeJoinMetrics, +} + +impl RightExistencePWMJStream { + pub(super) fn try_new( + schema: SchemaRef, + on_streamed: PhysicalExprRef, + join_type: JoinType, + operator: Operator, + streamed: SendableRecordBatchStream, + buffered_extreme_fut: OnceFut, + join_metrics: BuildProbeJoinMetrics, + ) -> Self { + Self { + schema, + on_streamed, + join_type, + operator, + streamed, + buffered_extreme_fut, + state: RightExistencePWMJStreamState::WaitBufferedExtreme, + buffered_extreme: None, + join_metrics, + } + } + + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>> { + loop { + return match self.state { + RightExistencePWMJStreamState::WaitBufferedExtreme => { + handle_state!(ready!(self.collect_buffered_extreme(cx))) + } + RightExistencePWMJStreamState::ScanStreamBatches => { + handle_state!(ready!(self.scan_stream_batch(cx))) + } + RightExistencePWMJStreamState::Completed => Poll::Ready(None), + }; + } + } + + /// Picks up the buffered side's reduced key, which is all this join type needs from it. + fn collect_buffered_extreme( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>>> { + let build_timer = self.join_metrics.build_time.timer(); + let buffered_extreme = ready!(self.buffered_extreme_fut.get_shared(cx))?; + build_timer.done(); + + // Null exactly when no buffered key is non-null, and NULLs match nothing. Cloned out by + // value -- it is one row -- so the shared state is not kept alive by this stream. + let extreme = buffered_extreme.extreme(); + self.buffered_extreme = + (extreme.null_count() == 0).then(|| Scalar::new(Arc::clone(extreme))); + + // With no non-null buffered key nothing matches, so `RightSemi` outputs nothing + // and does not need to read a single streamed batch. `RightAnti` still has to, + // since it outputs all of them. + self.state = match (&self.buffered_extreme, self.join_type) { + (None, JoinType::RightSemi) => { + let streamed_schema = self.streamed.schema(); + self.streamed = Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + RightExistencePWMJStreamState::Completed + } + _ => RightExistencePWMJStreamState::ScanStreamBatches, + }; + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Fetches one streamed batch and emits the rows it contributes, if any. + fn scan_stream_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>>> { + match ready!(self.streamed.poll_next_unpin(cx)) { + None => self.state = RightExistencePWMJStreamState::Completed, + Some(Ok(batch)) => { + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(batch.num_rows()); + + let output = self.filter_streamed_batch(&batch)?; + if output.num_rows() > 0 { + return Poll::Ready(Ok(StatefulStreamResult::Ready(Some(output)))); + } + // Nothing survived; take the next batch rather than yielding an empty one. + } + Some(Err(err)) => return Poll::Ready(Err(err)), + } + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Keeps the streamed rows that have a buffered match (`RightSemi`) or that have none + /// (`RightAnti`), by comparing each against the single buffered extreme. + fn filter_streamed_batch(&self, batch: &RecordBatch) -> Result { + let columns = match &self.buffered_extreme { + // No non-null buffered key, so no streamed row matches and `RightAnti` keeps + // the batch whole. `RightSemi` never gets here: it completed without reading + // the streamed side. + None => batch.columns().to_vec(), + Some(extreme) => { + let stream_values = self + .on_streamed + .evaluate(batch)? + .into_array(batch.num_rows())?; + + // `extreme` is the buffered key, so it goes on the left of the operator, + // matching the `buffered OP streamed` orientation of the predicate. + let matched = match self.operator { + Operator::Lt => lt(extreme, &stream_values), + Operator::LtEq => lt_eq(extreme, &stream_values), + Operator::Gt => gt(extreme, &stream_values), + Operator::GtEq => gt_eq(extreme, &stream_values), + other => { + return internal_err!( + "PiecewiseMergeJoin should not contain operator, {other}" + ); + } + }?; + + let predicate = match self.join_type { + // A NULL streamed key compares NULL rather than false, and `filter` + // already treats NULL as "not selected" -- which is what a + // non-matching row is. + JoinType::RightSemi => matched, + // Anti needs the complement, so those NULLs have to be folded into + // false first: `not(NULL)` is NULL, which would drop a row that + // matched nothing. + _ => not(&boolean_mask_from_filter(&matched))?, + }; + + filter_record_batch(batch, &predicate)?.columns().to_vec() + } + }; + + // Right existence joins output the streamed columns only. The streamed child's + // schema is field-for-field equal to the join's own, but rebuild against the + // latter so the stream's declared schema is what it yields. + Ok(RecordBatch::try_new(Arc::clone(&self.schema), columns)?) + } +} + +impl RecordBatchStream for RightExistencePWMJStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for RightExistencePWMJStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + // `record_poll` fills in `output_rows` and `end_time`; `elapsed_compute` is handled + // by `BuildProbeJoinMetrics::drop`. + let poll = self.poll_next_impl(cx); + self.join_metrics.baseline.record_poll(poll) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sorts::sort::SortExec; + use crate::{ + ExecutionPlan, ExecutionPlanProperties, common, joins::PiecewiseMergeJoinExec, + test::TestMemoryExec, + }; + use arrow::array::Int32Array; + use arrow::compute::SortOptions; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::test_util::batches_to_string; + use datafusion_execution::TaskContext; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; + use insta::assert_snapshot; + + // Coverage for right existence joins also lives in `pwmj.slt` (both correlation + // orientations, all four operators, NULLs, key types) and in the differential fuzz test + // `fuzz_pwmj_existence_matches_nested_loop`, which checks them against + // `NestedLoopJoinExec` over randomized inputs. The tests here pin the parts SQL cannot + // observe: which streamed batches were read, and which partition emitted them. + + fn kv_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("k", DataType::Int32, true), + ])) + } + + fn kv_batch(rows: &[(i32, Option)]) -> RecordBatch { + let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); + let keys: Vec> = rows.iter().map(|(_, k)| *k).collect(); + RecordBatch::try_new( + kv_schema(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(keys)), + ], + ) + .unwrap() + } + + fn kv_exec(partitions: &[Vec]) -> Arc { + TestMemoryExec::try_new_exec(partitions, kv_schema(), None).unwrap() + } + + /// Right existence joins declare no input ordering requirement, so the buffered sides + /// below are deliberately **unsorted**: only their min/max matters. That is also why + /// building the exec directly is faithful here, where the other streams would need a + /// `SortExec` the tests have to supply by hand. + fn join( + buffered: Arc, + streamed: Arc, + operator: Operator, + join_type: JoinType, + ) -> Result { + let on = ( + Arc::new(Column::new_with_schema("k", &buffered.schema())?) as _, + Arc::new(Column::new_with_schema("k", &streamed.schema())?) as _, + ); + let probe = PiecewiseMergeJoinExec::try_new( + Arc::clone(&buffered), + Arc::clone(&streamed), + on.clone(), + operator, + join_type, + 1, + )?; + assert!( + probe.required_input_ordering().iter().all(Option::is_none), + "right existence joins must not require an input ordering" + ); + assert!( + probe + .input_distribution_requirements() + .per_child_distributions() + .all(|d| matches!( + d, + datafusion_physical_expr::Distribution::UnspecifiedDistribution + )), + "right existence joins must not require the buffered side coalesced" + ); + assert_eq!( + probe.benefits_from_input_partitioning(), + vec![false, true], + "the buffered side must not be fanned out just to fold it" + ); + PiecewiseMergeJoinExec::try_new(buffered, streamed, on, operator, join_type, 1) + } + + fn input_batches(join: &PiecewiseMergeJoinExec) -> usize { + join.metrics() + .unwrap() + .sum_by_name("input_batches") + .expect("input_batches metric") + .as_usize() + } + + /// `RightSemi` keeps the streamed rows with at least one buffered match, and outputs + /// only the streamed columns. `>` needs the buffered maximum, 5, so of the streamed keys + /// {4, 5, 6} only 4 has some buffered key above it. + #[tokio::test] + async fn join_right_semi() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(5)), (2, Some(1)), (3, Some(2))])]]), + kv_exec(&[vec![kv_batch(&[ + (10, Some(4)), + (20, Some(5)), + (30, Some(6)), + ])]]), + Operator::Gt, + JoinType::RightSemi, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+ + | id | k | + +----+---+ + | 10 | 4 | + +----+---+ + "); + Ok(()) + } + + /// `RightAnti` is the complement: the streamed rows with no buffered match. + #[tokio::test] + async fn join_right_anti() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(5)), (2, Some(1)), (3, Some(2))])]]), + kv_exec(&[vec![kv_batch(&[ + (10, Some(4)), + (20, Some(5)), + (30, Some(6)), + ])]]), + Operator::Gt, + JoinType::RightAnti, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+ + | id | k | + +----+---+ + | 20 | 5 | + | 30 | 6 | + +----+---+ + "); + Ok(()) + } + + /// A NULL key satisfies no comparison predicate, on either side. + /// + /// The buffered side is sorted descending with NULLs first, as `<` requires, so the key + /// this operator needs -- the minimum, 5 -- is the last row and the NULL is nowhere near + /// it. On the streamed side the NULL-keyed row matches nothing, so `RightSemi` drops it + /// and `RightAnti` keeps it. + #[tokio::test] + async fn null_keys_match_nothing() -> Result<()> { + let buffered = + || kv_exec(&[vec![kv_batch(&[(1, Some(9)), (2, None), (3, Some(5))])]]); + let streamed = || { + kv_exec(&[vec![kv_batch(&[ + (10, Some(5)), + (20, Some(6)), + (30, None), + (40, Some(100)), + ])]]) + }; + + let semi = join(buffered(), streamed(), Operator::Lt, JoinType::RightSemi)?; + let batches = + common::collect(semi.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+-----+ + | id | k | + +----+-----+ + | 20 | 6 | + | 40 | 100 | + +----+-----+ + "); + + let anti = join(buffered(), streamed(), Operator::Lt, JoinType::RightAnti)?; + let batches = + common::collect(anti.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+ + | id | k | + +----+---+ + | 10 | 5 | + | 30 | | + +----+---+ + "); + Ok(()) + } + + /// With no non-null buffered key nothing can match, so `RightSemi` is empty however many + /// streamed rows there are -- and it does not read a single one of them. Asserted through + /// the `input_batches` metric, which SQL cannot observe. + #[tokio::test] + async fn all_null_buffered_side_makes_right_semi_read_nothing() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, None), (2, None)])]]), + kv_exec(&[vec![ + kv_batch(&[(10, Some(4))]), + kv_batch(&[(20, Some(5))]), + kv_batch(&[(30, Some(6))]), + ]]), + Operator::Gt, + JoinType::RightSemi, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + assert_eq!(input_batches(&join), 0, "no streamed batch should be read"); + Ok(()) + } + + /// The same buffered side, here empty rather than all-NULL, sends every streamed row to + /// `RightAnti` instead. It has to read them all: they are the output. + #[tokio::test] + async fn empty_buffered_side_passes_right_anti_through() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[])]]), + kv_exec(&[vec![ + kv_batch(&[(10, Some(4))]), + kv_batch(&[(20, None)]), + kv_batch(&[(30, Some(6))]), + ]]), + Operator::Gt, + JoinType::RightAnti, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+ + | id | k | + +----+---+ + | 10 | 4 | + | 20 | | + | 30 | 6 | + +----+---+ + "); + assert_eq!(input_batches(&join), 3); + Ok(()) + } + + /// The streamed side's ordering survives this join -- one output batch per streamed batch, in + /// order, with rows only removed -- so `maintains_input_order` claims it and the operator + /// advertises the streamed child's ordering as its own. That lets a downstream operator skip + /// a re-sort, which means a wrong claim here would produce wrong results, not just a slow + /// plan. The runtime side of it is pinned by the row-order snapshots above. + #[test] + fn output_ordering_follows_the_streamed_side() -> Result<()> { + let streamed_input = kv_exec(&[vec![kv_batch(&[(20, Some(4)), (10, Some(9))])]]); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new_with_schema("id", &streamed_input.schema())?), + SortOptions::new(true, false), + )]) + .unwrap(); + let streamed = Arc::new(SortExec::new(ordering.clone(), streamed_input)); + + for join_type in [JoinType::RightSemi, JoinType::RightAnti] { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(5))])]]), + Arc::clone(&streamed) as _, + Operator::Gt, + join_type, + )?; + assert_eq!( + join.properties().output_ordering(), + Some(&ordering), + "{join_type} should advertise the streamed side's ordering" + ); + } + + // The buffered side contributes no column, so its ordering must not leak out: a buffered + // child sorted on `k` cannot make the join claim anything. + let buffered_ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("k", 1)), + SortOptions::new(false, true), + )]) + .unwrap(); + let sorted_buffered = Arc::new(SortExec::new( + buffered_ordering, + kv_exec(&[vec![kv_batch(&[(1, Some(5))])]]), + )); + let join = join( + sorted_buffered as _, + kv_exec(&[vec![kv_batch(&[(20, Some(4))])]]), + Operator::Gt, + JoinType::RightSemi, + )?; + assert_eq!(join.properties().output_ordering(), None); + Ok(()) + } + + /// A zero-row streamed batch still reaches the comparison kernel, which yields an empty + /// mask rather than erroring, so the batch is filtered away and the surrounding batches + /// are unaffected. Covered for both join types since only anti also runs `not` over it. + #[tokio::test] + async fn empty_streamed_batch_is_skipped() -> Result<()> { + for (join_type, expected) in [(JoinType::RightSemi, 1), (JoinType::RightAnti, 1)] + { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(5)), (2, Some(1))])]]), + kv_exec(&[vec![ + kv_batch(&[]), + kv_batch(&[(10, Some(4)), (20, Some(9))]), + kv_batch(&[]), + ]]), + Operator::Gt, + join_type, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + // Buffered maximum is 5, so `>` keeps 4 for semi and 9 for anti -- one row each, + // and no empty batch in between. + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + expected, + "{join_type}" + ); + assert!( + batches.iter().all(|b| b.num_rows() > 0), + "{join_type} should not yield an empty batch" + ); + assert_eq!(input_batches(&join), 3, "{join_type}"); + } + Ok(()) + } + + /// The buffered side is folded to one key as it streams in and never materialized, so this + /// join runs in a memory pool far smaller than that side. 20k buffered rows is >80 KB of + /// Int32 keys alone and ~160 KB of batches; the pool here is 64 KB, which the collecting + /// path cannot fit -- it reserves every batch, then the concatenation, then the key array. + /// The fold needs only the one batch's key array it is reducing (~8 KB), which it does + /// reserve, so the pool is sized above that and well below the collecting path. + /// + /// This is the assertion that pins the retained buffered state at one row. `build_mem_used` + /// is checked too, since a regression to collecting would show up there even under a + /// generous pool. + #[tokio::test] + async fn buffered_side_is_folded_not_collected() -> Result<()> { + let buffered: Vec = (0..10) + .map(|b| { + kv_batch( + &(0..2000) + .map(|i| (i, Some(b * 2000 + i))) + .collect::>(), + ) + }) + .collect(); + + let join = join( + kv_exec(&[buffered]), + kv_exec(&[vec![kv_batch(&[(1, Some(0)), (2, Some(19_999))])]]), + Operator::Gt, + JoinType::RightSemi, + )?; + + let task_ctx = Arc::new( + TaskContext::default().with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024))) + .build()?, + )), + ); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // Buffered maximum is 19999, so `>` keeps only the streamed 0. + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+ + | id | k | + +----+---+ + | 1 | 0 | + +----+---+ + "); + + // What is held is one row, not 20k. + let build_mem_used = join + .metrics() + .unwrap() + .sum_by_name("build_mem_used") + .expect("build_mem_used metric") + .as_usize(); + assert!( + build_mem_used < 1024, + "buffered state should be a single key, got {build_mem_used} bytes" + ); + Ok(()) + } + + /// The key array a partition is reducing is transient but as wide as the batch, and it is + /// reserved for as long as it is held. A pool that cannot fit one of them fails the join + /// rather than allocating outside the accounting -- which is what keeps the memory claim + /// honest, since without the reservation this fold would simply succeed. + #[tokio::test] + async fn folding_reserves_the_batch_key_array() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch( + &(0..2000).map(|i| (i, Some(i))).collect::>(), + )]]), + kv_exec(&[vec![kv_batch(&[(1, Some(0))])]]), + Operator::Gt, + JoinType::RightSemi, + )?; + + let task_ctx = Arc::new( + TaskContext::default().with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024))) + .build()?, + )), + ); + let err = common::collect(join.execute(0, task_ctx)?) + .await + .expect_err("a 2000-row key array should not fit a 1 KB pool"); + assert!( + err.to_string().contains("PiecewiseMergeJoinBufferedFold"), + "expected the fold's reservation to fail, got: {err}" + ); + Ok(()) + } + + /// A min/max combines across partitions, so the buffered side is not required to be + /// coalesced -- and every partition of it has to be consumed. The deciding key here is in + /// the **last** buffered partition, so reading only partition 0 (which the other streams do, + /// since they require `SinglePartition`) silently loses it and changes the answer. + #[tokio::test] + async fn every_buffered_partition_is_folded() -> Result<()> { + // Buffered maxima per partition: 1, 3, 50. Only the last admits the streamed 40. + let buffered = vec![ + vec![kv_batch(&[(1, Some(0)), (2, Some(1))])], + vec![kv_batch(&[(3, Some(3))])], + vec![kv_batch(&[(4, Some(2)), (5, Some(50))])], + ]; + let streamed = || kv_exec(&[vec![kv_batch(&[(10, Some(2)), (20, Some(40))])]]); + + let semi = join( + kv_exec(&buffered), + streamed(), + Operator::Gt, + JoinType::RightSemi, + )?; + assert_eq!(semi.buffered().output_partitioning().partition_count(), 3); + let batches = + common::collect(semi.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+ + | id | k | + +----+----+ + | 10 | 2 | + | 20 | 40 | + +----+----+ + "); + + let anti = join( + kv_exec(&buffered), + streamed(), + Operator::Gt, + JoinType::RightAnti, + )?; + let batches = + common::collect(anti.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + + // All three partitions were read, not just the first. + assert_eq!(input_batches(&semi), 1, "streamed batches"); + let build_batches = semi + .metrics() + .unwrap() + .sum_by_name("build_input_batches") + .expect("build_input_batches metric") + .as_usize(); + assert_eq!( + build_batches, 3, + "every buffered partition should be folded" + ); + Ok(()) + } + + /// Unlike `LeftSemi`/`LeftAnti`, where one elected partition emits everything at the end, + /// each streamed partition here emits its own rows as it reads them -- so the join's N + /// advertised output partitions are all live. + #[tokio::test] + async fn every_streamed_partition_emits() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(5)), (2, Some(1))])]]), + kv_exec(&[ + vec![kv_batch(&[(10, Some(4)), (20, Some(9))])], + vec![kv_batch(&[(30, Some(0)), (40, Some(7))])], + ]), + Operator::Gt, + JoinType::RightSemi, + )?; + + assert_eq!(join.properties().output_partitioning().partition_count(), 2); + + let task_ctx = Arc::new(TaskContext::default()); + let mut per_partition = Vec::new(); + for partition in 0..2 { + let stream = join.execute(partition, Arc::clone(&task_ctx))?; + per_partition.push(common::collect(stream).await?); + } + + // Buffered maximum is 5: partition 0 keeps 4, partition 1 keeps 0. + for (partition, batches) in per_partition.iter().enumerate() { + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + 1, + "partition {partition} should have emitted its own row" + ); + } + let out = arrow::compute::concat_batches( + &join.schema(), + per_partition.iter().flatten(), + )?; + assert_snapshot!(batches_to_string(&[out]), @r" + +----+---+ + | id | k | + +----+---+ + | 10 | 4 | + | 30 | 0 | + +----+---+ + "); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs index 5093be0ca19be..604be024f6e6e 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs @@ -17,12 +17,14 @@ use datafusion_expr::JoinType; -// Returns boolean for whether the join is a right existence join -pub(super) fn is_right_existence_join(join_type: JoinType) -> bool { - matches!( - join_type, - JoinType::RightAnti | JoinType::RightSemi | JoinType::RightMark - ) +// Returns boolean for whether the join is a right existence join served by +// `RightExistencePWMJStream`, which reads nothing but a single min/max off the buffered side. +// +// `RightMark` is deliberately excluded even though it is a right existence join: it needs the +// buffered side walked in order and an extra boolean column, so it must not inherit this +// stream's relaxed input requirements if the `try_new` gate is ever loosened. +pub(super) fn is_supported_right_existence_join(join_type: JoinType) -> bool { + matches!(join_type, JoinType::RightSemi | JoinType::RightAnti) } // Returns boolean for whether the join is an existence join @@ -38,12 +40,16 @@ pub(super) fn is_existence_join(join_type: JoinType) -> bool { ) } -// Returns boolean for whether the join is a left existence join that is currently -// supported by `PiecewiseMergeJoin`. These do not require swapping the inputs: the -// marked (left) side is already the buffered side, so `ExistencePWMJStream` can track the -// matched suffix and slice the buffered batch at its start. +// Returns boolean for whether the join is an existence join that is currently supported by +// `PiecewiseMergeJoin`, which is every one of them except the Mark joins pub(super) fn is_supported_existence_join(join_type: JoinType) -> bool { - matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti) + matches!( + join_type, + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + ) } // Returns boolean to check if the join type needs to record @@ -55,17 +61,13 @@ pub(super) fn need_produce_result_in_final(join_type: JoinType) -> bool { // Returns boolean for whether or not we need to build the buffered side // bitmap for marking matched rows on the buffered side. // -// `LeftSemi`/`LeftAnti` are absent on purpose: `ExistencePWMJStream` only ever marks a -// contiguous suffix of the buffered side, so it tracks the boundary as a single index -// (`BufferedSideData::existence_min_marked`) and needs no bitmap. +// The Semi/Anti joins are absent on purpose. `LeftSemi`/`LeftAnti` only ever mark a +// contiguous suffix of the buffered side, so they track the boundary as a single index +// (`BufferedSideData::existence_min_marked`); `RightSemi`/`RightAnti` mark the streamed +// side and never touch the buffered side at all. pub(super) fn build_visited_indices_map(join_type: JoinType) -> bool { matches!( join_type, - JoinType::Full - | JoinType::Left - | JoinType::RightAnti - | JoinType::RightSemi - | JoinType::LeftMark - | JoinType::RightMark + JoinType::Full | JoinType::Left | JoinType::LeftMark | JoinType::RightMark ) } diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index ecf056560e015..d5740e63389fd 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1920,6 +1920,20 @@ pub(crate) fn symmetric_join_output_partitioning( Ok(result) } +/// Convert a boolean filter array into a unified mask bitmap. +/// +/// Caution: The filter result is NOT a bitmap; it contains true/false/null values. +/// For example, `1 < NULL` evaluates to NULL. Therefore, we must combine (AND) +/// the boolean array with its null bitmap to construct a unified bitmap. +#[inline] +pub(crate) fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray { + let (values, nulls) = filter_arr.clone().into_parts(); + match nulls { + Some(nulls) => BooleanArray::new(nulls.inner() & &values, None), + None => BooleanArray::new(values, None), + } +} + pub(crate) fn asymmetric_join_output_partitioning( left: &Arc, right: &Arc, diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index f71a1b667243e..b3a44b7a513be 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -1237,5 +1237,442 @@ physical_plan 05)--------DataSourceExec: partitions=1, partition_sizes=[1] 06)------DataSourceExec: partitions=1, partition_sizes=[1] +# ------------------------------------------------------------------ +# Right existence joins (RightSemi / RightAnti) via PiecewiseMergeJoin +# ------------------------------------------------------------------ +# These mark the streamed (right) side, which for a single range predicate is decided by one +# buffered key: the minimum for `<`/`<=`, the maximum for `>`/`>=`. So unlike LeftSemi/LeftAnti +# there is no watermark and no final pass -- each streamed batch is filtered and emitted as it +# arrives -- and no sort on either side, since a min/max is `O(B)` from any order. `EXISTS` +# always decorrelates to a *Left* existence join, so these are reached through the explicit +# join syntax instead. + +statement ok +CREATE TABLE rex_l(v INT); + +statement ok +INSERT INTO rex_l VALUES (2), (4); + +statement ok +CREATE TABLE rex_r(id INT, v INT); + +statement ok +INSERT INTO rex_r VALUES (1, 1), (2, 2), (3, 4), (4, 5), (5, NULL); + +# `>` : the deciding buffered key is max = 4. 4>{1,2}; 4 is not > 4 or 5; NULL matches nothing. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +1 +2 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +3 +4 +5 + +# `>=` additionally admits r.v = 4 via the equal buffered value, so it must differ from `>`. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v >= r.v ORDER BY 1; +---- +1 +2 +3 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN rex_r r ON l.v >= r.v ORDER BY 1; +---- +4 +5 + +# `<` : the deciding key flips to min = 2. 2<{4,5}; 2 is not < 1 or 2. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v < r.v ORDER BY 1; +---- +3 +4 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN rex_r r ON l.v < r.v ORDER BY 1; +---- +1 +2 +5 + +# `<=` additionally admits r.v = 2, so it must differ from `<`. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v <= r.v ORDER BY 1; +---- +2 +3 +4 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN rex_r r ON l.v <= r.v ORDER BY 1; +---- +1 +5 + +# **Neither** side is sorted: `required_input_ordering` is empty for both children, so the plan +# has no `SortExec` at all. The buffered side only needs its min/max, which is `O(B)` from any +# order -- every other PWMJ join type walks it and does require a global sort. +query TT +EXPLAIN SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +logical_plan +01)Sort: r.id ASC NULLS LAST +02)--Projection: r.id +03)----RightSemi Join: Filter: l.v > r.v +04)------SubqueryAlias: l +05)--------TableScan: rex_l projection=[v] +06)------SubqueryAlias: r +07)--------TableScan: rex_r projection=[id, v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=RightSemi, on=(v > v) +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] + +# A min/max combines across partitions, so the buffered side carries no `SinglePartition` +# requirement and every partition of it is folded. A `UNION ALL` gives it two, and the deciding +# key -- max = 7 -- is in the second one, so folding only the first would return {1,2} here. +statement ok +CREATE TABLE rex_l2(v INT); + +statement ok +INSERT INTO rex_l2 VALUES (3), (7); + +query I +SELECT r.id FROM (SELECT v FROM rex_l UNION ALL SELECT v FROM rex_l2) l +RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +1 +2 +3 +4 + +query I +SELECT r.id FROM (SELECT v FROM rex_l UNION ALL SELECT v FROM rex_l2) l +RIGHT ANTI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +5 + +# `UnionExec` feeds the join directly with both partitions -- no `CoalescePartitionsExec` +# funnelling every buffered row through one thread, which is what `SinglePartition` would force. +# Nor is a `RepartitionExec` added to split it: `benefits_from_input_partitioning` declines that +# for this side, since folding a batch is one linear scan and not worth a channel hop. +query TT +EXPLAIN SELECT r.id FROM (SELECT v FROM rex_l UNION ALL SELECT v FROM rex_l2) l +RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +logical_plan +01)Sort: r.id ASC NULLS LAST +02)--Projection: r.id +03)----RightSemi Join: Filter: l.v > r.v +04)------SubqueryAlias: l +05)--------Union +06)----------TableScan: rex_l projection=[v] +07)----------TableScan: rex_l2 projection=[v] +08)------SubqueryAlias: r +09)--------TableScan: rex_r projection=[id, v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=RightSemi, on=(v > v) +04)------UnionExec +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--------DataSourceExec: partitions=1, partition_sizes=[1] +07)------DataSourceExec: partitions=1, partition_sizes=[1] + +# An empty buffered side has no key at all, so nothing matches: RightSemi is empty and +# RightAnti keeps every streamed row, NULL-keyed ones included. +statement ok +CREATE TABLE rex_empty(v INT); + +query I +SELECT r.id FROM rex_empty l RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- + +query I +SELECT r.id FROM rex_empty l RIGHT ANTI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +1 +2 +3 +4 +5 + +# The same holds when the buffered side is present but every key is NULL, since a NULL key +# matches nothing either. +statement ok +CREATE TABLE rex_null_l(v INT); + +statement ok +INSERT INTO rex_null_l VALUES (NULL), (NULL); + +query I +SELECT r.id FROM rex_null_l l RIGHT SEMI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- + +query I +SELECT r.id FROM rex_null_l l RIGHT ANTI JOIN rex_r r ON l.v > r.v ORDER BY 1; +---- +1 +2 +3 +4 +5 + +# An empty streamed side has nothing to filter, either way. +statement ok +CREATE TABLE rex_empty_r(id INT, v INT); + +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN rex_empty_r r ON l.v > r.v ORDER BY 1; +---- + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN rex_empty_r r ON l.v > r.v ORDER BY 1; +---- + +# `count(*)` needs no column from the join, but the streamed side still has to project the +# join key, so the output is never zero-column -- which `RecordBatch::try_new` would reject. +query I +SELECT count(*) FROM rex_l l RIGHT SEMI JOIN rex_r r ON l.v > r.v; +---- +2 + +# Dictionary-encoded keys compare the buffered extreme as a dictionary `Scalar`, which is a +# different arrow path from the primitive one above. +statement ok +CREATE TABLE rex_dict_l AS + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS v FROM (VALUES ('c'), (NULL)); + +statement ok +CREATE TABLE rex_dict_r AS + SELECT column1 AS id, arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +# `>` : the only buffered key is 'c', so only 'a' has something above it. +query I +SELECT r.id FROM rex_dict_l l RIGHT SEMI JOIN rex_dict_r r ON l.v > r.v ORDER BY 1; +---- +1 + +query I +SELECT r.id FROM rex_dict_l l RIGHT ANTI JOIN rex_dict_r r ON l.v > r.v ORDER BY 1; +---- +2 +3 +4 + +# Float keys are the case where reducing the buffered side and comparing against the result must +# agree on an ordering. Both use a *total* order (arrow's `cmp` kernels and `min`/`max` compare +# floats with `total_cmp`), under which NaN is the largest value rather than incomparable. So the +# buffered maximum below is NaN, and id 5 matches through `NaN > 2.0`; were the comparison IEEE +# while the reduction stayed total, that would be false and no other buffered key is above 2.0, +# so the row would be lost. id 4 is excluded for the mirror reason: nothing is above NaN. +statement ok +CREATE TABLE rex_f_l(v DOUBLE); + +statement ok +INSERT INTO rex_f_l VALUES (1.0), (arrow_cast('NaN', 'Float64')); + +statement ok +CREATE TABLE rex_f_r(id INT, v DOUBLE); + +statement ok +INSERT INTO rex_f_r VALUES + (1, arrow_cast('-Inf', 'Float64')), (2, 0.0), (3, -0.0), + (4, arrow_cast('NaN', 'Float64')), (5, 2.0); + +query I +SELECT r.id FROM rex_f_l l RIGHT SEMI JOIN rex_f_r r ON l.v > r.v ORDER BY 1; +---- +1 +2 +3 +5 + +query I +SELECT r.id FROM rex_f_l l RIGHT ANTI JOIN rex_f_r r ON l.v > r.v ORDER BY 1; +---- +4 + +# No `SortExec` on either side for a float key either -- the operator is chosen for these types, +# not silently replaced by a NestedLoopJoin that would answer the queries above identically. +query TT +EXPLAIN SELECT r.id FROM rex_f_l l RIGHT SEMI JOIN rex_f_r r ON l.v > r.v ORDER BY 1; +---- +logical_plan +01)Sort: r.id ASC NULLS LAST +02)--Projection: r.id +03)----RightSemi Join: Filter: l.v > r.v +04)------SubqueryAlias: l +05)--------TableScan: rex_f_l projection=[v] +06)------SubqueryAlias: r +07)--------TableScan: rex_f_r projection=[id, v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=RightSemi, on=(v > v) +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] + +# String keys reduce through arrow's string kernels rather than the primitive ones: `Utf8` via +# `min_string`/`max_string` and `Utf8View` via their view equivalents. Same shape as the +# dictionary case above -- buffered 'c', so only 'a' has something above it. +statement ok +CREATE TABLE rex_s_l AS + SELECT arrow_cast(column1, 'Utf8') AS v FROM (VALUES ('c'), (NULL)); + +statement ok +CREATE TABLE rex_s_r AS + SELECT column1 AS id, arrow_cast(column2, 'Utf8') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +query I +SELECT r.id FROM rex_s_l l RIGHT SEMI JOIN rex_s_r r ON l.v > r.v ORDER BY 1; +---- +1 + +query I +SELECT r.id FROM rex_s_l l RIGHT ANTI JOIN rex_s_r r ON l.v > r.v ORDER BY 1; +---- +2 +3 +4 + +statement ok +CREATE TABLE rex_sv_l AS + SELECT arrow_cast(column1, 'Utf8View') AS v FROM (VALUES ('c'), (NULL)); + +statement ok +CREATE TABLE rex_sv_r AS + SELECT column1 AS id, arrow_cast(column2, 'Utf8View') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +query I +SELECT r.id FROM rex_sv_l l RIGHT SEMI JOIN rex_sv_r r ON l.v > r.v ORDER BY 1; +---- +1 + +query I +SELECT r.id FROM rex_sv_l l RIGHT ANTI JOIN rex_sv_r r ON l.v > r.v ORDER BY 1; +---- +2 +3 +4 + +# A temporal key checks that the reduced extreme keeps its own type: a `Date32` minimum that came +# back as the underlying `Int32` would fail the comparison against a `Date32` streamed column +# rather than answer it wrongly. +statement ok +CREATE TABLE rex_dt_l(v DATE); + +statement ok +INSERT INTO rex_dt_l VALUES (DATE '2024-03-01'), (NULL); + +statement ok +CREATE TABLE rex_dt_r(id INT, v DATE); + +statement ok +INSERT INTO rex_dt_r VALUES + (1, DATE '2024-01-15'), (2, DATE '2024-03-01'), (3, DATE '2024-06-30'), (4, NULL); + +query I +SELECT r.id FROM rex_dt_l l RIGHT SEMI JOIN rex_dt_r r ON l.v > r.v ORDER BY 1; +---- +1 + +query I +SELECT r.id FROM rex_dt_l l RIGHT ANTI JOIN rex_dt_r r ON l.v > r.v ORDER BY 1; +---- +2 +3 +4 + +# Every streamed partition emits its own rows here, rather than one elected partition emitting +# everything at the end as LeftSemi/LeftAnti do. The filter lets the streamed side be +# repartitioned, and the EXPLAIN below asserts that it really is. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN (SELECT * FROM rex_r WHERE id > 0) r ON l.v > r.v ORDER BY 1; +---- +1 +2 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN (SELECT * FROM rex_r WHERE id > 0) r ON l.v > r.v ORDER BY 1; +---- +3 +4 +5 + +query TT +EXPLAIN SELECT r.id FROM rex_l l RIGHT SEMI JOIN (SELECT * FROM rex_r WHERE id > 0) r ON l.v > r.v ORDER BY 1; +---- +logical_plan +01)Sort: r.id ASC NULLS LAST +02)--Projection: r.id +03)----RightSemi Join: Filter: l.v > r.v +04)------SubqueryAlias: l +05)--------TableScan: rex_l projection=[v] +06)------SubqueryAlias: r +07)--------Filter: rex_r.id > Int32(0) +08)----------TableScan: rex_r projection=[id, v] +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] +02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[id@0 as id] +04)------PiecewiseMergeJoin: operator=Gt, join_type=RightSemi, on=(v > v) +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +07)----------FilterExec: id@0 > 0 +08)------------DataSourceExec: partitions=1, partition_sizes=[1] + +# The streamed side's ordering survives this join -- one output batch per streamed batch, in +# order, with rows only removed -- so `maintains_input_order` claims it. Here that claim is +# load-bearing rather than merely advertised: the `LIMIT` pins an ordering below the join that the +# optimizer cannot drop, the outer `ORDER BY` requires that same ordering above it, and the plan +# below shows no `SortExec` above the join -- it was removed because the join carries the streamed +# side's ordering through. A wrong claim would surface as wrong row order here, not a slow plan. +query I +SELECT r.id FROM rex_l l RIGHT SEMI JOIN (SELECT * FROM rex_r ORDER BY id DESC LIMIT 5) r ON l.v > r.v +ORDER BY r.id DESC; +---- +2 +1 + +query I +SELECT r.id FROM rex_l l RIGHT ANTI JOIN (SELECT * FROM rex_r ORDER BY id DESC LIMIT 5) r ON l.v > r.v +ORDER BY r.id DESC; +---- +5 +4 +3 + +query TT +EXPLAIN SELECT r.id FROM rex_l l RIGHT SEMI JOIN (SELECT * FROM rex_r ORDER BY id DESC LIMIT 5) r ON l.v > r.v +ORDER BY r.id DESC; +---- +logical_plan +01)Sort: r.id DESC NULLS FIRST +02)--Projection: r.id +03)----RightSemi Join: Filter: l.v > r.v +04)------SubqueryAlias: l +05)--------TableScan: rex_l projection=[v] +06)------SubqueryAlias: r +07)--------Sort: rex_r.id DESC NULLS FIRST, fetch=5 +08)----------TableScan: rex_r projection=[id, v] +physical_plan +01)ProjectionExec: expr=[id@0 as id] +02)--PiecewiseMergeJoin: operator=Gt, join_type=RightSemi, on=(v > v) +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----SortExec: TopK(fetch=5), expr=[id@0 DESC], preserve_partitioning=[false] +05)------DataSourceExec: partitions=1, partition_sizes=[1] + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false;