jonathanc-n commented on code in PR #23741: URL: https://github.com/apache/datafusion/pull/23741#discussion_r3688199160
########## datafusion/physical-plan/src/joins/ie_join/exec.rs: ########## @@ -0,0 +1,509 @@ +// 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. + +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::compute::concat_batches; +use arrow::record_batch::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; +use datafusion_common::{ + JoinType, NullEquality, Result, Statistics, internal_err, plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::Operator; +use datafusion_physical_expr::equivalence::join_equivalence_properties; +use datafusion_physical_expr::{Distribution, PhysicalExprRef}; +use datafusion_physical_expr_common::physical_expr::{fmt_sql, is_volatile}; +use futures::{FutureExt, StreamExt}; + +use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::joins::ie_join::algorithm::IEJoinData; +use crate::joins::ie_join::stream::{IEJoinMetrics, IEJoinStream}; +use crate::joins::utils::{ + ColumnIndex, JoinFilter, build_join_schema, check_join_is_valid, + estimate_join_statistics, symmetric_join_output_partitioning, +}; +use crate::joins::{JoinOn, JoinOnRef}; +use crate::spill::get_record_batch_memory_size; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, check_if_same_properties, +}; + +/// One of the two range predicates that drive [`IEJoinExec`]. +#[derive(Debug, Clone)] +pub struct IEJoinCondition { + left: PhysicalExprRef, + right: PhysicalExprRef, + operator: Operator, +} + +impl IEJoinCondition { + /// Creates a normalized `left OP right` range condition. + pub fn new( + left: PhysicalExprRef, + right: PhysicalExprRef, + operator: Operator, + ) -> Self { + Self { + left, + right, + operator, + } + } + + /// Expression evaluated against the left input. + pub fn left(&self) -> &PhysicalExprRef { + &self.left + } + + /// Expression evaluated against the right input. + pub fn right(&self) -> &PhysicalExprRef { + &self.right + } + + /// Range comparison operator. + pub fn operator(&self) -> Operator { + self.operator + } +} + +/// An inequality join using two sorted orders, a permutation, and a bitmap. +/// +/// With equality keys, both inputs are co-partitioned by those keys. Inside +/// each input partition, rows are grouped by equality-key hash before the two +/// range predicates are evaluated. Hash collisions are checked with the full +/// equality expressions, so they cannot affect correctness. +/// +/// Both input partitions are materialized. Candidate pairs are emitted in +/// bounded batches and the bitmap scan retains its cursor across polls, so the +/// output cardinality never determines intermediate memory usage. +#[derive(Debug)] +pub struct IEJoinExec { + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + conditions: [IEJoinCondition; 2], + filter: Option<JoinFilter>, + join_type: JoinType, + null_equality: NullEquality, + schema: SchemaRef, + column_indices: Vec<ColumnIndex>, + metrics: crate::metrics::ExecutionPlanMetricsSet, + cache: Arc<PlanProperties>, +} + +impl IEJoinExec { + /// Creates an inner IEJoin from optional equality keys and exactly two + /// normalized range conditions. Each condition's left and right physical + /// expressions must have identical data types. SQL planning inserts any + /// required casts; programmatic callers must insert them explicitly. + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + conditions: [IEJoinCondition; 2], + filter: Option<JoinFilter>, + join_type: JoinType, + null_equality: NullEquality, + ) -> Result<Self> { + if join_type != JoinType::Inner { + return plan_err!( + "IEJoinExec currently supports Inner joins, got {join_type}" + ); + } + for condition in &conditions { + if is_volatile(condition.left()) || is_volatile(condition.right()) { + return plan_err!( + "IEJoinExec range conditions must not contain volatile expressions" + ); + } + if !matches!( + condition.operator, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "IEJoinExec requires <, <=, >, or >=, got {}", + condition.operator + ); + } + let left_type = condition.left.data_type(&left.schema())?; + let right_type = condition.right.data_type(&right.schema())?; + if left_type != right_type { + return plan_err!( + "IEJoinExec condition types differ: {left_type} and {right_type}; physical expressions must be explicitly coerced to identical types" + ); + } + } + + check_join_is_valid(&left.schema(), &right.schema(), &on)?; + let (schema, column_indices) = + build_join_schema(&left.schema(), &right.schema(), &join_type); + let schema = Arc::new(schema); + let cache = Arc::new(Self::compute_properties( + &left, &right, &schema, join_type, &on, + )?); + + Ok(Self { + left, + right, + on, + conditions, + filter, + join_type, + null_equality, + schema, + column_indices, + metrics: Default::default(), + cache, + }) + } + + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + pub fn on(&self) -> JoinOnRef<'_> { + &self.on + } + + pub fn conditions(&self) -> &[IEJoinCondition; 2] { + &self.conditions + } + + pub fn filter(&self) -> Option<&JoinFilter> { + self.filter.as_ref() + } + + pub fn join_type(&self) -> JoinType { + self.join_type + } + + pub fn null_equality(&self) -> NullEquality { + self.null_equality + } + + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: &SchemaRef, + join_type: JoinType, + on: JoinOnRef<'_>, + ) -> Result<PlanProperties> { + let eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + Arc::clone(schema), + &[false, false], + None, + on, + )?; + let output_partitioning = if on.is_empty() { + Partitioning::UnknownPartitioning(1) + } else { + symmetric_join_output_partitioning(left, right, &join_type)? + }; + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Final, + boundedness_from_children([left, right]), Review Comment: boundedness_from_children + emissiontype::final means an unbounded input would buffer forever. I think we can consider rejecting unbounded inputs ########## datafusion/physical-plan/src/joins/ie_join/stream.rs: ########## @@ -0,0 +1,429 @@ +// 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. + +use std::cmp::Ordering; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, ready}; + +use arrow::array::{UInt32Array, UInt64Array}; +use arrow::record_batch::RecordBatch; +use arrow_ord::ord::{DynComparator, make_comparator}; +use arrow_schema::{SchemaRef, SortOptions}; +use datafusion_common::{JoinSide, JoinType, NullEquality, Result}; +use datafusion_expr::Operator; +use futures::future::BoxFuture; +use futures::{FutureExt, Stream}; + +use crate::RecordBatchStream; +use crate::joins::JoinOn; +use crate::joins::ie_join::algorithm::{ActiveBitmap, IEJoinData, RightGroup}; +use crate::joins::ie_join::exec::IEJoinCondition; +use crate::joins::utils::{ + ColumnIndex, JoinFilter, apply_join_filter_to_indices, build_batch_from_indices, + equal_rows_arr, +}; +use crate::metrics::{Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time}; + +#[derive(Clone)] +pub(super) struct IEJoinMetrics { + pub load_time: Time, + pub join_time: Time, + pub left_input_batches: Count, + pub right_input_batches: Count, + pub left_input_rows: Count, + pub right_input_rows: Count, + pub candidate_rows: Count, + pub output_batches: Count, + pub output_rows: Count, + pub peak_mem_used: Gauge, +} + +impl IEJoinMetrics { + pub fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + load_time: MetricBuilder::new(metrics).subset_time("load_time", partition), + join_time: MetricBuilder::new(metrics).subset_time("join_time", partition), + left_input_batches: MetricBuilder::new(metrics) + .counter("left_input_batches", partition), + right_input_batches: MetricBuilder::new(metrics) + .counter("right_input_batches", partition), + left_input_rows: MetricBuilder::new(metrics) + .counter("left_input_rows", partition), + right_input_rows: MetricBuilder::new(metrics) + .counter("right_input_rows", partition), + candidate_rows: MetricBuilder::new(metrics) + .counter("candidate_rows", partition), + output_batches: MetricBuilder::new(metrics) + .counter("output_batches", partition), + output_rows: MetricBuilder::new(metrics).counter("output_rows", partition), + peak_mem_used: MetricBuilder::new(metrics) + .peak_memory_usage("peak_mem_used", partition), + } + } +} + +pub(super) struct IEJoinStream { + schema: SchemaRef, + column_indices: Vec<ColumnIndex>, + conditions: [IEJoinCondition; 2], + on: JoinOn, + filter: Option<JoinFilter>, + join_type: JoinType, + null_equality: NullEquality, + batch_size: usize, + state: StreamState, + metrics: IEJoinMetrics, +} + +enum StreamState { + Loading(BoxFuture<'static, Result<IEJoinData>>), + Joining(Box<JoinState>), + Done, +} + +struct CurrentLeft { + row: u32, + boundary: usize, + scan_position: usize, +} + +struct JoinState { + data: IEJoinData, + first_comparator: DynComparator, + second_comparator: DynComparator, + active: ActiveBitmap, + active_memory: usize, + current_hash: Option<u64>, + current_group: Option<RightGroup>, + right_group_index: usize, + right_second_cursor: usize, + left_position: usize, + current_left: Option<CurrentLeft>, +} + +impl IEJoinStream { + #[expect(clippy::too_many_arguments)] + pub fn new( + schema: SchemaRef, + column_indices: Vec<ColumnIndex>, + conditions: [IEJoinCondition; 2], + on: JoinOn, + filter: Option<JoinFilter>, + join_type: JoinType, + null_equality: NullEquality, + batch_size: usize, + load: BoxFuture<'static, Result<IEJoinData>>, + metrics: IEJoinMetrics, + ) -> Self { + Self { + schema, + column_indices, + conditions, + on, + filter, + join_type, + null_equality, + batch_size: batch_size.max(1), + state: StreamState::Loading(load), + metrics, + } + } + + fn poll_next_impl( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + match &mut self.state { + StreamState::Loading(load) => { + let data = match ready!(load.poll_unpin(cx)) { + Ok(data) => data, + Err(error) => { + self.state = StreamState::Done; + return Poll::Ready(Some(Err(error))); + } + }; + match JoinState::try_new(data, &self.metrics.peak_mem_used) { + Ok(state) => self.state = StreamState::Joining(Box::new(state)), + Err(error) => { + self.state = StreamState::Done; + return Poll::Ready(Some(Err(error))); + } + } + } + StreamState::Joining(state) => { + let timer = self.metrics.join_time.timer(); + let result = state.next_batch( + &self.schema, + &self.column_indices, + &self.conditions, + &self.on, + self.filter.as_ref(), + self.join_type, + self.null_equality, + self.batch_size, + &self.metrics, + ); + timer.done(); + match result { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => { + self.state = StreamState::Done; + return Poll::Ready(None); + } + Err(error) => { + self.state = StreamState::Done; + return Poll::Ready(Some(Err(error))); + } + } + } + StreamState::Done => return Poll::Ready(None), + } + } + } +} + +impl Stream for IEJoinStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +impl RecordBatchStream for IEJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl JoinState { + fn try_new(data: IEJoinData, peak_mem_used: &Gauge) -> Result<Self> { + let first_comparator = make_comparator( + data.left_keys[0].as_ref(), + data.right_keys[0].as_ref(), + SortOptions::default(), + )?; + let second_comparator = make_comparator( + data.left_keys[1].as_ref(), + data.right_keys[1].as_ref(), + SortOptions::default(), + )?; + let max_group = data + .right_groups + .iter() + .map(|group| group.first.len()) + .max() + .unwrap_or(0); + let active = ActiveBitmap::new(max_group); + let active_memory = active.memory_size(); + data.reservation.try_grow(active_memory)?; + peak_mem_used.set_max(data.reservation.size()); + + Ok(Self { + data, + first_comparator, + second_comparator, + active, + active_memory, + current_hash: None, + current_group: None, + right_group_index: 0, + right_second_cursor: 0, + left_position: 0, + current_left: None, + }) + } + + #[expect(clippy::too_many_arguments)] + fn next_batch( + &mut self, + schema: &SchemaRef, + column_indices: &[ColumnIndex], + conditions: &[IEJoinCondition; 2], + on: &JoinOn, + filter: Option<&JoinFilter>, + join_type: JoinType, + null_equality: NullEquality, + batch_size: usize, + metrics: &IEJoinMetrics, + ) -> Result<Option<RecordBatch>> { + loop { + let mut left_indices = Vec::with_capacity(batch_size); + let mut right_indices = Vec::with_capacity(batch_size); + + while left_indices.len() < batch_size { + if self.current_left.is_none() && !self.start_next_left(conditions) { + break; + } + let current = self.current_left.as_mut().expect("set above"); + let Some(position) = self + .active + .next_set(current.scan_position, current.boundary) + else { + self.current_left = None; + continue; + }; + current.scan_position = position + 1; + let group = self + .current_group + .as_ref() + .expect("group set with left row"); + let right_row = self.data.right_by_first[group.first.start + position]; + left_indices.push(current.row as u64); + right_indices.push(right_row); + } + + if left_indices.is_empty() { + return Ok(None); + } + metrics.candidate_rows.add(left_indices.len()); + let mut left_indices = UInt64Array::from(left_indices); + let mut right_indices = UInt32Array::from(right_indices); + + if !on.is_empty() { + (left_indices, right_indices) = equal_rows_arr( + &left_indices, + &right_indices, + &self.data.left_equality_keys, + &self.data.right_equality_keys, + null_equality, + )?; + } + if let Some(filter) = filter { + (left_indices, right_indices) = apply_join_filter_to_indices( + &self.data.left_batch, + &self.data.right_batch, + left_indices, + right_indices, + filter, + JoinSide::Left, + Some(batch_size), + join_type, + )?; + } + if left_indices.is_empty() { Review Comment: I think the candidates here satisfy both range predicates, and interval-overlap joins can generate o(n^2) of them. When equal_rows_arr or filter is very selective, this continue re enters the loop and a single poll_next can scan the entire candidate space without yielding I think it can be fixed by bounding work by candidates examined per poll (not just rows emitted) -- 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]
