From e546de5510de599c91f756bef4bf45243d93011a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:05:01 +0800 Subject: [PATCH 1/4] 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 2/4] 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 ed2ce92b0a7f822da9dd80c1b886ad7dc124a636 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 14:11:56 +0800 Subject: [PATCH 3/4] feat: serialize ASOF join plans --- .../proto-models/proto/datafusion.proto | 31 + .../proto-models/src/generated/pbjson.rs | 530 ++++++++++++++++++ .../proto-models/src/generated/prost.rs | 83 ++- .../proto/src/logical_plan/from_proto.rs | 16 + datafusion/proto/src/logical_plan/mod.rs | 121 +++- datafusion/proto/src/logical_plan/to_proto.rs | 20 +- datafusion/proto/src/physical_plan/mod.rs | 140 ++++- .../tests/cases/roundtrip_logical_plan.rs | 40 +- .../tests/cases/roundtrip_physical_plan.rs | 37 +- 9 files changed, 998 insertions(+), 20 deletions(-) diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 205cf89abed1b..f61a59f058cbd 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -63,6 +63,7 @@ message LogicalPlanNode { CteWorkTableScanNode cte_work_table_scan = 32; DmlNode dml = 33; EmptyTableScanNode empty_table_scan = 34; + AsOfJoinNode as_of_join = 35; } } @@ -276,6 +277,25 @@ message JoinNode { bool null_aware = 9; } +enum AsOfMatchOperator { + AS_OF_MATCH_OPERATOR_UNSPECIFIED = 0; + AS_OF_MATCH_OPERATOR_LT = 1; + AS_OF_MATCH_OPERATOR_LT_EQ = 2; + AS_OF_MATCH_OPERATOR_GT = 3; + AS_OF_MATCH_OPERATOR_GT_EQ = 4; +} + +message AsOfJoinNode { + LogicalPlanNode left = 1; + LogicalPlanNode right = 2; + repeated LogicalExprNode left_join_key = 3; + repeated LogicalExprNode right_join_key = 4; + LogicalExprNode left_match_expr = 5; + LogicalExprNode right_match_expr = 6; + AsOfMatchOperator match_operator = 7; + datafusion_common.JoinConstraint join_constraint = 8; +} + message DistinctNode { LogicalPlanNode input = 1; } @@ -881,6 +901,7 @@ message PhysicalPlanNode { BufferExecNode buffer = 37; ArrowScanExecNode arrow_scan = 38; ScalarSubqueryExecNode scalar_subquery = 39; + AsOfJoinExecNode as_of_join = 40; } } @@ -1637,6 +1658,16 @@ message SortMergeJoinExecNode { datafusion_common.NullEquality null_equality = 7; } +message AsOfJoinExecNode { + PhysicalPlanNode left = 1; + PhysicalPlanNode right = 2; + repeated JoinOn on = 3; + PhysicalExprNode left_match_expr = 4; + PhysicalExprNode right_match_expr = 5; + AsOfMatchOperator match_operator = 6; + repeated uint32 right_output_indices = 7; +} + message AsyncFuncExecNode { PhysicalPlanNode input = 1; repeated PhysicalExprNode async_exprs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d23f8eee5fd2c..d1f5c8d6985fa 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1520,6 +1520,508 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode { deserializer.deserialize_struct("datafusion.ArrowScanExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for AsOfJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.on.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if !self.right_output_indices.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.on.is_empty() { + struct_ser.serialize_field("on", &self.on)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if !self.right_output_indices.is_empty() { + struct_ser.serialize_field("rightOutputIndices", &self.right_output_indices)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "on", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "right_output_indices", + "rightOutputIndices", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + On, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + RightOutputIndices, + } + 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 { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "on" => Ok(GeneratedField::On), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "rightOutputIndices" | "right_output_indices" => Ok(GeneratedField::RightOutputIndices), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut on__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut right_output_indices__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::RightOutputIndices => { + if right_output_indices__.is_some() { + return Err(serde::de::Error::duplicate_field("rightOutputIndices")); + } + right_output_indices__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(AsOfJoinExecNode { + left: left__, + right: right__, + on: on__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + right_output_indices: right_output_indices__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfJoinNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.left_join_key.is_empty() { + len += 1; + } + if !self.right_join_key.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if self.join_constraint != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.left_join_key.is_empty() { + struct_ser.serialize_field("leftJoinKey", &self.left_join_key)?; + } + if !self.right_join_key.is_empty() { + struct_ser.serialize_field("rightJoinKey", &self.right_join_key)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if self.join_constraint != 0 { + let v = super::datafusion_common::JoinConstraint::try_from(self.join_constraint) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_constraint)))?; + struct_ser.serialize_field("joinConstraint", &v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "left_join_key", + "leftJoinKey", + "right_join_key", + "rightJoinKey", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "join_constraint", + "joinConstraint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + LeftJoinKey, + RightJoinKey, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + JoinConstraint, + } + 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 { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "leftJoinKey" | "left_join_key" => Ok(GeneratedField::LeftJoinKey), + "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "joinConstraint" | "join_constraint" => Ok(GeneratedField::JoinConstraint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut left_join_key__ = None; + let mut right_join_key__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut join_constraint__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::LeftJoinKey => { + if left_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("leftJoinKey")); + } + left_join_key__ = Some(map_.next_value()?); + } + GeneratedField::RightJoinKey => { + if right_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("rightJoinKey")); + } + right_join_key__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::JoinConstraint => { + if join_constraint__.is_some() { + return Err(serde::de::Error::duplicate_field("joinConstraint")); + } + join_constraint__ = Some(map_.next_value::()? as i32); + } + } + } + Ok(AsOfJoinNode { + left: left__, + right: right__, + left_join_key: left_join_key__.unwrap_or_default(), + right_join_key: right_join_key__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + join_constraint: join_constraint__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfMatchOperator { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for AsOfMatchOperator { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + "AS_OF_MATCH_OPERATOR_LT", + "AS_OF_MATCH_OPERATOR_LT_EQ", + "AS_OF_MATCH_OPERATOR_GT", + "AS_OF_MATCH_OPERATOR_GT_EQ", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = AsOfMatchOperator; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Ok(AsOfMatchOperator::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Ok(AsOfMatchOperator::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Ok(AsOfMatchOperator::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Ok(AsOfMatchOperator::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Ok(AsOfMatchOperator::GtEq), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for AsyncFuncExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -13674,6 +14176,9 @@ impl serde::Serialize for LogicalPlanNode { logical_plan_node::LogicalPlanType::EmptyTableScan(v) => { struct_ser.serialize_field("emptyTableScan", v)?; } + logical_plan_node::LogicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -13735,6 +14240,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "dml", "empty_table_scan", "emptyTableScan", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -13772,6 +14279,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { CteWorkTableScan, Dml, EmptyTableScan, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -13826,6 +14334,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "cteWorkTableScan" | "cte_work_table_scan" => Ok(GeneratedField::CteWorkTableScan), "dml" => Ok(GeneratedField::Dml), "emptyTableScan" | "empty_table_scan" => Ok(GeneratedField::EmptyTableScan), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14077,6 +14586,13 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { return Err(serde::de::Error::duplicate_field("emptyTableScan")); } logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::EmptyTableScan) +; + } + GeneratedField::AsOfJoin => { + if logical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::AsOfJoin) ; } } @@ -20331,6 +20847,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::ScalarSubquery(v) => { struct_ser.serialize_field("scalarSubquery", v)?; } + physical_plan_node::PhysicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -20403,6 +20922,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "arrowScan", "scalar_subquery", "scalarSubquery", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -20445,6 +20966,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { Buffer, ArrowScan, ScalarSubquery, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20504,6 +21026,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "buffer" => Ok(GeneratedField::Buffer), "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -20790,6 +21313,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("scalarSubquery")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::ScalarSubquery) +; + } + GeneratedField::AsOfJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::AsOfJoin) ; } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 6baabbf37a41c..0ece21da6dcb8 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -5,7 +5,7 @@ pub struct LogicalPlanNode { #[prost( oneof = "logical_plan_node::LogicalPlanType", - tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" + tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35" )] pub logical_plan_type: ::core::option::Option, } @@ -79,6 +79,8 @@ pub mod logical_plan_node { Dml(::prost::alloc::boxed::Box), #[prost(message, tag = "34")] EmptyTableScan(super::EmptyTableScanNode), + #[prost(message, tag = "35")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -421,6 +423,29 @@ pub struct JoinNode { pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub left_join_key: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "4")] + pub right_join_key: ::prost::alloc::vec::Vec, + #[prost(message, optional, boxed, tag = "5")] + pub left_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(message, optional, boxed, tag = "6")] + pub right_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(enumeration = "AsOfMatchOperator", tag = "7")] + pub match_operator: i32, + #[prost(enumeration = "super::datafusion_common::JoinConstraint", tag = "8")] + pub join_constraint: i32, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1292,7 +1317,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40" )] pub physical_plan_type: ::core::option::Option, } @@ -1378,6 +1403,8 @@ pub mod physical_plan_node { ArrowScan(super::ArrowScanExecNode), #[prost(message, tag = "39")] ScalarSubquery(::prost::alloc::boxed::Box), + #[prost(message, tag = "40")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2473,6 +2500,23 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub on: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "4")] + pub left_match_expr: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub right_match_expr: ::core::option::Option, + #[prost(enumeration = "AsOfMatchOperator", tag = "6")] + pub match_operator: i32, + #[prost(uint32, repeated, tag = "7")] + pub right_output_indices: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -2504,6 +2548,41 @@ pub struct PhysicalScalarSubqueryExprNode { #[prost(uint32, tag = "3")] pub index: u32, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum AsOfMatchOperator { + Unspecified = 0, + Lt = 1, + LtEq = 2, + Gt = 3, + GtEq = 4, +} +impl AsOfMatchOperator { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Some(Self::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Some(Self::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Some(Self::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Some(Self::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Some(Self::GtEq), + _ => None, + } + } +} /// Identifies a built-in file format supported by DataFusion. /// Used by DefaultLogicalExtensionCodec to serialize/deserialize /// FileFormatFactory instances (e.g. in CopyTo plans). diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 6d9a73e06ff45..8e6fe2f3d9c29 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -234,6 +234,22 @@ impl FromProto for JoinConstraint { } } +impl TryFromProto for Operator { + type Error = Error; + + fn try_from_proto(value: protobuf::AsOfMatchOperator) -> Result { + match value { + protobuf::AsOfMatchOperator::Lt => Ok(Self::Lt), + protobuf::AsOfMatchOperator::LtEq => Ok(Self::LtEq), + protobuf::AsOfMatchOperator::Gt => Ok(Self::Gt), + protobuf::AsOfMatchOperator::GtEq => Ok(Self::GtEq), + protobuf::AsOfMatchOperator::Unspecified => Err(Error::General( + "ASOF match operator must be specified".to_string(), + )), + } + } +} + impl FromProto for NullEquality { fn from_proto(t: protobuf::NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 353d44dc3536e..0fa9d501cd804 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -60,17 +60,17 @@ use datafusion_datasource_json::file_format::{ use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, Operator, RangePartitioning, RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, }; use datafusion_expr::{ DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ - Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection, - Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window, - builder::project, + Aggregate, AsOfJoin, AsOfMatch, CreateCatalog, CreateCatalogSchema, + CreateExternalTable, CreateView, DdlStatement, Distinct, EmptyRelation, + Extension, Join, Prepare, Projection, Repartition, Sort, SubqueryAlias, + TableScan, TableScanBuilder, Values, Window, builder::project, }, }; use datafusion_proto_common::protobuf_common; @@ -1033,6 +1033,61 @@ impl AsLogicalPlan for LogicalPlanNode { join.null_aware, )?)) } + LogicalPlanType::AsOfJoin(join) => { + let left_keys = + from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; + let right_keys = + from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + if left_keys.len() != right_keys.len() { + return Err(proto_error(format!( + "Received an AsOfJoinNode with left_join_key and right_join_key of different lengths: {} and {}", + left_keys.len(), + right_keys.len() + ))); + } + let left_match = from_proto::parse_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode left_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let right_match = from_proto::parse_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode right_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from( + join.match_operator, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + let join_constraint = protobuf::JoinConstraint::try_from( + join.join_constraint, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF JoinConstraint {}", + join.join_constraint + )) + })?; + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + left_keys.into_iter().zip(right_keys).collect(), + AsOfMatch::new(left_match, op, right_match), + JoinConstraint::from_proto(join_constraint), + )?)) + } LogicalPlanType::Union(union) => { assert_or_internal_err!( union.inputs.len() >= 2, @@ -1685,6 +1740,59 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + .. + }) => { + let left = LogicalPlanNode::try_from_logical_plan( + left.as_ref(), + extension_codec, + )?; + let right = LogicalPlanNode::try_from_logical_plan( + right.as_ref(), + extension_codec, + )?; + let (left_join_key, right_join_key) = on + .iter() + .map(|(left, right)| { + Ok(( + serialize_expr(left, extension_codec)?, + serialize_expr(right, extension_codec)?, + )) + }) + .collect::, ToProtoError>>()? + .into_iter() + .unzip(); + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(match_condition.op)?; + Ok(LogicalPlanNode { + logical_plan_type: Some(LogicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + left_join_key, + right_join_key, + left_match_expr: Some(Box::new(serialize_expr( + &match_condition.left, + extension_codec, + )?)), + right_match_expr: Some(Box::new(serialize_expr( + &match_condition.right, + extension_codec, + )?)), + match_operator: match_operator.into(), + join_constraint: protobuf::JoinConstraint::from_proto( + *join_constraint, + ) + .into(), + }, + ))), + }) + } LogicalPlan::Subquery(subquery) => { // Serialize the inner subquery plan directly — the // LogicalPlan::Subquery wrapper is reconstructed during @@ -2190,9 +2298,6 @@ 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/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 23ce254e99a40..1fc0a03944c71 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -32,8 +32,8 @@ use datafusion_expr::expr::{ }; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ - Expr, JoinConstraint, JoinType, SortExpr, TryCast, WindowFrame, WindowFrameBound, - WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, + Expr, JoinConstraint, JoinType, Operator, SortExpr, TryCast, WindowFrame, + WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, logical_plan::StringifiedPlan, }; @@ -782,6 +782,22 @@ impl FromProto for protobuf::JoinConstraint { } } +impl TryFromProto for protobuf::AsOfMatchOperator { + type Error = Error; + + fn try_from_proto(value: Operator) -> Result { + match value { + Operator::Lt => Ok(Self::Lt), + Operator::LtEq => Ok(Self::LtEq), + Operator::Gt => Ok(Self::Gt), + Operator::GtEq => Ok(Self::GtEq), + op => Err(Error::General(format!( + "Unsupported ASOF match operator {op}" + ))), + } + } +} + impl FromProto for protobuf::NullEquality { fn from_proto(t: NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index cea334e42aace..2c27fd23f1019 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -54,7 +54,7 @@ use datafusion_datasource_parquet::source::ParquetSource; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, Operator, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; @@ -83,8 +83,8 @@ use datafusion_physical_plan::expressions::PhysicalSortExpr; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -865,6 +865,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::SortMergeJoin(_) => { SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } + PhysicalPlanType::AsOfJoin(asof_join) => { + self.try_into_asof_join(asof_join, ctx, proto_converter) + } PhysicalPlanType::AsyncFunc(async_func) => { self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) } @@ -949,6 +952,14 @@ pub trait PhysicalPlanNodeExt: Sized { ); } + if let Some(exec) = plan.downcast_ref::() { + return protobuf::PhysicalPlanNode::try_from_asof_join_exec( + exec, + codec, + proto_converter, + ); + } + if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_cross_join_exec( exec, @@ -2561,6 +2572,73 @@ pub trait PhysicalPlanNodeExt: Sized { SortMergeJoinExec::try_from_proto(&node, &decode_ctx) } + fn try_into_asof_join( + &self, + join: &protobuf::AsOfJoinExecNode, + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let left = into_physical_plan(&join.left, ctx, proto_converter)?; + let right = into_physical_plan(&join.right, ctx, proto_converter)?; + let on = join + .on + .iter() + .map(|pair| { + let left_expr = proto_converter.proto_to_physical_expr( + pair.left.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode left equality expression is missing", + ) + })?, + left.schema().as_ref(), + ctx, + )?; + let right_expr = proto_converter.proto_to_physical_expr( + pair.right.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode right equality expression is missing", + ) + })?, + right.schema().as_ref(), + ctx, + )?; + Ok((left_expr, right_expr)) + }) + .collect::>()?; + let left_match = proto_converter.proto_to_physical_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode left_match_expr is missing") + })?, + left.schema().as_ref(), + ctx, + )?; + let right_match = proto_converter.proto_to_physical_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode right_match_expr is missing") + })?, + right.schema().as_ref(), + ctx, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from(join.match_operator) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, op, right_match), + join.right_output_indices + .iter() + .map(|index| *index as usize) + .collect(), + )?)) + } + fn try_into_generate_series_physical_plan( &self, generate_series: &protobuf::GenerateSeriesNode, @@ -3138,6 +3216,62 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + fn try_from_asof_join_exec( + exec: &AsOfJoinExec, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.left()), + codec, + proto_converter, + )?; + let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.right()), + codec, + proto_converter, + )?; + let on = exec + .on() + .iter() + .map(|(left, right)| { + Ok::<_, DataFusionError>(protobuf::JoinOn { + left: Some(proto_converter.physical_expr_to_proto(left, codec)?), + right: Some(proto_converter.physical_expr_to_proto(right, codec)?), + }) + }) + .collect::>()?; + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(exec.match_condition().op)?; + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + left_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().left, + codec, + )?, + ), + right_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().right, + codec, + )?, + ), + match_operator: match_operator.into(), + right_output_indices: exec + .right_output_indices() + .iter() + .map(|index| *index as u32) + .collect(), + }, + ))), + }) + } + fn try_from_cross_join_exec( exec: &CrossJoinExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 74f7253386764..f6113a807e8a7 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -74,8 +74,9 @@ use datafusion_common::format::{ }; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, - TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, + Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, + SplitPoint, TableReference, internal_datafusion_err, internal_err, not_impl_err, + plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -90,7 +91,7 @@ use datafusion_expr::logical_plan::{ ExplainOption, Extension, UserDefinedLogicalNodeCore, }; use datafusion_expr::{ - Accumulator, AggregateUDF, ColumnarValue, DmlStatement, ExprFunctionExt, + Accumulator, AggregateUDF, AsOfMatch, ColumnarValue, DmlStatement, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, @@ -3773,6 +3774,39 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_asof_join() -> Result<()> { + let ctx = SessionContext::new(); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + ctx.register_table("trades", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("prices", Arc::new(EmptyTable::new(right_schema)))?; + + let left = ctx.table("trades").await?.into_optimized_plan()?; + let right = ctx.table("prices").await?.into_optimized_plan()?; + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + let plan = LogicalPlanBuilder::from(left.clone()) + .asof_join_using( + right.clone(), + vec![Column::from_name("symbol")], + AsOfMatch::new(col("trades.ts"), op, col("prices.ts")), + )? + .build()?; + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{round_trip:?}")); + } + Ok(()) +} + // Single column, single split point range partitioning #[tokio::test] async fn roundtrip_range_partitioning_single_col() -> Result<()> { diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 2671f4d0152f7..682b8471320ab 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -72,8 +72,8 @@ use datafusion::physical_plan::expressions::{ }; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, HashJoinExec, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; @@ -478,6 +478,39 @@ fn roundtrip_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_asof_join() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_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 on = vec![( + Arc::new(Column::new("symbol", 0)) as _, + Arc::new(Column::new("symbol", 0)) as _, + )]; + + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + roundtrip_test(Arc::new(AsOfJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&left_schema))), + Arc::new(EmptyExec::new(Arc::clone(&right_schema))), + on.clone(), + AsOfMatchExpr::new( + Arc::new(Column::new("ts", 1)), + op, + Arc::new(Column::new("ts", 1)), + ), + vec![2], + )?))?; + } + Ok(()) +} + #[test] fn roundtrip_nested_loop_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); From d36fd277c1deb42d44b9d445c09795ed3a84ab2c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 14:39:34 +0800 Subject: [PATCH 4/4] fix: keep generated protobuf code in sync --- datafusion/proto-models/src/generated/pbjson.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d1f5c8d6985fa..eb2f554095b2d 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1701,7 +1701,7 @@ impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { if right_output_indices__.is_some() { return Err(serde::de::Error::duplicate_field("rightOutputIndices")); } - right_output_indices__ = + right_output_indices__ = Some(map_.next_value::>>()? .into_iter().map(|x| x.0).collect()) ;