From e546de5510de599c91f756bef4bf45243d93011a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:05:01 +0800 Subject: [PATCH 01/10] feat: add ASOF join physical operator --- .../physical-plan/src/joins/asof_join.rs | 1419 +++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 2 + 2 files changed, 1421 insertions(+) create mode 100644 datafusion/physical-plan/src/joins/asof_join.rs diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs new file mode 100644 index 0000000000000..5945c8d750716 --- /dev/null +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -0,0 +1,1419 @@ +// 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. + +//! Ordered, left-preserving ASOF join execution. + +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; +use datafusion_common::utils::{ + compare_rows, get_row_at_idx, normalize_float_zero_scalar, +}; +use datafusion_common::{ + ColumnStatistics, JoinType, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{Partitioning, PhysicalSortExpr}; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{StreamExt, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::joins::utils::{JoinOn, build_join_schema}; +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricsSet, RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, + check_if_same_properties, +}; + +/// 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 sort-merge ASOF join that emits exactly one row for every left row. +#[derive(Debug, Clone)] +pub struct AsOfJoinExec { + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + schema: SchemaRef, + metrics: ExecutionPlanMetricsSet, + left_ordering: LexOrdering, + right_ordering: LexOrdering, + cache: Arc, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + pub fn try_new( + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + ) -> Result { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "AsOfJoinExec requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + if left.boundedness().is_unbounded() || right.boundedness().is_unbounded() { + return plan_err!("AsOfJoinExec requires bounded inputs"); + } + if is_volatile(&match_condition.left) || is_volatile(&match_condition.right) { + return plan_err!("AsOfJoinExec match expression must be deterministic"); + } + if on + .iter() + .any(|(left, right)| is_volatile(left) || is_volatile(right)) + { + return plan_err!("AsOfJoinExec equality expressions must be deterministic"); + } + + let left_schema = left.schema(); + let right_schema = right.schema(); + validate_expr_side(&match_condition.left, &left_schema, "left match")?; + validate_expr_side(&match_condition.right, &right_schema, "right match")?; + for (left_expr, right_expr) in &on { + validate_expr_side(left_expr, &left_schema, "left equality")?; + validate_expr_side(right_expr, &right_schema, "right equality")?; + let left_type = left_expr.data_type(&left_schema)?; + let right_type = right_expr.data_type(&right_schema)?; + if left_type != right_type { + return plan_err!( + "AsOfJoinExec equality expression types differ: {left_type} and {right_type}" + ); + } + if !datafusion_expr::utils::can_hash(&left_type) { + return plan_err!( + "AsOfJoinExec equality expressions have unsupported hash type {left_type}" + ); + } + } + let left_match_type = match_condition.left.data_type(&left_schema)?; + let right_match_type = match_condition.right.data_type(&right_schema)?; + if left_match_type != right_match_type { + return plan_err!( + "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" + ); + } + if let Some(index) = right_output_indices + .iter() + .find(|index| **index >= right_schema.fields().len()) + { + return plan_err!( + "AsOfJoinExec right output index {index} is outside schema with {} fields", + right_schema.fields().len() + ); + } + if !right_output_indices + .windows(2) + .all(|pair| pair[0] < pair[1]) + { + return plan_err!( + "AsOfJoinExec right output indices must be strictly increasing" + ); + } + + let schema = + build_output_schema(&left_schema, &right_schema, &right_output_indices); + 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::>(); + 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::>(); + 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, &schema, on.is_empty())?); + + Ok(Self { + left, + right, + on, + match_condition, + right_output_indices, + schema, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + cache, + }) + } + + fn compute_properties( + left: &Arc, + schema: &SchemaRef, + single_partition: bool, + ) -> Result { + 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 eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); + let output_partitioning = if single_partition { + Partitioning::UnknownPartitioning(1) + } else { + left.output_partitioning() + .project(&mapping, input_eq_properties) + }; + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + + /// Equality expressions. + pub fn on(&self) -> &JoinOn { + &self.on + } + + /// Ordered match expression. + pub fn match_condition(&self) -> &AsOfMatchExpr { + &self.match_condition + } + + /// Indices of right input columns emitted after the left columns. + pub fn right_output_indices(&self) -> &[usize] { + &self.right_output_indices + } + + /// Left input. + pub fn left(&self) -> &Arc { + &self.left + } + + /// Right input. + pub fn right(&self) -> &Arc { + &self.right + } +} + +fn build_output_schema( + left: &SchemaRef, + right: &SchemaRef, + right_output_indices: &[usize], +) -> SchemaRef { + let full_schema = build_join_schema(left, right, &JoinType::Left).0; + let left_len = left.fields().len(); + let fields = full_schema + .fields() + .iter() + .take(left_len) + .cloned() + .chain( + right_output_indices + .iter() + .map(|index| Arc::clone(&full_schema.fields()[left_len + *index])), + ) + .collect::>(); + Arc::new(Schema::new_with_metadata( + fields, + full_schema.metadata().clone(), + )) +} + +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::>() + .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()) + ); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]", + Self::static_name(), + on, + match_condition + ), + 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 { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.on.is_empty() { + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) + } else { + let (left, right) = self + .on + .iter() + .map(|(left, right)| (Arc::clone(left), Arc::clone(right))) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left), + Distribution::KeyPartitioned(right), + ]) + } + } + + fn required_input_ordering(&self) -> Vec> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec { + vec![true, false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + check_if_same_properties!(self, children); + match &children[..] { + [left, right] => Ok(Arc::new(Self::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.match_condition.clone(), + self.right_output_indices.clone(), + )?)), + _ => internal_err!("AsOfJoinExec requires two children"), + } + } + + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + assert_eq_or_internal_err!( + children.len(), + 2, + "AsOfJoinExec requires two children" + ); + let left = children.remove(0); + let right = children.remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&self) + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + left_partitions, + right_partitions, + "AsOfJoinExec partition count mismatch: {left_partitions} != {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let right_stream = self.right.execute(partition, Arc::clone(&context))?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let state = AsOfJoinStreamState::new( + Arc::clone(&self.schema), + InputCursor::new( + left_stream, + left_keys, + Arc::clone(&self.match_condition.left), + ), + InputCursor::new( + right_stream, + right_keys, + Arc::clone(&self.match_condition.right), + ), + self.match_condition.op, + self.right_output_indices.clone(), + context.session_config().batch_size(), + AsOfJoinMetrics::new(partition, &self.metrics), + ); + let stream = stream::try_unfold(state, |mut state| async move { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, state))), + None => Ok(None), + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition), ChildStats::Skip] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let left = &input_stats[0]; + let mut column_statistics = left.column_statistics.clone(); + column_statistics.truncate(self.left.schema().fields().len()); + column_statistics.resize_with( + self.left.schema().fields().len(), + ColumnStatistics::new_unknown, + ); + column_statistics.extend( + self.right_output_indices + .iter() + .map(|_| ColumnStatistics::new_unknown()), + ); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec, + _config: &ConfigOptions, + ) -> Result { + let left_indices = (0..self.left.schema().fields().len()).collect::>(); + let left = ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + left_indices, + &self.left, + )?; + let right = ChildFilterDescription::all_unsupported(&parent_filters); + Ok(FilterDescription::new().with_child(left).with_child(right)) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result>> { + Ok(FilterPushdownPropagation::if_any(child_pushdown_result)) + } +} + +#[derive(Clone)] +struct Candidate { + batch: Arc, + row: usize, + group: Vec, +} + +struct InputCursor { + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + batch: Option>, + key_arrays: Vec, + match_array: Option, + row: usize, + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: vec![], + match_array: None, + row: 0, + eof: false, + } + } + + async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result { + loop { + if let Some(batch) = &self.batch + && self.row < batch.num_rows() + { + return Ok(true); + } + self.batch = None; + self.key_arrays.clear(); + 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(); + self.key_arrays = self + .key_exprs + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::>()?; + self.match_array = Some( + self.match_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?, + ); + self.batch = Some(batch); + } + } + + fn group(&self) -> Result> { + get_row_at_idx(&self.key_arrays, self.row) + .map(|row| row.into_iter().map(normalize_float_zero_scalar).collect()) + } + + fn match_value(&self) -> Result { + 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, 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; + } +} + +struct AsOfJoinMetrics { + baseline: BaselineMetrics, + matched_rows: Count, + unmatched_left_rows: Count, +} + +impl AsOfJoinMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + matched_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("matched_rows", partition), + unmatched_left_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("unmatched_left_rows", partition), + } + } +} + +#[derive(Default)] +struct PendingRows { + sources: Vec>, + source_by_ptr: HashMap, + indices: Vec>, +} + +impl PendingRows { + fn len(&self) -> usize { + self.indices.len() + } + + fn is_empty(&self) -> bool { + self.indices.is_empty() + } + + fn push(&mut self, batch: Arc, row: usize) { + let ptr = Arc::as_ptr(&batch) as usize; + let source = *self.source_by_ptr.entry(ptr).or_insert_with(|| { + let source = self.sources.len(); + self.sources.push(batch); + source + }); + self.indices.push(Some((source, row))); + } + + fn push_null(&mut self) { + self.indices.push(None); + } + + fn materialize_column( + &self, + source_column: usize, + data_type: &arrow::datatypes::DataType, + ) -> Result { + if self.indices.is_empty() { + return internal_err!("ASOF output materialization has no pending rows"); + } + + if self.sources.len() == 1 + && self.indices.iter().all(Option::is_some) + && let Some((0, first_row)) = self.indices[0] + && self + .indices + .iter() + .enumerate() + .all(|(offset, index)| *index == Some((0, first_row + offset))) + { + return Ok(self.sources[0] + .column(source_column) + .slice(first_row, self.indices.len())); + } + + let has_null = self.indices.iter().any(Option::is_none); + let null_array = has_null.then(|| new_null_array(data_type, 1)); + let mut source_arrays: Vec<&dyn Array> = + Vec::with_capacity(self.sources.len() + usize::from(has_null)); + if let Some(null_array) = &null_array { + source_arrays.push(null_array.as_ref()); + } + source_arrays.extend( + self.sources + .iter() + .map(|batch| batch.column(source_column).as_ref()), + ); + let source_offset = usize::from(has_null); + let interleave_indices = self + .indices + .iter() + .map(|index| match index { + Some((source, row)) => (source + source_offset, *row), + None => (0, 0), + }) + .collect::>(); + interleave(&source_arrays, &interleave_indices).map_err(Into::into) + } + + fn clear(&mut self) { + self.sources.clear(); + self.source_by_ptr.clear(); + self.indices.clear(); + } +} + +struct AsOfJoinStreamState { + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + candidate: Option, + group_sort_options: Vec, + pending_left: PendingRows, + pending_right: PendingRows, + batch_size: usize, + metrics: AsOfJoinMetrics, +} + +impl AsOfJoinStreamState { + fn new( + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + batch_size: usize, + metrics: AsOfJoinMetrics, + ) -> Self { + let group_sort_options = vec![ + SortOptions { + descending: false, + nulls_first: true, + }; + left.key_exprs.len() + ]; + Self { + pending_left: PendingRows::default(), + pending_right: PendingRows::default(), + schema, + left, + right, + op, + right_output_indices, + candidate: None, + group_sort_options, + batch_size: batch_size.max(1), + metrics, + } + } + + async fn next_batch(&mut self) -> Result> { + loop { + if self.pending_left.len() >= self.batch_size { + return self.flush().map(Some); + } + if !self + .left + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + if !self.pending_left.is_empty() { + return self.flush().map(Some); + } + self.metrics.baseline.done(); + return Ok(None); + } + + let (left_group, left_match) = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + (self.left.group()?, self.left.match_value()?) + }; + if left_match.is_null() || left_group.iter().any(ScalarValue::is_null) { + self.candidate = None; + self.push_current_left(None)?; + self.left.advance(); + continue; + } + let candidate_is_other_group = if let Some(candidate) = &self.candidate { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + compare_rows(&candidate.group, &left_group, &self.group_sort_options)? + != Ordering::Equal + } else { + false + }; + if candidate_is_other_group { + self.candidate = None; + } + + loop { + if !self + .right + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + break; + } + let action = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_group = self.right.group()?; + if right_group.iter().any(ScalarValue::is_null) { + RightAction::Advance + } else { + match compare_rows( + &right_group, + &left_group, + &self.group_sort_options, + )? { + Ordering::Less => RightAction::Advance, + Ordering::Greater => RightAction::Stop, + Ordering::Equal => { + let right_match = self.right.match_value()?; + if right_match.is_null() { + RightAction::Advance + } else if is_eligible(self.op, &left_match, &right_match)? + { + let (batch, row) = self.right.batch_row()?; + RightAction::Candidate(Candidate { + batch, + row, + group: right_group, + }) + } else { + RightAction::Stop + } + } + } + } + }; + match action { + RightAction::Advance => self.right.advance(), + RightAction::Candidate(candidate) => { + self.candidate = Some(candidate); + self.right.advance(); + } + RightAction::Stop => break, + } + } + + self.push_current_left(self.candidate.clone())?; + self.left.advance(); + } + } + + fn push_current_left(&mut self, candidate: Option) -> Result<()> { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let (left_batch, left_row) = self.left.batch_row()?; + self.pending_left.push(left_batch, left_row); + match candidate { + Some(candidate) => { + if !self.right_output_indices.is_empty() { + self.pending_right.push(candidate.batch, candidate.row); + } + self.metrics.matched_rows.add(1); + } + None => { + if !self.right_output_indices.is_empty() { + self.pending_right.push_null(); + } + self.metrics.unmatched_left_rows.add(1); + } + } + Ok(()) + } + + fn flush(&mut self) -> Result { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let left_len = self.schema.fields().len() - self.right_output_indices.len(); + let mut arrays = Vec::with_capacity(self.schema.fields().len()); + for index in 0..left_len { + arrays.push( + self.pending_left + .materialize_column(index, self.schema.field(index).data_type())?, + ); + } + for (offset, source_index) in self.right_output_indices.iter().enumerate() { + arrays.push(self.pending_right.materialize_column( + *source_index, + self.schema.field(left_len + offset).data_type(), + )?); + } + self.pending_left.clear(); + self.pending_right.clear(); + let batch = RecordBatch::try_new(Arc::clone(&self.schema), arrays)?; + (&batch).record_output(&self.metrics.baseline); + Ok(batch) + } +} + +fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Result<()> { + let columns = collect_columns(expr); + if columns.is_empty() { + return plan_err!("AsOfJoinExec {name} expression must reference its input"); + } + if let Some(column) = columns.iter().find(|column| { + schema + .fields() + .get(column.index()) + .is_none_or(|field| field.name() != column.name()) + }) { + return plan_err!( + "AsOfJoinExec {name} expression references column {column} outside its input" + ); + } + Ok(()) +} + +enum RightAction { + Advance, + Candidate(Candidate), + Stop, +} + +fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result { + let ordering = right.try_cmp(left)?; + Ok(match op { + Operator::Gt => ordering == Ordering::Less, + Operator::GtEq => ordering != Ordering::Greater, + Operator::Lt => ordering == Ordering::Greater, + Operator::LtEq => ordering != Ordering::Less, + _ => false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::test::TestMemoryExec; + use arrow::array::{ + DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Int8Type}; + use datafusion_execution::config::SessionConfig; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr_common::metrics::MetricValue; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + struct VolatileExpr; + + impl std::fmt::Display for VolatileExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile") + } + } + + impl PhysicalExpr for VolatileExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Int64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, _batch: &RecordBatch) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn is_volatile_node(&self) -> bool { + true + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile()") + } + } + + fn make_batch( + schema: &SchemaRef, + keys: Vec>, + times: Vec>, + values: Vec, + ) -> Result { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from(keys)), + Arc::new(Int64Array::from(times)), + Arc::new(Int32Array::from(values)), + ], + ) + .map_err(Into::into) + } + + fn test_exec() -> Result> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let left_batches = vec![ + RecordBatch::new_empty(Arc::clone(&left_schema)), + make_batch(&left_schema, vec![None], vec![Some(3)], vec![0])?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![None, Some(1)], + vec![1, 2], + )?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(4), Some(7)], + vec![3, 4], + )?, + make_batch( + &left_schema, + vec![Some("B"), Some("C")], + vec![Some(2), Some(3)], + vec![5, 6], + )?, + ]; + let left = TestMemoryExec::try_new_exec( + &[left_batches], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let right_batches = vec![ + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![None, Some("A")], + vec![Some(2), None], + vec![999, 777], + )?, + make_batch(&right_schema, vec![Some("A")], vec![Some(2)], vec![20])?, + make_batch(&right_schema, vec![Some("A")], vec![Some(4)], vec![40])?, + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![Some("A"), Some("B")], + vec![Some(6), Some(1)], + vec![60, 101], + )?, + ]; + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let on: JoinOn = vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )]; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?)) + } + + #[test] + fn eligibility_matches_public_semantics() -> Result<()> { + let left = ScalarValue::Int64(Some(10)); + let lower = ScalarValue::Int64(Some(9)); + let equal = ScalarValue::Int64(Some(10)); + let higher = ScalarValue::Int64(Some(11)); + assert!(is_eligible(Operator::Gt, &left, &lower)?); + assert!(!is_eligible(Operator::Gt, &left, &equal)?); + assert!(is_eligible(Operator::GtEq, &left, &equal)?); + assert!(is_eligible(Operator::Lt, &left, &higher)?); + assert!(!is_eligible(Operator::Lt, &left, &equal)?); + assert!(is_eligible(Operator::LtEq, &left, &equal)?); + Ok(()) + } + + #[tokio::test] + async fn state_survives_empty_input_batches_and_output_flushes() -> Result<()> { + let exec = test_exec()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let batches = collect(Arc::clone(&exec) as _, context).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 2, 2, 1] + ); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!( + ids, + vec![ + Some(0), + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + ] + ); + assert_eq!( + prices, + vec![None, None, None, Some(40), Some(60), Some(101), None] + ); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!(metrics.output_rows(), Some(7)); + assert_eq!( + metrics + .sum_by_name("matched_rows") + .map(|value| value.as_usize()), + Some(3) + ); + assert_eq!( + metrics + .sum_by_name("unmatched_left_rows") + .map(|value| value.as_usize()), + Some(4) + ); + assert!(metrics.elapsed_compute().is_some()); + assert!( + metrics.iter().any(|metric| { + matches!(metric.value(), MetricValue::ElapsedCompute(_)) + }) + ); + Ok(()) + } + + #[tokio::test] + async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut left_payload = StringDictionaryBuilder::::new(); + for _ in 0..129 { + left_payload.append_value("left"); + } + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(StringArray::from(vec!["A"; 129])), + Arc::new(Int64Array::from_iter_values(-1..128)), + Arc::new(left_payload.finish()), + ], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut right_payload = StringDictionaryBuilder::::new(); + right_payload.append_value("right"); + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(Int64Array::from(vec![0])), + Arc::new(right_payload.finish()), + ], + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(256)), + ); + let batches = collect(exec, context).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 129); + assert_eq!(batches[0].column(2).data_type(), &dictionary_type); + assert_eq!(batches[0].column(3).data_type(), &dictionary_type); + + let right_output = batches[0] + .column(3) + .as_any() + .downcast_ref::>() + .expect("right output must remain Dictionary(Int8, Utf8)"); + assert!(right_output.is_null(0)); + assert_eq!(right_output.null_count(), 1); + let values = right_output + .values() + .as_any() + .downcast_ref::() + .expect("dictionary values must be Utf8"); + for row in 1..129 { + assert_eq!( + values.value(right_output.keys().value(row) as usize), + "right" + ); + } + Ok(()) + } + + #[test] + fn rejects_volatile_physical_expressions() -> Result<()> { + let exec = test_exec()?; + let volatile = Arc::new(VolatileExpr) as PhysicalExprRef; + let match_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + exec.on().clone(), + AsOfMatchExpr::new( + Arc::clone(&volatile), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + ) + .expect_err("volatile match expression must be rejected"); + assert!(match_error.to_string().contains("must be deterministic")); + + let equality_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![(volatile, Arc::new(PhysicalColumn::new("key", 0)))], + exec.match_condition().clone(), + vec![2], + ) + .expect_err("volatile equality expression must be rejected"); + assert!(equality_error.to_string().contains("must be deterministic")); + Ok(()) + } + + #[test] + fn properties_and_statistics_follow_left_preserving_contract() -> Result<()> { + let exec = test_exec()?; + let exec_plan: Arc = Arc::clone(&exec) as _; + assert_eq!(exec.maintains_input_order(), vec![true, false]); + assert_eq!(exec_plan.pipeline_behavior(), EmissionType::Incremental); + assert_eq!(exec_plan.boundedness(), Boundedness::Bounded); + assert!(matches!( + &exec.input_distribution_requirements().into_per_child()[..], + [ + Distribution::KeyPartitioned(_), + Distribution::KeyPartitioned(_) + ] + )); + for ordering in exec.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 2); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + assert_eq!( + requirement[1].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + } + + let no_keys: Arc = Arc::new(AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::Lt, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + assert_eq!(no_keys.output_partitioning().partition_count(), 1); + assert!(matches!( + &no_keys.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + for ordering in no_keys.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 1); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: true, + nulls_first: true, + }) + ); + } + + let mut key_stats = ColumnStatistics::new_unknown(); + key_stats.null_count = Precision::Exact(1); + key_stats.distinct_count = Precision::Exact(4); + let mut ts_stats = ColumnStatistics::new_unknown(); + ts_stats.min_value = Precision::Exact(ScalarValue::Int64(Some(1))); + ts_stats.max_value = Precision::Exact(ScalarValue::Int64(Some(7))); + let mut id_stats = ColumnStatistics::new_unknown(); + id_stats.null_count = Precision::Exact(0); + id_stats.distinct_count = Precision::Exact(7); + let left_column_statistics = vec![key_stats, ts_stats, id_stats]; + let left_stats = Arc::new(Statistics { + num_rows: Precision::Exact(7), + total_byte_size: Precision::Exact(128), + column_statistics: left_column_statistics.clone(), + }); + let right_stats = Arc::new(Statistics::new_unknown(&exec.right().schema())); + let stats = exec + .statistics_from_inputs(&[left_stats, right_stats], &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(7)); + assert_eq!(stats.total_byte_size, Precision::Absent); + assert_eq!(stats.column_statistics.len(), 4); + assert_eq!( + &stats.column_statistics[..3], + left_column_statistics.as_slice() + ); + assert_eq!(stats.column_statistics[3], ColumnStatistics::new_unknown()); + assert_eq!( + exec.child_stats_requests(None), + vec![ChildStats::At(None), ChildStats::Skip] + ); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index bbb25dda65165..820f60b09b3ba 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -18,6 +18,7 @@ //! DataFusion Join implementations use arrow::array::BooleanBufferBuilder; +pub use asof_join::{AsOfJoinExec, AsOfMatchExpr}; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ @@ -29,6 +30,7 @@ use parking_lot::Mutex; pub use piecewise_merge_join::PiecewiseMergeJoinExec; pub use sort_merge_join::SortMergeJoinExec; pub use symmetric_hash_join::SymmetricHashJoinExec; +mod asof_join; pub mod chain; mod cross_join; mod hash_join; From 43110a3260548ad87b73fb5670c05c24839444bb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:38:22 +0800 Subject: [PATCH 02/10] feat: broadcast ASOF right input --- .../physical-plan/src/joins/asof_join.rs | 339 ++++++++++++++---- 1 file changed, 273 insertions(+), 66 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 5945c8d750716..b55a21260de8f 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -15,7 +15,42 @@ // specific language governing permissions and limitations // under the License. -//! Ordered, left-preserving ASOF join execution. +//! 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: +//! +//! ```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: +//! +//! ```text +//! AsOfJoinExec +//! SortExec(left equality keys, left match key) +//! RepartitionExec(RoundRobinBatch) +//! left +//! SortExec(right equality keys, right match key) +//! CoalescePartitionsExec +//! right +//! ``` +//! +//! 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. Each left partition owns its cursors, +//! equality-group state, and current candidate, while the collected right +//! batches are immutable and shared. +//! +//! 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, so a +//! repartitioned streaming mode remains a useful future alternative for large +//! right inputs. use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -35,27 +70,30 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, plan_err, }; 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; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{Partitioning, PhysicalSortExpr}; use datafusion_physical_expr_common::physical_expr::{ PhysicalExprRef, fmt_sql, is_volatile, }; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; -use futures::{StreamExt, stream}; +use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; -use crate::joins::utils::{JoinOn, build_join_schema}; +use crate::joins::utils::{JoinOn, OnceAsync, build_join_schema}; +use crate::memory::MemoryStream; use crate::metrics::{ - BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - MetricsSet, RecordOutput, Time, + BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, + MetricCategory, MetricsSet, RecordOutput, Time, }; +use crate::spill::get_record_batch_memory_size; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ @@ -82,8 +120,8 @@ impl AsOfMatchExpr { } } -/// A sort-merge ASOF join that emits exactly one row for every left row. -#[derive(Debug, Clone)] +/// A broadcast sort-merge ASOF join that emits one row for every left row. +#[derive(Debug)] pub struct AsOfJoinExec { left: Arc, right: Arc, @@ -94,6 +132,7 @@ pub struct AsOfJoinExec { metrics: ExecutionPlanMetricsSet, left_ordering: LexOrdering, right_ordering: LexOrdering, + right_fut: OnceAsync, cache: Arc, } @@ -216,7 +255,7 @@ impl AsOfJoinExec { "ASOF right ordering must not be empty" ) })?; - let cache = Arc::new(Self::compute_properties(&left, &schema, on.is_empty())?); + let cache = Arc::new(Self::compute_properties(&left, &schema)?); Ok(Self { left, @@ -228,6 +267,7 @@ impl AsOfJoinExec { metrics: ExecutionPlanMetricsSet::new(), left_ordering, right_ordering, + right_fut: Default::default(), cache, }) } @@ -235,7 +275,6 @@ impl AsOfJoinExec { fn compute_properties( left: &Arc, schema: &SchemaRef, - single_partition: bool, ) -> Result { let left_schema = left.schema(); let mapping = ProjectionMapping::try_new( @@ -254,12 +293,9 @@ impl AsOfJoinExec { )?; let input_eq_properties = left.equivalence_properties(); let eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); - let output_partitioning = if single_partition { - Partitioning::UnknownPartitioning(1) - } else { - left.output_partitioning() - .project(&mapping, input_eq_properties) - }; + let output_partitioning = left + .output_partitioning() + .project(&mapping, input_eq_properties); Ok(PlanProperties::new( eq_properties, output_partitioning, @@ -364,22 +400,10 @@ impl ExecutionPlan for AsOfJoinExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - if self.on.is_empty() { - InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - Distribution::SinglePartition, - ]) - } else { - let (left, right) = self - .on - .iter() - .map(|(left, right)| (Arc::clone(left), Arc::clone(right))) - .unzip(); - InputDistributionRequirements::co_partitioned(vec![ - Distribution::KeyPartitioned(left), - Distribution::KeyPartitioned(right), - ]) - } + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition, + ]) } fn required_input_ordering(&self) -> Vec> { @@ -428,8 +452,15 @@ impl ExecutionPlan for AsOfJoinExec { Ok(Arc::new(Self { left, right, + on: self.on.clone(), + match_condition: self.match_condition.clone(), + right_output_indices: self.right_output_indices.clone(), + schema: Arc::clone(&self.schema), metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&self) + left_ordering: self.left_ordering.clone(), + right_ordering: self.right_ordering.clone(), + right_fut: Default::default(), + cache: Arc::clone(&self.cache), })) } @@ -438,41 +469,62 @@ impl ExecutionPlan for AsOfJoinExec { partition: usize, context: Arc, ) -> Result { - let left_partitions = self.left.output_partitioning().partition_count(); let right_partitions = self.right.output_partitioning().partition_count(); assert_eq_or_internal_err!( - left_partitions, right_partitions, - "AsOfJoinExec partition count mismatch: {left_partitions} != {right_partitions}" + 1, + "AsOfJoinExec requires one right partition, found {right_partitions}" ); let left_stream = self.left.execute(partition, Arc::clone(&context))?; - let right_stream = self.right.execute(partition, Arc::clone(&context))?; - let (left_keys, right_keys) = self.on.iter().cloned().unzip(); - let state = AsOfJoinStreamState::new( - Arc::clone(&self.schema), - InputCursor::new( - left_stream, - left_keys, - Arc::clone(&self.match_condition.left), - ), - InputCursor::new( + 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, - right_keys, - Arc::clone(&self.match_condition.right), - ), - self.match_condition.op, - self.right_output_indices.clone(), - context.session_config().batch_size(), - AsOfJoinMetrics::new(partition, &self.metrics), - ); - let stream = stream::try_unfold(state, |mut state| async move { - match state.next_batch().await? { - Some(batch) => Ok(Some((batch, state))), - None => Ok(None), - } - }); + reservation, + build_metrics, + )) + })?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let output_schema = Arc::clone(&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 right_output_indices = self.right_output_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, + right_output_indices, + batch_size, + metrics, + ); + 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::(Box::pin( + RecordBatchStreamAdapter::new(stream_schema, stream), + )) + }) + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&self.schema), + output_schema, stream, ))) } @@ -535,6 +587,47 @@ impl ExecutionPlan for AsOfJoinExec { } } +struct BroadcastRightInput { + schema: SchemaRef, + batches: Vec, + _reservation: MemoryReservation, +} + +impl BroadcastRightInput { + fn stream(&self) -> Result { + 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 { + let schema = input.schema(); + let batches = input + .try_fold(Vec::new(), |mut batches, batch| { + let batch_size = get_record_batch_memory_size(&batch); + futures::future::ready(reservation.try_grow(batch_size).map(|_| { + metrics.build_mem_used.add(batch_size); + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + batches.push(batch); + batches + })) + }) + .await?; + Ok(BroadcastRightInput { + schema, + batches, + _reservation: reservation, + }) +} + #[derive(Clone)] struct Candidate { batch: Arc, @@ -632,10 +725,14 @@ impl InputCursor { } } +#[derive(Clone)] struct AsOfJoinMetrics { baseline: BaselineMetrics, matched_rows: Count, unmatched_left_rows: Count, + build_input_batches: Count, + build_input_rows: Count, + build_mem_used: Gauge, } impl AsOfJoinMetrics { @@ -648,6 +745,14 @@ impl AsOfJoinMetrics { unmatched_left_rows: MetricBuilder::new(metrics) .with_category(MetricCategory::Rows) .counter("unmatched_left_rows", partition), + build_input_batches: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_batches", partition), + build_input_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_rows", partition), + build_mem_used: MetricBuilder::new(metrics) + .peak_memory_usage("build_mem_used", partition), } } } @@ -958,8 +1063,8 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< #[cfg(test)] mod tests { use super::*; - use crate::collect; use crate::test::TestMemoryExec; + use crate::{collect, collect_partitioned}; use arrow::array::{ DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, }; @@ -1200,6 +1305,102 @@ mod tests { Ok(()) } + #[tokio::test] + async fn broadcasts_right_input_to_all_left_partitions() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("id", DataType::Int32, false), + ])); + let left = TestMemoryExec::try_new_exec( + &[ + vec![make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(1), Some(4)], + vec![0, 1], + )?], + vec![make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(2), Some(5)], + vec![2, 3], + )?], + ], + Arc::clone(&left_schema), + None, + )?; + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ])); + let right = TestMemoryExec::try_new_exec( + &[vec![ + make_batch(&right_schema, vec![Some("A")], vec![Some(1)], vec![10])?, + make_batch(&right_schema, vec![Some("A")], vec![Some(3)], vec![30])?, + ]], + Arc::clone(&right_schema), + None, + )?; + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + assert_eq!(exec.properties().output_partitioning().partition_count(), 2); + assert!(matches!( + &exec.input_distribution_requirements().into_per_child()[..], + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] + )); + + let partitions = collect_partitioned( + Arc::clone(&exec) as Arc, + Arc::new(TaskContext::default()), + ) + .await?; + assert_eq!(partitions.len(), 2); + for batches in partitions { + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(10), Some(30)]); + } + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!( + metrics + .sum_by_name("build_input_batches") + .map(|value| value.as_usize()), + Some(2) + ); + assert_eq!( + metrics + .sum_by_name("build_input_rows") + .map(|value| value.as_usize()), + Some(2) + ); + assert_eq!(metrics.output_rows(), Some(4)); + Ok(()) + } + #[tokio::test] async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { let dictionary_type = @@ -1333,8 +1534,8 @@ mod tests { assert!(matches!( &exec.input_distribution_requirements().into_per_child()[..], [ - Distribution::KeyPartitioned(_), - Distribution::KeyPartitioned(_) + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition ] )); for ordering in exec.required_input_ordering() { @@ -1367,10 +1568,16 @@ mod tests { ), vec![2], )?); - assert_eq!(no_keys.output_partitioning().partition_count(), 1); + assert_eq!( + no_keys.output_partitioning().partition_count(), + exec.left().output_partitioning().partition_count() + ); assert!(matches!( &no_keys.input_distribution_requirements().into_per_child()[..], - [Distribution::SinglePartition, Distribution::SinglePartition] + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] )); for ordering in no_keys.required_input_ordering() { let requirement = ordering.expect("ASOF ordering is required").into_single(); From 339d7857107c2e87bbffe87d89d2c647dbc45c63 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 01:40:50 +0800 Subject: [PATCH 03/10] fix: account shared ASOF build buffers once --- .../physical-plan/src/joins/asof_join.rs | 90 ++++++++++++++++++- 1 file changed, 88 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index b55a21260de8f..2daa5e5c9fe67 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -62,6 +62,7 @@ use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::{ compare_rows, get_row_at_idx, normalize_float_zero_scalar, }; @@ -93,7 +94,6 @@ use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricCategory, MetricsSet, RecordOutput, Time, }; -use crate::spill::get_record_batch_memory_size; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ @@ -609,9 +609,10 @@ async fn collect_right_input( metrics: AsOfJoinMetrics, ) -> Result { let schema = input.schema(); + let mut memory_counter = RecordBatchMemoryCounter::new(); let batches = input .try_fold(Vec::new(), |mut batches, batch| { - let batch_size = get_record_batch_memory_size(&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); metrics.build_input_batches.add(1); @@ -1070,6 +1071,7 @@ mod tests { }; use arrow::datatypes::{DataType, Field, Int8Type}; use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::metrics::MetricValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -1401,6 +1403,90 @@ mod tests { Ok(()) } + #[tokio::test] + async fn shared_right_buffers_are_reserved_once() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("id", DataType::Int32, false), + ])); + let left = TestMemoryExec::try_new_exec( + &[vec![make_batch( + &left_schema, + vec![Some("A")], + vec![Some(4095)], + vec![0], + )?]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ])); + let row_count = 4096; + let parent = make_batch( + &right_schema, + vec![Some("A"); row_count], + (0..row_count).map(|value| Some(value as i64)).collect(), + (0..row_count as i32).collect(), + )?; + let mut memory_counter = RecordBatchMemoryCounter::new(); + let retained_size = memory_counter.count_batch(&parent); + let right_batches = (0..16) + .map(|index| parent.slice(index * 256, 256)) + .collect(); + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(retained_size, 1.0) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + + let batches = collect(Arc::clone(&exec) as _, context).await?; + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(4095)]); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!( + metrics + .sum_by_name("build_mem_used") + .map(|value| value.as_usize()), + Some(retained_size) + ); + Ok(()) + } + #[tokio::test] async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { let dictionary_type = From 34552d158eef2737c9022de570dc6af5d54e5aa2 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:50:01 +0800 Subject: [PATCH 04/10] perf: compare ASOF keys without scalar materialization --- .../physical-plan/src/joins/asof_join.rs | 178 ++++++++++++------ 1 file changed, 124 insertions(+), 54 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 2daa5e5c9fe67..6da25bf50c26a 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -58,16 +58,15 @@ use std::fmt::Formatter; use std::sync::Arc; use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::buffer::NullBuffer; use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; -use datafusion_common::utils::{ - compare_rows, get_row_at_idx, normalize_float_zero_scalar, -}; +use datafusion_common::utils::normalize_float_zero_scalar; use datafusion_common::{ - ColumnStatistics, JoinType, Result, ScalarValue, Statistics, + ColumnStatistics, JoinType, NullEquality, Result, ScalarValue, Statistics, assert_eq_or_internal_err, internal_err, plan_err, }; use datafusion_execution::TaskContext; @@ -88,7 +87,9 @@ use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; -use crate::joins::utils::{JoinOn, OnceAsync, build_join_schema}; +use crate::joins::utils::{ + JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, matchable_join_keys, +}; use crate::memory::MemoryStream; use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, @@ -633,7 +634,8 @@ async fn collect_right_input( struct Candidate { batch: Arc, row: usize, - group: Vec, + key_arrays: Arc<[ArrayRef]>, + key_batch_id: usize, } struct InputCursor { @@ -641,8 +643,10 @@ struct InputCursor { key_exprs: Vec, match_expr: PhysicalExprRef, batch: Option>, - key_arrays: Vec, + key_arrays: Arc<[ArrayRef]>, + key_validity: Option, match_array: Option, + key_batch_id: usize, row: usize, eof: bool, } @@ -658,8 +662,10 @@ impl InputCursor { key_exprs, match_expr, batch: None, - key_arrays: vec![], + key_arrays: Arc::from([]), + key_validity: None, match_array: None, + key_batch_id: 0, row: 0, eof: false, } @@ -673,7 +679,8 @@ impl InputCursor { return Ok(true); } self.batch = None; - self.key_arrays.clear(); + self.key_arrays = Arc::from([]); + self.key_validity = None; self.match_array = None; self.row = 0; if self.eof { @@ -688,23 +695,28 @@ impl InputCursor { } let batch = Arc::new(batch); let _timer = elapsed_compute.timer(); - self.key_arrays = self + let key_arrays = self .key_exprs .iter() .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) - .collect::>()?; + .collect::>>()?; + 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(&self) -> Result> { - get_row_at_idx(&self.key_arrays, self.row) - .map(|row| row.into_iter().map(normalize_float_zero_scalar).collect()) + 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 { @@ -850,6 +862,8 @@ struct AsOfJoinStreamState { right_output_indices: Vec, candidate: Option, group_sort_options: Vec, + input_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + candidate_group_comparator: Option<(usize, usize, JoinKeyComparator)>, pending_left: PendingRows, pending_right: PendingRows, batch_size: usize, @@ -883,11 +897,76 @@ impl AsOfJoinStreamState { right_output_indices, candidate: None, group_sort_options, + input_group_comparator: None, + candidate_group_comparator: None, batch_size: batch_size.max(1), metrics, } } + fn compare_input_groups(&mut self) -> Result { + if self.group_sort_options.is_empty() { + return Ok(Ordering::Equal); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_batch_id = self.right.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .input_group_comparator + .as_ref() + .is_none_or(|(right, left, _)| { + *right != right_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + self.right.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.input_group_comparator = + Some((right_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .input_group_comparator + .as_ref() + .expect("ASOF input group comparator must be initialized"); + Ok(comparator.compare(self.right.row, self.left.row)) + } + + fn candidate_is_other_group(&mut self) -> Result { + let Some(candidate) = &self.candidate else { + return Ok(false); + }; + if self.group_sort_options.is_empty() { + return Ok(false); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let candidate_batch_id = candidate.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .candidate_group_comparator + .as_ref() + .is_none_or(|(candidate, left, _)| { + *candidate != candidate_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + candidate.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.candidate_group_comparator = + Some((candidate_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .candidate_group_comparator + .as_ref() + .expect("ASOF candidate group comparator must be initialized"); + Ok(comparator.compare(candidate.row, self.left.row) != Ordering::Equal) + } + async fn next_batch(&mut self) -> Result> { loop { if self.pending_left.len() >= self.batch_size { @@ -905,25 +984,20 @@ impl AsOfJoinStreamState { return Ok(None); } - let (left_group, left_match) = { + let left_match = { let _timer = self.metrics.baseline.elapsed_compute().timer(); - (self.left.group()?, self.left.match_value()?) + self.left.match_value()? }; - if left_match.is_null() || left_group.iter().any(ScalarValue::is_null) { + if left_match.is_null() || self.left.group_has_null() { self.candidate = None; + self.candidate_group_comparator = None; self.push_current_left(None)?; self.left.advance(); continue; } - let candidate_is_other_group = if let Some(candidate) = &self.candidate { - let _timer = self.metrics.baseline.elapsed_compute().timer(); - compare_rows(&candidate.group, &left_group, &self.group_sort_options)? - != Ordering::Equal - } else { - false - }; - if candidate_is_other_group { + if self.candidate_is_other_group()? { self.candidate = None; + self.candidate_group_comparator = None; } loop { @@ -934,34 +1008,27 @@ impl AsOfJoinStreamState { { break; } - let action = { - let _timer = self.metrics.baseline.elapsed_compute().timer(); - let right_group = self.right.group()?; - if right_group.iter().any(ScalarValue::is_null) { - RightAction::Advance - } else { - match compare_rows( - &right_group, - &left_group, - &self.group_sort_options, - )? { - Ordering::Less => RightAction::Advance, - Ordering::Greater => RightAction::Stop, - Ordering::Equal => { - let right_match = self.right.match_value()?; - if right_match.is_null() { - RightAction::Advance - } else if is_eligible(self.op, &left_match, &right_match)? - { - let (batch, row) = self.right.batch_row()?; - RightAction::Candidate(Candidate { - batch, - row, - group: right_group, - }) - } else { - RightAction::Stop - } + let action = if self.right.group_has_null() { + RightAction::Advance + } else { + match self.compare_input_groups()? { + Ordering::Less => RightAction::Advance, + Ordering::Greater => RightAction::Stop, + Ordering::Equal => { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_match = self.right.match_value()?; + if right_match.is_null() { + RightAction::Advance + } else if is_eligible(self.op, &left_match, &right_match)? { + let (batch, row) = self.right.batch_row()?; + RightAction::Candidate(Candidate { + batch, + row, + key_arrays: Arc::clone(&self.right.key_arrays), + key_batch_id: self.right.key_batch_id, + }) + } else { + RightAction::Stop } } } @@ -969,6 +1036,9 @@ impl AsOfJoinStreamState { match action { RightAction::Advance => self.right.advance(), RightAction::Candidate(candidate) => { + // Replacing the candidate selects the nearest eligible row. + // Equal match values have no secondary ordering, so which + // tied row wins is intentionally nondeterministic. self.candidate = Some(candidate); self.right.advance(); } @@ -1057,7 +1127,7 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< Operator::GtEq => ordering != Ordering::Greater, Operator::Lt => ordering == Ordering::Greater, Operator::LtEq => ordering != Ordering::Less, - _ => false, + _ => unreachable!("ASOF match operator is validated by try_new"), }) } From a127baaf87dc5440326976b889b6f1bbc006d3c5 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 2 Aug 2026 01:03:36 +0800 Subject: [PATCH 05/10] refactor: simplify ASOF physical operator --- .../physical-plan/src/joins/asof_join.rs | 353 ++++++++++-------- 1 file changed, 203 insertions(+), 150 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 6da25bf50c26a..a28f375c97945 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -19,7 +19,7 @@ //! //! 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: +//! one ordered comparison. This follows Snowflake's [ASOF JOIN] semantics: //! //! ```text //! left.ts >= right.ts => greatest eligible right.ts @@ -28,32 +28,44 @@ //! //! 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: +//! monotonic scan over the ordered right input. //! -//! ```text -//! AsOfJoinExec -//! SortExec(left equality keys, left match key) -//! RepartitionExec(RoundRobinBatch) -//! left -//! SortExec(right equality keys, right match key) -//! CoalescePartitionsExec -//! right -//! ``` +//! [`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. Each left partition owns its cursors, -//! equality-group state, and current candidate, while the collected right -//! batches are immutable and shared. +//! 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. //! //! 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, so a -//! repartitioned streaming mode remains a useful future alternative for large -//! right inputs. +//! 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, HashSet}; +use std::collections::HashMap; use std::fmt::Formatter; use std::sync::Arc; @@ -61,7 +73,6 @@ use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; use arrow::buffer::NullBuffer; use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; -use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::normalize_float_zero_scalar; @@ -83,10 +94,6 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequiremen use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; use crate::execution_plan::{Boundedness, EmissionType}; -use crate::filter_pushdown::{ - ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, - FilterPushdownPropagation, -}; use crate::joins::utils::{ JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, matchable_join_keys, }; @@ -128,17 +135,26 @@ pub struct AsOfJoinExec { right: Arc, on: JoinOn, match_condition: AsOfMatchExpr, + /// Sorted, unique indices of right columns appended after all left columns. right_output_indices: Vec, schema: SchemaRef, 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, cache: Arc, } 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. + /// Right output indices must be in bounds, sorted, and unique. pub fn try_new( left: Arc, right: Arc, @@ -146,73 +162,15 @@ impl AsOfJoinExec { match_condition: AsOfMatchExpr, right_output_indices: Vec, ) -> Result { - if !matches!( - match_condition.op, - Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq - ) { - return plan_err!( - "AsOfJoinExec requires <, <=, >, or >=, found {}", - match_condition.op - ); - } - if left.boundedness().is_unbounded() || right.boundedness().is_unbounded() { - return plan_err!("AsOfJoinExec requires bounded inputs"); - } - if is_volatile(&match_condition.left) || is_volatile(&match_condition.right) { - return plan_err!("AsOfJoinExec match expression must be deterministic"); - } - if on - .iter() - .any(|(left, right)| is_volatile(left) || is_volatile(right)) - { - return plan_err!("AsOfJoinExec equality expressions must be deterministic"); - } - + validate_asof_join( + left.as_ref(), + right.as_ref(), + &on, + &match_condition, + &right_output_indices, + )?; let left_schema = left.schema(); let right_schema = right.schema(); - validate_expr_side(&match_condition.left, &left_schema, "left match")?; - validate_expr_side(&match_condition.right, &right_schema, "right match")?; - for (left_expr, right_expr) in &on { - validate_expr_side(left_expr, &left_schema, "left equality")?; - validate_expr_side(right_expr, &right_schema, "right equality")?; - let left_type = left_expr.data_type(&left_schema)?; - let right_type = right_expr.data_type(&right_schema)?; - if left_type != right_type { - return plan_err!( - "AsOfJoinExec equality expression types differ: {left_type} and {right_type}" - ); - } - if !datafusion_expr::utils::can_hash(&left_type) { - return plan_err!( - "AsOfJoinExec equality expressions have unsupported hash type {left_type}" - ); - } - } - let left_match_type = match_condition.left.data_type(&left_schema)?; - let right_match_type = match_condition.right.data_type(&right_schema)?; - if left_match_type != right_match_type { - return plan_err!( - "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" - ); - } - if let Some(index) = right_output_indices - .iter() - .find(|index| **index >= right_schema.fields().len()) - { - return plan_err!( - "AsOfJoinExec right output index {index} is outside schema with {} fields", - right_schema.fields().len() - ); - } - if !right_output_indices - .windows(2) - .all(|pair| pair[0] < pair[1]) - { - return plan_err!( - "AsOfJoinExec right output indices must be strictly increasing" - ); - } - let schema = build_output_schema(&left_schema, &right_schema, &right_output_indices); let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); @@ -304,31 +262,6 @@ impl AsOfJoinExec { Boundedness::Bounded, )) } - - /// Equality expressions. - pub fn on(&self) -> &JoinOn { - &self.on - } - - /// Ordered match expression. - pub fn match_condition(&self) -> &AsOfMatchExpr { - &self.match_condition - } - - /// Indices of right input columns emitted after the left columns. - pub fn right_output_indices(&self) -> &[usize] { - &self.right_output_indices - } - - /// Left input. - pub fn left(&self) -> &Arc { - &self.left - } - - /// Right input. - pub fn right(&self) -> &Arc { - &self.right - } } fn build_output_schema( @@ -401,6 +334,8 @@ impl ExecutionPlan for AsOfJoinExec { } 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. InputDistributionRequirements::new(vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, @@ -415,7 +350,7 @@ impl ExecutionPlan for AsOfJoinExec { } fn maintains_input_order(&self) -> Vec { - vec![true, false] + vec![false, false] } fn children(&self) -> Vec<&Arc> { @@ -543,6 +478,8 @@ impl ExecutionPlan for AsOfJoinExec { input_stats: &[Arc], _args: &StatisticsArgs, ) -> Result> { + // 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 mut column_statistics = left.column_statistics.clone(); column_statistics.truncate(self.left.schema().fields().len()); @@ -561,36 +498,15 @@ impl ExecutionPlan for AsOfJoinExec { column_statistics, })) } - - fn gather_filters_for_pushdown( - &self, - _phase: FilterPushdownPhase, - parent_filters: Vec, - _config: &ConfigOptions, - ) -> Result { - let left_indices = (0..self.left.schema().fields().len()).collect::>(); - let left = ChildFilterDescription::from_child_with_allowed_indices( - &parent_filters, - left_indices, - &self.left, - )?; - let right = ChildFilterDescription::all_unsupported(&parent_filters); - Ok(FilterDescription::new().with_child(left).with_child(right)) - } - - fn handle_child_pushdown_result( - &self, - _phase: FilterPushdownPhase, - child_pushdown_result: ChildPushdownResult, - _config: &ConfigOptions, - ) -> Result>> { - Ok(FilterPushdownPropagation::if_any(child_pushdown_result)) - } } +/// 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, + /// Holds the memory-pool reservation for as long as the batches are shared. _reservation: MemoryReservation, } @@ -632,22 +548,40 @@ async fn collect_right_input( #[derive(Clone)] struct Candidate { + /// Right batch containing the nearest eligible row. batch: Arc, + /// 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, + /// Ordered match expression evaluated for each batch. match_expr: PhysicalExprRef, + /// Current non-empty batch. batch: Option>, + /// Evaluated equality-key arrays for `batch`. key_arrays: Arc<[ArrayRef]>, + /// Rows whose equality keys are all non-NULL. key_validity: Option, + /// Evaluated match values for `batch`. match_array: Option, + /// 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, } @@ -770,10 +704,18 @@ impl AsOfJoinMetrics { } } +/// 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>, + /// Maps an `Arc` pointer to its index in `sources`. source_by_ptr: HashMap, + /// Per-output-row `(source, row)` references or NULL padding. indices: Vec>, } @@ -854,18 +796,36 @@ impl PendingRows { } } +/// Per-left-partition state for the monotonic ASOF scan. +/// +/// For left rows `(A, 4), (A, 7)` and right rows `(A, 2), (A, 6)`, the +/// candidate advances from `(A, 2)` to `(A, 6)` without rewinding the right +/// cursor. Cursors and the candidate survive input batch changes and output +/// flushes; a change of equality group clears the candidate before reuse. struct AsOfJoinStreamState { + /// Output schema used when pending row references are materialized. schema: SchemaRef, + /// Cursor over the current left partition. left: InputCursor, + /// Independent cursor over the shared, ordered right input. right: InputCursor, + /// Validated ordered match operator. op: Operator, + /// Right columns appended to each output row. right_output_indices: Vec, + /// Nearest eligible right row for the current equality group. candidate: Option, + /// Equality-key ordering shared by the comparator caches. group_sort_options: Vec, + /// Cached comparator for the current right and left input batches. input_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + /// Cached comparator for the candidate and current left batches. candidate_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + /// Left row references accumulated for the next output batch. pending_left: PendingRows, + /// Matched right row references, aligned with `pending_left`. pending_right: PendingRows, + /// Maximum number of pending rows before an output flush. batch_size: usize, metrics: AsOfJoinMetrics, } @@ -967,6 +927,15 @@ impl AsOfJoinStreamState { Ok(comparator.compare(candidate.row, self.left.row) != Ordering::Equal) } + /// Produces the next output batch without resetting the merge state. + /// + /// Each left row first validates its equality group, then advances the right + /// cursor while right groups sort before it or right match values remain + /// eligible. The last eligible right row becomes the candidate. Empty input + /// batches are skipped. Right EOF preserves that candidate for later left + /// rows in the same group; left EOF flushes the final pending rows. NULL keys + /// and group changes clear the candidate, while output flushes only clear + /// pending row references. async fn next_batch(&mut self) -> Result> { loop { if self.pending_left.len() >= self.batch_size { @@ -1072,6 +1041,8 @@ impl AsOfJoinStreamState { Ok(()) } + /// Materializes pending row references while preserving both cursors and the + /// current equality-group candidate for the next output batch. fn flush(&mut self) -> Result { let _timer = self.metrics.baseline.elapsed_compute().timer(); let left_len = self.schema.fields().len() - self.right_output_indices.len(); @@ -1096,6 +1067,83 @@ impl AsOfJoinStreamState { } } +/// Validates all invariants required by the forward-only ASOF state machine. +fn validate_asof_join( + left: &dyn ExecutionPlan, + right: &dyn ExecutionPlan, + on: &JoinOn, + match_condition: &AsOfMatchExpr, + right_output_indices: &[usize], +) -> Result<()> { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "AsOfJoinExec requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + if left.boundedness().is_unbounded() || right.boundedness().is_unbounded() { + return plan_err!("AsOfJoinExec requires bounded inputs"); + } + if is_volatile(&match_condition.left) || is_volatile(&match_condition.right) { + return plan_err!("AsOfJoinExec match expression must be deterministic"); + } + if on + .iter() + .any(|(left, right)| is_volatile(left) || is_volatile(right)) + { + return plan_err!("AsOfJoinExec equality expressions must be deterministic"); + } + + let left_schema = left.schema(); + let right_schema = right.schema(); + validate_expr_side(&match_condition.left, &left_schema, "left match")?; + validate_expr_side(&match_condition.right, &right_schema, "right match")?; + for (left_expr, right_expr) in on { + validate_expr_side(left_expr, &left_schema, "left equality")?; + validate_expr_side(right_expr, &right_schema, "right equality")?; + let left_type = left_expr.data_type(&left_schema)?; + let right_type = right_expr.data_type(&right_schema)?; + if left_type != right_type { + return plan_err!( + "AsOfJoinExec equality expression types differ: {left_type} and {right_type}" + ); + } + if !datafusion_expr::utils::can_hash(&left_type) { + return plan_err!( + "AsOfJoinExec equality expressions have unsupported hash type {left_type}" + ); + } + } + let left_match_type = match_condition.left.data_type(&left_schema)?; + let right_match_type = match_condition.right.data_type(&right_schema)?; + if left_match_type != right_match_type { + return plan_err!( + "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" + ); + } + if let Some(index) = right_output_indices + .iter() + .find(|index| **index >= right_schema.fields().len()) + { + return plan_err!( + "AsOfJoinExec right output index {index} is outside schema with {} fields", + right_schema.fields().len() + ); + } + if !right_output_indices + .windows(2) + .all(|pair| pair[0] < pair[1]) + { + return plan_err!( + "AsOfJoinExec right output indices must be strictly increasing" + ); + } + Ok(()) +} + fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Result<()> { let columns = collect_columns(expr); if columns.is_empty() { @@ -1133,6 +1181,11 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< #[cfg(test)] mod tests { + // These tests cover physical-only contracts that SQL logic tests cannot + // observe, including batch-boundary state, shared build memory, Arrow type + // preservation, and execution properties. End-to-end SQL semantics live in + // the dependent SQL layer. + use super::*; use crate::test::TestMemoryExec; use crate::{collect, collect_partitioned}; @@ -1655,9 +1708,9 @@ mod tests { let exec = test_exec()?; let volatile = Arc::new(VolatileExpr) as PhysicalExprRef; let match_error = AsOfJoinExec::try_new( - Arc::clone(exec.left()), - Arc::clone(exec.right()), - exec.on().clone(), + Arc::clone(&exec.left), + Arc::clone(&exec.right), + exec.on.clone(), AsOfMatchExpr::new( Arc::clone(&volatile), Operator::GtEq, @@ -1669,10 +1722,10 @@ mod tests { assert!(match_error.to_string().contains("must be deterministic")); let equality_error = AsOfJoinExec::try_new( - Arc::clone(exec.left()), - Arc::clone(exec.right()), + Arc::clone(&exec.left), + Arc::clone(&exec.right), vec![(volatile, Arc::new(PhysicalColumn::new("key", 0)))], - exec.match_condition().clone(), + exec.match_condition.clone(), vec![2], ) .expect_err("volatile equality expression must be rejected"); @@ -1684,7 +1737,7 @@ mod tests { fn properties_and_statistics_follow_left_preserving_contract() -> Result<()> { let exec = test_exec()?; let exec_plan: Arc = Arc::clone(&exec) as _; - assert_eq!(exec.maintains_input_order(), vec![true, false]); + assert_eq!(exec.maintains_input_order(), vec![false, false]); assert_eq!(exec_plan.pipeline_behavior(), EmissionType::Incremental); assert_eq!(exec_plan.boundedness(), Boundedness::Bounded); assert!(matches!( @@ -1714,8 +1767,8 @@ mod tests { } let no_keys: Arc = Arc::new(AsOfJoinExec::try_new( - Arc::clone(exec.left()), - Arc::clone(exec.right()), + Arc::clone(&exec.left), + Arc::clone(&exec.right), vec![], AsOfMatchExpr::new( Arc::new(PhysicalColumn::new("ts", 1)), @@ -1726,7 +1779,7 @@ mod tests { )?); assert_eq!( no_keys.output_partitioning().partition_count(), - exec.left().output_partitioning().partition_count() + exec.left.output_partitioning().partition_count() ); assert!(matches!( &no_keys.input_distribution_requirements().into_per_child()[..], @@ -1762,7 +1815,7 @@ mod tests { total_byte_size: Precision::Exact(128), column_statistics: left_column_statistics.clone(), }); - let right_stats = Arc::new(Statistics::new_unknown(&exec.right().schema())); + let right_stats = Arc::new(Statistics::new_unknown(&exec.right.schema())); let stats = exec .statistics_from_inputs(&[left_stats, right_stats], &StatisticsArgs::new())?; assert_eq!(stats.num_rows, Precision::Exact(7)); From 450bffb5faaf751b2818a2566d53054867d06515 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 11 Aug 2026 17:59:52 +0800 Subject: [PATCH 06/10] Update datafusion/physical-plan/src/joins/asof_join.rs Co-authored-by: Yongting You <2010youy01@gmail.com> --- datafusion/physical-plan/src/joins/asof_join.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index a28f375c97945..a5ae7d4e7551e 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -1526,6 +1526,8 @@ mod tests { Ok(()) } + // Ensure the build-side memory usage equals the sum of all build-side input + // batches, verifying that the build-side buffer is shared. #[tokio::test] async fn shared_right_buffers_are_reserved_once() -> Result<()> { let left_schema = Arc::new(Schema::new(vec![ From 4aa77dc8640850c13df6d825f8910410d49c689e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 11 Aug 2026 19:19:05 +0800 Subject: [PATCH 07/10] refactor: align ASOF physical operator with join conventions --- .../physical-plan/src/joins/asof_join.rs | 762 ++++++++---------- 1 file changed, 317 insertions(+), 445 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index a5ae7d4e7551e..51c1d520cd670 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -54,6 +54,7 @@ //! //! 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 @@ -69,7 +70,7 @@ use std::collections::HashMap; use std::fmt::Formatter; use std::sync::Arc; -use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}; use arrow::buffer::NullBuffer; use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; @@ -77,15 +78,15 @@ use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::normalize_float_zero_scalar; use datafusion_common::{ - ColumnStatistics, JoinType, NullEquality, Result, ScalarValue, Statistics, - assert_eq_or_internal_err, internal_err, plan_err, + 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; +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, @@ -95,12 +96,13 @@ use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::joins::utils::{ - JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, matchable_join_keys, + ColumnIndex, JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, + matchable_join_keys, }; use crate::memory::MemoryStream; use crate::metrics::{ - BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, - MetricCategory, MetricsSet, RecordOutput, Time, + BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, + RecordOutput, Time, }; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; @@ -135,9 +137,12 @@ pub struct AsOfJoinExec { right: Arc, on: JoinOn, match_condition: AsOfMatchExpr, - /// Sorted, unique indices of right columns appended after all left columns. - right_output_indices: Vec, - schema: SchemaRef, + /// Unprojected left-join schema used to interpret `projection`. + join_schema: SchemaRef, + /// Information of index and left/right placement of columns. + column_indices: Vec, + /// Optional indices into the full left-then-right join schema. + projection: Option, metrics: ExecutionPlanMetricsSet, /// Required ordering for each left partition. left_ordering: LexOrdering, @@ -154,25 +159,21 @@ impl AsOfJoinExec { /// 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. - /// Right output indices must be in bounds, sorted, and unique. + /// Projection indices refer to the full left-then-right join schema. pub fn try_new( left: Arc, right: Arc, on: JoinOn, match_condition: AsOfMatchExpr, - right_output_indices: Vec, + projection: Option>, ) -> Result { - validate_asof_join( - left.as_ref(), - right.as_ref(), - &on, - &match_condition, - &right_output_indices, - )?; + validate_asof_join(left.as_ref(), right.as_ref(), &on, &match_condition)?; let left_schema = left.schema(); let right_schema = right.schema(); - let schema = - build_output_schema(&left_schema, &right_schema, &right_output_indices); + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &JoinType::Left); + let join_schema = Arc::new(join_schema); + let projection: Option = projection.map(Into::into); let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); let equality_options = SortOptions { descending: false, @@ -214,15 +215,20 @@ impl AsOfJoinExec { "ASOF right ordering must not be empty" ) })?; - let cache = Arc::new(Self::compute_properties(&left, &schema)?); + let cache = Arc::new(Self::compute_properties( + &left, + &join_schema, + projection.as_deref(), + )?); Ok(Self { left, right, on, match_condition, - right_output_indices, - schema, + join_schema, + column_indices, + projection, metrics: ExecutionPlanMetricsSet::new(), left_ordering, right_ordering, @@ -233,7 +239,8 @@ impl AsOfJoinExec { fn compute_properties( left: &Arc, - schema: &SchemaRef, + join_schema: &SchemaRef, + projection: Option<&[usize]>, ) -> Result { let left_schema = left.schema(); let mapping = ProjectionMapping::try_new( @@ -251,10 +258,19 @@ impl AsOfJoinExec { &left_schema, )?; let input_eq_properties = left.equivalence_properties(); - let eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); - let output_partitioning = left + 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, @@ -264,30 +280,6 @@ impl AsOfJoinExec { } } -fn build_output_schema( - left: &SchemaRef, - right: &SchemaRef, - right_output_indices: &[usize], -) -> SchemaRef { - let full_schema = build_join_schema(left, right, &JoinType::Left).0; - let left_len = left.fields().len(); - let fields = full_schema - .fields() - .iter() - .take(left_len) - .cloned() - .chain( - right_output_indices - .iter() - .map(|index| Arc::clone(&full_schema.fields()[left_len + *index])), - ) - .collect::>(); - Arc::new(Schema::new_with_metadata( - fields, - full_schema.metadata().clone(), - )) -} - impl DisplayAs for AsOfJoinExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { let on = self @@ -304,13 +296,32 @@ impl DisplayAs for AsOfJoinExec { 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::>() + .join(", ") + ) + }) + .unwrap_or_default(); match t { DisplayFormatType::Default | DisplayFormatType::Verbose => write!( f, - "{}: on=[{}], match=[{}]", + "{}: on=[{}], match=[{}]{}", Self::static_name(), on, - match_condition + match_condition, + projection ), DisplayFormatType::TreeRender => { writeln!(f, "on={on}")?; @@ -336,6 +347,9 @@ impl ExecutionPlan for AsOfJoinExec { 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, @@ -368,7 +382,7 @@ impl ExecutionPlan for AsOfJoinExec { Arc::clone(right), self.on.clone(), self.match_condition.clone(), - self.right_output_indices.clone(), + self.projection.as_deref().map(<[usize]>::to_vec), )?)), _ => internal_err!("AsOfJoinExec requires two children"), } @@ -390,8 +404,9 @@ impl ExecutionPlan for AsOfJoinExec { right, on: self.on.clone(), match_condition: self.match_condition.clone(), - right_output_indices: self.right_output_indices.clone(), - schema: Arc::clone(&self.schema), + 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(), @@ -425,12 +440,18 @@ impl ExecutionPlan for AsOfJoinExec { )) })?; let (left_keys, right_keys) = self.on.iter().cloned().unzip(); - let output_schema = Arc::clone(&self.schema); + 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 right_output_indices = self.right_output_indices.clone(); + 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; @@ -441,10 +462,12 @@ impl ExecutionPlan for AsOfJoinExec { InputCursor::new(left_stream, left_keys, left_match), InputCursor::new(right_stream, right_keys, right_match), match_op, - right_output_indices, + 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 { @@ -481,17 +504,24 @@ impl ExecutionPlan for AsOfJoinExec { // 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 mut column_statistics = left.column_statistics.clone(); - column_statistics.truncate(self.left.schema().fields().len()); - column_statistics.resize_with( - self.left.schema().fields().len(), - ColumnStatistics::new_unknown, - ); - column_statistics.extend( - self.right_output_indices + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection .iter() - .map(|_| ColumnStatistics::new_unknown()), - ); + .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, @@ -532,8 +562,6 @@ async fn collect_right_input( let batch_size = memory_counter.count_batch(&batch); futures::future::ready(reservation.try_grow(batch_size).map(|_| { metrics.build_mem_used.add(batch_size); - metrics.build_input_batches.add(1); - metrics.build_input_rows.add(batch.num_rows()); batches.push(batch); batches })) @@ -546,6 +574,11 @@ async fn collect_right_input( }) } +/// 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. @@ -674,11 +707,12 @@ impl InputCursor { #[derive(Clone)] struct AsOfJoinMetrics { + /// Standard output-row and elapsed-compute metrics. baseline: BaselineMetrics, - matched_rows: Count, - unmatched_left_rows: Count, - build_input_batches: Count, - build_input_rows: Count, + /// 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, } @@ -686,18 +720,6 @@ impl AsOfJoinMetrics { fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { Self { baseline: BaselineMetrics::new(metrics, partition), - matched_rows: MetricBuilder::new(metrics) - .with_category(MetricCategory::Rows) - .counter("matched_rows", partition), - unmatched_left_rows: MetricBuilder::new(metrics) - .with_category(MetricCategory::Rows) - .counter("unmatched_left_rows", partition), - build_input_batches: MetricBuilder::new(metrics) - .with_category(MetricCategory::Rows) - .counter("build_input_batches", partition), - build_input_rows: MetricBuilder::new(metrics) - .with_category(MetricCategory::Rows) - .counter("build_input_rows", partition), build_mem_used: MetricBuilder::new(metrics) .peak_memory_usage("build_mem_used", partition), } @@ -811,8 +833,10 @@ struct AsOfJoinStreamState { right: InputCursor, /// Validated ordered match operator. op: Operator, - /// Right columns appended to each output row. - right_output_indices: Vec, + /// Projected output columns and their input sides. + column_indices: Vec, + /// Whether any projected column needs a right row reference. + projects_right: bool, /// Nearest eligible right row for the current equality group. candidate: Option, /// Equality-key ordering shared by the comparator caches. @@ -836,7 +860,7 @@ impl AsOfJoinStreamState { left: InputCursor, right: InputCursor, op: Operator, - right_output_indices: Vec, + column_indices: Vec, batch_size: usize, metrics: AsOfJoinMetrics, ) -> Self { @@ -854,7 +878,10 @@ impl AsOfJoinStreamState { left, right, op, - right_output_indices, + projects_right: column_indices + .iter() + .any(|column| column.side == JoinSide::Right), + column_indices, candidate: None, group_sort_options, input_group_comparator: None, @@ -936,6 +963,17 @@ impl AsOfJoinStreamState { /// rows in the same group; left EOF flushes the final pending rows. NULL keys /// and group changes clear the candidate, while output flushes only clear /// pending row references. + /// + /// ```text + /// while the output batch is not full: + /// load the current left row, or flush/finish at left EOF + /// if its match or equality key is NULL, emit it unmatched and advance left + /// clear the candidate if the left equality group changed + /// while the current right row is before the left group or is eligible: + /// remember the nearest eligible row and advance right + /// emit the left row with the candidate (or NULLs), then advance left + /// flush pending rows without resetting either cursor or the candidate + /// ``` async fn next_batch(&mut self) -> Result> { loop { if self.pending_left.len() >= self.batch_size { @@ -1026,16 +1064,14 @@ impl AsOfJoinStreamState { self.pending_left.push(left_batch, left_row); match candidate { Some(candidate) => { - if !self.right_output_indices.is_empty() { + if self.projects_right { self.pending_right.push(candidate.batch, candidate.row); } - self.metrics.matched_rows.add(1); } None => { - if !self.right_output_indices.is_empty() { + if self.projects_right { self.pending_right.push_null(); } - self.metrics.unmatched_left_rows.add(1); } } Ok(()) @@ -1045,23 +1081,26 @@ impl AsOfJoinStreamState { /// current equality-group candidate for the next output batch. fn flush(&mut self) -> Result { let _timer = self.metrics.baseline.elapsed_compute().timer(); - let left_len = self.schema.fields().len() - self.right_output_indices.len(); + let row_count = self.pending_left.len(); let mut arrays = Vec::with_capacity(self.schema.fields().len()); - for index in 0..left_len { - arrays.push( - self.pending_left - .materialize_column(index, self.schema.field(index).data_type())?, - ); - } - for (offset, source_index) in self.right_output_indices.iter().enumerate() { - arrays.push(self.pending_right.materialize_column( - *source_index, - self.schema.field(left_len + offset).data_type(), - )?); + for (field, column) in self.schema.fields().iter().zip(&self.column_indices) { + let pending = match column.side { + JoinSide::Left => &self.pending_left, + JoinSide::Right => &self.pending_right, + JoinSide::None => { + return internal_err!("ASOF projection cannot contain a mark column"); + } + }; + arrays.push(pending.materialize_column(column.index, field.data_type())?); } self.pending_left.clear(); self.pending_right.clear(); - let batch = RecordBatch::try_new(Arc::clone(&self.schema), arrays)?; + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&self.schema), + arrays, + &options, + )?; (&batch).record_output(&self.metrics.baseline); Ok(batch) } @@ -1073,7 +1112,6 @@ fn validate_asof_join( right: &dyn ExecutionPlan, on: &JoinOn, match_condition: &AsOfMatchExpr, - right_output_indices: &[usize], ) -> Result<()> { if !matches!( match_condition.op, @@ -1124,23 +1162,6 @@ fn validate_asof_join( "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" ); } - if let Some(index) = right_output_indices - .iter() - .find(|index| **index >= right_schema.fields().len()) - { - return plan_err!( - "AsOfJoinExec right output index {index} is outside schema with {} fields", - right_schema.fields().len() - ); - } - if !right_output_indices - .windows(2) - .all(|pair| pair[0] < pair[1]) - { - return plan_err!( - "AsOfJoinExec right output indices must be strictly increasing" - ); - } Ok(()) } @@ -1181,23 +1202,21 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< #[cfg(test)] mod tests { - // These tests cover physical-only contracts that SQL logic tests cannot - // observe, including batch-boundary state, shared build memory, Arrow type - // preservation, and execution properties. End-to-end SQL semantics live in - // the dependent SQL layer. + // Keep physical tests focused on basic executor results, batch-boundary + // state, shared build memory, and constructor/statistics contracts. use super::*; + use crate::collect; use crate::test::TestMemoryExec; - use crate::{collect, collect_partitioned}; - use arrow::array::{ - DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, - }; - use arrow::datatypes::{DataType, Field, Int8Type}; + use arrow::array::{Int32Array, Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::ColumnarValue; - use datafusion_physical_expr_common::metrics::MetricValue; + use datafusion_physical_expr::expressions::BinaryExpr; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use insta::assert_snapshot; #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct VolatileExpr; @@ -1334,27 +1353,12 @@ mod tests { Operator::GtEq, Arc::new(PhysicalColumn::new("ts", 1)), ), - vec![2], + Some(vec![0, 1, 2, 5]), )?)) } - #[test] - fn eligibility_matches_public_semantics() -> Result<()> { - let left = ScalarValue::Int64(Some(10)); - let lower = ScalarValue::Int64(Some(9)); - let equal = ScalarValue::Int64(Some(10)); - let higher = ScalarValue::Int64(Some(11)); - assert!(is_eligible(Operator::Gt, &left, &lower)?); - assert!(!is_eligible(Operator::Gt, &left, &equal)?); - assert!(is_eligible(Operator::GtEq, &left, &equal)?); - assert!(is_eligible(Operator::Lt, &left, &higher)?); - assert!(!is_eligible(Operator::Lt, &left, &equal)?); - assert!(is_eligible(Operator::LtEq, &left, &equal)?); - Ok(()) - } - #[tokio::test] - async fn state_survives_empty_input_batches_and_output_flushes() -> Result<()> { + async fn simple_query() -> Result<()> { let exec = test_exec()?; let context = Arc::new( TaskContext::default() @@ -1368,161 +1372,186 @@ mod tests { .collect::>(), vec![2, 2, 2, 1] ); - let ids = batches - .iter() - .flat_map(|batch| { - batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - }) - .collect::>(); - let prices = batches - .iter() - .flat_map(|batch| { - batch - .column(3) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - }) - .collect::>(); - assert_eq!( - ids, - vec![ - Some(0), - Some(1), - Some(2), - Some(3), - Some(4), - Some(5), - Some(6), - ] - ); - assert_eq!( - prices, - vec![None, None, None, Some(40), Some(60), Some(101), None] - ); + assert_snapshot!(batches_to_sort_string(&batches), @r" + +-----+----+----+-------+ + | key | ts | id | price | + +-----+----+----+-------+ + | | 3 | 0 | | + | A | | 1 | | + | A | 1 | 2 | | + | A | 4 | 3 | 40 | + | A | 7 | 4 | 60 | + | B | 2 | 5 | 101 | + | C | 3 | 6 | | + +-----+----+----+-------+ + "); let metrics = exec.metrics().expect("ASOF metrics must be present"); assert_eq!(metrics.output_rows(), Some(7)); - assert_eq!( - metrics - .sum_by_name("matched_rows") - .map(|value| value.as_usize()), - Some(3) - ); - assert_eq!( - metrics - .sum_by_name("unmatched_left_rows") - .map(|value| value.as_usize()), - Some(4) - ); assert!(metrics.elapsed_compute().is_some()); - assert!( - metrics.iter().any(|metric| { - matches!(metric.value(), MetricValue::ElapsedCompute(_)) - }) - ); Ok(()) } + fn exec_without_equality_keys( + left_times: Vec, + right_times: Vec, + op: Operator, + ) -> Result> { + let left_batch = RecordBatch::try_from_iter(vec![ + ( + "ts", + Arc::new(Int64Array::from(left_times.clone())) as ArrayRef, + ), + ( + "id", + Arc::new(Int32Array::from( + left_times + .into_iter() + .map(|value| value as i32) + .collect::>(), + )) as ArrayRef, + ), + ])?; + let left_schema = left_batch.schema(); + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_batch = RecordBatch::try_from_iter(vec![ + ( + "ts", + Arc::new(Int64Array::from(right_times.clone())) as ArrayRef, + ), + ( + "price", + Arc::new(Int32Array::from( + right_times + .into_iter() + .map(|value| value as i32 * 10) + .collect::>(), + )) as ArrayRef, + ), + ])?; + let right_schema = right_batch.schema(); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 0)), + op, + Arc::new(PhysicalColumn::new("ts", 0)), + ), + Some(vec![1, 3]), + )?)) + } + #[tokio::test] - async fn broadcasts_right_input_to_all_left_partitions() -> Result<()> { - let left_schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("ts", DataType::Int64, false), - Field::new("id", DataType::Int32, false), - ])); + async fn comparison_directions_without_equality_keys() -> Result<()> { + let predecessor = + exec_without_equality_keys(vec![1, 4, 7], vec![2, 4, 6], Operator::GtEq)?; + let predecessor = collect(predecessor, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&predecessor), @r" + +----+-------+ + | id | price | + +----+-------+ + | 1 | | + | 4 | 40 | + | 7 | 60 | + +----+-------+ + "); + + let successor = + exec_without_equality_keys(vec![7, 4, 1], vec![6, 4, 2], Operator::Lt)?; + let successor = collect(successor, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&successor), @r" + +----+-------+ + | id | price | + +----+-------+ + | 1 | 20 | + | 4 | 60 | + | 7 | | + +----+-------+ + "); + Ok(()) + } + + #[tokio::test] + async fn complex_equality_and_match_expressions() -> Result<()> { + let left_batch = RecordBatch::try_from_iter(vec![ + ("g1", Arc::new(Int64Array::from(vec![0, 1])) as ArrayRef), + ("g2", Arc::new(Int64Array::from(vec![1, 1])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![4, 4])) as ArrayRef), + ("offset", Arc::new(Int64Array::from(vec![1, 0])) as ArrayRef), + ("id", Arc::new(Int64Array::from(vec![10, 20])) as ArrayRef), + ])?; + let left_schema = left_batch.schema(); let left = TestMemoryExec::try_new_exec( - &[ - vec![make_batch( - &left_schema, - vec![Some("A"), Some("A")], - vec![Some(1), Some(4)], - vec![0, 1], - )?], - vec![make_batch( - &left_schema, - vec![Some("A"), Some("A")], - vec![Some(2), Some(5)], - vec![2, 3], - )?], - ], + &[vec![left_batch]], Arc::clone(&left_schema), None, )?; - let right_schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("ts", DataType::Int64, false), - Field::new("price", DataType::Int32, false), - ])); + + let right_batch = RecordBatch::try_from_iter(vec![ + ("g1", Arc::new(Int64Array::from(vec![0, 0, 1])) as ArrayRef), + ("g2", Arc::new(Int64Array::from(vec![1, 1, 1])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![2, 5, 3])) as ArrayRef), + ( + "price", + Arc::new(Int64Array::from(vec![12, 15, 23])) as ArrayRef, + ), + ])?; + let right_schema = right_batch.schema(); let right = TestMemoryExec::try_new_exec( - &[vec![ - make_batch(&right_schema, vec![Some("A")], vec![Some(1)], vec![10])?, - make_batch(&right_schema, vec![Some("A")], vec![Some(3)], vec![30])?, - ]], + &[vec![right_batch]], Arc::clone(&right_schema), None, )?; + + let left_group = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("g1", 0)), + Operator::Plus, + Arc::new(PhysicalColumn::new("g2", 1)), + )); + let right_group = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("g1", 0)), + Operator::Plus, + Arc::new(PhysicalColumn::new("g2", 1)), + )); + let left_match = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("ts", 2)), + Operator::Plus, + Arc::new(PhysicalColumn::new("offset", 3)), + )); let exec = Arc::new(AsOfJoinExec::try_new( left, right, - vec![], + vec![(left_group, right_group)], AsOfMatchExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), + left_match, Operator::GtEq, - Arc::new(PhysicalColumn::new("ts", 1)), + Arc::new(PhysicalColumn::new("ts", 2)), ), - vec![2], + Some(vec![4, 8]), )?); - assert_eq!(exec.properties().output_partitioning().partition_count(), 2); - assert!(matches!( - &exec.input_distribution_requirements().into_per_child()[..], - [ - Distribution::UnspecifiedDistribution, - Distribution::SinglePartition - ] - )); - let partitions = collect_partitioned( - Arc::clone(&exec) as Arc, - Arc::new(TaskContext::default()), - ) - .await?; - assert_eq!(partitions.len(), 2); - for batches in partitions { - let prices = batches - .iter() - .flat_map(|batch| { - batch - .column(3) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - }) - .collect::>(); - assert_eq!(prices, vec![Some(10), Some(30)]); - } - - let metrics = exec.metrics().expect("ASOF metrics must be present"); - assert_eq!( - metrics - .sum_by_name("build_input_batches") - .map(|value| value.as_usize()), - Some(2) - ); - assert_eq!( - metrics - .sum_by_name("build_input_rows") - .map(|value| value.as_usize()), - Some(2) - ); - assert_eq!(metrics.output_rows(), Some(4)); + let batches = collect(exec, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+ + | id | price | + +----+-------+ + | 10 | 15 | + | 20 | 23 | + +----+-------+ + "); Ok(()) } @@ -1581,7 +1610,7 @@ mod tests { Operator::GtEq, Arc::new(PhysicalColumn::new("ts", 1)), ), - vec![2], + Some(vec![0, 1, 2, 5]), )?); let runtime = RuntimeEnvBuilder::new() .with_memory_limit(retained_size, 1.0) @@ -1612,99 +1641,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { - let dictionary_type = - DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); - let left_schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("ts", DataType::Int64, false), - Field::new("payload", dictionary_type.clone(), false), - ])); - let mut left_payload = StringDictionaryBuilder::::new(); - for _ in 0..129 { - left_payload.append_value("left"); - } - let left_batch = RecordBatch::try_new( - Arc::clone(&left_schema), - vec![ - Arc::new(StringArray::from(vec!["A"; 129])), - Arc::new(Int64Array::from_iter_values(-1..128)), - Arc::new(left_payload.finish()), - ], - )?; - let left = TestMemoryExec::try_new_exec( - &[vec![left_batch]], - Arc::clone(&left_schema), - None, - )?; - - let right_schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Utf8, false), - Field::new("ts", DataType::Int64, false), - Field::new("payload", dictionary_type.clone(), false), - ])); - let mut right_payload = StringDictionaryBuilder::::new(); - right_payload.append_value("right"); - let right_batch = RecordBatch::try_new( - Arc::clone(&right_schema), - vec![ - Arc::new(StringArray::from(vec!["A"])), - Arc::new(Int64Array::from(vec![0])), - Arc::new(right_payload.finish()), - ], - )?; - let right = TestMemoryExec::try_new_exec( - &[vec![right_batch]], - Arc::clone(&right_schema), - None, - )?; - - let exec = Arc::new(AsOfJoinExec::try_new( - left, - right, - vec![( - Arc::new(PhysicalColumn::new("key", 0)), - Arc::new(PhysicalColumn::new("key", 0)), - )], - AsOfMatchExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), - Operator::GtEq, - Arc::new(PhysicalColumn::new("ts", 1)), - ), - vec![2], - )?); - let context = Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().with_batch_size(256)), - ); - let batches = collect(exec, context).await?; - assert_eq!(batches.len(), 1); - assert_eq!(batches[0].num_rows(), 129); - assert_eq!(batches[0].column(2).data_type(), &dictionary_type); - assert_eq!(batches[0].column(3).data_type(), &dictionary_type); - - let right_output = batches[0] - .column(3) - .as_any() - .downcast_ref::>() - .expect("right output must remain Dictionary(Int8, Utf8)"); - assert!(right_output.is_null(0)); - assert_eq!(right_output.null_count(), 1); - let values = right_output - .values() - .as_any() - .downcast_ref::() - .expect("dictionary values must be Utf8"); - for row in 1..129 { - assert_eq!( - values.value(right_output.keys().value(row) as usize), - "right" - ); - } - Ok(()) - } - #[test] fn rejects_volatile_physical_expressions() -> Result<()> { let exec = test_exec()?; @@ -1718,7 +1654,7 @@ mod tests { Operator::GtEq, Arc::new(PhysicalColumn::new("ts", 1)), ), - vec![2], + Some(vec![0, 1, 2, 5]), ) .expect_err("volatile match expression must be rejected"); assert!(match_error.to_string().contains("must be deterministic")); @@ -1728,7 +1664,7 @@ mod tests { Arc::clone(&exec.right), vec![(volatile, Arc::new(PhysicalColumn::new("key", 0)))], exec.match_condition.clone(), - vec![2], + Some(vec![0, 1, 2, 5]), ) .expect_err("volatile equality expression must be rejected"); assert!(equality_error.to_string().contains("must be deterministic")); @@ -1736,72 +1672,8 @@ mod tests { } #[test] - fn properties_and_statistics_follow_left_preserving_contract() -> Result<()> { + fn statistics_follow_left_preserving_contract() -> Result<()> { let exec = test_exec()?; - let exec_plan: Arc = Arc::clone(&exec) as _; - assert_eq!(exec.maintains_input_order(), vec![false, false]); - assert_eq!(exec_plan.pipeline_behavior(), EmissionType::Incremental); - assert_eq!(exec_plan.boundedness(), Boundedness::Bounded); - assert!(matches!( - &exec.input_distribution_requirements().into_per_child()[..], - [ - Distribution::UnspecifiedDistribution, - Distribution::SinglePartition - ] - )); - for ordering in exec.required_input_ordering() { - let requirement = ordering.expect("ASOF ordering is required").into_single(); - assert_eq!(requirement.len(), 2); - assert_eq!( - requirement[0].options, - Some(SortOptions { - descending: false, - nulls_first: true, - }) - ); - assert_eq!( - requirement[1].options, - Some(SortOptions { - descending: false, - nulls_first: true, - }) - ); - } - - let no_keys: Arc = Arc::new(AsOfJoinExec::try_new( - Arc::clone(&exec.left), - Arc::clone(&exec.right), - vec![], - AsOfMatchExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), - Operator::Lt, - Arc::new(PhysicalColumn::new("ts", 1)), - ), - vec![2], - )?); - assert_eq!( - no_keys.output_partitioning().partition_count(), - exec.left.output_partitioning().partition_count() - ); - assert!(matches!( - &no_keys.input_distribution_requirements().into_per_child()[..], - [ - Distribution::UnspecifiedDistribution, - Distribution::SinglePartition - ] - )); - for ordering in no_keys.required_input_ordering() { - let requirement = ordering.expect("ASOF ordering is required").into_single(); - assert_eq!(requirement.len(), 1); - assert_eq!( - requirement[0].options, - Some(SortOptions { - descending: true, - nulls_first: true, - }) - ); - } - let mut key_stats = ColumnStatistics::new_unknown(); key_stats.null_count = Precision::Exact(1); key_stats.distinct_count = Precision::Exact(4); From 00ae56b6ca74bdee1cc311c7692c52434a9328a9 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 11 Aug 2026 23:38:35 +0800 Subject: [PATCH 08/10] fix: expose ASOF physical expressions --- datafusion/physical-plan/src/joins/asof_join.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 51c1d520cd670..fbb9fd9d77572 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -75,6 +75,7 @@ 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::{ @@ -371,6 +372,17 @@ impl ExecutionPlan for AsOfJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + 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 with_new_children( self: Arc, children: Vec>, From 52484024d0e089466e5e395bebbe846bac1cb824 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 15 Aug 2026 00:04:33 +0800 Subject: [PATCH 09/10] fix: reject floating ASOF equality keys --- .../physical-plan/src/joins/asof_join.rs | 120 ++++++++++++------ 1 file changed, 84 insertions(+), 36 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index fbb9fd9d77572..c7c068e494465 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -108,9 +108,9 @@ use crate::metrics::{ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, }; /// Physical ordered comparison for an ASOF join. @@ -159,8 +159,10 @@ impl AsOfJoinExec { /// /// 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. - /// Projection indices refer to the full left-then-right join schema. + /// 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, right: Arc, @@ -383,48 +385,57 @@ impl ExecutionPlan for AsOfJoinExec { ) } - fn with_new_children( + fn replace_children( self: Arc, - children: Vec>, + mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(Self::try_new( - Arc::clone(left), - Arc::clone(right), + 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), )?)), - _ => internal_err!("AsOfJoinExec requires two children"), } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 2, - "AsOfJoinExec requires two children" - ); - let left = children.remove(0); - let right = children.remove(0); - 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), - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -1166,6 +1177,11 @@ fn validate_asof_join( "AsOfJoinExec equality expressions have unsupported hash type {left_type}" ); } + if left_type.is_floating() { + return plan_err!( + "AsOfJoinExec equality expressions do not support floating-point type {left_type}" + ); + } } let left_match_type = match_condition.left.data_type(&left_schema)?; let right_match_type = match_condition.right.data_type(&right_schema)?; @@ -1226,7 +1242,7 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::ColumnarValue; - use datafusion_physical_expr::expressions::BinaryExpr; + use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use insta::assert_snapshot; @@ -1683,6 +1699,38 @@ mod tests { Ok(()) } + #[test] + fn rejects_floating_equality_expressions() -> Result<()> { + let exec = test_exec()?; + for data_type in [DataType::Float16, DataType::Float32, DataType::Float64] { + let left = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + data_type.clone(), + None, + )); + let right = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + data_type.clone(), + None, + )); + let error = AsOfJoinExec::try_new( + Arc::clone(&exec.left), + Arc::clone(&exec.right), + vec![(left, right)], + exec.match_condition.clone(), + Some(vec![0, 1, 2, 5]), + ) + .expect_err("floating equality expressions must be rejected"); + assert!( + error.to_string().contains(&format!( + "equality expressions do not support floating-point type {data_type}" + )), + "unexpected error: {error}" + ); + } + Ok(()) + } + #[test] fn statistics_follow_left_preserving_contract() -> Result<()> { let exec = test_exec()?; From 5201c34ae2cefb382f4b7e877ed46f7283a52e1c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 16 Aug 2026 23:47:49 +0800 Subject: [PATCH 10/10] refactor: make ASOF join state a stream --- .../physical-plan/src/joins/asof_join.rs | 165 ++++++++++-------- 1 file changed, 88 insertions(+), 77 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index c7c068e494465..206520f901c26 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -54,7 +54,7 @@ //! //! 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`. +//! The key state-machine entry point is [`AsOfJoinStream::poll_next_impl`]. //! //! 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 @@ -68,7 +68,9 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::fmt::Formatter; +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}; use arrow::buffer::NullBuffer; @@ -93,7 +95,7 @@ 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 futures::{Stream, StreamExt, TryStreamExt, future::poll_fn, ready, stream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::joins::utils::{ @@ -110,7 +112,8 @@ use crate::stream::RecordBatchStreamAdapter; use crate::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, - ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + validate_child_count, }; /// Physical ordered comparison for an ASOF join. @@ -480,7 +483,7 @@ impl ExecutionPlan for AsOfJoinExec { 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( + let stream = AsOfJoinStream::new( Arc::clone(&stream_schema), InputCursor::new(left_stream, left_keys, left_match), InputCursor::new(right_stream, right_keys, right_match), @@ -488,20 +491,10 @@ impl ExecutionPlan for AsOfJoinExec { 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), - } - }, + right_input, ); Ok::(Box::pin( - RecordBatchStreamAdapter::new(stream_schema, stream), + stream, )) }) .try_flatten(); @@ -661,12 +654,16 @@ impl InputCursor { } } - async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result { + fn poll_ensure_row( + &mut self, + cx: &mut Context<'_>, + elapsed_compute: &Time, + ) -> Poll> { loop { if let Some(batch) = &self.batch && self.row < batch.num_rows() { - return Ok(true); + return Poll::Ready(Ok(true)); } self.batch = None; self.key_arrays = Arc::from([]); @@ -674,11 +671,11 @@ impl InputCursor { self.match_array = None; self.row = 0; if self.eof { - return Ok(false); + return Poll::Ready(Ok(false)); } - let Some(batch) = self.stream.next().await.transpose()? else { + let Some(batch) = ready!(self.stream.poll_next_unpin(cx)).transpose()? else { self.eof = true; - return Ok(false); + return Poll::Ready(Ok(false)); }; if batch.num_rows() == 0 { continue; @@ -756,7 +753,8 @@ impl AsOfJoinMetrics { /// the first source batch, a NULL, and row 0 from the second source batch. #[derive(Default)] struct PendingRows { - /// Distinct source batches referenced by `indices`. + /// Distinct source batches referenced by `indices`. `Arc` keeps per-row + /// clones O(1) and provides stable identity for deduplication. sources: Vec>, /// Maps an `Arc` pointer to its index in `sources`. source_by_ptr: HashMap, @@ -847,13 +845,15 @@ impl PendingRows { /// candidate advances from `(A, 2)` to `(A, 6)` without rewinding the right /// cursor. Cursors and the candidate survive input batch changes and output /// flushes; a change of equality group clears the candidate before reuse. -struct AsOfJoinStreamState { +struct AsOfJoinStream { /// Output schema used when pending row references are materialized. schema: SchemaRef, /// Cursor over the current left partition. left: InputCursor, /// Independent cursor over the shared, ordered right input. right: InputCursor, + /// Retains the shared right batches and their memory reservation. + _right_input: Arc, /// Validated ordered match operator. op: Operator, /// Projected output columns and their input sides. @@ -877,7 +877,8 @@ struct AsOfJoinStreamState { metrics: AsOfJoinMetrics, } -impl AsOfJoinStreamState { +impl AsOfJoinStream { + #[expect(clippy::too_many_arguments)] fn new( schema: SchemaRef, left: InputCursor, @@ -886,6 +887,7 @@ impl AsOfJoinStreamState { column_indices: Vec, batch_size: usize, metrics: AsOfJoinMetrics, + right_input: Arc, ) -> Self { let group_sort_options = vec![ SortOptions { @@ -900,6 +902,7 @@ impl AsOfJoinStreamState { schema, left, right, + _right_input: right_input, op, projects_right: column_indices .iter() @@ -997,21 +1000,23 @@ impl AsOfJoinStreamState { /// emit the left row with the candidate (or NULLs), then advance left /// flush pending rows without resetting either cursor or the candidate /// ``` - async fn next_batch(&mut self) -> Result> { + fn poll_next_impl( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { loop { if self.pending_left.len() >= self.batch_size { - return self.flush().map(Some); + return Poll::Ready(Some(self.flush())); } - if !self - .left - .ensure_row(self.metrics.baseline.elapsed_compute()) - .await? - { + if !ready!( + self.left + .poll_ensure_row(cx, self.metrics.baseline.elapsed_compute()) + )? { if !self.pending_left.is_empty() { - return self.flush().map(Some); + return Poll::Ready(Some(self.flush())); } self.metrics.baseline.done(); - return Ok(None); + return Poll::Ready(None); } let left_match = { @@ -1031,49 +1036,44 @@ impl AsOfJoinStreamState { } loop { - if !self - .right - .ensure_row(self.metrics.baseline.elapsed_compute()) - .await? - { + if !ready!( + self.right + .poll_ensure_row(cx, self.metrics.baseline.elapsed_compute()) + )? { break; } - let action = if self.right.group_has_null() { - RightAction::Advance - } else { - match self.compare_input_groups()? { - Ordering::Less => RightAction::Advance, - Ordering::Greater => RightAction::Stop, - Ordering::Equal => { - let _timer = self.metrics.baseline.elapsed_compute().timer(); - let right_match = self.right.match_value()?; - if right_match.is_null() { - RightAction::Advance - } else if is_eligible(self.op, &left_match, &right_match)? { - let (batch, row) = self.right.batch_row()?; - RightAction::Candidate(Candidate { - batch, - row, - key_arrays: Arc::clone(&self.right.key_arrays), - key_batch_id: self.right.key_batch_id, - }) - } else { - RightAction::Stop - } - } - } - }; - match action { - RightAction::Advance => self.right.advance(), - RightAction::Candidate(candidate) => { - // Replacing the candidate selects the nearest eligible row. - // Equal match values have no secondary ordering, so which - // tied row wins is intentionally nondeterministic. - self.candidate = Some(candidate); + if self.right.group_has_null() { + self.right.advance(); + continue; + } + match self.compare_input_groups()? { + Ordering::Less => { self.right.advance(); + continue; } - RightAction::Stop => break, + Ordering::Greater => break, + Ordering::Equal => {} + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_match = self.right.match_value()?; + if right_match.is_null() { + self.right.advance(); + continue; + } + if !is_eligible(self.op, &left_match, &right_match)? { + break; } + let (batch, row) = self.right.batch_row()?; + // Replacing the candidate selects the nearest eligible row. + // Equal match values have no secondary ordering, so which tied + // row wins is intentionally nondeterministic. + self.candidate = Some(Candidate { + batch, + row, + key_arrays: Arc::clone(&self.right.key_arrays), + key_batch_id: self.right.key_batch_id, + }); + self.right.advance(); } self.push_current_left(self.candidate.clone())?; @@ -1129,6 +1129,23 @@ impl AsOfJoinStreamState { } } +impl RecordBatchStream for AsOfJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for AsOfJoinStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + self.poll_next_impl(cx) + } +} + /// Validates all invariants required by the forward-only ASOF state machine. fn validate_asof_join( left: &dyn ExecutionPlan, @@ -1211,12 +1228,6 @@ fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Re Ok(()) } -enum RightAction { - Advance, - Candidate(Candidate), - Stop, -} - fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result { let ordering = right.try_cmp(left)?; Ok(match op {