From e546de5510de599c91f756bef4bf45243d93011a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:05:01 +0800 Subject: [PATCH 01/20] 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 bd58e1debfe4622e93cb9fb52d08c35abad06799 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:28:02 +0800 Subject: [PATCH 02/20] feat: add ASOF join logical semantics --- datafusion/core/src/physical_planner.rs | 69 ++++- .../physical_optimizer/filter_pushdown.rs | 42 +++ datafusion/expr/src/logical_plan/builder.rs | 107 ++++++- datafusion/expr/src/logical_plan/display.rs | 23 +- datafusion/expr/src/logical_plan/mod.rs | 16 +- datafusion/expr/src/logical_plan/plan.rs | 279 +++++++++++++++++- datafusion/expr/src/logical_plan/tree_node.rs | 57 +++- .../optimizer/src/analyzer/type_coercion.rs | 37 ++- .../optimizer/src/common_subexpr_eliminate.rs | 1 + .../optimizer/src/optimize_projections/mod.rs | 43 ++- datafusion/optimizer/src/optimizer.rs | 5 + datafusion/optimizer/src/push_down_filter.rs | 29 ++ datafusion/proto/src/logical_plan/mod.rs | 3 + datafusion/sql/src/unparser/plan.rs | 1 + .../src/logical_plan/producer/rel/mod.rs | 3 + datafusion/substrait/tests/cases/serialize.rs | 27 +- 16 files changed, 701 insertions(+), 41 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aef8036c749a8..58b4da0f91a94 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -43,7 +43,8 @@ use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::filter::FilterExecBuilder; use crate::physical_plan::joins::utils as join_utils; use crate::physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -93,8 +94,8 @@ use datafusion_expr::physical_planning_context::{ use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, - FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, - WindowFrame, WindowFrameBound, WriteOp, + FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType, + StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, @@ -1860,6 +1861,67 @@ impl DefaultPhysicalPlanner { join } } + LogicalPlan::AsOfJoin(join) => { + let [physical_left, physical_right] = children.two()?; + let join_on = join + .on + .iter() + .map(|(left, right)| { + Ok(( + create_physical_expr( + left, + join.left.schema(), + execution_props, + planning_ctx, + )?, + create_physical_expr( + right, + join.right.schema(), + execution_props, + planning_ctx, + )?, + )) + }) + .collect::>()?; + let match_condition = AsOfMatchExpr::new( + create_physical_expr( + &join.match_condition.left, + join.left.schema(), + execution_props, + planning_ctx, + )?, + join.match_condition.op, + create_physical_expr( + &join.match_condition.right, + join.right.schema(), + execution_props, + planning_ctx, + )?, + ); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect(); + Arc::new(AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + match_condition, + right_output_indices, + )?) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, @@ -2359,6 +2421,7 @@ fn extract_dml_filters( | LogicalPlan::Sort(_) | LogicalPlan::Union(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Window(_) diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 909b80cadaae3..2f613264f45d9 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -55,6 +55,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{AsOfJoinExec, AsOfMatchExpr}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -123,6 +124,47 @@ fn test_pushdown_volatile_functions_not_allowed() { ); } +#[test] +fn test_asof_join_pushes_only_left_filters() { + let left = TestScanBuilder::new(schema()).with_support(true).build(); + let right = TestScanBuilder::new(schema()).with_support(true).build(); + let join = Arc::new( + AsOfJoinExec::try_new( + left, + right, + vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())], + AsOfMatchExpr::new( + col("b", &schema()).unwrap(), + Operator::GtEq, + col("b", &schema()).unwrap(), + ), + vec![0, 1, 2], + ) + .unwrap(), + ); + let left_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let right_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 5)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter)); + let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap()); + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); + let formatted = format_plan_for_test(&optimized); + + assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}"); + assert!(formatted.contains("AsOfJoinExec:"), "{formatted}"); + assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}"); + assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}"); +} + /// Show that we can use config options to determine how to do pushdown. #[test] fn test_pushdown_into_scan_with_config_options() { diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 2ecb12c30afad..d5dda08bf6a45 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -31,10 +31,10 @@ use crate::expr_rewriter::{ rewrite_sort_cols_by_aggs, }; use crate::logical_plan::{ - Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare, - Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest, - Values, Window, + Aggregate, Analyze, AsOfJoin, AsOfMatch, Distinct, DistinctOn, EmptyRelation, + Explain, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, + PlanType, Prepare, Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, + Union, Unnest, Values, Window, }; use crate::select_expr::SelectExpr; use crate::utils::{ @@ -1007,6 +1007,68 @@ impl LogicalPlanBuilder { ) } + /// Apply a left-preserving ASOF join using equality expressions and one + /// ordered match condition. + pub fn asof_join( + self, + right: LogicalPlan, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + ) -> Result { + self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On) + } + + /// Apply a left-preserving ASOF join using `USING` equality keys. + pub fn asof_join_using( + self, + right: LogicalPlan, + using_keys: Vec, + match_condition: AsOfMatch, + ) -> Result { + let on = using_keys + .into_iter() + .map(|key| { + let left = Self::normalize(&self.plan, key.clone())?; + let right = Self::normalize(&right, key)?; + Ok((Expr::Column(left), Expr::Column(right))) + }) + .collect::>()?; + self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using) + } + + fn asof_join_with_constraint( + self, + right: LogicalPlan, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + join_constraint: JoinConstraint, + ) -> Result { + let normalize = |expr, schema: &DFSchema| { + normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[]) + }; + let on = on + .into_iter() + .map(|(left, right_expr)| { + Ok(( + normalize(left, self.plan.schema())?, + normalize(right_expr, right.schema())?, + )) + }) + .collect::>()?; + let match_condition = AsOfMatch { + left: normalize(match_condition.left, self.plan.schema())?, + op: match_condition.op, + right: normalize(match_condition.right, right.schema())?, + }; + Ok(Self::new(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + self.plan, + Arc::new(right), + on, + match_condition, + join_constraint, + )?))) + } + pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result { if column.relation.is_some() { // column is already normalized @@ -1776,6 +1838,43 @@ pub fn build_join_schema( dfschema.with_functional_dependencies(func_dependencies) } +/// Creates the schema for a left-preserving ASOF join. +/// +/// `ON` emits all left fields followed by nullable right fields. `USING` emits +/// each equality key once by omitting the corresponding right field. +pub fn build_asof_join_schema( + left: &DFSchema, + right: &DFSchema, + on: &[(Expr, Expr)], + join_constraint: JoinConstraint, +) -> Result { + let omitted_right_indices = if join_constraint == JoinConstraint::Using { + on.iter() + .map(|(_, right_expr)| { + let column = right_expr.get_as_join_column().ok_or_else(|| { + plan_datafusion_err!("ASOF USING keys must be columns") + })?; + right.index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + + let full_schema = build_join_schema(left, right, &JoinType::Left)?; + let left_len = left.fields().len(); + let fields = full_schema + .iter() + .enumerate() + .filter(|(index, _)| { + *index < left_len || !omitted_right_indices.contains(&(*index - left_len)) + }) + .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field))) + .collect(); + DFSchema::new_with_metadata(fields, full_schema.metadata().clone())? + .with_functional_dependencies(left.functional_dependencies().clone()) +} + /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise /// conflict with the columns from the other. /// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 09f41c94f64fa..c5cd003d1f34e 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -21,10 +21,10 @@ use std::collections::HashMap; use std::fmt; use crate::{ - Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, - Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, - Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values, - Window, expr_vec_fmt, + Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, + Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, + Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, + Values, Window, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -493,6 +493,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Filter": format!("{}", filter_expr) }) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let join_expr: Vec = + on.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + json!({ + "Node Type": "AsOf Join", + "Join Constraint": format!("{join_constraint:?}"), + "Join Keys": join_expr.join(", "), + "Match Condition": match_condition.to_string(), + }) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 4766c3f33379f..98113d12c1b4a 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -28,8 +28,8 @@ pub mod tree_node; pub use builder::{ LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE, - build_join_schema, requalify_sides_if_needed, table_scan, union, - wrap_projection_for_join_if_necessary, + build_asof_join_schema, build_join_schema, requalify_sides_if_needed, table_scan, + union, wrap_projection_for_join_if_necessary, }; pub use ddl::{ CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction, @@ -41,12 +41,12 @@ pub use dml::{ WriteOp, }; pub use plan::{ - Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, - EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, - Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, - Unnest, Values, Window, projection_schema, + Aggregate, Analyze, AsOfJoin, AsOfMatch, ColumnUnnestList, DescribeTable, Distinct, + DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, + Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, + Projection, RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, + StringifiedPlan, Subquery, SubqueryAlias, TableScan, TableScanBuilder, + ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9cfab21a0395e..cac596b8bd58a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,13 +41,15 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement}; use crate::utils::{ - enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, - grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, + enumerate_grouping_sets, expr_to_columns, exprlist_to_fields, + find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, + merge_schema, split_conjunction, }; use crate::{ BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet, LogicalPlanBuilder, Operator, Prepare, TableProviderFilterPushDown, TableSource, - WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed, + WindowFunctionDefinition, build_asof_join_schema, build_join_schema, expr_vec_fmt, + requalify_sides_if_needed, }; use crate::statistics::StatisticsRequest; @@ -238,6 +240,9 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -341,6 +346,7 @@ impl LogicalPlan { LogicalPlan::Aggregate(Aggregate { schema, .. }) => schema, LogicalPlan::Sort(Sort { input, .. }) => input.schema(), LogicalPlan::Join(Join { schema, .. }) => schema, + LogicalPlan::AsOfJoin(AsOfJoin { schema, .. }) => schema, LogicalPlan::Repartition(Repartition { input, .. }) => input.schema(), LogicalPlan::Limit(Limit { input, .. }) => input.schema(), LogicalPlan::Statement(statement) => statement.schema(), @@ -369,7 +375,8 @@ impl LogicalPlan { | LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Unnest(_) - | LogicalPlan::Join(_) => self + | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) => self .inputs() .iter() .map(|input| input.schema().as_ref()) @@ -459,6 +466,9 @@ impl LogicalPlan { LogicalPlan::Aggregate(Aggregate { input, .. }) => vec![input], LogicalPlan::Sort(Sort { input, .. }) => vec![input], LogicalPlan::Join(Join { left, right, .. }) => vec![left, right], + LogicalPlan::AsOfJoin(AsOfJoin { left, right, .. }) => { + vec![left, right] + } LogicalPlan::Limit(Limit { input, .. }) => vec![input], LogicalPlan::Subquery(Subquery { subquery, .. }) => vec![subquery], LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => vec![input], @@ -494,12 +504,20 @@ impl LogicalPlan { let mut using_columns: Vec> = vec![]; self.apply_with_subqueries(|plan| { - if let LogicalPlan::Join(Join { - join_constraint: JoinConstraint::Using, - on, - .. - }) = plan - { + let on = match plan { + LogicalPlan::Join(Join { + join_constraint: JoinConstraint::Using, + on, + .. + }) + | LogicalPlan::AsOfJoin(AsOfJoin { + join_constraint: JoinConstraint::Using, + on, + .. + }) => Some(on), + _ => None, + }; + if let Some(on) = on { // The join keys in using-join must be columns. let columns = on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| { @@ -567,6 +585,7 @@ impl LogicalPlan { right.head_output_expr() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => left.head_output_expr(), LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => { static_term.head_output_expr() } @@ -690,6 +709,26 @@ impl LogicalPlan { null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema: _, + }) => Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + left, + right, + on.into_iter() + .map(|(left, right)| (left.unalias(), right.unalias())) + .collect(), + AsOfMatch { + left: match_condition.left.unalias(), + op: match_condition.op, + right: match_condition.right.unalias(), + }, + join_constraint, + )?)), LogicalPlan::Subquery(_) => Ok(self), LogicalPlan::SubqueryAlias(SubqueryAlias { input, @@ -985,6 +1024,45 @@ impl LogicalPlan { null_aware: *null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let (left, right) = self.only_two_inputs(inputs)?; + let expected = on.len() * 2 + 2; + assert_eq_or_internal_err!( + expected, + expr.len(), + "Invalid number of new ASOF join expressions: expected {}, got {}", + expected, + expr.len() + ); + + let mut iter = expr.into_iter(); + let mut new_on = Vec::with_capacity(on.len()); + for _ in 0..on.len() { + let left = iter.next().expect("expression count checked").unalias(); + let right = iter.next().expect("expression count checked").unalias(); + new_on.push((left, right)); + } + let match_left = iter.next().expect("expression count checked").unalias(); + let match_right = + iter.next().expect("expression count checked").unalias(); + + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + new_on, + AsOfMatch { + left: match_left, + op: match_condition.op, + right: match_right, + }, + *join_constraint, + )?)) + } LogicalPlan::Subquery(Subquery { outer_ref_columns, spans, @@ -1410,6 +1488,7 @@ impl LogicalPlan { right.max_rows() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => left.max_rows(), LogicalPlan::Repartition(Repartition { input, .. }) => input.max_rows(), LogicalPlan::Union(Union { inputs, .. }) => { inputs.iter().try_fold(0usize, |mut acc, plan| { @@ -1460,6 +1539,7 @@ impl LogicalPlan { LogicalPlan::Window(_) => Ok(None), LogicalPlan::Aggregate(_) => Ok(None), LogicalPlan::Join(_) => Ok(None), + LogicalPlan::AsOfJoin(_) => Ok(None), LogicalPlan::Repartition(_) => Ok(None), LogicalPlan::Union(_) => Ok(None), LogicalPlan::EmptyRelation(_) => Ok(None), @@ -1498,6 +1578,7 @@ impl LogicalPlan { LogicalPlan::Window(_) => Ok(None), LogicalPlan::Aggregate(_) => Ok(None), LogicalPlan::Join(_) => Ok(None), + LogicalPlan::AsOfJoin(_) => Ok(None), LogicalPlan::Repartition(_) => Ok(None), LogicalPlan::Union(_) => Ok(None), LogicalPlan::EmptyRelation(_) => Ok(None), @@ -2126,6 +2207,25 @@ impl LogicalPlan { } } } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let equality = on + .iter() + .map(|(left, right)| format!("{left} = {right}")) + .join(", "); + write!( + f, + "AsOf Join: match=[{match_condition}], constraint={join_constraint:?}" + )?; + if !equality.is_empty() { + write!(f, ", on=[{equality}]")?; + } + Ok(()) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. @@ -4237,6 +4337,165 @@ pub struct Join { pub null_aware: bool, } +/// The ordered comparison used by an [`AsOfJoin`]. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct AsOfMatch { + /// Expression evaluated against the left input. + pub left: Expr, + /// One of [`Operator::Lt`], [`Operator::LtEq`], [`Operator::Gt`], or + /// [`Operator::GtEq`]. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: Expr, +} + +impl AsOfMatch { + /// Creates an ordered ASOF match condition. + pub fn new(left: Expr, op: Operator, right: Expr) -> Self { + Self { left, op, right } + } +} + +impl Display for AsOfMatch { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{} {} {}", self.left, self.op, self.right) + } +} + +/// Match each left row with at most one ordered row from the right input. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AsOfJoin { + /// Left input. Every left row is preserved exactly once. + pub left: Arc, + /// Right input. + pub right: Arc, + /// Equality clauses expressed as pairs of left and right expressions. + pub on: Vec<(Expr, Expr)>, + /// Ordered match condition. + pub match_condition: Box, + /// Whether equality keys came from `ON` or `USING`. + pub join_constraint: JoinConstraint, + /// Output schema. + pub schema: DFSchemaRef, +} + +impl AsOfJoin { + /// Creates an ASOF join and validates its logical contract. + pub fn try_new( + left: Arc, + right: Arc, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + join_constraint: JoinConstraint, + ) -> Result { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + + Self::validate_side(&match_condition.left, left.schema(), "left match")?; + Self::validate_side(&match_condition.right, right.schema(), "right match")?; + if match_condition.left.is_volatile() || match_condition.right.is_volatile() { + return plan_err!("ASOF MATCH_CONDITION must be deterministic"); + } + + let left_type = match_condition.left.get_type(left.schema())?; + let right_type = match_condition.right.get_type(right.schema())?; + if crate::type_coercion::binary::comparison_coercion(&left_type, &right_type) + .is_none() + { + return plan_err!( + "ASOF match expressions have incompatible types {left_type} and {right_type}" + ); + } + + for (left_expr, right_expr) in &on { + Self::validate_side(left_expr, left.schema(), "left equality")?; + Self::validate_side(right_expr, right.schema(), "right equality")?; + if left_expr.is_volatile() || right_expr.is_volatile() { + return plan_err!("ASOF equality expressions must be deterministic"); + } + let left_type = left_expr.get_type(left.schema())?; + let right_type = right_expr.get_type(right.schema())?; + let Some(common_type) = crate::type_coercion::binary::comparison_coercion( + &left_type, + &right_type, + ) else { + return plan_err!( + "ASOF equality expressions have incompatible types {left_type} and {right_type}" + ); + }; + if !crate::utils::can_hash(&common_type) { + return plan_err!( + "ASOF equality expressions have unsupported hash type {common_type}" + ); + } + } + + if join_constraint == JoinConstraint::Using + && on.iter().any(|(left, right)| { + left.get_as_join_column().is_none() + || right.get_as_join_column().is_none() + }) + { + return plan_err!("ASOF USING keys must be columns"); + } + + let schema = + build_asof_join_schema(left.schema(), right.schema(), &on, join_constraint)?; + Ok(Self { + left, + right, + on, + match_condition: Box::new(match_condition), + join_constraint, + schema: Arc::new(schema), + }) + } + + fn validate_side(expr: &Expr, schema: &DFSchema, name: &str) -> Result<()> { + let mut columns = HashSet::new(); + expr_to_columns(expr, &mut columns)?; + if columns.is_empty() { + return plan_err!("ASOF {name} expression must reference its input"); + } + if let Some(column) = columns + .iter() + .find(|column| !schema.is_column_from_schema(column)) + { + return plan_err!( + "ASOF {name} expression references column {column} outside its input" + ); + } + Ok(()) + } +} + +impl PartialOrd for AsOfJoin { + fn partial_cmp(&self, other: &Self) -> Option { + ( + &self.left, + &self.right, + &self.on, + &self.match_condition, + &self.join_constraint, + ) + .partial_cmp(&( + &other.left, + &other.right, + &other.on, + &other.match_condition, + &other.join_constraint, + )) + .filter(|cmp| *cmp != Ordering::Equal || self == other) + } +} + impl Join { /// Creates a new Join operator with automatically computed schema. /// diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c10ac92eef4f5..daf31489b08c7 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -41,11 +41,11 @@ use std::sync::Arc; use crate::logical_plan::plan::RangePartitioning; use crate::{ - Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, - DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, - LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort, - Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode, - Values, Window, builder::unnest_with_options, dml::CopyTo, + Aggregate, Analyze, AsOfJoin, AsOfMatch, CreateMemoryTable, CreateView, DdlStatement, + Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, + Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, + Sort, Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, + UserDefinedLogicalNode, Values, Window, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -150,6 +150,23 @@ impl TreeNode for LogicalPlan { null_aware, }) }), + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) => (left, right).map_elements(f)?.update_data(|(left, right)| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) + }), LogicalPlan::Limit(Limit { skip, fetch, input }) => input .map_elements(f)? .update_data(|input| LogicalPlan::Limit(Limit { skip, fetch, input })), @@ -447,6 +464,13 @@ impl LogicalPlan { LogicalPlan::Join(Join { on, filter, .. }) => { (on, filter).apply_ref_elements(f) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + .. + }) => on.apply_elements(&mut f)?.visit_sibling(|| { + (&match_condition.left, &match_condition.right).apply_ref_elements(&mut f) + }), LogicalPlan::Sort(Sort { expr, .. }) => expr.apply_elements(f), LogicalPlan::Extension(extension) => { // would be nice to avoid this copy -- maybe can @@ -610,6 +634,29 @@ impl LogicalPlan { null_aware, }) }), + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) => (on, (match_condition.left, match_condition.right)) + .map_elements(f)? + .update_data(|(on, (left_match, right_match))| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition: Box::new(AsOfMatch { + left: left_match, + op: match_condition.op, + right: right_match, + }), + join_constraint, + schema, + }) + }), LogicalPlan::Sort(Sort { expr, input, fetch }) => expr .map_elements(f)? .update_data(|expr| LogicalPlan::Sort(Sort { expr, input, fetch })), diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index afd4e980b5424..5fb234ace4ac7 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,9 +57,9 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + AsOfJoin, AsOfMatch, Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -175,6 +175,7 @@ impl<'a> TypeCoercionRewriter<'a> { pub fn coerce_plan(&mut self, plan: LogicalPlan) -> Result { match plan { LogicalPlan::Join(join) => self.coerce_join(join), + LogicalPlan::AsOfJoin(join) => self.coerce_asof_join(join), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), _ => Ok(plan), @@ -218,6 +219,36 @@ impl<'a> TypeCoercionRewriter<'a> { Ok(LogicalPlan::Join(join)) } + /// Coerce ASOF equality and ordered match expressions across input schemas. + pub fn coerce_asof_join(&mut self, mut join: AsOfJoin) -> Result { + join.on = join + .on + .into_iter() + .map(|(left, right)| { + self.coerce_binary_op( + left, + join.left.schema(), + Operator::Eq, + right, + join.right.schema(), + ) + }) + .collect::>()?; + let (left, right) = self.coerce_binary_op( + join.match_condition.left, + join.left.schema(), + join.match_condition.op, + join.match_condition.right, + join.right.schema(), + )?; + join.match_condition = Box::new(AsOfMatch { + left, + op: join.match_condition.op, + right, + }); + Ok(LogicalPlan::AsOfJoin(join)) + } + /// Coerce the union’s inputs to a common schema compatible with all inputs. /// This occurs after wildcard expansion and the coercion of the input expressions. pub fn coerce_union(union_plan: Union) -> Result { diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 2775d62144c56..1f1bcbe60c53f 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -566,6 +566,7 @@ impl OptimizerRule for CommonSubexprEliminate { LogicalPlan::Window(window) => self.try_optimize_window(window, config)?, LogicalPlan::Aggregate(agg) => self.try_optimize_aggregate(agg, config)?, LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 80aceb8cad44c..da5c715d280cb 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,8 +24,9 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, - get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, + Column, DFSchema, HashMap, JoinConstraint, JoinType, Result, + assert_eq_or_internal_err, get_required_group_by_exprs_indices, + internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; use datafusion_expr::{ @@ -407,6 +408,44 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + std::collections::HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect::>(); + let mut left_required = Vec::new(); + let mut right_required = Vec::new(); + for index in indices.indices() { + if *index < left_len { + left_required.push(*index); + } else if let Some(right_index) = + right_output_indices.get(*index - left_len) + { + right_required.push(*right_index); + } + } + let left_indices = RequiredIndices::new_from_indices(left_required) + .with_plan_exprs(&plan, join.left.schema())?; + let right_indices = RequiredIndices::new_from_indices(right_required) + .with_plan_exprs(&plan, join.right.schema())?; + vec![ + left_indices.with_projection_beneficial(), + right_indices.with_projection_beneficial(), + ] + } // these nodes are explicitly rewritten in the match statement above LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index db7ad8475273a..0abb09bea8768 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -411,6 +411,11 @@ fn map_children_mut Result>( let r = f(Arc::make_mut(right))?; l || r } + LogicalPlan::AsOfJoin(join) => { + let l = f(Arc::make_mut(&mut join.left))?; + let r = f(Arc::make_mut(&mut join.right))?; + l || r + } LogicalPlan::Union(Union { inputs, .. }) => { let mut changed = false; for input in inputs { diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index f30b1187b7bca..44170d6702eb2 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1104,6 +1104,35 @@ impl OptimizerRule for PushDownFilter { result.map_data(|plan| Ok(with_filters(keep_predicates, plan))) } LogicalPlan::Join(join) => push_down_join(join, Some(filter.predicate)), + LogicalPlan::AsOfJoin(mut join) => { + let (push, keep): (Vec<_>, Vec<_>) = + split_conjunction_owned(filter.predicate) + .into_iter() + .partition(|predicate| { + !predicate.is_volatile() + && predicate.column_refs().iter().all(|column| { + join.left.schema().is_column_from_schema(column) + }) + }); + if push.is_empty() { + let Some(predicate) = conjunction(keep) else { + return internal_err!("ASOF join filter predicates are empty"); + }; + filter.predicate = predicate; + filter.input = Arc::new(LogicalPlan::AsOfJoin(join)); + Ok(Transformed::no(LogicalPlan::Filter(filter))) + } else { + let Some(predicate) = conjunction(push) else { + return internal_err!("ASOF join push-down predicates are empty"); + }; + join.left = + Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.left))); + Ok(Transformed::yes(with_filters( + keep, + LogicalPlan::AsOfJoin(join), + ))) + } + } LogicalPlan::TableScan(mut scan) => { let filter_predicates = split_conjunction(&filter.predicate); // Filters containing scalar subqueries cannot be pushed to diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 732676a3c0a0f..353d44dc3536e 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -2190,6 +2190,9 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlan::DescribeTable(_) => Err(proto_error( "LogicalPlan serde is not yet implemented for DescribeTable", )), + LogicalPlan::AsOfJoin(_) => Err(proto_error( + "LogicalPlan serde is not yet implemented for AsOfJoin", + )), LogicalPlan::RecursiveQuery(recursive) => { let static_term = LogicalPlanNode::try_from_logical_plan( recursive.static_term.as_ref(), diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 5eef9b82d975e..036e4b34f3ac2 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -135,6 +135,7 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs index c3599a2635ffa..15f59919a2a95 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs @@ -51,6 +51,9 @@ pub fn to_substrait_rel( LogicalPlan::Aggregate(plan) => producer.handle_aggregate(plan), LogicalPlan::Sort(plan) => producer.handle_sort(plan), LogicalPlan::Join(plan) => producer.handle_join(plan), + LogicalPlan::AsOfJoin(plan) => { + not_impl_err!("Substrait ASOF join is not supported: {plan:?}")? + } LogicalPlan::Repartition(plan) => producer.handle_repartition(plan), LogicalPlan::Union(plan) => producer.handle_union(plan), LogicalPlan::TableScan(plan) => producer.handle_table_scan(plan), diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 1981ef66db377..75e1fe251ac1b 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -18,7 +18,7 @@ #[cfg(test)] mod tests { use datafusion::datasource::provider_as_source; - use datafusion::logical_expr::LogicalPlanBuilder; + use datafusion::logical_expr::{AsOfMatch, LogicalPlanBuilder, Operator}; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::serializer; @@ -27,7 +27,7 @@ mod tests { use datafusion::prelude::*; use insta::assert_snapshot; - use std::fs; + use std::{fs, sync::Arc}; use substrait::proto::expression::field_reference::{ReferenceType, RootType}; use substrait::proto::expression::reference_segment; use substrait::proto::expression::{ReferenceSegment, RexType}; @@ -103,6 +103,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn asof_join_fails_closed_until_substrait_has_an_extension() -> Result<()> { + let ctx = create_context().await?; + let table = provider_as_source(ctx.table_provider("data").await?); + let left = LogicalPlanBuilder::scan("l", Arc::clone(&table), None)?.build()?; + let right = LogicalPlanBuilder::scan("r", table, None)?.build()?; + let plan = LogicalPlanBuilder::from(left) + .asof_join( + right, + vec![(col("l.b"), col("r.b"))], + AsOfMatch::new(col("l.a"), Operator::GtEq, col("r.a")), + )? + .build()?; + let error = to_substrait_plan(&plan, &ctx.state()) + .expect_err("ASOF must not be lowered to a generic Substrait join"); + assert!( + error + .to_string() + .contains("Substrait ASOF join is not supported") + ); + Ok(()) + } + #[tokio::test] async fn include_remaps_for_projects() -> Result<()> { let ctx = create_context().await?; From fe3b1827c9154d2a136299f7f4955c06bf500aba Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:45:16 +0800 Subject: [PATCH 03/20] feat: support ASOF JOIN SQL --- datafusion/core/tests/sql/joins.rs | 344 ++++++++++++++++++ datafusion/sql/src/relation/join.rs | 136 ++++++- datafusion/sql/src/unparser/plan.rs | 245 ++++++++++++- .../sqllogictest/test_files/asof_join.slt | 184 ++++++++++ docs/source/user-guide/sql/select.md | 34 ++ 5 files changed, 923 insertions(+), 20 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/asof_join.slt diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 7c0e89ee96418..15db970ad3283 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -20,6 +20,8 @@ use insta::assert_snapshot; use datafusion::assert_batches_eq; use datafusion::catalog::MemTable; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; +use datafusion::physical_plan::joins::AsOfJoinExec; +use datafusion::physical_plan::{Distribution, ExecutionPlanProperties}; use datafusion::test_util::register_unbounded_file_with_ordering; use datafusion_sql::unparser::plan_to_sql; @@ -299,3 +301,345 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } + +async fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { + let trades_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("trade_id", DataType::Int32, false), + ])); + let trades = vec![ + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(7), Some(2), Some(3)])), + Arc::new(Int32Array::from(vec![3, 4, 6])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(1), Some(4), Some(8)])), + Arc::new(Int32Array::from(vec![1, 2, 5])), + ], + )?, + ]; + ctx.register_table( + "trades", + Arc::new(MemTable::try_new(trades_schema, vec![trades])?), + )?; + + let prices_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let prices = vec![ + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(6), Some(1), Some(2)])), + Arc::new(Int32Array::from(vec![60, 101, 999])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(2), Some(4), Some(6)])), + Arc::new(Int32Array::from(vec![20, 40, 106])), + ], + )?, + ]; + ctx.register_table( + "prices", + Arc::new(MemTable::try_new(prices_schema, vec![prices])?), + )?; + Ok(()) +} + +fn find_asof_exec(plan: &Arc) -> Option> { + if plan.downcast_ref::().is_some() { + return Some(Arc::clone(plan)); + } + plan.children().into_iter().find_map(find_asof_exec) +} + +#[tokio::test] +async fn asof_join_all_match_directions_across_batches() -> Result<()> { + let config = SessionConfig::new() + .with_batch_size(2) + .with_target_partitions(2); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx).await?; + + for (op, expected) in [ + ( + ">=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + ">", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 20 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 40 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 60 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ] { + let batches = ctx + .sql(&format!( + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts {op} p.ts) \ + ON t.symbol = p.symbol ORDER BY t.trade_id" + )) + .await? + .collect() + .await?; + assert_batches_eq!(expected, &batches); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_coerces_equality_and_match_types() -> Result<()> { + let ctx = SessionContext::new(); + let batches = ctx + .sql( + "SELECT t.id, p.price \ + FROM (VALUES (CAST(1 AS INT), CAST(4 AS INT), 7)) t(k, ts, id) \ + ASOF JOIN \ + (VALUES (CAST(1 AS BIGINT), CAST(2 AS BIGINT), 20)) p(k, ts, price) \ + MATCH_CONDITION (t.ts >= p.ts) ON t.k = p.k", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----+-------+", + "| id | price |", + "+----+-------+", + "| 7 | 20 |", + "+----+-------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx).await?; + let df = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts)", + ) + .await?; + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert_contains!(sql.as_str(), "ASOF JOIN"); + assert!(!sql.contains(" ON "), "unexpected equality clause: {sql}"); + ctx.sql(&sql).await?; + let plan = df.create_physical_plan().await?; + let asof = find_asof_exec(&plan).expect("physical ASOF join must be present"); + assert_eq!(asof.output_partitioning().partition_count(), 1); + assert!(asof.output_ordering().is_some()); + assert!(matches!( + &asof.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); + Ok(()) +} + +#[tokio::test] +async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + let batches = ctx + .sql( + "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON t.symbol = p.symbol", + ) + .await? + .collect() + .await?; + let explain = arrow::util::pretty::pretty_format_batches(&batches)?.to_string(); + assert_contains!(explain.as_str(), "AsOf Join: match=[t.ts >= p.ts]"); + assert_contains!(explain.as_str(), "on=[t.symbol = p.symbol]"); + assert_contains!(explain.as_str(), "AsOfJoinExec:"); + assert_contains!(explain.as_str(), "on=[(symbol = symbol)]"); + assert_contains!(explain.as_str(), "match=[ts >= ts]"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result<()> { + let ctx = SessionContext::new(); + let tmp_dir = TempDir::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::UInt32, false), + Field::new("ts", DataType::UInt32, false), + ])); + let ordering = vec![vec![ + col("symbol").sort(true, true), + col("ts").sort(true, true), + ]]; + for table in ["left_stream", "right_stream"] { + let path = tmp_dir.path().join(format!("{table}.csv")); + File::create(&path)?; + register_unbounded_file_with_ordering( + &ctx, + Arc::clone(&schema), + &path, + table, + ordering.clone(), + )?; + } + let error = ctx + .sql( + "SELECT * FROM left_stream l ASOF JOIN right_stream r \ + MATCH_CONDITION (l.ts >= r.ts) ON l.symbol = r.symbol", + ) + .await? + .create_physical_plan() + .await + .expect_err("ASOF physical planning must reject unbounded inputs"); + assert_contains!(error.to_string(), "AsOfJoinExec requires bounded inputs"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + let df = ctx + .sql( + "SELECT * FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol)", + ) + .await?; + assert_eq!( + df.schema() + .fields() + .iter() + .map(|field| field.name()) + .collect::>(), + vec!["symbol", "ts", "trade_id", "ts", "price"] + ); + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert!(sql.contains("ASOF JOIN")); + assert!(sql.contains("MATCH_CONDITION")); + assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}"); + ctx.sql(&sql).await?; + Ok(()) +} + +#[tokio::test] +async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + for query in [ + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + "SELECT * FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + "SELECT t.trade_id, p.price FROM trades t \ + JOIN prices q ON t.symbol = q.symbol AND t.ts = q.ts \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON q.symbol = p.symbol ORDER BY t.trade_id", + "SELECT t.trade_id, q.trade_id FROM trades t \ + ASOF JOIN (prices p JOIN trades q \ + ON p.symbol = q.symbol AND p.ts = q.ts) \ + MATCH_CONDITION (t.ts >= q.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + ] { + let expected = ctx.sql(query).await?.collect().await?; + let plan = ctx.sql(query).await?.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + let actual = ctx.sql(&sql).await?.collect().await?; + assert_eq!( + datafusion_common::test_util::batches_to_string(&expected), + datafusion_common::test_util::batches_to_string(&actual), + "unparsed SQL changed ASOF candidate preselection: {sql}" + ); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_invalid_contracts() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + for sql in [ + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts = p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (p.ts >= t.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON t.symbol > p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (1 >= p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON 1 = 1", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts AND t.ts > p.ts) ON t.symbol = p.symbol", + ] { + assert!(ctx.sql(sql).await.is_err(), "query should fail: {sql}"); + } + Ok(()) +} diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b38099..70c9572598d1b 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,8 +16,13 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err}; -use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datafusion_common::{ + Column, DFSchema, Result, not_impl_err, plan_datafusion_err, plan_err, +}; +use datafusion_expr::utils::split_conjunction_owned; +use datafusion_expr::{ + AsOfMatch, BinaryExpr, Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, +}; use sqlparser::ast::{ Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, }; @@ -98,10 +103,129 @@ impl SqlToRel<'_, S> { JoinOperator::CrossJoin(JoinConstraint::None) => { self.parse_cross_join(left, right) } + JoinOperator::AsOf { + match_condition, + constraint, + } => self.parse_asof_join( + left, + right, + match_condition, + constraint, + planner_context, + ), other => not_impl_err!("Unsupported JOIN operator {other:?}"), } } + fn parse_asof_join( + &self, + left: LogicalPlan, + right: LogicalPlan, + sql_match_condition: sqlparser::ast::Expr, + constraint: JoinConstraint, + planner_context: &mut PlannerContext, + ) -> Result { + let join_schema = left.schema().join(right.schema())?; + let match_condition = + self.sql_to_expr(sql_match_condition, &join_schema, planner_context)?; + let Expr::BinaryExpr(BinaryExpr { + left: match_left, + op, + right: match_right, + }) = match_condition + else { + return plan_err!("ASOF MATCH_CONDITION must be a single comparison"); + }; + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {op}" + ); + } + if !expr_owned_by(&match_left, left.schema()) + || !expr_owned_by(&match_right, right.schema()) + { + return plan_err!( + "ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input" + ); + } + let match_condition = AsOfMatch::new(*match_left, op, *match_right); + + match constraint { + JoinConstraint::On(sql_on) => { + let on = self.sql_to_expr(sql_on, &join_schema, planner_context)?; + let on = split_conjunction_owned(on) + .into_iter() + .map(|predicate| { + let Expr::BinaryExpr(BinaryExpr { + left: on_left, + op: Operator::Eq, + right: on_right, + }) = predicate + else { + return plan_err!( + "ASOF ON accepts only equality conditions combined with AND" + ); + }; + if expr_owned_by(&on_left, left.schema()) + && expr_owned_by(&on_right, right.schema()) + { + Ok((*on_left, *on_right)) + } else if expr_owned_by(&on_right, left.schema()) + && expr_owned_by(&on_left, right.schema()) + { + Ok((*on_right, *on_left)) + } else { + plan_err!( + "Each ASOF equality condition must compare one left expression with one right expression" + ) + } + }) + .collect::>()?; + LogicalPlanBuilder::from(left) + .asof_join(right, on, match_condition)? + .build() + } + JoinConstraint::Using(object_names) => { + let keys = object_names + .into_iter() + .map(|object_name| { + let ObjectName(mut object_names) = object_name; + if object_names.len() != 1 { + return not_impl_err!( + "Invalid identifier in ASOF USING clause. Expected single identifier, got {}", + ObjectName(object_names) + ); + } + let id = object_names.swap_remove(0); + id.as_ident() + .ok_or_else(|| { + plan_datafusion_err!( + "Expected identifier in ASOF USING clause" + ) + }) + .map(|ident| { + Column::from_name( + self.ident_normalizer.normalize(ident.clone()), + ) + }) + }) + .collect::>>()?; + LogicalPlanBuilder::from(left) + .asof_join_using(right, keys, match_condition)? + .build() + } + JoinConstraint::None => LogicalPlanBuilder::from(left) + .asof_join(right, vec![], match_condition)? + .build(), + JoinConstraint::Natural => { + not_impl_err!("NATURAL ASOF JOIN is not supported") + } + } + } + fn parse_cross_join( &self, left: LogicalPlan, @@ -180,6 +304,14 @@ impl SqlToRel<'_, S> { } } +fn expr_owned_by(expr: &Expr, schema: &DFSchema) -> bool { + let columns = expr.column_refs(); + !columns.is_empty() + && columns + .iter() + .all(|column| schema.is_column_from_schema(column)) +} + /// Returns `true` if the given [`TableFactor`] is lateral. pub(crate) fn is_lateral(factor: &TableFactor) -> bool { match factor { diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 036e4b34f3ac2..d81cc735011d2 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -49,7 +49,7 @@ use datafusion_common::{ }; use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX}; use datafusion_expr::{ - Aggregate, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, + Aggregate, AsOfJoin, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, SkipType, Sort, SortExpr, TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; @@ -115,6 +115,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -135,7 +136,6 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) - | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } @@ -1299,11 +1299,8 @@ impl Unparser<'_> { let mut right_relation = RelationBuilder::default(); if already_projected - && let Some(nested_relation) = self - .qualified_passthrough_join_projection_to_nested_relation( - right_plan.as_ref(), - query, - )? + && let Some(nested_relation) = + self.join_input_to_nested_relation(right_plan.as_ref(), query)? { right_relation = nested_relation; } else { @@ -1437,6 +1434,9 @@ impl Unparser<'_> { Ok(()) } + LogicalPlan::AsOfJoin(join) => { + self.asof_join_to_sql(join, query, select, relation) + } LogicalPlan::SubqueryAlias(plan_alias) => { let (plan, mut columns) = subquery_alias_inner_query_and_columns(plan_alias); @@ -1734,6 +1734,145 @@ impl Unparser<'_> { } } + // Keep ASOF-specific locals out of the recursive plan unparser's stack frame. + #[inline(never)] + fn asof_join_to_sql( + &self, + join: &AsOfJoin, + query: &mut Option, + select: &mut SelectBuilder, + relation: &mut RelationBuilder, + ) -> Result<()> { + let already_projected = select.already_projected(); + let left_plan = + Self::unwrap_qualified_passthrough_join_projection(Arc::clone(&join.left)); + let inline_left_join = matches!(left_plan.as_ref(), LogicalPlan::Join(_)); + let left_projection = if already_projected { + None + } else if inline_left_join { + self.select_to_sql_recursively(left_plan.as_ref(), query, select, relation)?; + select.pop_projections(); + Some(self.derived_input_projection(join.left.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + let qualifier = self.derive_asof_input(join.left.as_ref(), relation)?; + Some(self.derived_input_projection(join.left.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively(join.left.as_ref(), query, select, relation)?; + Some(select.pop_projections()) + }; + if already_projected { + if inline_left_join { + self.select_to_sql_recursively( + left_plan.as_ref(), + query, + select, + relation, + )?; + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + self.derive_asof_input(join.left.as_ref(), relation)?; + } else { + self.select_to_sql_recursively( + join.left.as_ref(), + query, + select, + relation, + )?; + } + } + + let mut right_relation = RelationBuilder::default(); + let nested_right = + self.join_input_to_nested_relation(join.right.as_ref(), query)?; + let right_projection = if already_projected { + if let Some(nested_right) = nested_right { + right_relation = nested_right; + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + } + None + } else if let Some(nested_right) = nested_right { + right_relation = nested_right; + Some(self.derived_input_projection(join.right.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + let qualifier = + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + Some(self.derived_input_projection(join.right.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + Some(select.pop_projections()) + }; + let Ok(Some(relation)) = right_relation.build() else { + return internal_err!("Failed to build ASOF right relation"); + }; + let constraint = + self.join_constraint_to_sql(join.join_constraint, &join.on, None)?; + let match_condition = self.expr_to_sql(&Expr::BinaryExpr(BinaryExpr::new( + Box::new(join.match_condition.left.clone()), + join.match_condition.op, + Box::new(join.match_condition.right.clone()), + )))?; + let ast_join = ast::Join { + relation, + global: false, + join_operator: ast::JoinOperator::AsOf { + match_condition, + constraint, + }, + }; + let mut from = select + .pop_from() + .ok_or_else(|| internal_datafusion_err!("ASOF left relation is missing"))?; + from.push_join(ast_join); + select.push_from(from); + + if !already_projected { + let left_projection = left_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF left projection is missing") + })?; + let right_projection = right_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF right projection is missing") + })?; + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .filter_map(|(_, right)| right.get_as_join_column()) + .collect::>() + } else { + vec![] + }; + let right_projection = right_projection.into_iter().filter(|item| { + let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(ids)) = + item + else { + return true; + }; + let Some(name) = ids.last() else { + return true; + }; + !omitted_right.iter().any(|column| column.name == name.value) + }); + select.projection( + left_projection + .into_iter() + .chain(right_projection) + .collect(), + ); + } + Ok(()) + } + /// Walk through transparent nodes (SubqueryAlias) to find the inner /// Projection that feeds an Unnest node. /// @@ -2032,6 +2171,74 @@ impl Unparser<'_> { ) } + fn asof_input_requires_derived(plan: &LogicalPlan) -> bool { + let simple_scan = + |scan: &TableScan| scan.filters.is_empty() && scan.fetch.is_none(); + match plan { + LogicalPlan::TableScan(scan) => !simple_scan(scan), + LogicalPlan::SubqueryAlias(alias) => { + !matches!(alias.input.as_ref(), LogicalPlan::TableScan(scan) if simple_scan(scan)) + } + _ => true, + } + } + + fn derive_asof_input( + &self, + plan: &LogicalPlan, + relation: &mut RelationBuilder, + ) -> Result> { + if let LogicalPlan::SubqueryAlias(alias) = plan { + let (inner, columns) = subquery_alias_inner_query_and_columns(alias); + let table_alias = alias.alias.clone(); + if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() + { + let rewritten = + inject_column_aliases_into_subquery(inner.clone(), columns)?; + self.derive( + &rewritten, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), vec![])), + false, + )?; + } else { + self.derive( + inner, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), columns)), + false, + )?; + } + return Ok(Some(table_alias)); + } + + let qualifier = plan + .schema() + .iter() + .find_map(|(qualifier, _)| qualifier.cloned()); + let alias = qualifier + .as_ref() + .map(|qualifier| self.new_table_alias(qualifier.table().to_string(), vec![])); + self.derive(plan, relation, alias, false)?; + Ok(qualifier) + } + + fn derived_input_projection( + &self, + plan: &LogicalPlan, + qualifier: Option<&TableReference>, + ) -> Result> { + plan.schema() + .iter() + .map(|(field_qualifier, field)| { + self.select_item_to_sql(&Expr::Column(Column::new( + qualifier.cloned().or_else(|| field_qualifier.cloned()), + field.name(), + ))) + }) + .collect() + } + fn is_qualified_passthrough_projection(projection: &Projection) -> bool { projection .expr @@ -2052,26 +2259,28 @@ impl Unparser<'_> { } } - fn qualified_passthrough_join_projection_to_nested_relation( + fn join_input_to_nested_relation( &self, plan: &LogicalPlan, query: &mut Option, ) -> Result> { - let LogicalPlan::Projection(projection) = plan else { - return Ok(None); + let join_plan = match plan { + LogicalPlan::Join(_) => plan, + LogicalPlan::Projection(projection) + if matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + && Self::is_qualified_passthrough_projection(projection) => + { + projection.input.as_ref() + } + _ => return Ok(None), }; - if !matches!(projection.input.as_ref(), LogicalPlan::Join(_)) - || !Self::is_qualified_passthrough_projection(projection) - { - return Ok(None); - } let original_query = query.clone(); let mut nested_select = SelectBuilder::default(); nested_select.push_from(TableWithJoinsBuilder::default()); let mut nested_relation = RelationBuilder::default(); self.select_to_sql_recursively( - projection.input.as_ref(), + join_plan, query, &mut nested_select, &mut nested_relation, @@ -2082,11 +2291,11 @@ impl Unparser<'_> { } let Some(mut nested_from) = nested_select.pop_from() else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; nested_from.relation(nested_relation); let Some(table_with_joins) = nested_from.build()? else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; let mut relation = RelationBuilder::default(); diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt new file mode 100644 index 0000000000000..ec300fca74f24 --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,184 @@ +# 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. + +statement ok +CREATE TABLE asof_left(id INT, grp TEXT, ts INT) AS VALUES + (1, 'A', 1), + (2, 'A', 4), + (3, 'A', 7), + (4, 'B', 2), + (5, 'B', 8), + (6, NULL, 3), + (7, 'A', NULL); + +statement ok +CREATE TABLE asof_right(grp TEXT, ts INT, val TEXT) AS VALUES + ('A', 2, 'a2'), + ('A', 4, 'a4'), + ('A', 6, 'a6'), + ('B', 1, 'b1'), + ('B', 6, 'b6'), + (NULL, 2, 'null-group'), + ('A', NULL, 'null-ts'); + +# Inclusive predecessor per equality group. This also verifies unmatched left +# rows and NULL behavior for equality keys and ordered expressions. +query IIIT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 1 NULL NULL +2 4 4 a4 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Strict predecessor per equality group. +query IIIT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts > r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 1 NULL NULL +2 4 2 a2 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Inclusive successor per equality group. +query IIIT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts <= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 1 2 a2 +2 4 4 a4 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# Strict successor per equality group. +query IIIT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts < r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 1 2 a2 +2 4 6 a6 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# USING merges the equality key into one output column. +query TIIT +SELECT grp, l.id, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +A 1 NULL NULL +A 2 4 a4 +A 3 6 a6 +B 4 1 b1 +B 5 6 b6 +NULL 6 NULL NULL +A 7 NULL NULL + +# Equality keys are optional. +query IT +SELECT l.id, r.label +FROM (VALUES (1, 1), (2, 5), (3, CAST(NULL AS INT))) AS l(id, ts) +ASOF JOIN (VALUES (2, 'r2'), (4, 'r4')) AS r(ts, label) +MATCH_CONDITION (l.ts >= r.ts) +ORDER BY l.id; +---- +1 NULL +2 r4 +3 NULL + +query TT +EXPLAIN SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.val +02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +03)----SubqueryAlias: l +04)------TableScan: asof_left projection=[id, grp, ts] +05)----SubqueryAlias: r +06)------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, val@5 as val] +02)--AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts] +03)----SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts = r.ts) +ON l.grp = r.grp; + +query error ASOF MATCH_CONDITION left operand must reference only the left input +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (r.ts >= l.ts) +ON l.grp = r.grp; + +query error ASOF ON accepts only equality conditions combined with AND +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp > r.grp; + +query error ASOF MATCH_CONDITION must be a single comparison +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts) +ON l.grp = r.grp; diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index ea96f6ae4528d..50235c169fb20 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -296,6 +296,7 @@ SELECT a FROM table_name WHERE a > 10; ```text from_item [join_type] JOIN from_item [join_condition] +from_item ASOF JOIN from_item MATCH_CONDITION (condition) [join_condition] from_item CROSS JOIN from_item from_item NATURAL JOIN from_item from_item [join_type] JOIN LATERAL (query) AS alias [join_condition] @@ -377,6 +378,39 @@ SELECT * FROM x LEFT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ ``` +### ASOF JOIN + +An `ASOF JOIN` matches each left row with at most one right row according to an +ordered comparison. It preserves every left row and fills the right columns +with `NULL` when no right row matches. + +```sql +SELECT t.*, p.price +FROM trades AS t +ASOF JOIN prices AS p +MATCH_CONDITION (t.ts >= p.ts) +ON t.symbol = p.symbol; +``` + +`MATCH_CONDITION` must compare an expression from the left input with an +expression from the right input using one of the following operators: + +| Condition | Selected right row | +| --------- | ----------------------------------------- | +| `l >= r` | Greatest `r` less than or equal to `l` | +| `l > r` | Greatest `r` strictly less than `l` | +| `l <= r` | Smallest `r` greater than or equal to `l` | +| `l < r` | Smallest `r` strictly greater than `l` | + +An optional `ON` clause containing equality conditions combined with `AND`, or +a `USING` clause, divides rows into equality groups before the ordered match. +Without equality keys, all rows belong to one group and DataFusion executes the +join in a single partition. + +A `NULL` in either ordered expression or in any equality key does not match. +Both inputs must be bounded. If multiple right rows have the same equality keys +and ordered value, which tied row is selected is nondeterministic. + ### RIGHT OUTER JOIN The keywords `RIGHT JOIN` or `RIGHT OUTER JOIN` define a join that includes all rows from the right table even if there From 43110a3260548ad87b73fb5670c05c24839444bb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:38:22 +0800 Subject: [PATCH 04/20] 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 23e606017377bb212aaac6681877cc1bc68d2e9a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 17:04:40 +0800 Subject: [PATCH 05/20] test: cover broadcast ASOF SQL execution --- datafusion/core/tests/sql/joins.rs | 40 ++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 15db970ad3283..9668d90162818 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -302,7 +302,7 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } -async fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { +fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { let trades_schema = Arc::new(Schema::new(vec![ Field::new("symbol", DataType::Utf8, true), Field::new("ts", DataType::Int64, true), @@ -328,7 +328,10 @@ async fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { ]; ctx.register_table( "trades", - Arc::new(MemTable::try_new(trades_schema, vec![trades])?), + Arc::new(MemTable::try_new( + trades_schema, + trades.into_iter().map(|batch| vec![batch]).collect(), + )?), )?; let prices_schema = Arc::new(Schema::new(vec![ @@ -374,7 +377,7 @@ async fn asof_join_all_match_directions_across_batches() -> Result<()> { .with_batch_size(2) .with_target_partitions(2); let ctx = SessionContext::new_with_config(config); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; for (op, expected) in [ ( @@ -480,10 +483,10 @@ async fn asof_join_coerces_equality_and_match_types() -> Result<()> { } #[tokio::test] -async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { +async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> { let config = SessionConfig::new().with_target_partitions(4); let ctx = SessionContext::new_with_config(config); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let df = ctx .sql( "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices p \ @@ -496,11 +499,26 @@ async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { ctx.sql(&sql).await?; let plan = df.create_physical_plan().await?; let asof = find_asof_exec(&plan).expect("physical ASOF join must be present"); - assert_eq!(asof.output_partitioning().partition_count(), 1); + let output_partitions = asof.output_partitioning().partition_count(); + assert_eq!( + output_partitions, + asof.children()[0].output_partitioning().partition_count() + ); + assert!( + output_partitions > 1, + "ASOF join did not preserve left-side parallelism" + ); + assert_eq!( + asof.children()[1].output_partitioning().partition_count(), + 1 + ); assert!(asof.output_ordering().is_some()); assert!(matches!( &asof.input_distribution_requirements().into_per_child()[..], - [Distribution::SinglePartition, Distribution::SinglePartition] + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] )); let batches = collect(plan, ctx.task_ctx()).await?; assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); @@ -510,7 +528,7 @@ async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { #[tokio::test] async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let batches = ctx .sql( "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ @@ -568,7 +586,7 @@ async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result #[tokio::test] async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let df = ctx .sql( "SELECT * FROM trades t ASOF JOIN prices p \ @@ -594,7 +612,7 @@ async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { #[tokio::test] async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; for query in [ "SELECT t.trade_id, p.price FROM trades t \ ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ @@ -630,7 +648,7 @@ async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { #[tokio::test] async fn asof_join_rejects_invalid_contracts() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; for sql in [ "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts = p.ts) ON t.symbol = p.symbol", "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (p.ts >= t.ts) ON t.symbol = p.symbol", From 339d7857107c2e87bbffe87d89d2c647dbc45c63 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 01:40:50 +0800 Subject: [PATCH 06/20] 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 3a4f995a75a7a47f502dcb2b65c45e124c080914 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 01:52:40 +0800 Subject: [PATCH 07/20] fix: align ASOF logical contracts --- datafusion/core/src/physical_planner.rs | 22 +++--------- datafusion/expr/src/logical_plan/builder.rs | 36 +++---------------- datafusion/expr/src/logical_plan/plan.rs | 35 +++++++++++++++--- .../optimizer/src/optimize_projections/mod.rs | 27 +++----------- 4 files changed, 42 insertions(+), 78 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d852d56b1e52d..43ce2e6d71b33 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -94,8 +94,8 @@ use datafusion_expr::physical_planning_context::{ use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, - FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType, - StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, + FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, + WindowFrame, WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, @@ -1903,22 +1903,8 @@ impl DefaultPhysicalPlanner { planning_ctx, )?, ); - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .map(|(_, right)| { - let column = right.get_as_join_column().ok_or_else(|| { - internal_datafusion_err!("ASOF USING key is not a column") - })?; - join.right.schema().index_of_column(column) - }) - .collect::>>()? - } else { - HashSet::new() - }; - let right_output_indices = (0..join.right.schema().fields().len()) - .filter(|index| !omitted_right.contains(index)) - .collect(); + let right_output_indices = + (0..join.right.schema().fields().len()).collect(); Arc::new(AsOfJoinExec::try_new( physical_left, physical_right, diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 14603c3f62e24..67436332a123f 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1840,38 +1840,10 @@ pub fn build_join_schema( /// Creates the schema for a left-preserving ASOF join. /// -/// `ON` emits all left fields followed by nullable right fields. `USING` emits -/// each equality key once by omitting the corresponding right field. -pub fn build_asof_join_schema( - left: &DFSchema, - right: &DFSchema, - on: &[(Expr, Expr)], - join_constraint: JoinConstraint, -) -> Result { - let omitted_right_indices = if join_constraint == JoinConstraint::Using { - on.iter() - .map(|(_, right_expr)| { - let column = right_expr.get_as_join_column().ok_or_else(|| { - plan_datafusion_err!("ASOF USING keys must be columns") - })?; - right.index_of_column(column) - }) - .collect::>>()? - } else { - HashSet::new() - }; - - let full_schema = build_join_schema(left, right, &JoinType::Left)?; - let left_len = left.fields().len(); - let fields = full_schema - .iter() - .enumerate() - .filter(|(index, _)| { - *index < left_len || !omitted_right_indices.contains(&(*index - left_len)) - }) - .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field))) - .collect(); - DFSchema::new_with_metadata(fields, full_schema.metadata().clone())? +/// Both `ON` and `USING` preserve all qualified input fields. SQL wildcard +/// expansion handles the unqualified `USING` key as a single column. +pub fn build_asof_join_schema(left: &DFSchema, right: &DFSchema) -> Result { + build_join_schema(left, right, &JoinType::Left)? .with_functional_dependencies(left.functional_dependencies().clone()) } diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 5cee0ad97bee2..3635aeb83b6e0 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -240,9 +240,6 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), - /// Match each left row with at most one ordered row from the right input. - /// This is used to implement SQL `ASOF JOIN`. - AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -299,6 +296,9 @@ pub enum LogicalPlan { Unnest(Unnest), /// A variadic query (e.g. "Recursive CTEs") RecursiveQuery(RecursiveQuery), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), } impl Default for LogicalPlan { @@ -4455,8 +4455,7 @@ impl AsOfJoin { return plan_err!("ASOF USING keys must be columns"); } - let schema = - build_asof_join_schema(left.schema(), right.schema(), &on, join_constraint)?; + let schema = build_asof_join_schema(left.schema(), right.schema())?; Ok(Self { left, right, @@ -6867,6 +6866,32 @@ mod tests { Ok(()) } + #[test] + fn test_asof_using_preserves_qualified_keys() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("ts", DataType::Int64, false), + ]); + let left = Arc::new(table_scan(Some("t1"), &schema, None)?.build()?); + let right = Arc::new(table_scan(Some("t2"), &schema, None)?.build()?); + let join = AsOfJoin::try_new( + left, + right, + vec![(col("t1.id"), col("t2.id"))], + AsOfMatch::new(col("t1.ts"), Operator::GtEq, col("t2.ts")), + JoinConstraint::Using, + )?; + + assert_eq!(join.schema.fields().len(), 4); + assert_eq!( + join.schema + .index_of_column(&Column::from_qualified_name("t2.id"))?, + 2 + ); + assert!(join.schema.field(2).is_nullable()); + Ok(()) + } + #[test] fn test_join_try_new_schema_validation() -> Result<()> { let left_schema = Schema::new(vec![ diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index da5c715d280cb..74ce2ad374af8 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,9 +24,8 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinConstraint, JoinType, Result, - assert_eq_or_internal_err, get_required_group_by_exprs_indices, - internal_datafusion_err, internal_err, + Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, + get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; use datafusion_expr::{ @@ -410,31 +409,13 @@ fn optimize_projections( } LogicalPlan::AsOfJoin(join) => { let left_len = join.left.schema().fields().len(); - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .map(|(_, right)| { - let column = right.get_as_join_column().ok_or_else(|| { - internal_datafusion_err!("ASOF USING key is not a column") - })?; - join.right.schema().index_of_column(column) - }) - .collect::>>()? - } else { - std::collections::HashSet::new() - }; - let right_output_indices = (0..join.right.schema().fields().len()) - .filter(|index| !omitted_right.contains(index)) - .collect::>(); let mut left_required = Vec::new(); let mut right_required = Vec::new(); for index in indices.indices() { if *index < left_len { left_required.push(*index); - } else if let Some(right_index) = - right_output_indices.get(*index - left_len) - { - right_required.push(*right_index); + } else { + right_required.push(*index - left_len); } } let left_indices = RequiredIndices::new_from_indices(left_required) From fd8fddb87ca3e31485b0b778477f0b60984f390a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 02:09:23 +0800 Subject: [PATCH 08/20] fix: align ASOF USING and broadcast docs --- datafusion/core/tests/sql/joins.rs | 30 +++++++++++++++++-- datafusion/sql/src/unparser/plan.rs | 19 ------------ .../sqllogictest/test_files/asof_join.slt | 19 +++++++++++- docs/source/user-guide/sql/select.md | 10 +++++-- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 9668d90162818..2504c0077e6ff 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -584,7 +584,7 @@ async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result } #[tokio::test] -async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { +async fn asof_join_using_preserves_key_access_and_unparser_round_trips() -> Result<()> { let ctx = SessionContext::new(); register_asof_test_tables(&ctx)?; let df = ctx @@ -599,13 +599,39 @@ async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { .iter() .map(|field| field.name()) .collect::>(), - vec!["symbol", "ts", "trade_id", "ts", "price"] + vec!["ts", "trade_id", "symbol", "ts", "price"] ); let sql = plan_to_sql(df.logical_plan())?.to_string(); assert!(sql.contains("ASOF JOIN")); assert!(sql.contains("MATCH_CONDITION")); assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}"); ctx.sql(&sql).await?; + + let batches = ctx + .sql( + "SELECT t.trade_id, t.symbol AS left_symbol, p.symbol AS right_symbol \ + FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----------+-------------+--------------+", + "| trade_id | left_symbol | right_symbol |", + "+----------+-------------+--------------+", + "| 1 | A | |", + "| 2 | A | A |", + "| 3 | A | A |", + "| 4 | B | B |", + "| 5 | B | B |", + "| 6 | | |", + "+----------+-------------+--------------+", + ], + &batches + ); Ok(()) } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 12d9b65be1a7d..9c0e7786fd0e1 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -1885,25 +1885,6 @@ impl Unparser<'_> { let right_projection = right_projection.ok_or_else(|| { internal_datafusion_err!("ASOF right projection is missing") })?; - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .filter_map(|(_, right)| right.get_as_join_column()) - .collect::>() - } else { - vec![] - }; - let right_projection = right_projection.into_iter().filter(|item| { - let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(ids)) = - item - else { - return true; - }; - let Some(name) = ids.last() else { - return true; - }; - !omitted_right.iter().any(|column| column.name == name.value) - }); select.projection( left_projection .into_iter() diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index ec300fca74f24..62a9e3cf35ceb 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -104,7 +104,7 @@ ORDER BY l.id; 6 3 NULL NULL 7 NULL NULL NULL -# USING merges the equality key into one output column. +# USING exposes one unqualified equality key. query TIIT SELECT grp, l.id, r.ts, r.val FROM asof_left l @@ -121,6 +121,23 @@ B 5 6 b6 NULL 6 NULL NULL A 7 NULL NULL +# Both qualified equality keys remain addressable. +query ITT +SELECT l.id, l.grp, r.grp +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +1 A NULL +2 A A +3 A A +4 B B +5 B B +6 NULL NULL +7 A NULL + # Equality keys are optional. query IT SELECT l.id, r.label diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index 50235c169fb20..817f98a8b236f 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -404,8 +404,14 @@ expression from the right input using one of the following operators: An optional `ON` clause containing equality conditions combined with `AND`, or a `USING` clause, divides rows into equality groups before the ordered match. -Without equality keys, all rows belong to one group and DataFusion executes the -join in a single partition. +An unqualified `USING` key appears once in wildcard output, while both qualified +input keys remain addressable. + +Without equality keys, all rows belong to one group. The initial execution +strategy collects one ordered right partition and shares it across every left +partition, so output partitioning follows the left input. The complete right +input must fit in memory and may be scanned once per left partition; spilling +and repartitioned ASOF execution are not yet supported. A `NULL` in either ordered expression or in any equality key does not match. Both inputs must be bounded. If multiple right rows have the same equality keys From 34552d158eef2737c9022de570dc6af5d54e5aa2 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:50:01 +0800 Subject: [PATCH 09/20] 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 d618879332b4c0fea6fc981650541f3c2874cc55 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:50:13 +0800 Subject: [PATCH 10/20] test: cover temporal and multi-partition ASOF joins --- datafusion/core/tests/sql/joins.rs | 33 ++++++- .../sqllogictest/test_files/asof_join.slt | 94 +++++++++---------- 2 files changed, 75 insertions(+), 52 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 2504c0077e6ff..842f5690e3828 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -17,12 +17,12 @@ use insta::assert_snapshot; -use datafusion::assert_batches_eq; use datafusion::catalog::MemTable; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use datafusion::physical_plan::joins::AsOfJoinExec; use datafusion::physical_plan::{Distribution, ExecutionPlanProperties}; use datafusion::test_util::register_unbounded_file_with_ordering; +use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion_sql::unparser::plan_to_sql; use super::*; @@ -359,7 +359,10 @@ fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { ]; ctx.register_table( "prices", - Arc::new(MemTable::try_new(prices_schema, vec![prices])?), + Arc::new(MemTable::try_new( + prices_schema, + prices.into_iter().map(|batch| vec![batch]).collect(), + )?), )?; Ok(()) } @@ -483,13 +486,14 @@ async fn asof_join_coerces_equality_and_match_types() -> Result<()> { } #[tokio::test] -async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> { +async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { let config = SessionConfig::new().with_target_partitions(4); let ctx = SessionContext::new_with_config(config); register_asof_test_tables(&ctx)?; let df = ctx .sql( - "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices p \ + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN \ + (SELECT ts, price FROM prices WHERE symbol = 'A') p \ MATCH_CONDITION (t.ts >= p.ts)", ) .await?; @@ -512,6 +516,11 @@ async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> asof.children()[1].output_partitioning().partition_count(), 1 ); + let right_plan = displayable(asof.children()[1].as_ref()) + .indent(true) + .to_string(); + assert_contains!(right_plan.as_str(), "SortPreservingMergeExec"); + assert_contains!(right_plan.as_str(), "DataSourceExec: partitions=2"); assert!(asof.output_ordering().is_some()); assert!(matches!( &asof.input_distribution_requirements().into_per_child()[..], @@ -521,7 +530,21 @@ async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> ] )); let batches = collect(plan, ctx.task_ctx()).await?; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); + assert_batches_sorted_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 20 |", + "| 5 | 60 |", + "| 6 | 20 |", + "+----------+-------+", + ], + &batches + ); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index 62a9e3cf35ceb..46cf6cdf10bf7 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -16,28 +16,28 @@ # under the License. statement ok -CREATE TABLE asof_left(id INT, grp TEXT, ts INT) AS VALUES - (1, 'A', 1), - (2, 'A', 4), - (3, 'A', 7), - (4, 'B', 2), - (5, 'B', 8), - (6, NULL, 3), +CREATE TABLE asof_left(id INT, grp TEXT, ts TIMESTAMP) AS VALUES + (1, 'A', TIMESTAMP '2024-01-01 09:00:01'), + (2, 'A', TIMESTAMP '2024-01-01 09:00:04'), + (3, 'A', TIMESTAMP '2024-01-01 09:00:07'), + (4, 'B', TIMESTAMP '2024-01-01 09:00:02'), + (5, 'B', TIMESTAMP '2024-01-01 09:00:08'), + (6, NULL, TIMESTAMP '2024-01-01 09:00:03'), (7, 'A', NULL); statement ok -CREATE TABLE asof_right(grp TEXT, ts INT, val TEXT) AS VALUES - ('A', 2, 'a2'), - ('A', 4, 'a4'), - ('A', 6, 'a6'), - ('B', 1, 'b1'), - ('B', 6, 'b6'), - (NULL, 2, 'null-group'), +CREATE TABLE asof_right(grp TEXT, ts TIMESTAMP, val TEXT) AS VALUES + ('A', TIMESTAMP '2024-01-01 09:00:02', 'a2'), + ('A', TIMESTAMP '2024-01-01 09:00:04', 'a4'), + ('A', TIMESTAMP '2024-01-01 09:00:06', 'a6'), + ('B', TIMESTAMP '2024-01-01 09:00:01', 'b1'), + ('B', TIMESTAMP '2024-01-01 09:00:06', 'b6'), + (NULL, TIMESTAMP '2024-01-01 09:00:02', 'null-group'), ('A', NULL, 'null-ts'); # Inclusive predecessor per equality group. This also verifies unmatched left # rows and NULL behavior for equality keys and ordered expressions. -query IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -45,16 +45,16 @@ MATCH_CONDITION (l.ts >= r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 NULL NULL -2 4 4 a4 -3 7 6 a6 -4 2 1 b1 -5 8 6 b6 -6 3 NULL NULL +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL 7 NULL NULL NULL # Strict predecessor per equality group. -query IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -62,16 +62,16 @@ MATCH_CONDITION (l.ts > r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 NULL NULL -2 4 2 a2 -3 7 6 a6 -4 2 1 b1 -5 8 6 b6 -6 3 NULL NULL +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:02 a2 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL 7 NULL NULL NULL # Inclusive successor per equality group. -query IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -79,16 +79,16 @@ MATCH_CONDITION (l.ts <= r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 2 a2 -2 4 4 a4 -3 7 NULL NULL -4 2 6 b6 -5 8 NULL NULL -6 3 NULL NULL +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL 7 NULL NULL NULL # Strict successor per equality group. -query IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -96,16 +96,16 @@ MATCH_CONDITION (l.ts < r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 2 a2 -2 4 6 a6 -3 7 NULL NULL -4 2 6 b6 -5 8 NULL NULL -6 3 NULL NULL +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:06 a6 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL 7 NULL NULL NULL # USING exposes one unqualified equality key. -query TIIT +query TIPT SELECT grp, l.id, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -114,10 +114,10 @@ USING (grp) ORDER BY l.id; ---- A 1 NULL NULL -A 2 4 a4 -A 3 6 a6 -B 4 1 b1 -B 5 6 b6 +A 2 2024-01-01T09:00:04 a4 +A 3 2024-01-01T09:00:06 a6 +B 4 2024-01-01T09:00:01 b1 +B 5 2024-01-01T09:00:06 b6 NULL 6 NULL NULL A 7 NULL NULL From f157eaecd60702e0612c647a0c38cd95b2433814 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 2 Aug 2026 01:03:36 +0800 Subject: [PATCH 11/20] test: remove ASOF filter pushdown coverage --- .../physical_optimizer/filter_pushdown.rs | 42 ------------------- 1 file changed, 42 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 2f613264f45d9..909b80cadaae3 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -55,7 +55,6 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, - joins::{AsOfJoinExec, AsOfMatchExpr}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -124,47 +123,6 @@ fn test_pushdown_volatile_functions_not_allowed() { ); } -#[test] -fn test_asof_join_pushes_only_left_filters() { - let left = TestScanBuilder::new(schema()).with_support(true).build(); - let right = TestScanBuilder::new(schema()).with_support(true).build(); - let join = Arc::new( - AsOfJoinExec::try_new( - left, - right, - vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())], - AsOfMatchExpr::new( - col("b", &schema()).unwrap(), - Operator::GtEq, - col("b", &schema()).unwrap(), - ), - vec![0, 1, 2], - ) - .unwrap(), - ); - let left_filter = Arc::new(BinaryExpr::new( - Arc::new(Column::new("c", 2)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), - )) as Arc; - let right_filter = Arc::new(BinaryExpr::new( - Arc::new(Column::new("c", 5)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), - )) as Arc; - let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter)); - let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap()); - let mut config = ConfigOptions::default(); - config.execution.parquet.pushdown_filters = true; - let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); - let formatted = format_plan_for_test(&optimized); - - assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}"); - assert!(formatted.contains("AsOfJoinExec:"), "{formatted}"); - assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}"); - assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}"); -} - /// Show that we can use config options to determine how to do pushdown. #[test] fn test_pushdown_into_scan_with_config_options() { From a127baaf87dc5440326976b889b6f1bbc006d3c5 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 2 Aug 2026 01:03:36 +0800 Subject: [PATCH 12/20] 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 763a06ea2a2948a2778dda4bce78d532dc9df0bd Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 2 Aug 2026 01:03:37 +0800 Subject: [PATCH 13/20] test: expand ASOF SQL coverage --- .../sqllogictest/test_files/asof_join.slt | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index 46cf6cdf10bf7..daca2d13c4d81 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -150,6 +150,50 @@ ORDER BY l.id; 2 r4 3 NULL +# Multiple equality keys form one candidate group. +query IT +SELECT l.id, r.val +FROM (VALUES + (1, 'X', 'A', TIMESTAMP '2024-01-01 09:00:04'), + (2, 'Y', 'A', TIMESTAMP '2024-01-01 09:00:04') +) AS l(id, venue, grp, ts) +ASOF JOIN (VALUES + ('X', 'A', TIMESTAMP '2024-01-01 09:00:02', 'x-a2'), + ('Y', 'A', TIMESTAMP '2024-01-01 09:00:03', 'y-a3'), + ('X', 'B', TIMESTAMP '2024-01-01 09:00:04', 'x-b4') +) AS r(venue, grp, ts, val) +MATCH_CONDITION (l.ts >= r.ts) +ON l.venue = r.venue AND l.grp = r.grp +ORDER BY l.id; +---- +1 x-a2 +2 y-a3 + +# Candidate selection sees the right input after subquery filtering. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN (SELECT * FROM asof_right WHERE val <> 'a6') r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +WHERE l.id IN (3, 5) +ORDER BY l.id; +---- +3 a4 +5 b6 + +# Equality and match operands use the planner's common coercion types. +query II +SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +1 90 + query TT EXPLAIN SELECT l.id, r.val FROM asof_left l From 450bffb5faaf751b2818a2566d53054867d06515 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 11 Aug 2026 17:59:52 +0800 Subject: [PATCH 14/20] 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 15/20] 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 16/20] 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 bb58ee0c24f6ae458297010e6ac48415138c4fb0 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 11 Aug 2026 23:50:53 +0800 Subject: [PATCH 17/20] fix: align ASOF planner projection --- datafusion/core/src/physical_planner.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 424d3f0532bde..94300acfdbdd1 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1803,14 +1803,12 @@ impl DefaultPhysicalPlanner { planning_ctx, )?, ); - let right_output_indices = - (0..join.right.schema().fields().len()).collect(); Arc::new(AsOfJoinExec::try_new( physical_left, physical_right, join_on, match_condition, - right_output_indices, + None, )?) } LogicalPlan::RecursiveQuery(RecursiveQuery { From 52484024d0e089466e5e395bebbe846bac1cb824 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 15 Aug 2026 00:04:33 +0800 Subject: [PATCH 18/20] 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 e5703257251baccd32782ae658a51c993f4e4954 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 15 Aug 2026 00:22:27 +0800 Subject: [PATCH 19/20] fix: support floating ASOF equality keys --- .../physical-expr/src/expressions/mod.rs | 2 + .../src/expressions/normalize_float_zero.rs | 199 ++++++++++++++++++ .../physical-plan/src/joins/asof_join.rs | 153 +++++++++----- .../proto-models/proto/datafusion.proto | 5 + .../proto-models/src/generated/pbjson.rs | 105 +++++++++ .../proto-models/src/generated/prost.rs | 11 +- .../proto/src/physical_plan/from_proto.rs | 6 +- datafusion/proto/tests/cases/plans/sorts.rs | 18 +- 8 files changed, 449 insertions(+), 50 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/normalize_float_zero.rs diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 035dd5d5072b0..a4065e2aadb98 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -32,6 +32,7 @@ mod like; mod literal; mod negative; mod no_op; +mod normalize_float_zero; mod not; mod try_cast; mod unknown_column; @@ -57,6 +58,7 @@ pub use like::{LikeExpr, like}; pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; +pub use normalize_float_zero::NormalizeFloatZeroExpr; pub use not::{NotExpr, not}; pub use try_cast::{TryCastExpr, try_cast}; pub use unknown_column::UnKnownColumn; diff --git a/datafusion/physical-expr/src/expressions/normalize_float_zero.rs b/datafusion/physical-expr/src/expressions/normalize_float_zero.rs new file mode 100644 index 0000000000000..8e6335354deb6 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/normalize_float_zero.rs @@ -0,0 +1,199 @@ +// 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. + +//! Floating-point signed-zero normalization expression. + +use std::hash::Hash; +use std::sync::Arc; + +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; +use datafusion_expr::ColumnarValue; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::ExprProperties; + +use crate::PhysicalExpr; + +/// Replaces floating-point `-0.0` values with `+0.0`. +/// +/// Other values and data types are returned unchanged. This expression is +/// order-preserving but not strictly order-preserving because it collapses the +/// two signed-zero representations. +#[derive(Debug, Eq)] +pub struct NormalizeFloatZeroExpr { + arg: Arc, +} + +impl PartialEq for NormalizeFloatZeroExpr { + fn eq(&self, other: &Self) -> bool { + self.arg.eq(&other.arg) + } +} + +impl Hash for NormalizeFloatZeroExpr { + fn hash(&self, state: &mut H) { + self.arg.hash(state); + } +} + +impl NormalizeFloatZeroExpr { + /// Creates a signed-zero normalization expression. + pub fn new(arg: Arc) -> Self { + Self { arg } + } + + /// Returns the input expression. + pub fn arg(&self) -> &Arc { + &self.arg + } +} + +impl std::fmt::Display for NormalizeFloatZeroExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "normalize_float_zero({})", self.arg) + } +} + +impl PhysicalExpr for NormalizeFloatZeroExpr { + fn data_type(&self, input_schema: &Schema) -> Result { + self.arg.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.arg.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + Ok(match self.arg.evaluate(batch)? { + ColumnarValue::Array(array) => { + ColumnarValue::Array(normalize_float_zero(&array)) + } + ColumnarValue::Scalar(scalar) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(scalar)) + } + }) + } + + fn return_field(&self, input_schema: &Schema) -> Result { + self.arg.return_field(input_schema) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.arg] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } + + fn evaluate_bounds(&self, children: &[&Interval]) -> Result { + Ok(children[0].clone()) + } + + fn get_properties(&self, children: &[ExprProperties]) -> Result { + Ok(children[0].clone().with_strictly_order_preserving(false)) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "normalize_float_zero(")?; + self.arg.fmt_sql(f)?; + write!(f, ")") + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::NormalizeFloatZero( + Box::new(protobuf::PhysicalNormalizeFloatZeroNode { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }), + )), + })) + } +} + +#[cfg(feature = "proto")] +impl NormalizeFloatZeroExpr { + /// Reconstructs a [`NormalizeFloatZeroExpr`] from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::NormalizeFloatZero, + "NormalizeFloatZero", + ); + let arg = ctx.decode_required_expression( + node.expr.as_deref(), + "NormalizeFloatZeroExpr", + "expr", + )?; + Ok(Arc::new(Self::new(arg))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ArrayRef, AsArray, Float64Array}; + use datafusion_common::ScalarValue; + + use crate::expressions::{Column, Literal}; + + #[test] + fn normalizes_array_and_scalar_signed_zero() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Float64Array::from(vec![-0.0, 0.0, 1.0])) as ArrayRef, + )])?; + let expr = NormalizeFloatZeroExpr::new(Arc::new(Column::new("a", 0))); + let ColumnarValue::Array(array) = expr.evaluate(&batch)? else { + panic!("column evaluation must return an array"); + }; + let array = array.as_primitive::(); + assert_eq!(array.value(0).to_bits(), 0.0_f64.to_bits()); + assert_eq!(array.value(1).to_bits(), 0.0_f64.to_bits()); + assert_eq!(array.value(2), 1.0); + + let expr = NormalizeFloatZeroExpr::new(Arc::new(Literal::new( + ScalarValue::Float64(Some(-0.0)), + ))); + let ColumnarValue::Scalar(ScalarValue::Float64(Some(value))) = + expr.evaluate(&RecordBatch::new_empty(Arc::new(Schema::empty())))? + else { + panic!("literal evaluation must return a Float64 scalar"); + }; + assert_eq!(value.to_bits(), 0.0_f64.to_bits()); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index c7c068e494465..69a377fd352d3 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -86,7 +86,9 @@ 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::expressions::{ + Column as PhysicalColumn, NormalizeFloatZeroExpr, +}; use datafusion_physical_expr::projection::{ProjectionMapping, ProjectionRef}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr_common::physical_expr::{ @@ -159,10 +161,8 @@ 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; - /// 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. + /// input, and have matching input types. Equality types must support hashing. + /// Projection indices refer to the full left-then-right join schema. pub fn try_new( left: Arc, right: Arc, @@ -186,24 +186,33 @@ impl AsOfJoinExec { descending, nulls_first: true, }; - let mut left_sort_exprs = on - .iter() - .map(|(left, _)| PhysicalSortExpr { - expr: Arc::clone(left), + let mut left_sort_exprs = Vec::with_capacity(on.len() + 1); + let mut right_sort_exprs = Vec::with_capacity(on.len() + 1); + for (left, right) in &on { + let left_expr = if left.data_type(&left_schema)?.is_floating() { + Arc::new(NormalizeFloatZeroExpr::new(Arc::clone(left))) as PhysicalExprRef + } else { + Arc::clone(left) + }; + let right_expr = if right.data_type(&right_schema)?.is_floating() { + Arc::new(NormalizeFloatZeroExpr::new(Arc::clone(right))) + as PhysicalExprRef + } else { + Arc::clone(right) + }; + left_sort_exprs.push(PhysicalSortExpr { + expr: left_expr, options: equality_options, - }) - .collect::>(); + }); + right_sort_exprs.push(PhysicalSortExpr { + expr: right_expr, + options: equality_options, + }); + } 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, @@ -1177,11 +1186,6 @@ 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)?; @@ -1235,8 +1239,9 @@ mod tests { use super::*; use crate::collect; + use crate::sorts::sort::SortExec; use crate::test::TestMemoryExec; - use arrow::array::{Int32Array, Int64Array, StringArray}; + use arrow::array::{Float64Array, Int32Array, Int64Array, StringArray}; use arrow::datatypes::{DataType, Field}; use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::config::SessionConfig; @@ -1699,34 +1704,88 @@ mod tests { Ok(()) } - #[test] - fn rejects_floating_equality_expressions() -> Result<()> { - let exec = test_exec()?; + #[tokio::test] + async fn floating_equality_keys_treat_signed_zero_as_equal() -> Result<()> { + let left_batch = RecordBatch::try_from_iter(vec![ + ("key", Arc::new(Float64Array::from(vec![0.0])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![5])) as ArrayRef), + ("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef), + ])?; + let right_batch = RecordBatch::try_from_iter(vec![ + ( + "key", + Arc::new(Float64Array::from(vec![-0.0, 0.0])) as ArrayRef, + ), + ("ts", Arc::new(Int64Array::from(vec![10, 1])) as ArrayRef), + ( + "price", + Arc::new(Int32Array::from(vec![100, 10])) as ArrayRef, + ), + ])?; 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(), + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch.clone()]], + left_batch.schema(), None, - )); - let right = Arc::new(CastExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), - data_type.clone(), + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()]], + right_batch.schema(), 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}" + )?; + let on: JoinOn = vec![( + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("key", 0)), + data_type.clone(), + None, + )), + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("key", 0)), + data_type.clone(), + None, )), - "unexpected error: {error}" + )]; + let match_condition = AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), ); + let unsorted = AsOfJoinExec::try_new( + left, + right, + on.clone(), + match_condition.clone(), + Some(vec![2, 5]), + )?; + let left = Arc::new(SortExec::new( + unsorted.left_ordering.clone(), + Arc::clone(&unsorted.left), + )); + let right = Arc::new(SortExec::new( + unsorted.right_ordering.clone(), + Arc::clone(&unsorted.right), + )); + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + match_condition, + Some(vec![2, 5]), + )?); + + let batches = collect(exec, Arc::new(TaskContext::default())).await?; + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(10)], "data type: {data_type}"); } Ok(()) } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 43a90264c2b1f..8f103a0233a49 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1051,6 +1051,7 @@ message PhysicalExprNode { PhysicalLambdaExprNode lambda = 25; PhysicalLambdaVariableExprNode lambda_variable = 26; PhysicalRangeExprNode range_expr = 27; + PhysicalNormalizeFloatZeroNode normalize_float_zero = 28; } } @@ -1192,6 +1193,10 @@ message PhysicalNegativeNode { PhysicalExprNode expr = 1; } +message PhysicalNormalizeFloatZeroNode { + PhysicalExprNode expr = 1; +} + message PhysicalExtensionExprNode { bytes expr = 1; repeated PhysicalExprNode inputs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 908f9752b7f18..603b6291a6c60 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -18492,6 +18492,9 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::RangeExpr(v) => { struct_ser.serialize_field("rangeExpr", v)?; } + physical_expr_node::ExprType::NormalizeFloatZero(v) => { + struct_ser.serialize_field("normalizeFloatZero", v)?; + } } } struct_ser.end() @@ -18549,6 +18552,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable", "range_expr", "rangeExpr", + "normalize_float_zero", + "normalizeFloatZero", ]; #[allow(clippy::enum_variant_names)] @@ -18579,6 +18584,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { Lambda, LambdaVariable, RangeExpr, + NormalizeFloatZero, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -18626,6 +18632,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambda" => Ok(GeneratedField::Lambda), "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), "rangeExpr" | "range_expr" => Ok(GeneratedField::RangeExpr), + "normalizeFloatZero" | "normalize_float_zero" => Ok(GeneratedField::NormalizeFloatZero), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18830,6 +18837,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("rangeExpr")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::RangeExpr) +; + } + GeneratedField::NormalizeFloatZero => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("normalizeFloatZero")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::NormalizeFloatZero) ; } } @@ -20200,6 +20214,97 @@ impl<'de> serde::Deserialize<'de> for PhysicalNegativeNode { deserializer.deserialize_struct("datafusion.PhysicalNegativeNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalNormalizeFloatZeroNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.expr.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalNormalizeFloatZeroNode", len)?; + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalNormalizeFloatZeroNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "expr", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Expr, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "expr" => Ok(GeneratedField::Expr), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalNormalizeFloatZeroNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalNormalizeFloatZeroNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut expr__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + } + } + Ok(PhysicalNormalizeFloatZeroNode { + expr: expr__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalNormalizeFloatZeroNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalNot { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index ba00577ab9a1b..0b42ae7d8d110 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1557,7 +1557,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option, #[prost( oneof = "physical_expr_node::ExprType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28" )] pub expr_type: ::core::option::Option, } @@ -1622,6 +1622,10 @@ pub mod physical_expr_node { LambdaVariable(super::PhysicalLambdaVariableExprNode), #[prost(message, tag = "27")] RangeExpr(super::PhysicalRangeExprNode), + #[prost(message, tag = "28")] + NormalizeFloatZero( + ::prost::alloc::boxed::Box, + ), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1847,6 +1851,11 @@ pub struct PhysicalNegativeNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalNormalizeFloatZeroNode { + #[prost(message, optional, boxed, tag = "1")] + pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalExtensionExprNode { #[prost(bytes = "vec", tag = "1")] pub expr: ::prost::alloc::vec::Vec, diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 06105be806cfc..a50f0eae89ae3 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -36,7 +36,8 @@ use datafusion_physical_expr::{ }; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, - LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, + LikeExpr, Literal, NegativeExpr, NormalizeFloatZeroExpr, NotExpr, TryCastExpr, + UnKnownColumn, }; use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; @@ -287,6 +288,9 @@ pub fn parse_physical_expr_with_converter( ExprType::IsNotNullExpr(_) => IsNotNullExpr::try_from_proto(proto, &decode_ctx)?, ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::NormalizeFloatZero(_) => { + NormalizeFloatZeroExpr::try_from_proto(proto, &decode_ctx)? + } ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, diff --git a/datafusion/proto/tests/cases/plans/sorts.rs b/datafusion/proto/tests/cases/plans/sorts.rs index 1172775b1bad7..92a4d5173b944 100644 --- a/datafusion/proto/tests/cases/plans/sorts.rs +++ b/datafusion/proto/tests/cases/plans/sorts.rs @@ -22,7 +22,9 @@ use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::LexOrdering; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::physical_plan::expressions::{ + NormalizeFloatZeroExpr, PhysicalSortExpr, col, +}; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::prelude::SessionContext; @@ -61,6 +63,20 @@ fn roundtrip_sort() -> Result<()> { ))) } +#[test] +fn roundtrip_sort_with_normalized_float_zero() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); + let sort_exprs = [PhysicalSortExpr { + expr: Arc::new(NormalizeFloatZeroExpr::new(col("a", &schema)?)), + options: SortOptions::default(), + }] + .into(); + roundtrip_test(Arc::new(SortExec::new( + sort_exprs, + Arc::new(EmptyExec::new(schema)), + ))) +} + #[test] fn roundtrip_sort_preserve_partitioning() -> Result<()> { let field_a = Field::new("a", DataType::Boolean, false); From 5201c34ae2cefb382f4b7e877ed46f7283a52e1c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 16 Aug 2026 23:47:49 +0800 Subject: [PATCH 20/20] 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 {