comphead commented on code in PR #25584: URL: https://github.com/apache/datafusion/pull/25584#discussion_r4124469567
########## datafusion/physical-plan/src/joins/sort_merge_join/existence_summary.rs: ########## @@ -0,0 +1,688 @@ +// 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. + +//! Exact bounded summaries of an ordinary semi/anti join's residual predicate. +//! +//! For one equi-key group, existence of `outer != inner` needs at most one +//! non-null representative and a second-distinct-value bit. Existence of +//! `outer <[=] inner` needs the inner maximum, and `outer >[=] inner` needs its +//! minimum. Existence distributes over OR, but not AND: two comparisons may +//! have different witnesses. Each compiled disjunct therefore has at most one +//! cross-side comparison, with any total side-local guards attached to it. +//! +//! Results are non-null *existence* bits, not the residual's three-valued SQL +//! result. A null comparison is not a witness. The caller supplies only rows +//! from one matching equi-key group and implements semi/anti polarity itself. +//! Mark and null-aware joins must not use this interface. +//! +//! Input batches are borrowed only during update/evaluation. A representative +//! is copied after reserving its storage, never retained as a slice of a source +//! array. The same dedicated reservation must accompany updates and resets. +//! Memory admission failure is an execution error; after input has been +//! consumed it is not safe to fall back without an explicit replay path. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, RecordBatch}; +use arrow::compute::SortOptions; +use arrow::compute::kernels::cmp::{gt, gt_eq, lt, lt_eq, neq}; +use arrow::datatypes::{DataType, Schema}; +use arrow_ord::ord::make_comparator; +use datafusion_common::{JoinSide, Result, ScalarValue}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, +}; + +use crate::joins::utils::JoinFilter; + +type Expr = Arc<dyn PhysicalExpr>; + +// Bound compilation work, expression recursion and the number of retained +// representatives even if distributive normalization would grow exponentially. +const MAX_CLAUSES: usize = 32; +const MAX_DEPTH: usize = 64; + +#[derive(Debug, Clone)] +struct Comparison { + outer: Expr, + inner: Expr, + op: Operator, +} + +#[derive(Debug, Default)] +struct State { + /// Used for clauses with only side-local predicates. Even an outer TRUE + /// needs at least one qualifying inner row to witness existence. + present: bool, + representative: Option<ScalarValue>, + multiple: bool, + reserved: usize, +} + +impl State { + fn set_multiple(&mut self, reservation: &MemoryReservation) { + self.representative = None; + reservation.shrink(self.reserved); + *self = Self { + multiple: true, + ..Self::default() + }; + } + + /// Admit the simultaneous old and new copies before copying the candidate. + /// For ranges, retain it only if it improves the current extremum. + fn update_from_array( + &mut self, + array: &ArrayRef, + index: usize, + order: Option<Ordering>, + reservation: &MemoryReservation, + peak: &mut usize, + ) -> Result<()> { + let bytes = scalar_storage_size(array, index); + reservation.try_grow(bytes)?; + *peak = (*peak).max(reservation.size()); + let candidate = ScalarValue::try_from_array(array, index)?; + if let Some(previous) = &self.representative + && order.is_some_and(|order| candidate.partial_cmp(previous) != Some(order)) + { + drop(candidate); + reservation.shrink(bytes); + return Ok(()); + } + self.representative = Some(candidate); + reservation.shrink(self.reserved); + self.reserved = bytes; + Ok(()) + } +} + +#[derive(Debug, Default)] +struct Clause { + outer_guard: Option<Expr>, + inner_guard: Option<Expr>, + comparison: Option<Comparison>, + state: State, +} + +/// A compiled residual and the bounded state for its current inner key group. +#[derive(Debug)] +pub(super) struct ExistenceSummary { + clauses: Vec<Clause>, + peak: usize, +} + +impl ExistenceSummary { + /// Compile only exact, total expressions over supported identically typed + /// values. `None` means the ordinary residual path must be retained. + pub(super) fn try_new( + filter: &JoinFilter, + outer_is_left: bool, + outer_schema: &Schema, + inner_schema: &Schema, + ) -> Result<Option<Self>> { + let compiler = Compiler { + filter, + outer_is_left, + outer_schema, + inner_schema, + }; + Ok(compiler + .compile(filter.expression(), 0)? + .map(|clauses| Self { clauses, peak: 0 })) + } + + /// Incorporate one slice of the current inner group. The caller bounds + /// slice size for cancellation latency and continues draining the group + /// even when every not-equal summary has saturated. + pub(super) fn update( + &mut self, + inner: &RecordBatch, + reservation: &MemoryReservation, + ) -> Result<()> { + for clause in &mut self.clauses { + let guard = evaluate_guard(clause.inner_guard.as_ref(), inner)?; + let selected = |row| { + guard + .as_ref() + .is_none_or(|g| g.is_valid(row) && g.value(row)) + }; + let Some(comparison) = &clause.comparison else { + clause.state.present |= (0..inner.num_rows()).any(selected); + continue; + }; + if clause.state.multiple { + continue; + } + let values = comparison + .inner + .evaluate(inner)? + .into_array(inner.num_rows())?; + let Some(first) = + (0..values.len()).find(|&row| selected(row) && values.is_valid(row)) + else { + continue; + }; + if comparison.op == Operator::NotEq { + if clause.state.representative.is_none() { + clause.state.update_from_array( + &values, + first, + None, + reservation, + &mut self.peak, + )?; + } + let value = clause.state.representative.as_ref().unwrap(); + for row in first..values.len() { + if selected(row) + && values.is_valid(row) + && !value.eq_array(&values, row)? + { + clause.state.set_multiple(reservation); + break; + } + } + } else { + let order = if matches!(comparison.op, Operator::Lt | Operator::LtEq) { + Ordering::Greater + } else { + Ordering::Less + }; + let comparator = make_comparator( + values.as_ref(), + values.as_ref(), + SortOptions::default(), + )?; + let mut candidate = first; + for row in first + 1..values.len() { + if selected(row) + && values.is_valid(row) + && comparator(row, candidate) == order + { + candidate = row; + } + } + clause.state.update_from_array( + &values, + candidate, + Some(order), + reservation, + &mut self.peak, + )?; + } + } + Ok(()) + } + + /// Return non-null witness bits for an outer slice. Null outer values do + /// not match even when a not-equal summary contains two distinct values. + pub(super) fn evaluate(&self, outer: &RecordBatch) -> Result<BooleanArray> { + let mut matches = vec![false; outer.num_rows()]; Review Comment: The probe can follow PWMJ as well. `RightExistencePWMJStream` (`piecewise_merge_join/right_existence_join.rs:290-312`) calls `apply_cmp` with the extreme as a `ColumnarValue::Scalar`, and `boolean_mask_from_filter` (`joins/utils.rs:1987`) folds NULL into false. One mask per clause (comparison mask AND guard mask), OR-ed straight into `self.matched` with `apply_bitwise_binary_op` as `eval_filter_for_inner_slice` does, removes the loop over rows × clauses, the `Vec<bool>`, the `BooleanArray::from` repack and the second `Operator` table ending in `unreachable!`. `BinaryExpr` compares through `apply_cmp` too (`binary.rs:657`), so both paths would share one comparison implementation if more types are admitted later. This carries over unchanged to the min/max shape. ########## datafusion/physical-plan/src/joins/sort_merge_join/existence_summary.rs: ########## @@ -0,0 +1,688 @@ +// 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. + +//! Exact bounded summaries of an ordinary semi/anti join's residual predicate. +//! +//! For one equi-key group, existence of `outer != inner` needs at most one +//! non-null representative and a second-distinct-value bit. Existence of +//! `outer <[=] inner` needs the inner maximum, and `outer >[=] inner` needs its +//! minimum. Existence distributes over OR, but not AND: two comparisons may +//! have different witnesses. Each compiled disjunct therefore has at most one +//! cross-side comparison, with any total side-local guards attached to it. +//! +//! Results are non-null *existence* bits, not the residual's three-valued SQL +//! result. A null comparison is not a witness. The caller supplies only rows +//! from one matching equi-key group and implements semi/anti polarity itself. +//! Mark and null-aware joins must not use this interface. +//! +//! Input batches are borrowed only during update/evaluation. A representative +//! is copied after reserving its storage, never retained as a slice of a source +//! array. The same dedicated reservation must accompany updates and resets. +//! Memory admission failure is an execution error; after input has been +//! consumed it is not safe to fall back without an explicit replay path. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, RecordBatch}; +use arrow::compute::SortOptions; +use arrow::compute::kernels::cmp::{gt, gt_eq, lt, lt_eq, neq}; +use arrow::datatypes::{DataType, Schema}; +use arrow_ord::ord::make_comparator; +use datafusion_common::{JoinSide, Result, ScalarValue}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr, +}; + +use crate::joins::utils::JoinFilter; + +type Expr = Arc<dyn PhysicalExpr>; + +// Bound compilation work, expression recursion and the number of retained +// representatives even if distributive normalization would grow exponentially. +const MAX_CLAUSES: usize = 32; +const MAX_DEPTH: usize = 64; + +#[derive(Debug, Clone)] +struct Comparison { + outer: Expr, + inner: Expr, + op: Operator, +} + +#[derive(Debug, Default)] +struct State { Review Comment: The reduction already exists in `joins/`. `CollectLeftAccumulator` (`hash_join/exec.rs:2661`) evaluates an expression per batch and folds `MinAccumulator`/`MaxAccumulator` across batches, and PWMJ's `extreme_key` (`piecewise_merge_join/existence_join.rs:456`) uses `min_batch`/`max_batch`. Moving `CollectLeftAccumulator` to `joins/utils.rs` and feeding it the guard-filtered inner column would replace `State`, `update_from_array`, the `make_comparator` loop and `scalar_storage_size` (`Accumulator::size()` goes through `ScalarValue::size()`). With min and max available, `<>` needs no `multiple` flag (`x <> min OR x <> max`, as @jayzhan211 wrote). ########## datafusion/sqllogictest/test_files/sort_merge_join.slt: ########## @@ -971,6 +971,47 @@ WHERE t1_sorted.data < 0 ---- 100 +statement ok +SET datafusion.execution.enable_sort_merge_join_existence_summary = true; + +# OR clauses can have different witnesses. A null component is not itself +# a conflict, and a missing inner group must survive an anti join. +statement ok +CREATE TABLE identities(id BIGINT, k BIGINT, user_id VARCHAR, org_id VARCHAR) AS VALUES + (1, 1, 'alice', 'x'), (2, 2, 'alice', 'x'), + (3, 3, NULL, 'x'), (4, 4, 'alice', NULL), (5, 5, 'carol', 'x'); + +statement ok +CREATE TABLE observations(k BIGINT, user_id VARCHAR, org_id VARCHAR, active BOOLEAN) AS VALUES + (1, 'alice', 'x', true), (1, 'alice', 'x', true), + (1, 'bob', 'y', false), (2, 'alice', 'y', true), + (2, 'bob', 'x', true), (3, 'bob', NULL, true), (4, 'bob', NULL, true); + +query I rowsort +SELECT l.id FROM identities l LEFT ANTI JOIN observations r +ON l.k = r.k AND r.active AND (l.user_id <> r.user_id OR l.org_id <> r.org_id); Review Comment: `r.active` is a top-level ON conjunct, so `PushDownFilter` moves it into `observations` for LEFT SEMI and LEFT ANTI (`JoinType::on_lr_is_preserved`, see the `left_semi_join_with_filters` and `left_anti_join_with_filters` snapshots in `push_down_filter.rs`). Only the OR reaches the join, so none of these three queries exercises the inner-guard code. It also supports @jayzhan211's point that side-local guards rarely survive to the join. A guard only stays when nested under OR, e.g. `((r.active AND l.user_id <> r.user_id) OR l.org_id <> r.org_id)`, or as an outer-only conjunct of an anti join. ########## datafusion/physical-plan/src/joins/sort_merge_join/existence_summary_tests.rs: ########## @@ -0,0 +1,753 @@ +// 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. + +//! Differential execution tests for bounded semi/anti existence summaries. + +use std::sync::Arc; +use std::task::Context; + +use arrow::array::{ + Array, ArrayRef, BooleanArray, Date32Array, Decimal128Array, Int32Array, + LargeStringArray, RecordBatch, StringArray, StringViewArray, + TimestampMicrosecondArray, +}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::{ + DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, + assert_contains, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, Column, IsNullExpr, Literal, NotExpr, +}; + +use crate::joins::SortMergeJoinExec; +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::test::TestMemoryExec; +use crate::test::exec::MockExec; +use crate::{ExecutionPlan, PhysicalExpr, common}; + +const JOINS: [JoinType; 4] = [ + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, +]; +type Row = (Option<i32>, Option<i32>, Option<bool>); + +fn column(index: usize) -> Arc<dyn PhysicalExpr> { + Arc::new(Column::new( + ["left_value", "right_value", "left_guard", "right_guard"][index], + index, + )) +} + +fn binary( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, +) -> Arc<dyn PhysicalExpr> { + Arc::new(BinaryExpr::new(left, op, right)) +} + +fn filter(expr: Arc<dyn PhysicalExpr>, value_type: DataType) -> JoinFilter { + JoinFilter::new( + expr, + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("left_value", value_type.clone(), true), + Field::new("right_value", value_type, true), + Field::new("left_guard", DataType::Boolean, true), + Field::new("right_guard", DataType::Boolean, true), + ])), + ) +} + +fn comparison(op: Operator) -> JoinFilter { + filter(binary(column(0), op, column(1)), DataType::Int32) +} + +fn batch(rows: &[Row]) -> Result<RecordBatch> { + batch_values( + rows, + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.1))), + ) +} + +fn batch_values(rows: &[Row], values: ArrayRef) -> Result<RecordBatch> { + RecordBatch::try_from_iter(vec![ + ( + "key", + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.0))) as ArrayRef, + ), + ("value", values), + ( + "guard", + Arc::new(BooleanArray::from_iter(rows.iter().map(|r| r.2))), + ), + ( + "id", + Arc::new(Int32Array::from_iter_values(0..rows.len() as i32)), + ), + ]) + .map_err(Into::into) +} + +fn input(batch: &RecordBatch, chunk: usize) -> Result<Arc<dyn ExecutionPlan>> { + // Empty batches between slices exercise both empty input and non-zero array offsets. + let mut batches = vec![batch.slice(0, 0)]; + for offset in (0..batch.num_rows()).step_by(chunk) { + batches.push(batch.slice(offset, chunk.min(batch.num_rows() - offset))); + batches.push(batch.slice(offset, 0)); + } + Ok(TestMemoryExec::try_new_exec( + &[batches], + batch.schema(), + None, + )?) +} + +fn join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: JoinType, + filter: JoinFilter, + options: SortOptions, + nulls: NullEquality, +) -> Result<SortMergeJoinExec> { + SortMergeJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + Some(filter), + join_type, + vec![options], + nulls, + ) +} + +fn config(batch_size: usize, enabled: bool) -> SessionConfig { + let mut config = SessionConfig::new().with_batch_size(batch_size); + config + .options_mut() + .execution + .enable_sort_merge_join_existence_summary = enabled; + config +} + +fn context(batch_size: usize, enabled: bool) -> Arc<TaskContext> { + Arc::new(TaskContext::default().with_session_config(config(batch_size, enabled))) +} + +fn metric(plan: &dyn ExecutionPlan, name: &str) -> usize { + plan.metrics() + .unwrap() + .iter() + .filter(|metric| metric.value().name() == name) + .map(|metric| metric.value().as_usize()) + .sum() +} + +async fn ids(plan: &dyn ExecutionPlan, ctx: Arc<TaskContext>) -> Result<Vec<i32>> { + let mut ids = common::collect(plan.execute(0, ctx)?) + .await? + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap() + .values() + .to_vec() + }) + .collect::<Vec<_>>(); + ids.sort_unstable(); + Ok(ids) +} + +#[tokio::test] +async fn comparisons_match_generic_execution_across_groups_and_orientations() -> Result<()> +{ + let left = vec![ + (None, None, None), + (None, Some(4), Some(true)), + (Some(0), Some(7), Some(true)), + (Some(1), None, None), + (Some(1), Some(2), Some(true)), + (Some(1), Some(2), Some(false)), + (Some(1), Some(6), None), + (Some(2), Some(3), Some(true)), + (Some(4), Some(5), Some(false)), + ]; + let right = vec![ + (None, Some(1), Some(true)), + (Some(1), None, None), + (Some(1), Some(2), Some(false)), + (Some(1), Some(2), Some(true)), + (Some(1), Some(6), Some(true)), + (Some(2), None, None), + (Some(3), Some(4), Some(false)), + ]; + let mut cases = vec![]; + for op in [ + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + ] { + for kind in JOINS { + cases.push(( + op, + kind, + SortOptions::default(), + NullEquality::NullEqualsNothing, + )); + } + } + // Key ordering and null equality are independent of residual comparison. + cases.extend([ + ( + Operator::NotEq, + JoinType::LeftAnti, + SortOptions { + descending: false, + nulls_first: false, + }, + NullEquality::NullEqualsNull, + ), + ( + Operator::Lt, + JoinType::RightSemi, + SortOptions { + descending: true, + nulls_first: true, + }, + NullEquality::NullEqualsNull, + ), + ( + Operator::LtEq, + JoinType::LeftSemi, + SortOptions { + descending: true, + nulls_first: false, + }, + NullEquality::NullEqualsNothing, + ), + ]); + for (op, kind, options, nulls) in cases { + let sort = |rows: &[Row]| { + let mut rows = rows.to_vec(); + rows.sort_by(|a, b| { + let null_order = if options.nulls_first { + b.0.is_none().cmp(&a.0.is_none()) + } else { + a.0.is_none().cmp(&b.0.is_none()) + }; + null_order.then_with(|| { + if options.descending { + b.0.cmp(&a.0) + } else { + a.0.cmp(&b.0) + } + }) + }); + rows + }; + let (left, right) = (batch(&sort(&left))?, batch(&sort(&right))?); + let mut outputs = vec![]; + for enabled in [false, true] { + let plan = join( + input(&left, 2)?, + input(&right, 3)?, + kind, + comparison(op), + options, + nulls, + )?; + outputs.push(ids(&plan, context(2, enabled)).await?); + assert_eq!( + metric(&plan, "existence_summary_enabled"), + usize::from(enabled) + ); + if enabled { + assert!(metric(&plan, "existence_summary_inner_rows") > 0); + } else { + assert!(plan.metrics().unwrap().iter().all(|metric| { + !metric.value().name().starts_with("existence_summary_") + })); + } + } + assert_eq!( + outputs[0], outputs[1], + "{kind:?} {op:?} {options:?} {nulls:?}" + ); + if kind == JoinType::LeftSemi + && options == SortOptions::default() + && nulls == NullEquality::NullEqualsNothing + { + // Rows 4/5 equal the minimum; row 6 equals the maximum. + let expected = match op { + Operator::Lt => vec![4, 5], + Operator::Gt => vec![6], + Operator::NotEq | Operator::LtEq | Operator::GtEq => vec![4, 5, 6], + _ => unreachable!(), + }; + assert_eq!(outputs[1], expected, "{op:?}"); + } + } + Ok(()) +} + +#[tokio::test] +async fn guarded_or_preserves_anti_rows_and_requires_an_inner_witness() -> Result<()> { + let left = vec![ + (Some(0), Some(5), Some(true)), + (Some(1), Some(2), Some(true)), + (Some(1), Some(3), Some(false)), + (Some(1), None, None), + (Some(2), Some(7), Some(true)), + (Some(3), None, Some(true)), + ]; + let right = vec![ + (Some(1), Some(2), Some(false)), + (Some(1), Some(4), Some(true)), + (Some(1), None, None), + (Some(2), Some(7), None), + ]; + for kind in JOINS { + for outer_only_or in [false, true] { + let guarded = binary( + binary( + column(2), + Operator::And, + binary(column(0), Operator::NotEq, column(1)), + ), + Operator::And, + column(3), + ); + let expr = if outer_only_or { + binary(guarded, Operator::Or, Arc::new(IsNullExpr::new(column(0)))) + } else { + binary( + guarded, + Operator::Or, + binary(column(0), Operator::Lt, column(1)), + ) + }; + let expected = match (kind, outer_only_or) { + (JoinType::LeftSemi, false) => vec![1, 2], + (JoinType::LeftAnti, false) => vec![0, 3, 4, 5], + (JoinType::RightSemi, false) => vec![1], + (JoinType::RightAnti, false) => vec![0, 2, 3], + (JoinType::LeftSemi, true) => vec![1, 3], + (JoinType::LeftAnti, true) => vec![0, 2, 4, 5], + (JoinType::RightSemi, true) => vec![0, 1, 2], + (JoinType::RightAnti, true) => vec![3], + _ => unreachable!(), + }; + for enabled in [false, true] { + let plan = join( + input(&batch(&left)?, 1)?, + input(&batch(&right)?, 2)?, + kind, + filter(Arc::clone(&expr), DataType::Int32), + SortOptions::default(), + NullEquality::NullEqualsNothing, + )?; + assert_eq!( + ids(&plan, context(1, enabled)).await?, + expected, + "{kind:?} {outer_only_or} {enabled}" + ); + assert_eq!( + metric(&plan, "existence_summary_enabled"), + usize::from(enabled) + ); + } + } + } + Ok(()) +} + +#[tokio::test] +async fn strings_dates_timestamps_and_decimals_match_generic_execution() -> Result<()> { + let rows = vec![(Some(1), None, Some(true)); 5]; + let types: Vec<ArrayRef> = vec![ + Arc::new(StringArray::from(vec![ + None, + Some(""), + Some("é"), + Some("a"), + Some("é"), + ])), + Arc::new(LargeStringArray::from(vec![ + None, + Some(""), + Some("é"), + Some("a"), + Some("é"), + ])), + Arc::new(StringViewArray::from(vec![ + None, + Some(""), + Some("out-of-line-string-z"), + Some("out-of-line-string-a"), + Some("out-of-line-string-z"), + ])), + Arc::new(Date32Array::from(vec![ + None, + Some(-1), + Some(0), + Some(1), + Some(1), + ])), + Arc::new( + TimestampMicrosecondArray::from(vec![ + None, + Some(-100), + Some(0), + Some(100), + Some(100), + ]) + .with_timezone("UTC"), + ), + Arc::new( + Decimal128Array::from(vec![None, Some(-123), Some(0), Some(456), Some(456)]) + .with_precision_and_scale(20, 2)?, + ), + ]; + for values in types { + let data_type = values.data_type().clone(); + let batch = batch_values(&rows, values)?; + // NotEq uses scalar equality; ranges use array and scalar ordering. + for op in [Operator::NotEq, Operator::Lt] { + let mut outputs = vec![]; + for enabled in [false, true] { + let plan = join( + input(&batch, 2)?, + input(&batch, 3)?, + JoinType::LeftSemi, + filter(binary(column(0), op, column(1)), data_type.clone()), + SortOptions::default(), + NullEquality::NullEqualsNothing, + )?; + outputs.push(ids(&plan, context(2, enabled)).await?); + assert_eq!( + metric(&plan, "existence_summary_enabled"), + usize::from(enabled), + "{data_type:?} {op:?}" + ); + } + assert_eq!(outputs[0], outputs[1], "{data_type:?} {op:?}"); + } + } + Ok(()) +} + +#[tokio::test] +async fn independent_cross_side_witnesses_use_generic_fallback() -> Result<()> { + let values = batch(&[ + (Some(1), Some(0), Some(true)), + (Some(1), Some(5), None), + (Some(1), Some(10), Some(false)), + ])?; + // Min/max have separate witnesses for 5, but no row satisfies both clauses. + let expr = binary( + binary(column(0), Operator::Lt, column(1)), + Operator::And, + binary(column(0), Operator::Gt, column(1)), + ); + for enabled in [false, true] { + let plan = join( + input(&values, 1)?, + input(&values, 2)?, + JoinType::LeftSemi, + filter(Arc::clone(&expr), DataType::Int32), + SortOptions::default(), + NullEquality::NullEqualsNothing, + )?; + assert!(ids(&plan, context(2, enabled)).await?.is_empty()); + assert_eq!(metric(&plan, "existence_summary_enabled"), 0); + assert_eq!( + metric(&plan, "existence_summary_fallback"), + usize::from(enabled) + ); + } + Ok(()) +} + +#[tokio::test] +async fn normalized_nullable_strings_and_negated_equality_keep_sql_semantics() +-> Result<()> { + let left = batch_values( + &[(Some(1), None, Some(true)); 4], + Arc::new(StringArray::from(vec![ + None, + Some(""), + Some("a"), + Some("out-of-line-value"), + ])), + )?; + let right = batch_values( + &[(Some(1), None, Some(true)); 2], + Arc::new(StringArray::from(vec![None, Some("")])), + )?; + let normalize = |index| -> Result<Arc<dyn PhysicalExpr>> { + Ok(Arc::new(CaseExpr::try_new( + None, + vec![( + Arc::new(IsNullExpr::new(column(index))), + Arc::new(Literal::new(ScalarValue::Utf8(Some(String::new())))), + )], + Some(column(index)), + )?)) + }; + for kind in JOINS { + let expected = match kind { + JoinType::LeftSemi => vec![2, 3], + JoinType::LeftAnti => vec![0, 1], + JoinType::RightSemi => vec![0, 1], + JoinType::RightAnti => vec![], + _ => unreachable!(), + }; + for enabled in [false, true] { + let expr = Arc::new(NotExpr::new(binary( + normalize(0)?, + Operator::Eq, + normalize(1)?, + ))); + let plan = join( + input(&left, 1)?, + input(&right, 1)?, + kind, + filter(expr, DataType::Utf8), + SortOptions::default(), + NullEquality::NullEqualsNothing, + )?; + assert_eq!( + ids(&plan, context(1, enabled)).await?, + expected, + "{kind:?} {enabled}" + ); + assert_eq!( + metric(&plan, "existence_summary_enabled"), + usize::from(enabled) + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn empty_inputs_never_synthesize_a_witness() -> Result<()> { Review Comment: This never reaches the summary. With an empty inner side `join` goes straight to `drain_outer`, and with an empty outer side the loop never runs, so neither `summarize_inner_key_group` nor `evaluate` is called. It checks that the filter compiled, plus the existing drain behavior. The case in the comment needs a matching key whose inner rows all fail an inner guard, so that `present` stays false. ########## datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs: ########## @@ -912,6 +1003,122 @@ impl BitwiseSortMergeJoinStream { } } + /// Return a bounded slice length, yielding periodically even if every + /// child batch is immediately ready. Pausing the timer excludes scheduling + /// time from the join's own work. + async fn summary_slice_len(&mut self, remaining: usize) -> usize { + if self.existence_summary.as_ref().unwrap().rows_until_yield == 0 { + self.stop_join_time(); + tokio::task::yield_now().await; + self.start_join_time(); + self.existence_summary.as_mut().unwrap().rows_until_yield = + SUMMARY_WORK_BUDGET; + } + let state = self.existence_summary.as_mut().unwrap(); + let len = remaining.min(state.rows_until_yield); + state.rows_until_yield -= len; + len + } + + /// Consume the whole inner key group into owned summary state. Continue + /// draining after saturation so input errors and group boundaries retain + /// their ordinary behavior. + async fn summarize_inner_key_group(&mut self) -> Result<()> { Review Comment: Apart from the per-slice block, this is `buffer_inner_key_group` with renamed variables, and `process_summary_match_loop` repeats the tail of `process_filtered_match_loop`. If skipping the buffer is worth keeping, fold min/max at the per-slice point in `buffer_inner_key_group` (where it pushes into `inner_key_buffer`, around line 750) and probe where `process_filtered_match_loop` calls `process_key_match_with_filter`. That keeps one traversal and one reservation (`inner_buffer_size` / `try_resize_reservation`), and eligible groups never buffer or spill. Per-group fixed cost matters more here than the large-group benchmarks show. In TPC-H Q21 only `l_suppkey != l_suppkey` reaches this join (`tpch/plans/q21.slt.part:65-66`), and an order has 1 to 7 lineitems. That is closer to `small_group_no_witness` (1.6x) than to the 57x/106x cases. Each group here slices both batches over all columns, re-evaluates the operand and guard expressions, and allocates a `to_scalar()` array, a `Vec<bool>` and a `BooleanArray`. Evaluating the operand and guard columns once per input batch, as `evaluate_join_keys` does for keys, and passing row ranges should remove most of that. This is from reading the code, not profiling. ########## datafusion/physical-plan/src/joins/sort_merge_join/existence_summary_tests.rs: ########## @@ -0,0 +1,753 @@ +// 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. + +//! Differential execution tests for bounded semi/anti existence summaries. + +use std::sync::Arc; +use std::task::Context; + +use arrow::array::{ + Array, ArrayRef, BooleanArray, Date32Array, Decimal128Array, Int32Array, + LargeStringArray, RecordBatch, StringArray, StringViewArray, + TimestampMicrosecondArray, +}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::{ + DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, + assert_contains, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, Column, IsNullExpr, Literal, NotExpr, +}; + +use crate::joins::SortMergeJoinExec; +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::test::TestMemoryExec; +use crate::test::exec::MockExec; +use crate::{ExecutionPlan, PhysicalExpr, common}; + +const JOINS: [JoinType; 4] = [ + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, +]; +type Row = (Option<i32>, Option<i32>, Option<bool>); + +fn column(index: usize) -> Arc<dyn PhysicalExpr> { + Arc::new(Column::new( + ["left_value", "right_value", "left_guard", "right_guard"][index], + index, + )) +} + +fn binary( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, +) -> Arc<dyn PhysicalExpr> { + Arc::new(BinaryExpr::new(left, op, right)) +} + +fn filter(expr: Arc<dyn PhysicalExpr>, value_type: DataType) -> JoinFilter { + JoinFilter::new( + expr, + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("left_value", value_type.clone(), true), + Field::new("right_value", value_type, true), + Field::new("left_guard", DataType::Boolean, true), + Field::new("right_guard", DataType::Boolean, true), + ])), + ) +} + +fn comparison(op: Operator) -> JoinFilter { + filter(binary(column(0), op, column(1)), DataType::Int32) +} + +fn batch(rows: &[Row]) -> Result<RecordBatch> { + batch_values( + rows, + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.1))), + ) +} + +fn batch_values(rows: &[Row], values: ArrayRef) -> Result<RecordBatch> { + RecordBatch::try_from_iter(vec![ + ( + "key", + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.0))) as ArrayRef, + ), + ("value", values), + ( + "guard", + Arc::new(BooleanArray::from_iter(rows.iter().map(|r| r.2))), + ), + ( + "id", + Arc::new(Int32Array::from_iter_values(0..rows.len() as i32)), + ), + ]) + .map_err(Into::into) +} + +fn input(batch: &RecordBatch, chunk: usize) -> Result<Arc<dyn ExecutionPlan>> { + // Empty batches between slices exercise both empty input and non-zero array offsets. + let mut batches = vec![batch.slice(0, 0)]; + for offset in (0..batch.num_rows()).step_by(chunk) { + batches.push(batch.slice(offset, chunk.min(batch.num_rows() - offset))); + batches.push(batch.slice(offset, 0)); + } + Ok(TestMemoryExec::try_new_exec( + &[batches], + batch.schema(), + None, + )?) +} + +fn join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: JoinType, + filter: JoinFilter, + options: SortOptions, + nulls: NullEquality, +) -> Result<SortMergeJoinExec> { + SortMergeJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + Some(filter), + join_type, + vec![options], + nulls, + ) +} + +fn config(batch_size: usize, enabled: bool) -> SessionConfig { + let mut config = SessionConfig::new().with_batch_size(batch_size); + config + .options_mut() + .execution + .enable_sort_merge_join_existence_summary = enabled; + config +} + +fn context(batch_size: usize, enabled: bool) -> Arc<TaskContext> { + Arc::new(TaskContext::default().with_session_config(config(batch_size, enabled))) +} + +fn metric(plan: &dyn ExecutionPlan, name: &str) -> usize { + plan.metrics() + .unwrap() + .iter() + .filter(|metric| metric.value().name() == name) + .map(|metric| metric.value().as_usize()) + .sum() +} + +async fn ids(plan: &dyn ExecutionPlan, ctx: Arc<TaskContext>) -> Result<Vec<i32>> { + let mut ids = common::collect(plan.execute(0, ctx)?) + .await? + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap() + .values() + .to_vec() + }) + .collect::<Vec<_>>(); + ids.sort_unstable(); + Ok(ids) +} + +#[tokio::test] +async fn comparisons_match_generic_execution_across_groups_and_orientations() -> Result<()> Review Comment: Most of this file could be SLT, like the PWMJ existence coverage in `pwmj.slt`. `sort_merge_join_matrix.slt` already checks SMJ against HJ at `batch_size` 1, 2, 100 and 8192, but has no residual-filtered existence queries. The flag as a third `# configMatrix:` dimension plus residual EXISTS / NOT EXISTS / RIGHT SEMI / RIGHT ANTI queries (five operators, NULLs, OR, `arrow_cast` to LargeUtf8/Utf8View, dates, timestamps, decimals, and COALESCE, which simplifies to CASE) covers the operator × join type, type and multi-batch cases here. - The three sort-option / null-equality cases mostly re-test the shared `find_key_group_end` / `keys_match` comparators. - `independent_cross_side_witnesses_use_generic_fallback` repeats `sort_merge_join.slt:1005-1010`. - The activation test in `core/tests/sql/joins.rs` fits `sort_merge_join.slt` as `EXPLAIN ANALYZE` with `<slt:ignore>` around the counter, as `aggregate_memory_spill.slt` does for `spill_count`. - The memory, upstream-error and drop tests, and the `NotExpr` compile test, need to stay in Rust. The SQL simplifier rewrites `NOT (a = b)` into `a != b`, so SQL cannot reach that path. The remaining metric sums can use `MetricsSet::sum_by_name`. ########## datafusion/physical-plan/benches/sort_merge_join.rs: ########## @@ -248,5 +251,167 @@ fn bench_smj(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_smj); +/// Compare execution with summaries enabled and disabled in the same binary. +/// Inputs are already sorted: SQL versions of these EXISTS/NOT EXISTS queries +/// also measure sorting and depend on the optimizer's choice of join algorithm. +/// These cases isolate the residual semi/anti join, including output collection. +fn bench_existence_summary(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let mut group = c.benchmark_group("sort_merge_join_existence_summary"); + group.sample_size(10); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2)); + + // Vary group count, rows per side, residual selectivity, and join type. + // The early-witness case controls for the generic join's short circuit; + // the small-group case measures the cost of repeatedly resetting summaries. + for (name, groups, probe_rows, inner_rows, distinct, op, kind) in [ Review Comment: None of these cases shows the summary losing, which is the worst case @jayzhan211 asked to measure. `not_equal_early_witness` saturates, so both paths stop early. With a range predicate whose first inner row matches every outer row, the generic path stops after one filter evaluation, and `buffer_inner_key_group` only slices batches (`find_key_group_end` is a binary search), so the summary's scan of the whole group is extra work. Only Int64 `Column op Column` is measured, not strings, guards or OR. For end-to-end numbers, `benchmarks/src/smj.rs` Q11-Q13 and Q18 are LEFT SEMI `<>` queries with SMJ forced, and the runner reads config from the environment (`benchmarks/src/util/options.rs:81`). So `DATAFUSION_EXECUTION_ENABLE_SORT_MERGE_JOIN_EXISTENCE_SUMMARY=true ./bench.sh run smj` gives an on/off comparison with no new code. I expect those queries to qualify, since their `%` predicate should be pushed into the inner input, but I have not checked the plans. ########## datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs: ########## @@ -912,6 +1003,122 @@ impl BitwiseSortMergeJoinStream { } } + /// Return a bounded slice length, yielding periodically even if every + /// child batch is immediately ready. Pausing the timer excludes scheduling + /// time from the join's own work. + async fn summary_slice_len(&mut self, remaining: usize) -> usize { Review Comment: +1. The `ExecutionPlan::execute` docs reserve manual yields for work that is superlinear per batch, and the quadratic `process_key_match_with_filter` does not yield today, so this linear path should not need one. If a yield is ever needed, `tokio::task::coop::consume_budget()` (as used in `coop.rs`) follows the task budget instead of rescheduling every 1024 rows. `summary_work_yields_and_dropping_stream_releases_memory` would go with it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
