diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 757d540c77b86..c88c34645b6de 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1620,14 +1620,6 @@ impl DefaultPhysicalPlanner { Arc::new(CrossJoinExec::new(physical_left, physical_right)) } else if num_range_filters == 1 && total_filters == 1 - // 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::LeftMark | JoinType::RightMark - ) && session_state .config_options() .optimizer diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 91a26d083d504..f57bd7be5eccf 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -20,7 +20,9 @@ use std::time::SystemTime; use crate::fuzz_cases::join_fuzz::JoinTestType::{HjSmj, NljHj}; -use arrow::array::{Array, ArrayRef, BinaryArray, Float64Array, Int32Array}; +use arrow::array::{ + Array, ArrayRef, BinaryArray, BooleanArray, Float64Array, Int32Array, +}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; @@ -1492,7 +1494,7 @@ fn pwmj_plan( // 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, + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => left, _ => { let sort_options = match op { Operator::Lt | Operator::LtEq => SortOptions::new(true, true), @@ -1600,6 +1602,47 @@ async fn pwmj_collect_id_pairs( pairs } +/// Executes every output partition concurrently and returns a `LeftMark`/`RightMark` join's +/// rows as `(id, mark)` pairs, sorted so partition interleaving does not affect the +/// comparison. `mark` is documented to never be NULL (see `JoinType::LeftMark`), so it is +/// unwrapped rather than carried as `Option`. +async fn pwmj_collect_mark_pairs( + plan: Arc, + task_ctx: Arc, +) -> Vec<(i32, bool)> { + let streams = (0..plan.output_partitioning().partition_count()) + .map(|partition| plan.execute(partition, Arc::clone(&task_ctx)).unwrap()) + .collect::>(); + let per_partition = + futures::future::join_all(streams.into_iter().map(|stream| { + SpawnedTask::spawn(async move { common::collect(stream).await }) + })) + .await; + + let mut pairs = Vec::new(); + for batches in per_partition { + for batch in batches.unwrap().unwrap() { + let id = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + // `mark` is always the last column: it is appended after either side's own + // columns (see `build_join_schema`). + let mark = batch + .column(batch.num_columns() - 1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + pairs.push((id.value(row), mark.value(row))); + } + } + } + pairs.sort_unstable(); + pairs +} + /// Differential test for every join type `PiecewiseMergeJoin` supports, against a /// `NestedLoopJoin` oracle, run once over `i32` keys and once over `f64` keys via /// [`run_pwmj_fuzz`]. @@ -1637,6 +1680,8 @@ async fn run_pwmj_fuzz(task_ctx: &Arc(task_ctx: &Arc { + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { pwmj_parts_exec(&left_ids, &left_keys, buffered_nparts) } _ => pwmj_single_exec(&left_ids, &left_keys), }; + + // Mark joins output every row of their side with a `mark` column rather than + // a filtered subset, so they are compared as `(id, mark)` pairs instead of + // the `(left id, right id)` pairs every other join type here uses. + if matches!(join_type, JoinType::LeftMark | JoinType::RightMark) { + let got = pwmj_collect_mark_pairs( + pwmj_plan( + buffered, + pwmj_parts_exec(&right_ids, &right_keys, nparts), + op, + join_type, + ), + Arc::clone(task_ctx), + ) + .await; + let want = pwmj_collect_mark_pairs( + pwmj_nlj_oracle_plan::( + pwmj_single_exec(&left_ids, &left_keys), + pwmj_single_exec(&right_ids, &right_keys), + op, + join_type, + ), + Arc::clone(task_ctx), + ) + .await; + + assert_eq!( + got, + want, + "mismatch key_type={} seed={seed} op={op:?} join_type={join_type:?} \ + nparts={nparts} buffered_nparts={buffered_nparts} \ + left_keys={left_keys:?} right_keys={right_keys:?}", + std::any::type_name::(), + ); + continue; + } + let got = pwmj_collect_id_pairs( pwmj_plan( buffered, @@ -1699,7 +1781,7 @@ async fn run_pwmj_fuzz(task_ctx: &Arc(), ); } 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 07454dc50db9c..f2395211be861 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -21,7 +21,6 @@ use arrow::{ compute::{concat, concat_batches}, }; use arrow_schema::{SchemaRef, SortOptions}; -use datafusion_common::not_impl_err; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{JoinSide, Result, internal_datafusion_err, internal_err}; use datafusion_common_runtime::SpawnedTask; @@ -50,9 +49,7 @@ 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::{ - is_existence_join, is_supported_existence_join, is_supported_right_existence_join, -}; +use crate::joins::piecewise_merge_join::utils::is_right_existence_join; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; use crate::{ @@ -164,14 +161,15 @@ use crate::{ /// ``` /// /// ## Existence Joins (Semi, Anti, Mark) -/// 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. +/// Every Semi/Anti/Mark join is supported. The two sides are served by different streams, +/// because a single range predicate makes them different problems. /// -/// `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. +/// `LeftSemi`/`LeftAnti`/`LeftMark` 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, once every streamed partition has been consumed, either slices the +/// buffered batch at that index (`LeftSemi`/`LeftAnti`) or emits the whole batch with a `mark` +/// column built from it (`LeftMark`). /// /// ```text /// // Using the example of a less than `<` operation @@ -213,13 +211,14 @@ use crate::{ /// min value: 200 /// ``` /// -/// `RightSemi`/`RightAnti` mark the streamed (right) side, which is +/// `RightSemi`/`RightAnti`/`RightMark` 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. +/// against it and emitted straight away rather than at the end -- filtered for +/// `RightSemi`/`RightAnti`, kept whole with a `mark` column appended for `RightMark`. 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 @@ -230,8 +229,8 @@ use crate::{ /// 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 +/// Except for `RightSemi`/`RightAnti`/`RightMark`, the buffered side must be sorted ascending +/// for `Operator::Lt` (<) or `Operator::LtEq` (<=) and descending for `Operator::Gt` (>) or /// `Operator::GtEq` (>=). /// /// # Partitioning Logic @@ -240,10 +239,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`. +/// `RightSemi`/`RightAnti`/`RightMark` 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: @@ -320,14 +320,13 @@ impl PiecewiseMergeJoinExec { join_type: JoinType, num_partitions: usize, ) -> Result { - // 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" - ); - } - + // There is no `null_aware` parameter here, unlike `HashJoinExec::try_new`: + // `mark_streamed_batch`/`emit_matched` always build a non-nullable `mark` column. + // That is sound because `join_on` is guaranteed non-empty whenever `null_aware` is + // set (see `ExtractEquijoinPredicate`, which moves the hash-equality predicate that + // makes a mark join null-aware into `on`), and the PWMJ branch in + // `physical_planner.rs` is only reached when `join_on` is empty. + // // Take the operator and enforce a sort order on the streamed + buffered side based on // the operator type. let sort_options = match operator { @@ -467,11 +466,13 @@ impl PiecewiseMergeJoinExec { 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], + // *removed* by a filter (`RightSemi`/`RightAnti`) or left in place with a `mark` + // column appended (`RightMark`) -- neither reorders a surviving row relative to + // the others, so each output partition keeps its streamed partition's order. The + // buffered side contributes no output column of its own, hence `false` there. + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + vec![false, true] + } // Unlike the right existence joins above, output here is gated on a watermark // shared across every streamed partition (see `ExistencePWMJStream`) and emitted // only from whichever partition finishes last, so no streamed partition's output @@ -479,7 +480,6 @@ impl PiecewiseMergeJoinExec { JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { 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` @@ -555,7 +555,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { // 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) { + let buffered = if is_right_existence_join(self.join_type) { Distribution::UnspecifiedDistribution } else { Distribution::SinglePartition @@ -578,7 +578,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { // 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) + if is_right_existence_join(self.join_type) && let Some(buffered) = benefits.first_mut() { *buffered = false; @@ -594,7 +594,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { // // 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) { + if is_right_existence_join(self.join_type) { vec![None, None] } else { vec![ @@ -702,51 +702,54 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { let streamed_partitions = self.streamed.output_partitioning().partition_count(); let batch_size = context.session_config().batch_size(); - 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, - )) - })?; + // 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. `RightMark` decides its `mark` column with the exact same + // comparison as `RightSemi`/`RightAnti` -- it just keeps every row instead of + // filtering by it -- so it takes the same path. + if is_right_existence_join(self.join_type) { + // `∃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))?; + 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, - batch_size, - ))) - } - JoinType::LeftSemi | JoinType::LeftAnti => { + return Ok(Box::pin(RightExistencePWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + extreme_fut, + metrics, + batch_size, + ))); + } + + match self.join_type { + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { let buffered_side = self.buffered_side( &context, &on_buffered, @@ -767,10 +770,6 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { batch_size, ))) } - JoinType::LeftMark | JoinType::RightMark => internal_err!( - "PiecewiseMergeJoin does not support existence join {} (should have been rejected in try_new)", - self.join_type - ), JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { let buffered_side = self.buffered_side( &context, @@ -793,6 +792,12 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { batch_size, ))) } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + unreachable!( + "is_right_existence_join returned false for {:?}, which the `if` above handles", + self.join_type + ) + } } } 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 693e64a06f275..2ea33239c1106 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 @@ -17,11 +17,9 @@ //! PiecewiseMergeJoin stream specialized for existence joins. //! -//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `LeftSemi` or `LeftAnti`. -//! `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. +//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `LeftSemi`, `LeftAnti`, or +//! `LeftMark`. `RightSemi`/`RightAnti`/`RightMark` mark the streamed side instead and are +//! served by `RightExistencePWMJStream` (see `right_existence_join.rs`). //! //! # Motivation //! @@ -80,8 +78,11 @@ //! //! Once every streamed partition has been consumed, the last one to finish slices the //! buffered batch: `LeftSemi` takes `[min_marked, len)`, `LeftAnti` the complementary -//! prefix `[0, min_marked)`, which is where the null-keyed rows live. Only the buffered -//! (left) columns are produced. +//! prefix `[0, min_marked)`, which is where the null-keyed rows live. `LeftMark` takes +//! neither slice: every buffered row is preserved, with a `mark` column built from the same +//! watermark (`true` from `min_marked` on, `false` before it) appended instead of any row +//! being dropped. Only the buffered (left) columns, plus that `mark` column for `LeftMark`, +//! are produced. //! //! [`PiecewiseMergeJoinExec`]: super::PiecewiseMergeJoinExec @@ -90,7 +91,7 @@ use std::sync::Arc; use std::sync::atomic::Ordering as AtomicOrdering; use std::task::{Poll, ready}; -use arrow::array::{Array, ArrayRef, RecordBatch}; +use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch}; use arrow::compute::BatchCoalescer; use arrow_schema::{SchemaRef, SortOptions}; use datafusion_common::{NullEquality, Result, internal_err}; @@ -118,12 +119,13 @@ pub(super) enum ExistencePWMJStreamState { } pub(super) struct ExistencePWMJStream { - /// Output schema, which for `LeftSemi`/`LeftAnti` is the buffered side's schema + /// Output schema, which for `LeftSemi`/`LeftAnti` is the buffered side's schema, and for + /// `LeftMark` is that schema plus a trailing `mark` column 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, - /// `LeftSemi` or `LeftAnti` + /// `LeftSemi`, `LeftAnti`, or `LeftMark` join_type: JoinType, /// Comparison operator operator: Operator, @@ -389,7 +391,8 @@ impl ExistencePWMJStream { } /// Emits the existence result by slicing at the watermark: the marked buffered rows for - /// `LeftSemi`, the unmarked ones for `LeftAnti`. + /// `LeftSemi`, the unmarked ones for `LeftAnti`, or -- for `LeftMark` -- every buffered row + /// with a `mark` column built from the same watermark. fn emit_matched(&mut self) -> Result>> { let _join_timer = self.join_metrics.join_time.timer(); if !self.emitted { @@ -409,22 +412,47 @@ impl ExistencePWMJStream { .load(AtomicOrdering::SeqCst) .min(buffered_len); - let sliced = match self.join_type { + let (num_rows, columns) = match self.join_type { JoinType::LeftSemi => { - buffered_batch.slice(min_marked, buffered_len - min_marked) + let sliced = + buffered_batch.slice(min_marked, buffered_len - min_marked); + (sliced.num_rows(), sliced.columns().to_vec()) + } + // `LeftMark` keeps every buffered row -- nothing to slice -- and appends + // the watermark as a `mark` column instead of using it to drop rows: `false` + // for the unmatched prefix `[0, min_marked)`, `true` for the matched suffix + // `[min_marked, len)` -- the same split `LeftSemi`/`LeftAnti` slice the + // buffered batch on, just kept as one column instead of used to drop rows. + JoinType::LeftMark => { + let mut mark = BooleanBufferBuilder::new(buffered_len); + mark.append_n(min_marked, false); + mark.append_n(buffered_len - min_marked, true); + + let mut columns = buffered_batch.columns().to_vec(); + columns.push( + Arc::new(BooleanArray::new(mark.finish(), None)) as ArrayRef + ); + (buffered_len, columns) + } + // `LeftAnti`: the unmarked prefix, which includes every null-keyed row -- + // nulls sort first and the watermark never drops below the buffered null + // count. + JoinType::LeftAnti => { + let sliced = buffered_batch.slice(0, min_marked); + (sliced.num_rows(), sliced.columns().to_vec()) + } + other => { + return internal_err!( + "ExistencePWMJStream does not support join type {other:?}" + ); } - // The unmarked prefix, which includes every null-keyed row: nulls sort - // first and the watermark never drops below the buffered null count. - _ => buffered_batch.slice(0, min_marked), }; - if sliced.num_rows() > 0 { - // Existence joins output the buffered (left) columns only; rebuild against - // the join's own schema, which keeps the slice zero-copy. - let batch = RecordBatch::try_new( - Arc::clone(&self.schema), - sliced.columns().to_vec(), - )?; + if num_rows > 0 { + // Existence joins output the buffered (left) columns only (plus `mark` for + // `LeftMark`); rebuild against the join's own schema, which keeps a slice + // zero-copy. + let batch = RecordBatch::try_new(Arc::clone(&self.schema), columns)?; self.output_batches.push_batch(batch)?; self.output_batches.finish_buffered_batch()?; } @@ -569,6 +597,25 @@ mod tests { Ok(()) } + /// `LeftMark` keeps every buffered row -- unlike `LeftSemi`/`LeftAnti`, neither is + /// dropped -- and appends a `mark` column: `true` for exactly the rows `LeftSemi` would + /// have kept. + #[tokio::test] + async fn join_left_mark() -> Result<()> { + let batches = join_collect(JoinType::LeftMark).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+-------+ + | a1 | b1 | c1 | mark | + +----+----+----+-------+ + | 1 | 1 | 7 | false | + | 2 | 2 | 8 | false | + | 3 | 5 | 9 | true | + +----+----+----+-------+ + "); + Ok(()) + } + /// Once every markable buffered row is marked, no later streamed batch can lower the /// watermark, so the stream stops reading rather than scanning batches that provably /// cannot contribute. The first batch below is smaller than every buffered value, so it @@ -777,42 +824,4 @@ mod tests { assert_eq!(consumed, 1, "partition 1 should not have read any batch"); Ok(()) } - - /// 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( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 2, 5]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - for join_type in [JoinType::LeftMark, JoinType::RightMark] { - let err = PiecewiseMergeJoinExec::try_new( - Arc::clone(&left), - Arc::clone(&right), - on.clone(), - Operator::Gt, - join_type, - 1, - ) - .expect_err(&format!("{join_type} should be rejected")) - .to_string(); - assert!( - err.contains("not supported for PiecewiseMergeJoin"), - "unexpected error for {join_type}: {err}" - ); - } - Ok(()) - } } 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 index 483e828e4147e..23cc609009353 100644 --- 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 @@ -17,21 +17,20 @@ //! 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`. +//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `RightSemi`, `RightAnti`, +//! or `RightMark`. `LeftSemi`/`LeftAnti`/`LeftMark` are served by `ExistencePWMJStream` (see +//! `existence_join.rs`). //! //! # 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: +//! `LeftSemi`/`LeftAnti`/`LeftMark` 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`/`RightMark` 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 @@ -53,8 +52,9 @@ //! 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: +//! vectorized `cmp` kernel per batch, followed by a filter (`RightSemi`/`RightAnti`) or a +//! `mark` column appended to the batch unfiltered (`RightMark`). 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, @@ -64,9 +64,11 @@ //! 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`. +//! without reading the streamed side, and `RightAnti`/`RightMark` a passthrough of it (`mark` +//! false on every row for the latter). A NULL streamed key makes its comparison NULL rather +//! than false, which is "no match" — dropped by `RightSemi`, kept by `RightAnti`, and folded +//! to `false` rather than left NULL for `RightMark`, matching the `mark` column's documented +//! never-NULL semantics (see `JoinType::LeftMark`). //! //! Picking the extreme and comparing against it must agree on ordering. `min_batch`/`max_batch` //! order floats by `total_cmp` (`arrow-arith`'s `MinAccumulator` compares with @@ -75,13 +77,13 @@ //! actually evaluates to: any `BinaryExpr` comparison, including the one the //! `NestedLoopJoinExec` oracle in the differential fuzz test builds its filter from, normalizes //! `-0.0` to `+0.0` first (see `apply_cmp` in `datafusion-physical-expr-common`, which -//! [`filter_streamed_batch`](RightExistencePWMJStream::filter_streamed_batch) calls to compare -//! the extreme against each streamed batch). `apply_cmp` normalizes both operands itself, so -//! nothing here has to: the reduction can order `-0.0` below `+0.0` as `min`/`max` do, and the -//! comparison it feeds into still agrees with SQL semantics. `NaN` needs no such fix-up: every -//! kernel involved orders it as the maximum, so it is treated identically on both sides already. -//! `apply_cmp` also dispatches to the nested-aware comparator for List/Struct keys, which the -//! raw arrow `lt`/`lt_eq`/`gt`/`gt_eq` kernels reject outright. +//! [`matched_mask`](RightExistencePWMJStream::matched_mask) calls to compare the extreme +//! against each streamed batch). `apply_cmp` normalizes both operands itself, so nothing here +//! has to: the reduction can order `-0.0` below `+0.0` as `min`/`max` do, and the comparison it +//! feeds into still agrees with SQL semantics. `NaN` needs no such fix-up: every kernel involved +//! orders it as the maximum, so it is treated identically on both sides already. `apply_cmp` +//! also dispatches to the nested-aware comparator for List/Struct keys, which the raw arrow +//! `lt`/`lt_eq`/`gt`/`gt_eq` kernels reject outright. //! //! # Cost //! @@ -99,7 +101,8 @@ use std::sync::Arc; use std::task::{Poll, ready}; -use arrow::array::{AsArray, RecordBatch}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, RecordBatch}; +use arrow::buffer::BooleanBuffer; use arrow::compute::BatchCoalescer; use arrow::compute::filter_record_batch; use arrow::compute::kernels::boolean::not; @@ -122,7 +125,7 @@ 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. + /// (`RightAnti`) have a buffered match, or every row with `mark` appended (`RightMark`). ScanStreamBatches, /// Streamed side exhausted; drain whatever `output_batches` still holds. Draining, @@ -130,12 +133,13 @@ pub(super) enum RightExistencePWMJStreamState { } pub(super) struct RightExistencePWMJStream { - /// Output schema, which for `RightSemi`/`RightAnti` is the streamed side's schema + /// Output schema, which for `RightSemi`/`RightAnti` is the streamed side's schema, and for + /// `RightMark` is that schema plus a trailing `mark` column 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` + /// `RightSemi`, `RightAnti`, or `RightMark` join_type: JoinType, /// Comparison operator operator: Operator, @@ -214,13 +218,13 @@ impl RightExistencePWMJStream { // Null exactly when no buffered key is non-null, and NULLs match nothing. `apply_cmp` // normalizes `-0.0`/`+0.0` on both operands when it compares this against the streamed - // key in `filter_streamed_batch`, so no normalization is needed here. + // key in `matched_mask`, so no normalization is needed here. let extreme = ScalarValue::try_from_array(buffered_extreme.extreme(), 0)?; self.buffered_extreme = (!extreme.is_null()).then_some(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. + // and does not need to read a single streamed batch. `RightAnti`/`RightMark` still + // have to, since they output all of them (the latter with `mark` false throughout). self.state = match (&self.buffered_extreme, self.join_type) { (None, JoinType::RightSemi) => { let streamed_schema = self.streamed.schema(); @@ -233,7 +237,7 @@ impl RightExistencePWMJStream { Poll::Ready(Ok(StatefulStreamResult::Continue)) } - /// Fetches one streamed batch, filters it, and pushes the survivors into + /// Fetches one streamed batch, filters or marks it, and pushes the result into /// `output_batches`. Once the streamed side is exhausted, flushes the coalescer's /// partial batch and hands off to [`Self::drain_output`]. fn scan_stream_batch( @@ -253,7 +257,10 @@ impl RightExistencePWMJStream { self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); - let output = self.filter_streamed_batch(&batch)?; + let output = match self.join_type { + JoinType::RightMark => self.mark_streamed_batch(&batch)?, + _ => self.filter_streamed_batch(&batch)?, + }; if output.num_rows() > 0 { self.output_batches.push_batch(output)?; } @@ -277,43 +284,52 @@ impl RightExistencePWMJStream { } } + /// Compares a streamed batch's key against the single buffered extreme, on the same + /// `buffered OP streamed` orientation the predicate has. `None` when the buffered side has + /// no non-null key, in which case no streamed row can match anything. + fn matched_mask(&self, batch: &RecordBatch) -> Result> { + let Some(extreme) = &self.buffered_extreme else { + return Ok(None); + }; + + let stream_values = ColumnarValue::Array( + 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. `apply_cmp` dispatches to the + // nested-aware comparator for List/Struct keys, unlike the raw arrow + // `lt`/`lt_eq`/`gt`/`gt_eq` kernels, which reject those types outright. + let matched = apply_cmp( + self.operator, + &ColumnarValue::Scalar(extreme.clone()), + &stream_values, + )? + .into_array(batch.num_rows())?; + + Ok(Some(matched.as_boolean().clone())) + } + /// 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 { + let columns = match self.matched_mask(batch)? { // 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 = ColumnarValue::Array( - 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. `apply_cmp` - // dispatches to the nested-aware comparator for List/Struct keys, unlike the - // raw arrow `lt`/`lt_eq`/`gt`/`gt_eq` kernels this used to call directly, which - // reject those types outright. - let matched = apply_cmp( - self.operator, - &ColumnarValue::Scalar(extreme.clone()), - &stream_values, - )? - .into_array(batch.num_rows())?; - let matched = matched.as_boolean(); - + Some(matched) => { 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.clone(), + 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))?, + _ => not(&boolean_mask_from_filter(&matched))?, }; filter_record_batch(batch, &predicate)?.columns().to_vec() @@ -325,6 +341,29 @@ impl RightExistencePWMJStream { // latter so the stream's declared schema is what it yields. Ok(RecordBatch::try_new(Arc::clone(&self.schema), columns)?) } + + /// Keeps every streamed row -- nothing is filtered -- and appends the comparison against + /// the single buffered extreme as a `mark` column: `true` where some buffered row + /// matches, `false` everywhere else, including where the streamed key is NULL (the `mark` + /// column is documented never to be NULL; see `JoinType::LeftMark`). + fn mark_streamed_batch(&self, batch: &RecordBatch) -> Result { + let mark: ArrayRef = match self.matched_mask(batch)? { + None => Arc::new(BooleanArray::new( + BooleanBuffer::new_unset(batch.num_rows()), + None, + )), + Some(matched) => Arc::new(boolean_mask_from_filter(&matched)), + }; + + let mut columns = batch.columns().to_vec(); + columns.push(mark); + + // Right existence joins output the streamed columns (plus `mark` for `RightMark`) + // only. The streamed child's schema is field-for-field equal to the join's own for + // those columns, 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 { @@ -355,7 +394,9 @@ mod tests { ExecutionPlan, ExecutionPlanProperties, common, joins::PiecewiseMergeJoinExec, test::TestMemoryExec, }; - use arrow::array::{DictionaryArray, Int32Array, ListArray, StructArray}; + use arrow::array::{ + DictionaryArray, Float64Array, Int32Array, ListArray, StructArray, + }; use arrow::compute::SortOptions; use arrow::datatypes::Int32Type; use arrow_schema::{DataType, Field, Schema}; @@ -507,6 +548,106 @@ mod tests { Ok(()) } + /// `RightMark` keeps every streamed row -- unlike `RightSemi`/`RightAnti`, none is + /// dropped -- and appends a `mark` column: `true` for exactly the rows `RightSemi` would + /// have kept. + #[tokio::test] + async fn join_right_mark() -> 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::RightMark, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+-------+ + | id | k | mark | + +----+---+-------+ + | 10 | 4 | true | + | 20 | 5 | false | + | 30 | 6 | false | + +----+---+-------+ + "); + Ok(()) + } + + fn kv_float_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("k", DataType::Float64, true), + ])) + } + + fn kv_float_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_float_schema(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Float64Array::from(keys)), + ], + ) + .unwrap() + } + + fn kv_float_exec(partitions: &[Vec]) -> Arc { + TestMemoryExec::try_new_exec(partitions, kv_float_schema(), None).unwrap() + } + + /// `matched_mask` compares against the buffered extreme with `apply_cmp`, the same total + /// order (`f64::total_cmp`) `min`/`max` used to fold that extreme in the first place. Under + /// that order NaN is the largest float, so the buffered max here is NaN and every non-NaN + /// streamed key marks `true` -- `-Inf`, `0.0`, `-0.0`, and `2.0` are all `< NaN` -- while + /// NaN itself marks `false` (`NaN < NaN` is false under `total_cmp`, unlike an unordered + /// IEEE comparison, which would make it incomparable rather than false). `0.0` and `-0.0` + /// both marking `true` here also confirms they compare equal to each other and to every + /// ordinary float on this path, not just relative to a `-0.0` buffered key as in the + /// `LeftMark` case (see the `pwmj.slt` `-0.0` coverage under `LeftMark`). Same data as the + /// `RIGHT SEMI`/`RIGHT ANTI` NaN coverage in `pwmj.slt`, through `RightMark` instead. + #[tokio::test] + async fn join_right_mark_supports_nan_and_negative_zero() -> Result<()> { + let join = join( + kv_float_exec(&[vec![kv_float_batch(&[ + (1, Some(1.0)), + (2, Some(f64::NAN)), + ])]]), + kv_float_exec(&[vec![kv_float_batch(&[ + (1, Some(f64::NEG_INFINITY)), + (2, Some(0.0)), + (3, Some(-0.0)), + (4, Some(f64::NAN)), + (5, Some(2.0)), + ])]]), + Operator::Gt, + JoinType::RightMark, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+------+-------+ + | id | k | mark | + +----+------+-------+ + | 1 | -inf | true | + | 2 | 0.0 | true | + | 3 | -0.0 | true | + | 4 | NaN | false | + | 5 | 2.0 | true | + +----+------+-------+ + "); + Ok(()) + } + fn kv_list_schema() -> Arc { Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), @@ -692,8 +833,8 @@ mod tests { /// /// 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. + /// it. On the streamed side the NULL-keyed row matches nothing, so `RightSemi` drops it, + /// `RightAnti` keeps it, and `RightMark` keeps it with `mark` folded to `false`. #[tokio::test] async fn null_keys_match_nothing() -> Result<()> { let buffered = @@ -730,6 +871,27 @@ mod tests { | 30 | | +----+---+ "); + + // `RightMark`'s NULL-keyed row (id 30) takes a different code path from + // `all_null_buffered_side_marks_every_row_false`: the buffered extreme here is a real, + // non-null 5, so `matched_mask` returns `Some` and the NULL streamed key compares NULL + // rather than short-circuiting on a `None` extreme. `mark_streamed_batch` must still + // fold that NULL to `false` via `boolean_mask_from_filter`, the same as `RightAnti` + // folds it to "keep" above -- this is the one test that exercises that fold for `mark` + // outside of the differential fuzzer's random NULLs. + let mark = join(buffered(), streamed(), Operator::Lt, JoinType::RightMark)?; + let batches = + common::collect(mark.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+-----+-------+ + | id | k | mark | + +----+-----+-------+ + | 10 | 5 | false | + | 20 | 6 | true | + | 30 | | false | + | 40 | 100 | true | + +----+-----+-------+ + "); Ok(()) } @@ -788,6 +950,39 @@ mod tests { Ok(()) } + /// The same all-NULL buffered side as the `RightSemi` test above, but for `RightMark`: + /// nothing can ever match, so every streamed row's `mark` is `false` -- never NULL, even + /// though the streamed keys include one. `RightMark` still has to read every batch, since + /// (unlike `RightSemi`) it outputs one row per streamed row regardless of its mark. + #[tokio::test] + async fn all_null_buffered_side_marks_every_row_false() -> Result<()> { + let join = join( + kv_exec(&[vec![kv_batch(&[(1, None), (2, None)])]]), + kv_exec(&[vec![ + kv_batch(&[(10, Some(4))]), + kv_batch(&[(20, None)]), + kv_batch(&[(30, Some(6))]), + ]]), + Operator::Gt, + JoinType::RightMark, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+---+-------+ + | id | k | mark | + +----+---+-------+ + | 10 | 4 | false | + | 20 | | false | + | 30 | 6 | false | + +----+---+-------+ + "); + assert_eq!(input_batches(&join), 3, "RightMark must read every batch"); + 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 @@ -803,7 +998,11 @@ mod tests { .unwrap(); let streamed = Arc::new(SortExec::new(ordering.clone(), streamed_input)); - for join_type in [JoinType::RightSemi, JoinType::RightAnti] { + for join_type in [ + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightMark, + ] { let join = join( kv_exec(&[vec![kv_batch(&[(1, Some(5))])]]), Arc::clone(&streamed) as _, @@ -840,11 +1039,15 @@ mod tests { /// 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. + /// are unaffected. Covered for all three join types since only anti also runs `not` over + /// it, and mark keeps every row regardless. #[tokio::test] async fn empty_streamed_batch_is_skipped() -> Result<()> { - for (join_type, expected) in [(JoinType::RightSemi, 1), (JoinType::RightAnti, 1)] - { + for (join_type, expected) in [ + (JoinType::RightSemi, 1), + (JoinType::RightAnti, 1), + (JoinType::RightMark, 2), + ] { let join = join( kv_exec(&[vec![kv_batch(&[(1, Some(5)), (2, Some(1))])]]), kv_exec(&[vec![ @@ -859,8 +1062,8 @@ mod tests { 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. + // Buffered maximum is 5, so `>` keeps 4 for semi and 9 for anti -- one row each -- + // and `RightMark` keeps both, with no empty batch in between for any of them. assert_eq!( batches.iter().map(|b| b.num_rows()).sum::(), expected, @@ -977,6 +1180,74 @@ mod tests { Ok(()) } + /// The `RightMark` sibling of the test above: `mark_streamed_batch` never filters, so + /// every row (not a subset) has to stay in the streamed side's sorted order once the + /// coalescer repacks across the original `SortExec` batch boundaries, and each row's + /// `mark` has to travel with the right `id` through that repacking. + #[tokio::test] + async fn coalescing_does_not_reorder_a_sorted_streamed_side_for_mark() -> Result<()> { + let streamed_input = kv_exec(&[vec![kv_batch(&[ + (50, Some(5)), + (10, Some(199)), + (90, Some(6)), + (30, Some(200)), + (20, Some(7)), + (70, Some(8)), + (40, Some(199)), + (60, Some(9)), + ])]]); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new_with_schema("id", &streamed_input.schema())?), + SortOptions::new(false, false), + )]) + .unwrap(); + let streamed = Arc::new(SortExec::new(ordering, streamed_input)); + + // Buffered max is 100, so `mark` is true for k < 100 -- ids 10, 30, 40 (k in the + // 100s) are false, unevenly across the sorted id order, so the coalescer must not + // only preserve order but keep each id's mark paired with it while repacking. + let join = join( + kv_exec(&[vec![kv_batch(&[(1, Some(100))])]]), + streamed as _, + Operator::Gt, + JoinType::RightMark, + )?; + + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let ids: Vec = batches + .iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect(); + let marks: Vec = batches + .iter() + .flat_map(|b| b.column(2).as_boolean().values().iter()) + .collect(); + assert_eq!( + ids, + vec![10, 20, 30, 40, 50, 60, 70, 90], + "coalescing must not disturb the streamed side's sorted order, even though \ + RightMark keeps every row rather than a subset" + ); + assert_eq!( + marks, + vec![false, true, false, false, true, true, true, true], + "each id's mark must travel with it through the coalescer's repacking" + ); + assert_eq!( + batches.iter().map(|b| b.num_rows()).collect::>(), + vec![2, 2, 2, 2], + "8 rows packed to batch_size = 2 confirms the coalescer, not a pass-through, \ + produced this shape" + ); + 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 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 7fef429f54e40..7f63e7eac5f97 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs @@ -20,35 +20,13 @@ use datafusion_expr::JoinType; // 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 -pub(super) fn is_existence_join(join_type: JoinType) -> bool { - matches!( - join_type, - JoinType::LeftAnti - | JoinType::RightAnti - | JoinType::LeftSemi - | JoinType::RightSemi - | JoinType::LeftMark - | JoinType::RightMark - ) -} - -// 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 { +// `RightMark` belongs here too: deciding its mark column is the same one-key comparison as +// `RightSemi`/`RightAnti`, just kept instead of used to filter, so it needs no more of the +// buffered side than they do. +pub(super) fn is_right_existence_join(join_type: JoinType) -> bool { matches!( join_type, - JoinType::LeftSemi - | JoinType::LeftAnti - | JoinType::RightSemi - | JoinType::RightAnti + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark ) } diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index da0c0bd5733a1..6ff44227915db 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -607,11 +607,8 @@ fn piecewise_join( /// `PiecewiseMergeJoinExec` derives its schema, sort options, required input /// orderings and plan properties inside `try_new`, so only the six constructor /// arguments travel on the wire. Cover the full cartesian product of the two -/// enum-valued ones: every range operator against every supported join type -- -/// the four classic ones plus the four existence joins, whose output schema is -/// one side alone (buffered for `LeftSemi`/`LeftAnti`, streamed for -/// `RightSemi`/`RightAnti`). Mark joins are rejected by `try_new`, so this is -/// the complete set. +/// enum-valued ones: every range operator against every join type the operator +/// supports. #[test] fn roundtrip_piecewise_merge_join() -> Result<()> { let (schema_buffered, schema_streamed) = piecewise_schemas(); @@ -630,6 +627,8 @@ fn roundtrip_piecewise_merge_join() -> Result<()> { JoinType::LeftAnti, JoinType::RightSemi, JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, ] { let result = roundtrip_test_and_return( Arc::new(PiecewiseMergeJoinExec::try_new( @@ -654,15 +653,17 @@ fn roundtrip_piecewise_merge_join() -> Result<()> { // Both the name and the index have to survive on the correct side. assert_eq!(result.on.0.to_string(), "a@2"); assert_eq!(result.on.1.to_string(), "b@0"); - // The left existence joins output the buffered side alone (3 fields), - // the right existence joins the streamed side alone (2 fields), and the - // classic joins output both sides (3 + 2). The before/after comparison - // inside the helper already covers the schema; this pins the expected - // width absolutely, so each existence output contract is stated rather - // than merely preserved. + // The classic joins output both sides (3 + 2); the existence joins + // output only the side they mark (buffered's 3 fields for Left*, streamed's + // 2 for Right*), plus one more for a Mark join's `mark` column. The + // before/after comparison inside the helper already covers the schema; + // this pins the expected width absolutely, so the existence output + // contract is stated rather than merely preserved. let expected_fields = match join_type { JoinType::LeftSemi | JoinType::LeftAnti => 3, JoinType::RightSemi | JoinType::RightAnti => 2, + JoinType::LeftMark => 4, + JoinType::RightMark => 3, _ => 5, }; assert_eq!( @@ -881,13 +882,10 @@ fn piecewise_merge_join_rejects_bad_operator_on_the_wire() -> Result<()> { Ok(()) } -/// `join_type` is a proto3 enum, so any `i32` is representable on the wire -- including -/// the existence joins `try_new` still rejects (right-sided and mark) and values that -/// map to no `JoinType` at all. `try_to_proto` can emit none of these, but a payload -/// from a newer peer or a hand-built one can, and each must surface as an error -/// rather than a panic or a silently different operator. The right-sided cases matter -/// most: `try_new` derives *reversed* sort options for them, so accepting one would -/// build a plan whose buffered side is sorted the wrong way. +/// `join_type` is a proto3 enum, so any `i32` is representable on the wire, including values +/// that map to no `JoinType` at all. `try_to_proto` can emit none of these, but a payload +/// from a newer peer or a hand-built one can, and that must surface as an error rather than a +/// panic or a silently different operator. #[test] fn piecewise_merge_join_rejects_unsupported_join_type_on_the_wire() -> Result<()> { let schemas = piecewise_schemas(); @@ -904,46 +902,28 @@ fn piecewise_merge_join_rejects_unsupported_join_type_on_the_wire() -> Result<() let mut node = protobuf::PhysicalPlanNode::decode(valid.as_ref()) .expect("a plan just encoded by try_to_proto must decode as a PhysicalPlanNode"); - // (wire value, description, expected error fragment) - // RightSemi/RightAnti are supported (see `piecewise_merge_join_existence_wire_layout` - // and `right_existence_join.rs`); only the Mark joins and unknown variants remain - // unsupported on the wire. - let cases: [(i32, &str, &str); 3] = [ - ( - protobuf::JoinType::Leftmark as i32, - "LeftMark", - "Existence join LeftMark is currently not supported", - ), - ( - protobuf::JoinType::Rightmark as i32, - "RightMark", - "Existence join RightMark is currently not supported", - ), - // Past the last tag, so it maps to no variant. - (i32::MAX, "no variant", "unknown JoinType"), - ]; + // Past the last tag, so it maps to no variant. + let wire_value = i32::MAX; - for (wire_value, description, expected) in cases { - let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = - node.physical_plan_type.as_mut() - else { - panic!("expected a PiecewiseMergeJoin node"); - }; - join.join_type = wire_value; + let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = + node.physical_plan_type.as_mut() + else { + panic!("expected a PiecewiseMergeJoin node"); + }; + join.join_type = wire_value; - let Err(err) = physical_plan_from_bytes_with_proto_converter( - &node.encode_to_vec(), - ctx.task_ctx().as_ref(), - &codec, - &proto_converter, - ) else { - panic!("decoding must fail for join_type {description}"); - }; - assert!( - err.to_string().contains(expected), - "join_type {description}: expected an error containing {expected:?}, got: {err}" - ); - } + let Err(err) = physical_plan_from_bytes_with_proto_converter( + &node.encode_to_vec(), + ctx.task_ctx().as_ref(), + &codec, + &proto_converter, + ) else { + panic!("decoding must fail for an out-of-range join_type"); + }; + assert!( + err.to_string().contains("unknown JoinType"), + "expected an error containing \"unknown JoinType\", got: {err}" + ); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 03670c06fd4cd..a15f4a3b82410 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -503,6 +503,53 @@ ORDER BY 1 NULLS FIRST; NULL 1 +# ------------------------------------------------------------------ +# LeftMark via PiecewiseMergeJoin +# ------------------------------------------------------------------ + +# `LeftMark` has no SQL syntax of its own -- it is produced by decorrelating an EXISTS/IN +# subquery that appears inside a disjunction, since the outer filter then needs the +# subquery's match result (the `mark`) as a value rather than as a row filter. Same data and +# correlation as the plain-EXISTS `LeftSemi` case above: no `t1.t1_id` is over 100, so the +# `OR` degenerates to the mark alone and the result is identical -- 11 is the one row whose +# mark is `false` and is correctly excluded. +query I +SELECT t1.t1_id +FROM join_t1 t1 +WHERE t1.t1_id > 100 OR EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +22 +33 +44 + +query TT +EXPLAIN +SELECT t1.t1_id +FROM join_t1 t1 +WHERE t1.t1_id > 100 OR EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +logical_plan +01)Sort: t1.t1_id ASC NULLS LAST +02)--Projection: t1.t1_id +03)----Filter: t1.t1_id > Int32(100) OR __correlated_sq_1.mark +04)------LeftMark Join: Filter: t1.t1_id > __correlated_sq_1.t2_id +05)--------SubqueryAlias: t1 +06)----------TableScan: join_t1 projection=[t1_id] +07)--------SubqueryAlias: __correlated_sq_1 +08)----------SubqueryAlias: t2 +09)------------TableScan: join_t2 projection=[t2_id] +physical_plan +01)SortPreservingMergeExec: [t1_id@0 ASC NULLS LAST] +02)--SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----FilterExec: t1_id@0 > 100 OR mark@1, projection=[t1_id@0] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------PiecewiseMergeJoin: operator=Gt, join_type=LeftMark, on=(t1_id > t2_id) +06)----------SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] +07)------------DataSourceExec: partitions=1, partition_sizes=[1] +08)----------DataSourceExec: partitions=1, partition_sizes=[1] + # ------------------------------------------------------------------ # Existence joins: operator, NULL, empty-side and type coverage # ------------------------------------------------------------------ @@ -1079,6 +1126,51 @@ SELECT l.id FROM ex_multinull_l l WHERE NOT EXISTS (SELECT 1 FROM ex_multinull_r 3 4 +# Same multi-NULL buffered side, but through `LeftMark` instead of `LeftSemi`/`LeftAnti`: the +# `id > 100 OR EXISTS(...)` disjunction forces the decorrelation to produce a `mark` value the +# outer `Filter` reads, rather than filtering rows directly off the join. Both NULL-keyed rows +# (ids 1, 3) must still mark `false` under either operator -- a NULL buffered key can never +# land in the matched suffix `min_marked` starts at, the same watermark mechanism `LeftSemi` +# already relies on -- so these pin the same two outcomes the `EXISTS`/`NOT EXISTS` queries +# above already covered, just observed through `LeftMark`. +query I +SELECT l.id FROM ex_multinull_l l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v < r.v) ORDER BY 1; +---- +4 + +query I +SELECT l.id FROM ex_multinull_l l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v > r.v) ORDER BY 1; +---- +2 + +query TT +EXPLAIN SELECT l.id FROM ex_multinull_l l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v < r.v) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----Filter: l.id > Int32(100) OR __correlated_sq_1.mark +04)------Projection: l.id, __correlated_sq_1.mark +05)--------LeftMark Join: Filter: l.v < __correlated_sq_1.v +06)----------SubqueryAlias: l +07)------------TableScan: ex_multinull_l projection=[id, v] +08)----------SubqueryAlias: __correlated_sq_1 +09)------------SubqueryAlias: r +10)--------------TableScan: ex_multinull_r projection=[v] +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] +02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----FilterExec: id@0 > 100 OR mark@1, projection=[id@0] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------ProjectionExec: expr=[id@0 as id, mark@2 as mark] +06)----------PiecewiseMergeJoin: operator=Lt, join_type=LeftMark, on=(v < v) +07)------------SortExec: expr=[v@1 DESC], preserve_partitioning=[false] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +09)------------DataSourceExec: partitions=1, partition_sizes=[1] + # ------------------------------------------------------------------ # Existence joins: float keys and negative zero # ------------------------------------------------------------------ @@ -1118,6 +1210,49 @@ SELECT l.id FROM ex_fl l WHERE NOT EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r. ---- 2 +# Same `-0.0` normalization, but through `LeftMark` instead of `LeftSemi`/`LeftAnti`: putting +# the `EXISTS` in a disjunction with an always-false outer predicate (no `id` is over 100) +# forces the decorrelation to produce a `mark` value the outer `Filter` then reads, rather +# than filtering rows directly off the join. `<` must still mark id 1 (`v = -0.0`) false -- +# `-0.0 < 0.0` is false under SQL equality -- and `<=` must still mark it true, the same pair +# of outcomes the `EXISTS`/`NOT EXISTS` queries above already pinned for `LeftSemi`/`LeftAnti`. +query I +SELECT l.id FROM ex_fl l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_fr r WHERE l.v < r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_fl l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r.v) ORDER BY 1; +---- +1 + +query TT +EXPLAIN SELECT l.id FROM ex_fl l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r.v) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----Filter: l.id > Int32(100) OR __correlated_sq_1.mark +04)------Projection: l.id, __correlated_sq_1.mark +05)--------LeftMark Join: Filter: l.v <= __correlated_sq_1.v +06)----------SubqueryAlias: l +07)------------TableScan: ex_fl projection=[id, v] +08)----------SubqueryAlias: __correlated_sq_1 +09)------------SubqueryAlias: r +10)--------------TableScan: ex_fr projection=[v] +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] +02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----FilterExec: id@0 > 100 OR mark@1, projection=[id@0] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------ProjectionExec: expr=[id@0 as id, mark@2 as mark] +06)----------PiecewiseMergeJoin: operator=LtEq, join_type=LeftMark, on=(v <= v) +07)------------SortExec: expr=[v@1 DESC], preserve_partitioning=[false] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +09)------------DataSourceExec: partitions=1, partition_sizes=[1] + # ------------------------------------------------------------------ # Classic joins: final pass when the streamed side has fewer partitions than # `target_partitions` @@ -1236,6 +1371,42 @@ physical_plan 05)--------DataSourceExec: partitions=1, partition_sizes=[1] 06)------DataSourceExec: partitions=1, partition_sizes=[1] +# `LeftMark` shares the exact same buffered-side comparator (`JoinKeyComparator`, sorted by +# `required_input_ordering`) as `LeftSemi`/`LeftAnti` above -- only `emit_matched`'s final step +# differs (slice vs. `mark` column) -- so this pins that a dictionary key marks correctly +# through that shared path too. `>` : only 'e' (id 3) is above 'c', same outcome the `EXISTS` +# query above already pinned for `LeftSemi`. +query I +SELECT l.id FROM ex_dict_l l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 + +# ------------------------------------------------------------------ +# Existence joins: LeftMark with List/LargeList/FixedSizeList keys +# ------------------------------------------------------------------ + +statement ok +CREATE TABLE ex_list_l(id INT, v INT[]); + +statement ok +INSERT INTO ex_list_l VALUES (1, [1]), (2, [2]), (3, [4]), (4, [5]); + +statement ok +CREATE TABLE ex_list_r(v INT[]); + +statement ok +INSERT INTO ex_list_r VALUES ([2]), ([4]); + +# `>` : `EXISTS r WHERE l.v > r.v` holds iff `l.v` beats the smallest `r.v` ([2]), which only +# [4] (id 3) and [5] (id 4) do. +query I +SELECT l.id FROM ex_list_l l +WHERE l.id > 100 OR EXISTS (SELECT 1 FROM ex_list_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 +4 + # ------------------------------------------------------------------ # Right existence joins (RightSemi / RightAnti) via PiecewiseMergeJoin # ------------------------------------------------------------------