jayzhan211 commented on code in PR #23828: URL: https://github.com/apache/datafusion/pull/23828#discussion_r3791821850
########## datafusion/physical-plan/src/joins/asof_join.rs: ########## @@ -0,0 +1,1769 @@ +// 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. + +//! Broadcast, left-preserving ASOF join execution. +//! +//! An ASOF join emits exactly one output row for every left row. Within an +//! optional equality-key group, it selects the closest right row that satisfies +//! one ordered comparison. This follows Snowflake's [ASOF JOIN] semantics: +//! +//! ```text +//! left.ts >= right.ts => greatest eligible right.ts +//! left.ts <= right.ts => smallest eligible right.ts +//! ``` +//! +//! The right input is collected and shared by all output partitions. The left +//! input remains partitioned, and each partition performs an independent +//! monotonic scan over the ordered right input. +//! +//! [`AsOfJoinExec::input_distribution_requirements`] requires a single right +//! partition but leaves the left distribution unrestricted. +//! [`AsOfJoinExec::required_input_ordering`] requires both inputs to be ordered. +//! The physical optimizer satisfies these contracts by inserting operators such +//! as `RepartitionExec`, `SortExec`, `CoalescePartitionsExec`, or +//! `SortPreservingMergeExec`, depending on the input properties. The inserted +//! plan shape is therefore not fixed by this operator. +//! +//! Both inputs must be ordered by their equality keys followed by the match +//! key. For `<` and `<=`, the match ordering is reversed so all directions use +//! the same forward-only state machine. For example: +//! +//! ```text +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts >= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts ASC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts ASC NULLS FIRST] +//! +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts <= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts DESC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts DESC NULLS FIRST] +//! ``` +//! +//! Each left partition owns its cursors, equality-group state, and current +//! candidate, while the collected right batches are immutable and shared. +//! The key state-machine entry point is `AsOfJoinStreamState::next_batch`. +//! +//! This mode preserves probe-side parallelism when there are no equality keys +//! or when equality keys have low cardinality or skew. It retains the complete +//! right input in the memory pool and may scan it once per left partition. +//! Alternative strategies, including broadcasting the other side or +//! repartitioning both inputs, remain future work for other input-size and +//! key-distribution profiles. +//! +//! [ASOF JOIN]: https://docs.snowflake.com/en/sql-reference/constructs/asof-join + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}; +use arrow::buffer::NullBuffer; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; +use datafusion_common::utils::normalize_float_zero_scalar; +use datafusion_common::{ + ColumnStatistics, JoinSide, JoinType, NullEquality, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, project_schema, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::{ProjectionMapping, ProjectionRef}; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::joins::utils::{ + ColumnIndex, JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, + matchable_join_keys, +}; +use crate::memory::MemoryStream; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, + RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, +}; + +/// Physical ordered comparison for an ASOF join. +#[derive(Debug, Clone)] +pub struct AsOfMatchExpr { + /// Expression evaluated against the left input. + pub left: PhysicalExprRef, + /// Ordered comparison operator. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: PhysicalExprRef, +} + +impl AsOfMatchExpr { + /// Creates a physical ASOF match expression. + pub fn new(left: PhysicalExprRef, op: Operator, right: PhysicalExprRef) -> Self { + Self { left, op, right } + } +} + +/// A broadcast sort-merge ASOF join that emits one row for every left row. +#[derive(Debug)] +pub struct AsOfJoinExec { + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + match_condition: AsOfMatchExpr, + /// Unprojected left-join schema used to interpret `projection`. + join_schema: SchemaRef, + /// Information of index and left/right placement of columns. + column_indices: Vec<ColumnIndex>, + /// Optional indices into the full left-then-right join schema. + projection: Option<ProjectionRef>, + metrics: ExecutionPlanMetricsSet, + /// Required ordering for each left partition. + left_ordering: LexOrdering, + /// Required global ordering for the single right partition. + right_ordering: LexOrdering, + /// Shared collection future that materializes the right input only once. + right_fut: OnceAsync<BroadcastRightInput>, + cache: Arc<PlanProperties>, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + /// + /// The match operator must be `<`, `<=`, `>`, or `>=`. Equality and match + /// expressions must be deterministic, reference only their corresponding + /// input, and have matching input types. Equality types must support hashing; + /// floating-point equality keys are not supported because Arrow sorting + /// distinguishes signed zero while SQL equality does not. Projection indices + /// refer to the full left-then-right join schema. + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + match_condition: AsOfMatchExpr, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + validate_asof_join(left.as_ref(), right.as_ref(), &on, &match_condition)?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &JoinType::Left); + let join_schema = Arc::new(join_schema); + let projection: Option<ProjectionRef> = projection.map(Into::into); + let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); + let equality_options = SortOptions { + descending: false, + nulls_first: true, + }; + let match_options = SortOptions { + descending, + nulls_first: true, + }; + let mut left_sort_exprs = on + .iter() + .map(|(left, _)| PhysicalSortExpr { + expr: Arc::clone(left), + options: equality_options, + }) + .collect::<Vec<_>>(); + left_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.left), + options: match_options, + }); + let mut right_sort_exprs = on + .iter() + .map(|(_, right)| PhysicalSortExpr { + expr: Arc::clone(right), + options: equality_options, + }) + .collect::<Vec<_>>(); + right_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.right), + options: match_options, + }); + let left_ordering = LexOrdering::new(left_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF left ordering must not be empty" + ) + })?; + let right_ordering = LexOrdering::new(right_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF right ordering must not be empty" + ) + })?; + let cache = Arc::new(Self::compute_properties( + &left, + &join_schema, + projection.as_deref(), + )?); + + Ok(Self { + left, + right, + on, + match_condition, + join_schema, + column_indices, + projection, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + right_fut: Default::default(), + cache, + }) + } + + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + join_schema: &SchemaRef, + projection: Option<&[usize]>, + ) -> Result<PlanProperties> { + let left_schema = left.schema(); + let mapping = ProjectionMapping::try_new( + left_schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + ( + Arc::new(PhysicalColumn::new(field.name(), index)) + as PhysicalExprRef, + field.name().to_string(), + ) + }), + &left_schema, + )?; + let input_eq_properties = left.equivalence_properties(); + let mut eq_properties = + input_eq_properties.project(&mapping, Arc::clone(join_schema)); + let mut output_partitioning = left + .output_partitioning() + .project(&mapping, input_eq_properties); + if let Some(projection) = projection { + let projection_mapping = + ProjectionMapping::from_indices(projection, join_schema)?; + let output_schema = project_schema(join_schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, output_schema); + } + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + let on = self + .on + .iter() + .map(|(left, right)| { + format!("({} = {})", fmt_sql(left.as_ref()), fmt_sql(right.as_ref())) + }) + .collect::<Vec<_>>() + .join(", "); + let match_condition = format!( + "{} {} {}", + fmt_sql(self.match_condition.left.as_ref()), + self.match_condition.op, + fmt_sql(self.match_condition.right.as_ref()) + ); + let projection = self + .projection + .as_ref() + .map(|projection| { + format!( + ", projection=[{}]", + projection + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.field(*index).name(), + index + )) + .collect::<Vec<_>>() + .join(", ") + ) + }) + .unwrap_or_default(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]{}", + Self::static_name(), + on, + match_condition, + projection + ), + DisplayFormatType::TreeRender => { + writeln!(f, "on={on}")?; + writeln!(f, "match={match_condition}") + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + // Every left partition scans the complete broadcast right input, so + // equality keys do not require the inputs to be co-partitioned. + // `UnspecifiedDistribution` imposes no layout requirement; because this + // operator uses the default `benefits_from_input_partitioning`, the + // optimizer may still add round-robin repartitioning when it is useful. + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition, + ]) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![false, false] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + crate::apply_expression_roots( + join_keys.chain([&self.match_condition.left, &self.match_condition.right]), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + let left = children.swap_remove(0); + let right = children.swap_remove(0); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + left, + right, + on: self.on.clone(), + match_condition: self.match_condition.clone(), + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + metrics: ExecutionPlanMetricsSet::new(), + left_ordering: self.left_ordering.clone(), + right_ordering: self.right_ordering.clone(), + right_fut: Default::default(), + cache: Arc::clone(&self.cache), + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::try_new( + left, + right, + self.on.clone(), + self.match_condition.clone(), + self.projection.as_deref().map(<[usize]>::to_vec), + )?)), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + right_partitions, + 1, + "AsOfJoinExec requires one right partition, found {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let metrics = AsOfJoinMetrics::new(partition, &self.metrics); + let build_metrics = metrics.clone(); + let right_fut = self.right_fut.try_once(|| { + let right_stream = self.right.execute(0, Arc::clone(&context))?; + let reservation = + MemoryConsumer::new("AsOfJoinInput").register(context.memory_pool()); + Ok(collect_right_input( + right_stream, + reservation, + build_metrics, + )) + })?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let output_schema = self.schema(); + let stream_schema = Arc::clone(&output_schema); + let left_match = Arc::clone(&self.match_condition.left); + let right_match = Arc::clone(&self.match_condition.right); + let match_op = self.match_condition.op; + let column_indices = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let batch_size = context.session_config().batch_size(); + let stream = stream::once(async move { + let mut right_fut = right_fut; + let right_input = poll_fn(|cx| right_fut.get_shared(cx)).await?; + let right_stream = right_input.stream()?; + let state = AsOfJoinStreamState::new( + Arc::clone(&stream_schema), + InputCursor::new(left_stream, left_keys, left_match), + InputCursor::new(right_stream, right_keys, right_match), + match_op, + column_indices, + batch_size, + metrics, + ); + // `next_batch` is the key state-machine entry point. `try_unfold` + // preserves that state between emitted batches. + let stream = stream::try_unfold( + (state, right_input), + |(mut state, right_input)| async { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, (state, right_input)))), + None => Ok(None), + } + }, + ); + Ok::<SendableRecordBatchStream, datafusion_common::DataFusionError>(Box::pin( + RecordBatchStreamAdapter::new(stream_schema, stream), + )) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + output_schema, + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition), ChildStats::Skip] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + // The default is fully unknown, but ASOF emits exactly one output row + // per left row and preserves statistics for unmodified left columns. + let left = &input_stats[0]; + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let column_statistics = column_indices_after_projection + .iter() + .map(|column| match column.side { + JoinSide::Left => left + .column_statistics + .get(column.index) + .cloned() + .unwrap_or_else(ColumnStatistics::new_unknown), + JoinSide::Right | JoinSide::None => ColumnStatistics::new_unknown(), + }) + .collect(); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } +} + +/// Materialized right input shared by every left output partition. +struct BroadcastRightInput { + /// Schema retained even when the input has no batches. + schema: SchemaRef, + /// Ordered right batches; their buffers are shared without copying. + batches: Vec<RecordBatch>, + /// Holds the memory-pool reservation for as long as the batches are shared. + _reservation: MemoryReservation, +} + +impl BroadcastRightInput { + fn stream(&self) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(MemoryStream::try_new( + self.batches.clone(), + Arc::clone(&self.schema), + None, + )?)) + } +} + +async fn collect_right_input( + input: SendableRecordBatchStream, + reservation: MemoryReservation, + metrics: AsOfJoinMetrics, +) -> Result<BroadcastRightInput> { + let schema = input.schema(); + let mut memory_counter = RecordBatchMemoryCounter::new(); + let batches = input + .try_fold(Vec::new(), |mut batches, batch| { + let batch_size = memory_counter.count_batch(&batch); + futures::future::ready(reservation.try_grow(batch_size).map(|_| { + metrics.build_mem_used.add(batch_size); + batches.push(batch); + batches + })) + }) + .await?; + Ok(BroadcastRightInput { + schema, + batches, + _reservation: reservation, + }) +} + +/// Last eligible right row for the current left equality group. +/// +/// The row and its evaluated keys survive right batch changes and output +/// flushes. It belongs to the join state rather than `InputCursor` because its +/// validity also depends on the current left equality group. +#[derive(Clone)] +struct Candidate { + /// Right batch containing the nearest eligible row. + batch: Arc<RecordBatch>, + /// Row index within `batch`. + row: usize, + /// Evaluated equality keys retained when the right cursor changes batches. + key_arrays: Arc<[ArrayRef]>, + /// Identity used to invalidate the cached candidate/left comparator. + key_batch_id: usize, +} + +/// Cursor over one ordered input stream. +/// +/// Expressions are evaluated once per non-empty batch. `key_batch_id` changes +/// whenever a new batch is loaded so comparators cannot retain stale arrays. +struct InputCursor { + /// Remaining input batches. + stream: SendableRecordBatchStream, + /// Equality expressions evaluated for each batch. + key_exprs: Vec<PhysicalExprRef>, + /// Ordered match expression evaluated for each batch. + match_expr: PhysicalExprRef, + /// Current non-empty batch. + batch: Option<Arc<RecordBatch>>, + /// Evaluated equality-key arrays for `batch`. + key_arrays: Arc<[ArrayRef]>, + /// Rows whose equality keys are all non-NULL. + key_validity: Option<NullBuffer>, + /// Evaluated match values for `batch`. + match_array: Option<ArrayRef>, + /// Monotonic identity of the current key arrays. + key_batch_id: usize, + /// Current row within `batch`. + row: usize, + /// Whether the input stream has returned EOF. + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec<PhysicalExprRef>, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: Arc::from([]), + key_validity: None, + match_array: None, + key_batch_id: 0, + row: 0, + eof: false, + } + } + + async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result<bool> { + loop { + if let Some(batch) = &self.batch + && self.row < batch.num_rows() + { + return Ok(true); + } + self.batch = None; + self.key_arrays = Arc::from([]); + self.key_validity = None; + self.match_array = None; + self.row = 0; + if self.eof { + return Ok(false); + } + let Some(batch) = self.stream.next().await.transpose()? else { + self.eof = true; + return Ok(false); + }; + if batch.num_rows() == 0 { + continue; + } + let batch = Arc::new(batch); + let _timer = elapsed_compute.timer(); + let key_arrays = self + .key_exprs + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::<Result<Vec<_>>>()?; + self.key_validity = + matchable_join_keys(&key_arrays, NullEquality::NullEqualsNothing); + self.key_arrays = key_arrays.into(); + self.match_array = Some( + self.match_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?, + ); + self.key_batch_id += 1; + self.batch = Some(batch); + } + } + + fn group_has_null(&self) -> bool { + self.key_validity + .as_ref() + .is_some_and(|validity| validity.is_null(self.row)) + } + + fn match_value(&self) -> Result<ScalarValue> { + let array = self.match_array.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF match array is missing") + })?; + ScalarValue::try_from_array(array, self.row).map(normalize_float_zero_scalar) + } + + fn batch_row(&self) -> Result<(Arc<RecordBatch>, usize)> { + let batch = self.batch.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF input batch is missing") + })?; + Ok((Arc::clone(batch), self.row)) + } + + fn advance(&mut self) { + self.row += 1; + } +} + +#[derive(Clone)] +struct AsOfJoinMetrics { + /// Standard output-row and elapsed-compute metrics. + baseline: BaselineMetrics, + /// Peak bytes retained for the shared right input. + /// + /// `peak_memory_usage` records this as `MetricValue::PeakMemoryUsage`; `Gauge` + /// is the handle used to update that metric. + build_mem_used: Gauge, +} + +impl AsOfJoinMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + build_mem_used: MetricBuilder::new(metrics) + .peak_memory_usage("build_mem_used", partition), + } + } +} + +/// Row references accumulated for the next output batch. +/// +/// For right output, `None` represents NULL padding for an unmatched left row. +/// For example, indices `[Some((0, 2)), None, Some((1, 0))]` select row 2 from +/// the first source batch, a NULL, and row 0 from the second source batch. +#[derive(Default)] +struct PendingRows { + /// Distinct source batches referenced by `indices`. + sources: Vec<Arc<RecordBatch>>, Review Comment: Do we need `Vec<Arc<RecordBatch>>` and not just `Vec<RecordBatch>` ########## datafusion/physical-plan/src/joins/asof_join.rs: ########## @@ -0,0 +1,1769 @@ +// 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. + +//! Broadcast, left-preserving ASOF join execution. +//! +//! An ASOF join emits exactly one output row for every left row. Within an +//! optional equality-key group, it selects the closest right row that satisfies +//! one ordered comparison. This follows Snowflake's [ASOF JOIN] semantics: +//! +//! ```text +//! left.ts >= right.ts => greatest eligible right.ts +//! left.ts <= right.ts => smallest eligible right.ts +//! ``` +//! +//! The right input is collected and shared by all output partitions. The left +//! input remains partitioned, and each partition performs an independent +//! monotonic scan over the ordered right input. +//! +//! [`AsOfJoinExec::input_distribution_requirements`] requires a single right +//! partition but leaves the left distribution unrestricted. +//! [`AsOfJoinExec::required_input_ordering`] requires both inputs to be ordered. +//! The physical optimizer satisfies these contracts by inserting operators such +//! as `RepartitionExec`, `SortExec`, `CoalescePartitionsExec`, or +//! `SortPreservingMergeExec`, depending on the input properties. The inserted +//! plan shape is therefore not fixed by this operator. +//! +//! Both inputs must be ordered by their equality keys followed by the match +//! key. For `<` and `<=`, the match ordering is reversed so all directions use +//! the same forward-only state machine. For example: +//! +//! ```text +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts >= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts ASC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts ASC NULLS FIRST] +//! +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts <= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts DESC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts DESC NULLS FIRST] +//! ``` +//! +//! Each left partition owns its cursors, equality-group state, and current +//! candidate, while the collected right batches are immutable and shared. +//! The key state-machine entry point is `AsOfJoinStreamState::next_batch`. +//! +//! This mode preserves probe-side parallelism when there are no equality keys +//! or when equality keys have low cardinality or skew. It retains the complete +//! right input in the memory pool and may scan it once per left partition. +//! Alternative strategies, including broadcasting the other side or +//! repartitioning both inputs, remain future work for other input-size and +//! key-distribution profiles. +//! +//! [ASOF JOIN]: https://docs.snowflake.com/en/sql-reference/constructs/asof-join + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}; +use arrow::buffer::NullBuffer; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; +use datafusion_common::utils::normalize_float_zero_scalar; +use datafusion_common::{ + ColumnStatistics, JoinSide, JoinType, NullEquality, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, project_schema, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::{ProjectionMapping, ProjectionRef}; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::joins::utils::{ + ColumnIndex, JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, + matchable_join_keys, +}; +use crate::memory::MemoryStream; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, + RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, +}; + +/// Physical ordered comparison for an ASOF join. +#[derive(Debug, Clone)] +pub struct AsOfMatchExpr { + /// Expression evaluated against the left input. + pub left: PhysicalExprRef, + /// Ordered comparison operator. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: PhysicalExprRef, +} + +impl AsOfMatchExpr { + /// Creates a physical ASOF match expression. + pub fn new(left: PhysicalExprRef, op: Operator, right: PhysicalExprRef) -> Self { + Self { left, op, right } + } +} + +/// A broadcast sort-merge ASOF join that emits one row for every left row. +#[derive(Debug)] +pub struct AsOfJoinExec { + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + match_condition: AsOfMatchExpr, + /// Unprojected left-join schema used to interpret `projection`. + join_schema: SchemaRef, + /// Information of index and left/right placement of columns. + column_indices: Vec<ColumnIndex>, + /// Optional indices into the full left-then-right join schema. + projection: Option<ProjectionRef>, + metrics: ExecutionPlanMetricsSet, + /// Required ordering for each left partition. + left_ordering: LexOrdering, + /// Required global ordering for the single right partition. + right_ordering: LexOrdering, + /// Shared collection future that materializes the right input only once. + right_fut: OnceAsync<BroadcastRightInput>, + cache: Arc<PlanProperties>, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + /// + /// The match operator must be `<`, `<=`, `>`, or `>=`. Equality and match + /// expressions must be deterministic, reference only their corresponding + /// input, and have matching input types. Equality types must support hashing; + /// floating-point equality keys are not supported because Arrow sorting + /// distinguishes signed zero while SQL equality does not. Projection indices + /// refer to the full left-then-right join schema. + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + match_condition: AsOfMatchExpr, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + validate_asof_join(left.as_ref(), right.as_ref(), &on, &match_condition)?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &JoinType::Left); + let join_schema = Arc::new(join_schema); + let projection: Option<ProjectionRef> = projection.map(Into::into); + let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); + let equality_options = SortOptions { + descending: false, + nulls_first: true, + }; + let match_options = SortOptions { + descending, + nulls_first: true, + }; + let mut left_sort_exprs = on + .iter() + .map(|(left, _)| PhysicalSortExpr { + expr: Arc::clone(left), + options: equality_options, + }) + .collect::<Vec<_>>(); + left_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.left), + options: match_options, + }); + let mut right_sort_exprs = on + .iter() + .map(|(_, right)| PhysicalSortExpr { + expr: Arc::clone(right), + options: equality_options, + }) + .collect::<Vec<_>>(); + right_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.right), + options: match_options, + }); + let left_ordering = LexOrdering::new(left_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF left ordering must not be empty" + ) + })?; + let right_ordering = LexOrdering::new(right_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF right ordering must not be empty" + ) + })?; + let cache = Arc::new(Self::compute_properties( + &left, + &join_schema, + projection.as_deref(), + )?); + + Ok(Self { + left, + right, + on, + match_condition, + join_schema, + column_indices, + projection, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + right_fut: Default::default(), + cache, + }) + } + + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + join_schema: &SchemaRef, + projection: Option<&[usize]>, + ) -> Result<PlanProperties> { + let left_schema = left.schema(); + let mapping = ProjectionMapping::try_new( + left_schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + ( + Arc::new(PhysicalColumn::new(field.name(), index)) + as PhysicalExprRef, + field.name().to_string(), + ) + }), + &left_schema, + )?; + let input_eq_properties = left.equivalence_properties(); + let mut eq_properties = + input_eq_properties.project(&mapping, Arc::clone(join_schema)); + let mut output_partitioning = left + .output_partitioning() + .project(&mapping, input_eq_properties); + if let Some(projection) = projection { + let projection_mapping = + ProjectionMapping::from_indices(projection, join_schema)?; + let output_schema = project_schema(join_schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, output_schema); + } + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + let on = self + .on + .iter() + .map(|(left, right)| { + format!("({} = {})", fmt_sql(left.as_ref()), fmt_sql(right.as_ref())) + }) + .collect::<Vec<_>>() + .join(", "); + let match_condition = format!( + "{} {} {}", + fmt_sql(self.match_condition.left.as_ref()), + self.match_condition.op, + fmt_sql(self.match_condition.right.as_ref()) + ); + let projection = self + .projection + .as_ref() + .map(|projection| { + format!( + ", projection=[{}]", + projection + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.field(*index).name(), + index + )) + .collect::<Vec<_>>() + .join(", ") + ) + }) + .unwrap_or_default(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]{}", + Self::static_name(), + on, + match_condition, + projection + ), + DisplayFormatType::TreeRender => { + writeln!(f, "on={on}")?; + writeln!(f, "match={match_condition}") + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + // Every left partition scans the complete broadcast right input, so + // equality keys do not require the inputs to be co-partitioned. + // `UnspecifiedDistribution` imposes no layout requirement; because this + // operator uses the default `benefits_from_input_partitioning`, the + // optimizer may still add round-robin repartitioning when it is useful. + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition, + ]) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![false, false] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + crate::apply_expression_roots( + join_keys.chain([&self.match_condition.left, &self.match_condition.right]), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + let left = children.swap_remove(0); + let right = children.swap_remove(0); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + left, + right, + on: self.on.clone(), + match_condition: self.match_condition.clone(), + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + metrics: ExecutionPlanMetricsSet::new(), + left_ordering: self.left_ordering.clone(), + right_ordering: self.right_ordering.clone(), + right_fut: Default::default(), + cache: Arc::clone(&self.cache), + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::try_new( + left, + right, + self.on.clone(), + self.match_condition.clone(), + self.projection.as_deref().map(<[usize]>::to_vec), + )?)), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + right_partitions, + 1, + "AsOfJoinExec requires one right partition, found {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let metrics = AsOfJoinMetrics::new(partition, &self.metrics); + let build_metrics = metrics.clone(); + let right_fut = self.right_fut.try_once(|| { + let right_stream = self.right.execute(0, Arc::clone(&context))?; + let reservation = + MemoryConsumer::new("AsOfJoinInput").register(context.memory_pool()); + Ok(collect_right_input( + right_stream, + reservation, + build_metrics, + )) + })?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let output_schema = self.schema(); + let stream_schema = Arc::clone(&output_schema); + let left_match = Arc::clone(&self.match_condition.left); + let right_match = Arc::clone(&self.match_condition.right); + let match_op = self.match_condition.op; + let column_indices = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let batch_size = context.session_config().batch_size(); + let stream = stream::once(async move { + let mut right_fut = right_fut; + let right_input = poll_fn(|cx| right_fut.get_shared(cx)).await?; + let right_stream = right_input.stream()?; + let state = AsOfJoinStreamState::new( + Arc::clone(&stream_schema), + InputCursor::new(left_stream, left_keys, left_match), + InputCursor::new(right_stream, right_keys, right_match), + match_op, + column_indices, + batch_size, + metrics, + ); + // `next_batch` is the key state-machine entry point. `try_unfold` + // preserves that state between emitted batches. + let stream = stream::try_unfold( + (state, right_input), + |(mut state, right_input)| async { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, (state, right_input)))), + None => Ok(None), + } + }, + ); + Ok::<SendableRecordBatchStream, datafusion_common::DataFusionError>(Box::pin( + RecordBatchStreamAdapter::new(stream_schema, stream), + )) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + output_schema, + stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition), ChildStats::Skip] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + // The default is fully unknown, but ASOF emits exactly one output row + // per left row and preserves statistics for unmodified left columns. + let left = &input_stats[0]; + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let column_statistics = column_indices_after_projection + .iter() + .map(|column| match column.side { + JoinSide::Left => left + .column_statistics + .get(column.index) + .cloned() + .unwrap_or_else(ColumnStatistics::new_unknown), + JoinSide::Right | JoinSide::None => ColumnStatistics::new_unknown(), + }) + .collect(); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } +} + +/// Materialized right input shared by every left output partition. +struct BroadcastRightInput { + /// Schema retained even when the input has no batches. + schema: SchemaRef, + /// Ordered right batches; their buffers are shared without copying. + batches: Vec<RecordBatch>, + /// Holds the memory-pool reservation for as long as the batches are shared. + _reservation: MemoryReservation, +} + +impl BroadcastRightInput { + fn stream(&self) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(MemoryStream::try_new( + self.batches.clone(), + Arc::clone(&self.schema), + None, + )?)) + } +} + +async fn collect_right_input( + input: SendableRecordBatchStream, + reservation: MemoryReservation, + metrics: AsOfJoinMetrics, +) -> Result<BroadcastRightInput> { + let schema = input.schema(); + let mut memory_counter = RecordBatchMemoryCounter::new(); + let batches = input + .try_fold(Vec::new(), |mut batches, batch| { + let batch_size = memory_counter.count_batch(&batch); + futures::future::ready(reservation.try_grow(batch_size).map(|_| { + metrics.build_mem_used.add(batch_size); + batches.push(batch); + batches + })) + }) + .await?; + Ok(BroadcastRightInput { + schema, + batches, + _reservation: reservation, + }) +} + +/// Last eligible right row for the current left equality group. +/// +/// The row and its evaluated keys survive right batch changes and output +/// flushes. It belongs to the join state rather than `InputCursor` because its +/// validity also depends on the current left equality group. +#[derive(Clone)] +struct Candidate { + /// Right batch containing the nearest eligible row. + batch: Arc<RecordBatch>, + /// Row index within `batch`. + row: usize, + /// Evaluated equality keys retained when the right cursor changes batches. + key_arrays: Arc<[ArrayRef]>, + /// Identity used to invalidate the cached candidate/left comparator. + key_batch_id: usize, +} + +/// Cursor over one ordered input stream. +/// +/// Expressions are evaluated once per non-empty batch. `key_batch_id` changes +/// whenever a new batch is loaded so comparators cannot retain stale arrays. +struct InputCursor { + /// Remaining input batches. + stream: SendableRecordBatchStream, + /// Equality expressions evaluated for each batch. + key_exprs: Vec<PhysicalExprRef>, + /// Ordered match expression evaluated for each batch. + match_expr: PhysicalExprRef, + /// Current non-empty batch. + batch: Option<Arc<RecordBatch>>, + /// Evaluated equality-key arrays for `batch`. + key_arrays: Arc<[ArrayRef]>, + /// Rows whose equality keys are all non-NULL. + key_validity: Option<NullBuffer>, + /// Evaluated match values for `batch`. + match_array: Option<ArrayRef>, + /// Monotonic identity of the current key arrays. + key_batch_id: usize, + /// Current row within `batch`. + row: usize, + /// Whether the input stream has returned EOF. + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec<PhysicalExprRef>, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: Arc::from([]), + key_validity: None, + match_array: None, + key_batch_id: 0, + row: 0, + eof: false, + } + } + + async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result<bool> { Review Comment: It seems similar to `poll_next` method from Stream what we have for other join like HashJoinStream Could we have Stream for AsofJoinStream as well, or is there any reason that current implementation is better for AsofJoin? -- 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]
