From 7880ea2da8a4bd0e11d766a41cdfc64ce862a64c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 20 Aug 2026 01:46:26 +0000 Subject: [PATCH 1/3] feat: add ASOF join physical operator (#23828) ## Which issue does this PR close? - Part of #318. - Umbrella PR: #23738. - Follow-up for floating-point equality keys: #24375. ## Rationale for this change This is the first layer of the ASOF JOIN stack. It establishes a broadcast-based physical execution contract independently so later floating-point equality, logical-plan, SQL, DataFrame, and serialization changes can be reviewed as smaller follow-up PRs. The initial implementation deliberately favors the simpler broadcast design: the right input must fit in memory and each left partition scans the shared right-side batches. A repartitioned implementation can be evaluated separately without changing the ASOF semantics introduced here. Floating-point equality keys are rejected in this base layer because Arrow's required sort order distinguishes `-0.0` from `+0.0` while join equality does not. #24375 adds the required ordering normalization as an independently reviewable layer. ## What changes are included in this PR? - Add `AsOfJoinExec` for left-preserving, Snowflake-style ASOF semantics. - Coalesce and collect the ordered right input once, then share it across all left partitions. - Keep the left input partitioned so each partition can scan independently and preserve the left-side output partitioning. - Preserve merge state across input and output batch boundaries. - Reserve each retained Arrow buffer exactly once, including when right-side batches are zero-copy slices, and expose build, match, and output metrics. - Define output properties and statistics for the broadcast execution model. - Reject floating-point equality keys until #24375 supplies a sort/equality contract that handles signed zero correctly. - Add physical operator tests covering match directions, equality groups, batch boundaries, unmatched rows, invalid contracts, shared-buffer memory accounting, multi-partition broadcast execution, and float-key rejection. ## Are these changes tested? Yes: - `cargo fmt --all` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-physical-plan joins::asof_join --all-features` - Extended workspace tests from the contributor guide - FFI integration tests ## Are there any user-facing changes? This adds a new physical operator API. The base operator deliberately rejects floating-point equality keys; #24375 adds full Float16, Float32, and Float64 support. SQL and DataFrame APIs are left to later dependent PRs. --------- Co-authored-by: Yongting You <2010youy01@gmail.com> --- .../physical-plan/src/joins/asof_join.rs | 1782 +++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 2 + 2 files changed, 1784 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..22625afca1f49 --- /dev/null +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -0,0 +1,1782 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Broadcast, left-preserving ASOF join execution. +//! +//! An ASOF join emits exactly one output row for every left row. Within an +//! optional equality-key group, it selects the closest right row that satisfies +//! one ordered comparison. This follows Snowflake's [ASOF JOIN] semantics: +//! +//! ```text +//! left.ts >= right.ts => greatest eligible right.ts +//! left.ts <= right.ts => smallest eligible right.ts +//! ``` +//! +//! The right input is collected and shared by all output partitions. The left +//! input remains partitioned, and each partition performs an independent +//! monotonic scan over the ordered right input. +//! +//! [`AsOfJoinExec::input_distribution_requirements`] requires a single right +//! partition but leaves the left distribution unrestricted. +//! [`AsOfJoinExec::required_input_ordering`] requires both inputs to be ordered. +//! The physical optimizer satisfies these contracts by inserting operators such +//! as `RepartitionExec`, `SortExec`, `CoalescePartitionsExec`, or +//! `SortPreservingMergeExec`, depending on the input properties. The inserted +//! plan shape is therefore not fixed by this operator. +//! +//! Both inputs must be ordered by their equality keys followed by the match +//! key. For `<` and `<=`, the match ordering is reversed so all directions use +//! the same forward-only state machine. For example: +//! +//! ```text +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts >= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts ASC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts ASC NULLS FIRST] +//! +//! ON left.symbol = right.symbol MATCH_CONDITION(left.ts <= right.ts) +//! left: [left.symbol ASC NULLS FIRST, left.ts DESC NULLS FIRST] +//! right: [right.symbol ASC NULLS FIRST, right.ts DESC NULLS FIRST] +//! ``` +//! +//! Each left partition owns its cursors, equality-group state, and current +//! candidate, while the collected right batches are immutable and shared. +//! The key state-machine entry point is [`AsOfJoinStream::poll_next_impl`]. +//! +//! This mode preserves probe-side parallelism when there are no equality keys +//! or when equality keys have low cardinality or skew. It retains the complete +//! right input in the memory pool and may scan it once per left partition. +//! Alternative strategies, including broadcasting the other side or +//! repartitioning both inputs, remain future work for other input-size and +//! key-distribution profiles. +//! +//! [ASOF JOIN]: https://docs.snowflake.com/en/sql-reference/constructs/asof-join + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}; +use arrow::buffer::NullBuffer; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; +use datafusion_common::utils::normalize_float_zero_scalar; +use datafusion_common::{ + ColumnStatistics, JoinSide, JoinType, NullEquality, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, project_schema, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::{ProjectionMapping, ProjectionRef}; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{Stream, StreamExt, TryStreamExt, future::poll_fn, ready, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::joins::utils::{ + ColumnIndex, JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, + matchable_join_keys, +}; +use crate::memory::MemoryStream; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, + RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + validate_child_count, +}; + +/// Physical ordered comparison for an ASOF join. +#[derive(Debug, Clone)] +pub struct AsOfMatchExpr { + /// Expression evaluated against the left input. + pub left: PhysicalExprRef, + /// Ordered comparison operator. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: PhysicalExprRef, +} + +impl AsOfMatchExpr { + /// Creates a physical ASOF match expression. + pub fn new(left: PhysicalExprRef, op: Operator, right: PhysicalExprRef) -> Self { + Self { left, op, right } + } +} + +/// A broadcast sort-merge ASOF join that emits one row for every left row. +#[derive(Debug)] +pub struct AsOfJoinExec { + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + /// Unprojected left-join schema used to interpret `projection`. + join_schema: SchemaRef, + /// Information of index and left/right placement of columns. + column_indices: Vec, + /// Optional indices into the full left-then-right join schema. + projection: Option, + metrics: ExecutionPlanMetricsSet, + /// Required ordering for each left partition. + left_ordering: LexOrdering, + /// Required global ordering for the single right partition. + right_ordering: LexOrdering, + /// Shared collection future that materializes the right input only once. + right_fut: OnceAsync, + cache: Arc, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + /// + /// The match operator must be `<`, `<=`, `>`, or `>=`. Equality and match + /// expressions must be deterministic, reference only their corresponding + /// input, and have matching input types. Equality types must support hashing; + /// floating-point equality keys are not supported because Arrow sorting + /// distinguishes signed zero while SQL equality does not. Projection indices + /// refer to the full left-then-right join schema. + pub fn try_new( + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + projection: Option>, + ) -> Result { + validate_asof_join(left.as_ref(), right.as_ref(), &on, &match_condition)?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &JoinType::Left); + let join_schema = Arc::new(join_schema); + let projection: Option = projection.map(Into::into); + let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); + let equality_options = SortOptions { + descending: false, + nulls_first: true, + }; + let match_options = SortOptions { + descending, + nulls_first: true, + }; + let mut left_sort_exprs = on + .iter() + .map(|(left, _)| PhysicalSortExpr { + expr: Arc::clone(left), + options: equality_options, + }) + .collect::>(); + 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, + &join_schema, + projection.as_deref(), + )?); + + Ok(Self { + left, + right, + on, + match_condition, + join_schema, + column_indices, + projection, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + right_fut: Default::default(), + cache, + }) + } + + fn compute_properties( + left: &Arc, + join_schema: &SchemaRef, + projection: Option<&[usize]>, + ) -> 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 mut eq_properties = + input_eq_properties.project(&mapping, Arc::clone(join_schema)); + let mut output_partitioning = left + .output_partitioning() + .project(&mapping, input_eq_properties); + if let Some(projection) = projection { + let projection_mapping = + ProjectionMapping::from_indices(projection, join_schema)?; + let output_schema = project_schema(join_schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, output_schema); + } + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + let on = self + .on + .iter() + .map(|(left, right)| { + format!("({} = {})", fmt_sql(left.as_ref()), fmt_sql(right.as_ref())) + }) + .collect::>() + .join(", "); + let match_condition = format!( + "{} {} {}", + fmt_sql(self.match_condition.left.as_ref()), + self.match_condition.op, + fmt_sql(self.match_condition.right.as_ref()) + ); + let projection = self + .projection + .as_ref() + .map(|projection| { + format!( + ", projection=[{}]", + projection + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.field(*index).name(), + index + )) + .collect::>() + .join(", ") + ) + }) + .unwrap_or_default(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]{}", + Self::static_name(), + on, + match_condition, + projection + ), + DisplayFormatType::TreeRender => { + writeln!(f, "on={on}")?; + writeln!(f, "match={match_condition}") + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + // Every left partition scans the complete broadcast right input, so + // equality keys do not require the inputs to be co-partitioned. + // `UnspecifiedDistribution` imposes no layout requirement; because this + // operator uses the default `benefits_from_input_partitioning`, the + // optimizer may still add round-robin repartitioning when it is useful. + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition, + ]) + } + + fn required_input_ordering(&self) -> Vec> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec { + // ASOF emits exactly one row for each left row and never reorders the + // left input. The right input is scanned independently. + vec![true, false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + crate::apply_expression_roots( + join_keys.chain([&self.match_condition.left, &self.match_condition.right]), + f, + ) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + let left = children.swap_remove(0); + let right = children.swap_remove(0); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + left, + right, + on: self.on.clone(), + match_condition: self.match_condition.clone(), + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + metrics: ExecutionPlanMetricsSet::new(), + left_ordering: self.left_ordering.clone(), + right_ordering: self.right_ordering.clone(), + right_fut: Default::default(), + cache: Arc::clone(&self.cache), + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::try_new( + left, + right, + self.on.clone(), + self.match_condition.clone(), + self.projection.as_deref().map(<[usize]>::to_vec), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + right_partitions, + 1, + "AsOfJoinExec requires one right partition, found {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let metrics = AsOfJoinMetrics::new(partition, &self.metrics); + let build_metrics = metrics.clone(); + let right_fut = self.right_fut.try_once(|| { + let right_stream = self.right.execute(0, Arc::clone(&context))?; + let reservation = + MemoryConsumer::new("AsOfJoinInput").register(context.memory_pool()); + Ok(collect_right_input( + right_stream, + reservation, + build_metrics, + )) + })?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let output_schema = self.schema(); + let stream_schema = Arc::clone(&output_schema); + let left_match = Arc::clone(&self.match_condition.left); + let right_match = Arc::clone(&self.match_condition.right); + let match_op = self.match_condition.op; + let column_indices = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let batch_size = context.session_config().batch_size(); + let stream = stream::once(async move { + let mut right_fut = right_fut; + let right_input = poll_fn(|cx| right_fut.get_shared(cx)).await?; + let right_stream = right_input.stream()?; + let stream = AsOfJoinStream::new( + Arc::clone(&stream_schema), + InputCursor::new(left_stream, left_keys, left_match), + InputCursor::new(right_stream, right_keys, right_match), + match_op, + column_indices, + batch_size, + metrics, + right_input, + ); + Ok::(Box::pin( + stream, + )) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + output_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> { + // The default is fully unknown, but ASOF emits exactly one output row + // per left row and preserves statistics for unmodified left columns. + let left = &input_stats[0]; + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|index| self.column_indices[*index].clone()) + .collect(), + None => self.column_indices.clone(), + }; + let column_statistics = column_indices_after_projection + .iter() + .map(|column| match column.side { + JoinSide::Left => left + .column_statistics + .get(column.index) + .cloned() + .unwrap_or_else(ColumnStatistics::new_unknown), + JoinSide::Right | JoinSide::None => ColumnStatistics::new_unknown(), + }) + .collect(); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } +} + +/// Materialized right input shared by every left output partition. +struct BroadcastRightInput { + /// Schema retained even when the input has no batches. + schema: SchemaRef, + /// Ordered right batches; their buffers are shared without copying. + batches: Vec, + /// Holds the memory-pool reservation for as long as the batches are shared. + _reservation: MemoryReservation, +} + +impl BroadcastRightInput { + fn stream(&self) -> Result { + Ok(Box::pin(MemoryStream::try_new( + self.batches.clone(), + Arc::clone(&self.schema), + None, + )?)) + } +} + +async fn collect_right_input( + input: SendableRecordBatchStream, + reservation: MemoryReservation, + metrics: AsOfJoinMetrics, +) -> Result { + let schema = input.schema(); + let mut memory_counter = RecordBatchMemoryCounter::new(); + let batches = input + .try_fold(Vec::new(), |mut batches, batch| { + let batch_size = memory_counter.count_batch(&batch); + futures::future::ready(reservation.try_grow(batch_size).map(|_| { + metrics.build_mem_used.add(batch_size); + batches.push(batch); + batches + })) + }) + .await?; + Ok(BroadcastRightInput { + schema, + batches, + _reservation: reservation, + }) +} + +/// Last eligible right row for the current left equality group. +/// +/// The row and its evaluated keys survive right batch changes and output +/// flushes. It belongs to the join state rather than `InputCursor` because its +/// validity also depends on the current left equality group. +#[derive(Clone)] +struct Candidate { + /// Right batch containing the nearest eligible row. + batch: Arc, + /// Row index within `batch`. + row: usize, + /// Evaluated equality keys retained when the right cursor changes batches. + key_arrays: Arc<[ArrayRef]>, + /// Identity used to invalidate the cached candidate/left comparator. + key_batch_id: usize, +} + +/// Cursor over one ordered input stream. +/// +/// Expressions are evaluated once per non-empty batch. `key_batch_id` changes +/// whenever a new batch is loaded so comparators cannot retain stale arrays. +struct InputCursor { + /// Remaining input batches. + stream: SendableRecordBatchStream, + /// Equality expressions evaluated for each batch. + key_exprs: Vec, + /// Ordered match expression evaluated for each batch. + match_expr: PhysicalExprRef, + /// Current non-empty batch. + batch: Option>, + /// Evaluated equality-key arrays for `batch`. + key_arrays: Arc<[ArrayRef]>, + /// Rows whose equality keys are all non-NULL. + key_validity: Option, + /// Evaluated match values for `batch`. + match_array: Option, + /// Monotonic identity of the current key arrays. + key_batch_id: usize, + /// Current row within `batch`. + row: usize, + /// Whether the input stream has returned EOF. + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: Arc::from([]), + key_validity: None, + match_array: None, + key_batch_id: 0, + row: 0, + eof: false, + } + } + + fn poll_ensure_row( + &mut self, + cx: &mut Context<'_>, + elapsed_compute: &Time, + ) -> Poll> { + loop { + if let Some(batch) = &self.batch + && self.row < batch.num_rows() + { + return Poll::Ready(Ok(true)); + } + self.batch = None; + self.key_arrays = Arc::from([]); + self.key_validity = None; + self.match_array = None; + self.row = 0; + if self.eof { + return Poll::Ready(Ok(false)); + } + let Some(batch) = ready!(self.stream.poll_next_unpin(cx)).transpose()? else { + self.eof = true; + return Poll::Ready(Ok(false)); + }; + if batch.num_rows() == 0 { + continue; + } + let batch = Arc::new(batch); + let _timer = elapsed_compute.timer(); + let key_arrays = self + .key_exprs + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::>>()?; + self.key_validity = + matchable_join_keys(&key_arrays, NullEquality::NullEqualsNothing); + self.key_arrays = key_arrays.into(); + self.match_array = Some( + self.match_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?, + ); + self.key_batch_id += 1; + self.batch = Some(batch); + } + } + + fn group_has_null(&self) -> bool { + self.key_validity + .as_ref() + .is_some_and(|validity| validity.is_null(self.row)) + } + + fn match_value(&self) -> Result { + 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; + } +} + +#[derive(Clone)] +struct AsOfJoinMetrics { + /// Standard output-row and elapsed-compute metrics. + baseline: BaselineMetrics, + /// Peak bytes retained for the shared right input. + /// + /// `peak_memory_usage` records this as `MetricValue::PeakMemoryUsage`; `Gauge` + /// is the handle used to update that metric. + build_mem_used: Gauge, +} + +impl AsOfJoinMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + build_mem_used: MetricBuilder::new(metrics) + .peak_memory_usage("build_mem_used", partition), + } + } +} + +/// Row references accumulated for the next output batch. +/// +/// For right output, `None` represents NULL padding for an unmatched left row. +/// For example, indices `[Some((0, 2)), None, Some((1, 0))]` select row 2 from +/// the first source batch, a NULL, and row 0 from the second source batch. +#[derive(Default)] +struct PendingRows { + /// Distinct source batches referenced by `indices`. `Arc` keeps per-row + /// clones O(1) and provides stable identity for deduplication. + sources: Vec>, + /// Maps an `Arc` pointer to its index in `sources`. + source_by_ptr: HashMap, + /// Per-output-row `(source, row)` references or NULL padding. + 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(); + } +} + +/// Per-left-partition state for the monotonic ASOF scan. +/// +/// For left rows `(A, 4), (A, 7)` and right rows `(A, 2), (A, 6)`, the +/// candidate advances from `(A, 2)` to `(A, 6)` without rewinding the right +/// cursor. Cursors and the candidate survive input batch changes and output +/// flushes; a change of equality group clears the candidate before reuse. +struct AsOfJoinStream { + /// Output schema used when pending row references are materialized. + schema: SchemaRef, + /// Cursor over the current left partition. + left: InputCursor, + /// Independent cursor over the shared, ordered right input. + right: InputCursor, + /// Retains the shared right batches and their memory reservation. + _right_input: Arc, + /// Validated ordered match operator. + op: Operator, + /// Projected output columns and their input sides. + column_indices: Vec, + /// Whether any projected column needs a right row reference. + projects_right: bool, + /// Nearest eligible right row for the current equality group. + candidate: Option, + /// Equality-key ordering shared by the comparator caches. + group_sort_options: Vec, + /// Cached comparator for the current right and left input batches. + input_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + /// Cached comparator for the candidate and current left batches. + candidate_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + /// Left row references accumulated for the next output batch. + pending_left: PendingRows, + /// Matched right row references, aligned with `pending_left`. + pending_right: PendingRows, + /// Maximum number of pending rows before an output flush. + batch_size: usize, + metrics: AsOfJoinMetrics, +} + +impl AsOfJoinStream { + #[expect(clippy::too_many_arguments)] + fn new( + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + column_indices: Vec, + batch_size: usize, + metrics: AsOfJoinMetrics, + right_input: Arc, + ) -> 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, + _right_input: right_input, + op, + projects_right: column_indices + .iter() + .any(|column| column.side == JoinSide::Right), + column_indices, + candidate: None, + group_sort_options, + input_group_comparator: None, + candidate_group_comparator: None, + batch_size: batch_size.max(1), + metrics, + } + } + + fn compare_input_groups(&mut self) -> Result { + if self.group_sort_options.is_empty() { + return Ok(Ordering::Equal); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_batch_id = self.right.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .input_group_comparator + .as_ref() + .is_none_or(|(right, left, _)| { + *right != right_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + self.right.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.input_group_comparator = + Some((right_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .input_group_comparator + .as_ref() + .expect("ASOF input group comparator must be initialized"); + Ok(comparator.compare(self.right.row, self.left.row)) + } + + fn candidate_is_other_group(&mut self) -> Result { + let Some(candidate) = &self.candidate else { + return Ok(false); + }; + if self.group_sort_options.is_empty() { + return Ok(false); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let candidate_batch_id = candidate.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .candidate_group_comparator + .as_ref() + .is_none_or(|(candidate, left, _)| { + *candidate != candidate_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + candidate.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.candidate_group_comparator = + Some((candidate_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .candidate_group_comparator + .as_ref() + .expect("ASOF candidate group comparator must be initialized"); + Ok(comparator.compare(candidate.row, self.left.row) != Ordering::Equal) + } + + /// Produces the next output batch without resetting the merge state. + /// + /// Each left row first validates its equality group, then advances the right + /// cursor while right groups sort before it or right match values remain + /// eligible. The last eligible right row becomes the candidate. Empty input + /// batches are skipped. Right EOF preserves that candidate for later left + /// rows in the same group; left EOF flushes the final pending rows. NULL keys + /// and group changes clear the candidate, while output flushes only clear + /// pending row references. + /// + /// ```text + /// while the output batch is not full: + /// load the current left row, or flush/finish at left EOF + /// if its match or equality key is NULL, emit it unmatched and advance left + /// clear the candidate if the left equality group changed + /// while the current right row is before the left group or is eligible: + /// remember the nearest eligible row and advance right + /// emit the left row with the candidate (or NULLs), then advance left + /// flush pending rows without resetting either cursor or the candidate + /// ``` + fn poll_next_impl( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + loop { + if self.pending_left.len() >= self.batch_size { + return Poll::Ready(Some(self.flush())); + } + if !ready!( + self.left + .poll_ensure_row(cx, self.metrics.baseline.elapsed_compute()) + )? { + if !self.pending_left.is_empty() { + return Poll::Ready(Some(self.flush())); + } + self.metrics.baseline.done(); + return Poll::Ready(None); + } + + let left_match = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + self.left.match_value()? + }; + if left_match.is_null() || self.left.group_has_null() { + self.candidate = None; + self.candidate_group_comparator = None; + self.push_current_left(None)?; + self.left.advance(); + continue; + } + if self.candidate_is_other_group()? { + self.candidate = None; + self.candidate_group_comparator = None; + } + + loop { + if !ready!( + self.right + .poll_ensure_row(cx, self.metrics.baseline.elapsed_compute()) + )? { + break; + } + if self.right.group_has_null() { + self.right.advance(); + continue; + } + match self.compare_input_groups()? { + Ordering::Less => { + self.right.advance(); + continue; + } + Ordering::Greater => break, + Ordering::Equal => {} + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_match = self.right.match_value()?; + if right_match.is_null() { + self.right.advance(); + continue; + } + if !is_eligible(self.op, &left_match, &right_match)? { + break; + } + let (batch, row) = self.right.batch_row()?; + // Replacing the candidate selects the nearest eligible row. + // Equal match values have no secondary ordering, so which tied + // row wins is intentionally nondeterministic. + self.candidate = Some(Candidate { + batch, + row, + key_arrays: Arc::clone(&self.right.key_arrays), + key_batch_id: self.right.key_batch_id, + }); + self.right.advance(); + } + + self.push_current_left(self.candidate.clone())?; + 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.projects_right { + self.pending_right.push(candidate.batch, candidate.row); + } + } + None => { + if self.projects_right { + self.pending_right.push_null(); + } + } + } + Ok(()) + } + + /// Materializes pending row references while preserving both cursors and the + /// current equality-group candidate for the next output batch. + fn flush(&mut self) -> Result { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let row_count = self.pending_left.len(); + let mut arrays = Vec::with_capacity(self.schema.fields().len()); + for (field, column) in self.schema.fields().iter().zip(&self.column_indices) { + let pending = match column.side { + JoinSide::Left => &self.pending_left, + JoinSide::Right => &self.pending_right, + JoinSide::None => { + return internal_err!("ASOF projection cannot contain a mark column"); + } + }; + arrays.push(pending.materialize_column(column.index, field.data_type())?); + } + self.pending_left.clear(); + self.pending_right.clear(); + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&self.schema), + arrays, + &options, + )?; + (&batch).record_output(&self.metrics.baseline); + Ok(batch) + } +} + +impl RecordBatchStream for AsOfJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for AsOfJoinStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + self.poll_next_impl(cx) + } +} + +/// Validates all invariants required by the forward-only ASOF state machine. +fn validate_asof_join( + left: &dyn ExecutionPlan, + right: &dyn ExecutionPlan, + on: &JoinOn, + match_condition: &AsOfMatchExpr, +) -> 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}" + ); + } + if left_type.is_floating() { + return plan_err!( + "AsOfJoinExec equality expressions do not support floating-point type {left_type}" + ); + } + } + let left_match_type = match_condition.left.data_type(&left_schema)?; + let right_match_type = match_condition.right.data_type(&right_schema)?; + if left_match_type != right_match_type { + return plan_err!( + "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" + ); + } + Ok(()) +} + +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(()) +} + +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, + _ => unreachable!("ASOF match operator is validated by try_new"), + }) +} + +#[cfg(test)] +mod tests { + // Keep physical tests focused on basic executor results, batch-boundary + // state, shared build memory, and constructor/statistics contracts. + + use super::*; + use crate::collect; + use crate::test::TestMemoryExec; + use arrow::array::{Int32Array, Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::test_util::batches_to_sort_string; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr}; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use insta::assert_snapshot; + + #[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)), + ), + Some(vec![0, 1, 2, 5]), + )?)) + } + + #[tokio::test] + async fn simple_query() -> 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] + ); + assert_snapshot!(batches_to_sort_string(&batches), @r" + +-----+----+----+-------+ + | key | ts | id | price | + +-----+----+----+-------+ + | | 3 | 0 | | + | A | | 1 | | + | A | 1 | 2 | | + | A | 4 | 3 | 40 | + | A | 7 | 4 | 60 | + | B | 2 | 5 | 101 | + | C | 3 | 6 | | + +-----+----+----+-------+ + "); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!(metrics.output_rows(), Some(7)); + assert!(metrics.elapsed_compute().is_some()); + Ok(()) + } + + fn exec_without_equality_keys( + left_times: Vec, + right_times: Vec, + op: Operator, + ) -> Result> { + let left_batch = RecordBatch::try_from_iter(vec![ + ( + "ts", + Arc::new(Int64Array::from(left_times.clone())) as ArrayRef, + ), + ( + "id", + Arc::new(Int32Array::from( + left_times + .into_iter() + .map(|value| value as i32) + .collect::>(), + )) as ArrayRef, + ), + ])?; + let left_schema = left_batch.schema(); + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_batch = RecordBatch::try_from_iter(vec![ + ( + "ts", + Arc::new(Int64Array::from(right_times.clone())) as ArrayRef, + ), + ( + "price", + Arc::new(Int32Array::from( + right_times + .into_iter() + .map(|value| value as i32 * 10) + .collect::>(), + )) as ArrayRef, + ), + ])?; + let right_schema = right_batch.schema(); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 0)), + op, + Arc::new(PhysicalColumn::new("ts", 0)), + ), + Some(vec![1, 3]), + )?)) + } + + #[tokio::test] + async fn comparison_directions_without_equality_keys() -> Result<()> { + let predecessor = + exec_without_equality_keys(vec![1, 4, 7], vec![2, 4, 6], Operator::GtEq)?; + let predecessor = collect(predecessor, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&predecessor), @r" + +----+-------+ + | id | price | + +----+-------+ + | 1 | | + | 4 | 40 | + | 7 | 60 | + +----+-------+ + "); + + let successor = + exec_without_equality_keys(vec![7, 4, 1], vec![6, 4, 2], Operator::Lt)?; + let successor = collect(successor, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&successor), @r" + +----+-------+ + | id | price | + +----+-------+ + | 1 | 20 | + | 4 | 60 | + | 7 | | + +----+-------+ + "); + Ok(()) + } + + #[tokio::test] + async fn complex_equality_and_match_expressions() -> Result<()> { + let left_batch = RecordBatch::try_from_iter(vec![ + ("g1", Arc::new(Int64Array::from(vec![0, 1])) as ArrayRef), + ("g2", Arc::new(Int64Array::from(vec![1, 1])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![4, 4])) as ArrayRef), + ("offset", Arc::new(Int64Array::from(vec![1, 0])) as ArrayRef), + ("id", Arc::new(Int64Array::from(vec![10, 20])) as ArrayRef), + ])?; + let left_schema = left_batch.schema(); + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_batch = RecordBatch::try_from_iter(vec![ + ("g1", Arc::new(Int64Array::from(vec![0, 0, 1])) as ArrayRef), + ("g2", Arc::new(Int64Array::from(vec![1, 1, 1])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![2, 5, 3])) as ArrayRef), + ( + "price", + Arc::new(Int64Array::from(vec![12, 15, 23])) as ArrayRef, + ), + ])?; + let right_schema = right_batch.schema(); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + let left_group = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("g1", 0)), + Operator::Plus, + Arc::new(PhysicalColumn::new("g2", 1)), + )); + let right_group = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("g1", 0)), + Operator::Plus, + Arc::new(PhysicalColumn::new("g2", 1)), + )); + let left_match = Arc::new(BinaryExpr::new( + Arc::new(PhysicalColumn::new("ts", 2)), + Operator::Plus, + Arc::new(PhysicalColumn::new("offset", 3)), + )); + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![(left_group, right_group)], + AsOfMatchExpr::new( + left_match, + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 2)), + ), + Some(vec![4, 8]), + )?); + + let batches = collect(exec, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+ + | id | price | + +----+-------+ + | 10 | 15 | + | 20 | 23 | + +----+-------+ + "); + Ok(()) + } + + // Ensure the build-side memory usage equals the sum of all build-side input + // batches, verifying that the build-side buffer is shared. + #[tokio::test] + async fn shared_right_buffers_are_reserved_once() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("id", DataType::Int32, false), + ])); + let left = TestMemoryExec::try_new_exec( + &[vec![make_batch( + &left_schema, + vec![Some("A")], + vec![Some(4095)], + vec![0], + )?]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ])); + let row_count = 4096; + let parent = make_batch( + &right_schema, + vec![Some("A"); row_count], + (0..row_count).map(|value| Some(value as i64)).collect(), + (0..row_count as i32).collect(), + )?; + let mut memory_counter = RecordBatchMemoryCounter::new(); + let retained_size = memory_counter.count_batch(&parent); + let right_batches = (0..16) + .map(|index| parent.slice(index * 256, 256)) + .collect(); + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + Some(vec![0, 1, 2, 5]), + )?); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(retained_size, 1.0) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + + let batches = collect(Arc::clone(&exec) as _, context).await?; + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(4095)]); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!( + metrics + .sum_by_name("build_mem_used") + .map(|value| value.as_usize()), + Some(retained_size) + ); + Ok(()) + } + + #[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)), + ), + Some(vec![0, 1, 2, 5]), + ) + .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(), + Some(vec![0, 1, 2, 5]), + ) + .expect_err("volatile equality expression must be rejected"); + assert!(equality_error.to_string().contains("must be deterministic")); + Ok(()) + } + + #[test] + fn rejects_floating_equality_expressions() -> Result<()> { + let exec = test_exec()?; + for data_type in [DataType::Float16, DataType::Float32, DataType::Float64] { + let left = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + data_type.clone(), + None, + )); + let right = Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + data_type.clone(), + None, + )); + let error = AsOfJoinExec::try_new( + Arc::clone(&exec.left), + Arc::clone(&exec.right), + vec![(left, right)], + exec.match_condition.clone(), + Some(vec![0, 1, 2, 5]), + ) + .expect_err("floating equality expressions must be rejected"); + assert!( + error.to_string().contains(&format!( + "equality expressions do not support floating-point type {data_type}" + )), + "unexpected error: {error}" + ); + } + Ok(()) + } + + #[test] + fn statistics_follow_left_preserving_contract() -> Result<()> { + let exec = test_exec()?; + 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 e4f7e2e123e0e..ca0e883f7ea5e 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 e246fbd86f55ff7f3c80ab1da6f432bd47dc4ccb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 4 Sep 2026 01:36:31 +0000 Subject: [PATCH 2/3] feat: add ASOF join logical semantics (#23829) - Part of #318. - Umbrella PR: #23738. - Depends on #23828 (merged). This is the logical-planning layer of the ASOF JOIN stack. It defines the logical contract and planner behavior separately from the SQL frontend and serialization formats. logical layer. It no longer depends on the optional floating-point follow-up - Add `LogicalPlan::AsOfJoin`, `AsOfJoin`, and `AsOfMatch`. - Validate deterministic expressions, input ownership, supported match operators, equality-key types, and USING constraints. - Add `LogicalPlanBuilder` entry points and schema construction that preserves both qualified `USING` keys while exposing one unqualified wildcard key. - Integrate ASOF joins with tree transforms, display, type coercion, projection pruning, row bounds, and physical planning. - Plan logical ASOF joins to the broadcast-based `AsOfJoinExec` from - Fail closed at proto, SQL unparser, and Substrait boundaries until their owning stack layers add explicit support. - Defer ASOF-specific functional-dependency refinement to #24799 and filter pushdown to #24801 so each optimization can be reviewed independently. Yes: - `cargo fmt --all` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-expr min_rows_of_joins --all-features` - `cargo test -p datafusion-substrait asof_join_fails_closed_until_substrait_has_an_extension --all-features` - The extended workspace test command from the contributor guide This adds logical-plan and builder APIs for ASOF joins. SQL syntax, DataFrame APIs, and plan serialization are intentionally left to dependent stack PRs. Floating equality keys remain rejected by the merged physical operator unless the independent follow-up #24375 is also included. As with any new public `LogicalPlan` variant, downstream exhaustive matches must add an arm. The variant is appended so existing variants retain their `PartialOrd` ordering; maintainers should still treat the enum addition as a Rust source-compatibility break. This PR can be reviewed independently now that #23828 has merged. The optimization follow-ups #24799 and #24801 are not required by the core ASOF stack. --- datafusion/core/src/physical_planner.rs | 49 ++- datafusion/expr/src/logical_plan/builder.rs | 78 ++++- datafusion/expr/src/logical_plan/display.rs | 23 +- datafusion/expr/src/logical_plan/mod.rs | 16 +- datafusion/expr/src/logical_plan/plan.rs | 310 +++++++++++++++++- datafusion/expr/src/logical_plan/tree_node.rs | 58 +++- .../optimizer/src/analyzer/type_coercion.rs | 39 ++- .../optimizer/src/common_subexpr_eliminate.rs | 1 + .../optimizer/src/optimize_projections/mod.rs | 20 ++ datafusion/optimizer/src/optimizer.rs | 5 + .../physical-plan/src/joins/asof_join.rs | 4 + 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 +- 15 files changed, 598 insertions(+), 39 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 3c1e7b50780a5..3803936b50327 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -41,7 +41,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}; @@ -1765,6 +1766,51 @@ 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, + )?, + ); + Arc::new(AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + match_condition, + None, + )?) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, @@ -2264,6 +2310,7 @@ fn extract_dml_filters( | LogicalPlan::Sort(_) | LogicalPlan::Union(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Window(_) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 82734ebbcadd4..f348336e768ff 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,14 @@ pub fn build_join_schema( dfschema.with_functional_dependencies(func_dependencies) } +/// Creates the schema for a left-preserving ASOF join. +/// +/// Both `ON` and `USING` preserve all qualified input fields. SQL wildcard +/// expansion handles the unqualified `USING` key as a single column. +pub fn build_asof_join_schema(left: &DFSchema, right: &DFSchema) -> Result { + build_join_schema(left, right, &JoinType::Left) +} + /// (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 1a141ea52a13a..c45feaf35b46a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,14 +41,15 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement, WriteOp}; use crate::utils::{ - check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, - find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, - merge_schema, split_conjunction, + check_aggregate_and_window_nesting, 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; @@ -295,6 +296,9 @@ pub enum LogicalPlan { Unnest(Unnest), /// A variadic query (e.g. "Recursive CTEs") RecursiveQuery(RecursiveQuery), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), } impl Default for LogicalPlan { @@ -342,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(), @@ -370,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()) @@ -460,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], @@ -495,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)| { @@ -568,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() } @@ -691,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, @@ -994,6 +1032,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, @@ -1419,6 +1496,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| { @@ -1469,6 +1547,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), @@ -1507,6 +1586,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), @@ -2135,6 +2215,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, .. @@ -4255,6 +4354,169 @@ 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. + /// + /// This is the pre-coercion boundary. The physical ASOF constructor repeats + /// the shared operator, side-ownership, and determinism checks for direct + /// physical-plan callers and adds execution-only constraints. Keep the + /// shared checks aligned across both entry points. + 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())?; + 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. /// @@ -6617,6 +6879,32 @@ mod tests { Ok(()) } + #[test] + fn test_asof_using_preserves_qualified_keys() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("ts", DataType::Int64, false), + ]); + let left = Arc::new(table_scan(Some("t1"), &schema, None)?.build()?); + let right = Arc::new(table_scan(Some("t2"), &schema, None)?.build()?); + let join = AsOfJoin::try_new( + left, + right, + vec![(col("t1.id"), col("t2.id"))], + AsOfMatch::new(col("t1.ts"), Operator::GtEq, col("t2.ts")), + JoinConstraint::Using, + )?; + + assert_eq!(join.schema.fields().len(), 4); + assert_eq!( + join.schema + .index_of_column(&Column::from_qualified_name("t2.id"))?, + 2 + ); + assert!(join.schema.field(2).is_nullable()); + Ok(()) + } + #[test] fn test_join_try_new_schema_validation() -> Result<()> { let left_schema = Schema::new(vec![ diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c4c1d743b58b6..ee43666736fe6 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -41,11 +41,12 @@ 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, WriteOp, 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, WriteOp, builder::unnest_with_options, + dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -150,6 +151,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 +465,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 @@ -614,6 +639,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 d11c3e7435fde..6186cc5350bb2 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,10 +57,10 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, - Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, - WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, - lit, not, + AsOfJoin, AsOfMatch, Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, + LogicalPlan, Operator, Projection, Union, ValueOrLambda, WindowFrame, + WindowFrameBound, WindowFrameUnits, WriteOp, is_false, is_not_false, is_not_true, + is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -191,6 +191,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), LogicalPlan::Dml(dml) => self.coerce_dml(dml), @@ -284,6 +285,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 41d09db7c2bbe..5dfa9c61553ac 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..74ce2ad374af8 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -407,6 +407,26 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + 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 { + right_required.push(*index - left_len); + } + } + 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/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 22625afca1f49..d7570c080cf79 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -166,6 +166,10 @@ impl AsOfJoinExec { /// floating-point equality keys are not supported because Arrow sorting /// distinguishes signed zero while SQL equality does not. Projection indices /// refer to the full left-then-right join schema. + /// + /// The logical ASOF constructor validates the corresponding pre-coercion + /// contract. Keep the shared operator, side-ownership, and determinism checks + /// aligned across both public entry points. pub fn try_new( left: Arc, right: Arc, diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 2efff1cbde793..f01cb6b9a78fb 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -2212,6 +2212,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 f4b60176cfba9..5cf2e7a75e91c 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -198,6 +198,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 c340ae6001f1dc33d629c26e9baa2086a65cb13c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 9 Sep 2026 02:23:19 +0000 Subject: [PATCH 3/3] feat: support ASOF JOIN SQL (#23830) - Part of #318. - Umbrella PR: #23738. - Builds on #23829 and #23828, both merged. This is the SQL frontend layer of the ASOF JOIN stack. It adds syntax and unparsing on top of the merged physical and logical contracts. The PR now contains only the isolated SQL frontend diff. - Plan `ASOF JOIN ... MATCH_CONDITION (...)` with optional `ON` or `USING` equality keys. - Reject unsupported match shapes and non-equality `ON` predicates. - Unparse ASOF joins while preserving right-side candidate preselection and nested join scope. - Document the supported SQL syntax and semantics, including qualified `USING` keys, left-partitioned broadcast execution, the full-right memory requirement, repeated scans, and the absence of spill/repartitioned ASOF. - Add SQL integration and sqllogictest coverage for all four match directions, coercion, equality-free joins, USING, invalid contracts, EXPLAIN, boundedness, and optimized-plan round trips. - Verify the broadcast topology with a multi-partition left input: left partitioning is preserved while the right input is single-partitioned. Yes: - `cargo fmt --all` - `./ci/scripts/doc_prettier_check.sh --write --allow-dirty` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion --test core_integration asof --all-features` - `cargo test -p datafusion-sqllogictest --test sqllogictests --all-features -- asof_join` - The extended workspace test command from the contributor guide Users can express Snowflake-style ASOF joins in SQL with `MATCH_CONDITION`, optional equality keys, and `<`, `<=`, `>`, or `>=` match directions. With `USING`, wildcard output exposes one unqualified key while both qualified input keys remain addressable. The user guide also documents the initial broadcast strategy and its memory/no-spill limitations. stack and does not depend on the optional floating-point follow-up #24375. --- datafusion/core/tests/sql/joins.rs | 413 +++++++++++++++++- datafusion/expr/src/logical_plan/builder.rs | 154 ++++++- datafusion/expr/src/logical_plan/plan.rs | 19 + datafusion/sql/src/relation/join.rs | 67 ++- datafusion/sql/src/unparser/plan.rs | 236 +++++++++- datafusion/sql/tests/cases/plan_to_sql.rs | 90 ++++ .../sqllogictest/test_files/asof_join.slt | 287 ++++++++++++ datafusion/substrait/tests/cases/serialize.rs | 8 +- docs/source/user-guide/sql/select.md | 43 ++ 9 files changed, 1278 insertions(+), 39 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/asof_join.slt diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index f1ffd6c5b12d0..58dc5df9e2b6f 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -17,10 +17,12 @@ use insta::assert_snapshot; -use datafusion::assert_batches_eq; use datafusion::catalog::MemTable; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; +use datafusion::physical_plan::joins::AsOfJoinExec; +use datafusion::physical_plan::{Distribution, ExecutionPlanProperties}; use datafusion::test_util::register_unbounded_file_with_ordering; +use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion_sql::unparser::plan_to_sql; use super::*; @@ -326,3 +328,412 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } + +fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { + let trades_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("trade_id", DataType::Int32, false), + ])); + let trades = vec![ + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(7), Some(2), Some(3)])), + Arc::new(Int32Array::from(vec![3, 4, 6])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(1), Some(4), Some(8)])), + Arc::new(Int32Array::from(vec![1, 2, 5])), + ], + )?, + ]; + ctx.register_table( + "trades", + Arc::new(MemTable::try_new( + trades_schema, + trades.into_iter().map(|batch| vec![batch]).collect(), + )?), + )?; + + let prices_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let prices = vec![ + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(6), Some(1), Some(2)])), + Arc::new(Int32Array::from(vec![60, 101, 999])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(2), Some(4), Some(6)])), + Arc::new(Int32Array::from(vec![20, 40, 106])), + ], + )?, + ]; + ctx.register_table( + "prices", + Arc::new(MemTable::try_new( + prices_schema, + prices.into_iter().map(|batch| vec![batch]).collect(), + )?), + )?; + Ok(()) +} + +fn find_asof_exec(plan: &Arc) -> Option> { + if plan.downcast_ref::().is_some() { + return Some(Arc::clone(plan)); + } + plan.children().into_iter().find_map(find_asof_exec) +} + +#[tokio::test] +async fn asof_join_all_match_directions_across_batches() -> Result<()> { + let config = SessionConfig::new() + .with_batch_size(2) + .with_target_partitions(2); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx)?; + + for (op, expected) in [ + ( + ">=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + ">", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 20 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 40 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 60 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ] { + let batches = ctx + .sql(&format!( + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts {op} p.ts) \ + ON t.symbol = p.symbol ORDER BY t.trade_id" + )) + .await? + .collect() + .await?; + assert_batches_eq!(expected, &batches); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_coerces_equality_and_match_types() -> Result<()> { + let ctx = SessionContext::new(); + let batches = ctx + .sql( + "SELECT t.id, p.price \ + FROM (VALUES (CAST(1 AS INT), CAST(4 AS INT), 7)) t(k, ts, id) \ + ASOF JOIN \ + (VALUES (CAST(1 AS BIGINT), CAST(2 AS BIGINT), 20)) p(k, ts, price) \ + MATCH_CONDITION (t.ts >= p.ts) ON t.k = p.k", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----+-------+", + "| id | price |", + "+----+-------+", + "| 7 | 20 |", + "+----+-------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx)?; + let df = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN \ + (SELECT ts, price FROM prices WHERE symbol = 'A') p \ + MATCH_CONDITION (t.ts >= p.ts)", + ) + .await?; + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert_contains!(sql.as_str(), "ASOF JOIN"); + assert!(!sql.contains(" ON "), "unexpected equality clause: {sql}"); + ctx.sql(&sql).await?; + let plan = df.create_physical_plan().await?; + let asof = find_asof_exec(&plan).expect("physical ASOF join must be present"); + let output_partitions = asof.output_partitioning().partition_count(); + assert_eq!( + output_partitions, + asof.children()[0].output_partitioning().partition_count() + ); + assert!( + output_partitions > 1, + "ASOF join did not preserve left-side parallelism" + ); + assert_eq!( + asof.children()[1].output_partitioning().partition_count(), + 1 + ); + let right_plan = displayable(asof.children()[1].as_ref()) + .indent(true) + .to_string(); + assert_contains!(right_plan.as_str(), "SortPreservingMergeExec"); + assert_contains!(right_plan.as_str(), "DataSourceExec: partitions=2"); + assert!(asof.output_ordering().is_some()); + assert!(matches!( + &asof.input_distribution_requirements().into_per_child()[..], + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] + )); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_batches_sorted_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 20 |", + "| 5 | 60 |", + "| 6 | 20 |", + "+----------+-------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + let batches = ctx + .sql( + "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON t.symbol = p.symbol", + ) + .await? + .collect() + .await?; + let explain = arrow::util::pretty::pretty_format_batches(&batches)?.to_string(); + assert_contains!(explain.as_str(), "AsOf Join: match=[t.ts >= p.ts]"); + assert_contains!(explain.as_str(), "on=[t.symbol = p.symbol]"); + assert_contains!(explain.as_str(), "AsOfJoinExec:"); + assert_contains!(explain.as_str(), "on=[(symbol = symbol)]"); + assert_contains!(explain.as_str(), "match=[ts >= ts]"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result<()> { + let ctx = SessionContext::new(); + let tmp_dir = TempDir::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::UInt32, false), + Field::new("ts", DataType::UInt32, false), + ])); + let ordering = vec![vec![ + col("symbol").sort(true, true), + col("ts").sort(true, true), + ]]; + for table in ["left_stream", "right_stream"] { + let path = tmp_dir.path().join(format!("{table}.csv")); + File::create(&path)?; + register_unbounded_file_with_ordering( + &ctx, + Arc::clone(&schema), + &path, + table, + ordering.clone(), + )?; + } + let error = ctx + .sql( + "SELECT * FROM left_stream l ASOF JOIN right_stream r \ + MATCH_CONDITION (l.ts >= r.ts) ON l.symbol = r.symbol", + ) + .await? + .create_physical_plan() + .await + .expect_err("ASOF physical planning must reject unbounded inputs"); + assert_contains!(error.to_string(), "AsOfJoinExec requires bounded inputs"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_using_preserves_key_access_and_unparser_round_trips() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + let df = ctx + .sql( + "SELECT * FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol)", + ) + .await?; + assert_eq!( + df.schema() + .fields() + .iter() + .map(|field| field.name()) + .collect::>(), + vec!["ts", "trade_id", "symbol", "ts", "price"] + ); + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert!(sql.contains("ASOF JOIN")); + assert!(sql.contains("MATCH_CONDITION")); + assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}"); + ctx.sql(&sql).await?; + + let batches = ctx + .sql( + "SELECT t.trade_id, t.symbol AS left_symbol, p.symbol AS right_symbol \ + FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----------+-------------+--------------+", + "| trade_id | left_symbol | right_symbol |", + "+----------+-------------+--------------+", + "| 1 | A | |", + "| 2 | A | A |", + "| 3 | A | A |", + "| 4 | B | B |", + "| 5 | B | B |", + "| 6 | | |", + "+----------+-------------+--------------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + for query in [ + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + "SELECT * FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + "SELECT t.trade_id, p.price FROM trades t \ + JOIN prices q ON t.symbol = q.symbol AND t.ts = q.ts \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON q.symbol = p.symbol ORDER BY t.trade_id", + "SELECT t.trade_id, q.trade_id FROM trades t \ + ASOF JOIN (prices p JOIN trades q \ + ON p.symbol = q.symbol AND p.ts = q.ts) \ + MATCH_CONDITION (t.ts >= q.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + ] { + let expected = ctx.sql(query).await?.collect().await?; + let plan = ctx.sql(query).await?.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + let actual = ctx.sql(&sql).await?.collect().await?; + assert_eq!( + datafusion_common::test_util::batches_to_string(&expected), + datafusion_common::test_util::batches_to_string(&actual), + "unparsed SQL changed ASOF candidate preselection: {sql}" + ); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_invalid_contracts() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + for sql in [ + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts = p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (p.ts >= t.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON t.symbol > p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (1 >= p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON 1 = 1", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts AND t.ts > p.ts) ON t.symbol = p.symbol", + ] { + assert!(ctx.sql(sql).await.is_err(), "query should fail: {sql}"); + } + Ok(()) +} diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index f348336e768ff..666e60d3a3177 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -38,13 +38,15 @@ use crate::logical_plan::{ }; use crate::select_expr::SelectExpr; use crate::utils::{ - can_hash, columnize_expr, compare_sort_expr, expand_qualified_wildcard, - expand_wildcard, expr_to_columns, find_valid_equijoin_key_pair, - group_window_expr_by_sort_keys, + can_hash, check_all_columns_from_schema, columnize_expr, compare_sort_expr, + expand_qualified_wildcard, expand_wildcard, expr_to_columns, + find_valid_equijoin_key_pair, group_window_expr_by_sort_keys, + split_conjunction_owned, }; use crate::{ - DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, RecursiveQuery, - Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, binary_expr, lit, + BinaryExpr, DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, + RecursiveQuery, Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, + binary_expr, lit, }; use super::dml::InsertOp; @@ -1007,23 +1009,67 @@ impl LogicalPlanBuilder { ) } - /// Apply a left-preserving ASOF join using equality expressions and one - /// ordered match condition. - pub fn asof_join( + /// Apply a left-preserving ASOF join using an optional equality condition + /// and one ordered match condition. + /// + /// When present, `on_expr` must contain equality comparisons combined with + /// `AND`. Each comparison must have one operand that references only the + /// left input and one that references only `right`; their order does not + /// matter. `match_condition` must be a single `<`, `<=`, `>`, or `>=` + /// comparison whose left operand references only the left input and whose + /// right operand references only `right`. + pub fn asof_join_on( self, right: LogicalPlan, - on: Vec<(Expr, Expr)>, - match_condition: AsOfMatch, + on_expr: Option, + match_condition: Expr, ) -> Result { - self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On) + let on = on_expr + .into_iter() + .flat_map(split_conjunction_owned) + .map(|predicate| { + let Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::Eq, + right: right_expr, + }) = predicate + else { + return plan_err!( + "ASOF ON accepts only equality conditions combined with AND" + ); + }; + find_valid_equijoin_key_pair( + &left, + &right_expr, + self.plan.schema(), + right.schema(), + )? + .ok_or_else(|| { + plan_datafusion_err!( + "Each ASOF equality condition must compare one left expression with one right expression" + ) + }) + }) + .collect::>()?; + self.asof_join_with_constraint( + right, + on, + AsOfMatch::try_from(match_condition)?, + JoinConstraint::On, + ) } - /// Apply a left-preserving ASOF join using `USING` equality keys. + /// Apply a left-preserving ASOF join using `USING` equality keys and one + /// ordered match condition. + /// + /// Every key in `using_keys` must resolve in both inputs. + /// `match_condition` follows the same operand and operator requirements as + /// [`asof_join_on`](Self::asof_join_on). pub fn asof_join_using( self, right: LogicalPlan, using_keys: Vec, - match_condition: AsOfMatch, + match_condition: Expr, ) -> Result { let on = using_keys .into_iter() @@ -1033,7 +1079,12 @@ impl LogicalPlanBuilder { Ok((Expr::Column(left), Expr::Column(right))) }) .collect::>()?; - self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using) + self.asof_join_with_constraint( + right, + on, + AsOfMatch::try_from(match_condition)?, + JoinConstraint::Using, + ) } fn asof_join_with_constraint( @@ -1043,6 +1094,17 @@ impl LogicalPlanBuilder { match_condition: AsOfMatch, join_constraint: JoinConstraint, ) -> Result { + let left_columns = match_condition.left.column_refs(); + let right_columns = match_condition.right.column_refs(); + if left_columns.is_empty() + || right_columns.is_empty() + || !check_all_columns_from_schema(&left_columns, self.plan.schema())? + || !check_all_columns_from_schema(&right_columns, right.schema())? + { + return plan_err!( + "ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input" + ); + } let normalize = |expr, schema: &DFSchema| { normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[]) }; @@ -2971,6 +3033,70 @@ mod tests { Ok(()) } + #[test] + fn asof_join_on_extracts_and_validates_conditions() -> Result<()> { + let values = vec![vec![lit(1), lit(2)]]; + let left = LogicalPlanBuilder::values(values.clone())? + .alias("l")? + .build()?; + let right = LogicalPlanBuilder::values(values)?.alias("r")?.build()?; + + let plan = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some( + col("r.column1") + .eq(col("l.column1")) + .and(col("l.column2").eq(col("r.column2"))), + ), + col("l.column2").gt_eq(col("r.column2")), + )? + .build()?; + let LogicalPlan::AsOfJoin(join) = plan else { + panic!("expected ASOF join") + }; + assert_eq!( + join.on, + vec![ + (col("l.column1"), col("r.column1")), + (col("l.column2"), col("r.column2")), + ] + ); + assert_eq!( + join.match_condition.as_ref(), + &AsOfMatch::new(col("l.column2"), Operator::GtEq, col("r.column2")) + ); + + let invalid_on = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some(col("l.column1").gt(col("r.column1"))), + col("l.column2").gt_eq(col("r.column2")), + ) + .expect_err("non-equality ASOF ON should fail"); + assert_snapshot!(invalid_on.strip_backtrace(), @r#"Error during planning: ASOF ON accepts only equality conditions combined with AND"#); + + let invalid_match = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some(col("l.column1").eq(col("r.column1"))), + col("l.column2").eq(col("r.column2")), + ) + .expect_err("equality ASOF MATCH_CONDITION should fail"); + assert_snapshot!(invalid_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION requires <, <=, >, or >=, found ="#); + + let reversed_match = LogicalPlanBuilder::from(left) + .asof_join_on( + right, + Some(col("l.column1").eq(col("r.column1"))), + col("r.column2").gt_eq(col("l.column2")), + ) + .expect_err("reversed ASOF MATCH_CONDITION should fail"); + assert_snapshot!(reversed_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input"#); + + Ok(()) + } + #[test] fn plan_builder_from_logical_plan() -> Result<()> { let plan = diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c45feaf35b46a..fa8ceb9d826b7 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -4373,6 +4373,25 @@ impl AsOfMatch { } } +impl TryFrom for AsOfMatch { + type Error = DataFusionError; + + fn try_from(condition: Expr) -> Result { + let Expr::BinaryExpr(BinaryExpr { left, op, right }) = condition else { + return plan_err!("ASOF MATCH_CONDITION must be a single comparison"); + }; + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {op}" + ); + } + Ok(Self::new(*left, op, *right)) + } +} + impl Display for AsOfMatch { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { write!(f, "{} {} {}", self.left, self.op, self.right) diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b38099..73c2f05504269 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,7 +16,7 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err}; +use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err, plan_err}; use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{ Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, @@ -98,10 +98,75 @@ impl SqlToRel<'_, S> { JoinOperator::CrossJoin(JoinConstraint::None) => { self.parse_cross_join(left, right) } + JoinOperator::AsOf { + match_condition, + constraint, + } => self.parse_asof_join( + left, + right, + match_condition, + constraint, + planner_context, + ), other => not_impl_err!("Unsupported JOIN operator {other:?}"), } } + fn parse_asof_join( + &self, + left: LogicalPlan, + right: LogicalPlan, + sql_match_condition: sqlparser::ast::Expr, + constraint: JoinConstraint, + planner_context: &mut PlannerContext, + ) -> Result { + let join_schema = left.schema().join(right.schema())?; + let match_condition = + self.sql_to_expr(sql_match_condition, &join_schema, planner_context)?; + + match constraint { + JoinConstraint::On(sql_on) => { + let on = self.sql_to_expr(sql_on, &join_schema, planner_context)?; + LogicalPlanBuilder::from(left) + .asof_join_on(right, Some(on), match_condition)? + .build() + } + JoinConstraint::Using(object_names) => { + let keys = object_names + .into_iter() + .map(|object_name| { + let ObjectName(mut object_names) = object_name; + if object_names.len() != 1 { + return not_impl_err!( + "Invalid identifier in ASOF USING clause. Expected single identifier, got {}", + ObjectName(object_names) + ); + } + let id = object_names.swap_remove(0); + id.as_ident() + .ok_or_else(|| { + plan_datafusion_err!( + "Expected identifier in ASOF USING clause" + ) + }) + .map(|ident| { + Column::from_name( + self.ident_normalizer.normalize(ident.clone()), + ) + }) + }) + .collect::>>()?; + LogicalPlanBuilder::from(left) + .asof_join_using(right, keys, match_condition)? + .build() + } + JoinConstraint::None => LogicalPlanBuilder::from(left) + .asof_join_on(right, None, match_condition)? + .build(), + JoinConstraint::Natural => plan_err!("NATURAL ASOF JOIN is not supported"), + } + } + fn parse_cross_join( &self, left: LogicalPlan, diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 5cf2e7a75e91c..cf8bf1a03c8f6 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -49,7 +49,7 @@ use datafusion_common::{ }; use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX}; use datafusion_expr::{ - Aggregate, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, + Aggregate, AsOfJoin, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, SkipType, Sort, SortExpr, TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; @@ -178,6 +178,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -198,7 +199,6 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) - | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } @@ -1340,11 +1340,8 @@ impl Unparser<'_> { let mut right_relation = RelationBuilder::default(); if already_projected - && let Some(nested_relation) = self - .qualified_passthrough_join_projection_to_nested_relation( - right_plan.as_ref(), - query, - )? + && let Some(nested_relation) = + self.join_input_to_nested_relation(right_plan.as_ref(), query)? { right_relation = nested_relation; } else { @@ -1478,6 +1475,9 @@ impl Unparser<'_> { Ok(()) } + LogicalPlan::AsOfJoin(join) => { + self.asof_join_to_sql(join, query, select, relation) + } LogicalPlan::SubqueryAlias(plan_alias) => { let (plan, mut columns) = subquery_alias_inner_query_and_columns(plan_alias); @@ -1775,6 +1775,129 @@ impl Unparser<'_> { } } + // Keep ASOF-specific locals out of the recursive plan unparser's stack frame. + #[inline(never)] + fn asof_join_to_sql( + &self, + join: &AsOfJoin, + query: &mut Option, + select: &mut SelectBuilder, + relation: &mut RelationBuilder, + ) -> Result<()> { + let already_projected = select.already_projected(); + let left_plan = + Self::unwrap_qualified_passthrough_join_projection(Arc::clone(&join.left)); + let inline_left_join = matches!( + left_plan.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ); + let left_projection = if already_projected { + None + } else if inline_left_join { + self.select_to_sql_recursively(left_plan.as_ref(), query, select, relation)?; + select.pop_projections(); + Some(self.derived_input_projection(join.left.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + let qualifier = self.derive_asof_input(join.left.as_ref(), relation)?; + Some(self.derived_input_projection(join.left.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively(join.left.as_ref(), query, select, relation)?; + Some(select.pop_projections()) + }; + if already_projected { + if inline_left_join { + self.select_to_sql_recursively( + left_plan.as_ref(), + query, + select, + relation, + )?; + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + self.derive_asof_input(join.left.as_ref(), relation)?; + } else { + self.select_to_sql_recursively( + join.left.as_ref(), + query, + select, + relation, + )?; + } + } + + let mut right_relation = RelationBuilder::default(); + let nested_right = + self.join_input_to_nested_relation(join.right.as_ref(), query)?; + let right_projection = if already_projected { + if let Some(nested_right) = nested_right { + right_relation = nested_right; + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + } + None + } else if let Some(nested_right) = nested_right { + right_relation = nested_right; + Some(self.derived_input_projection(join.right.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + let qualifier = + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + Some(self.derived_input_projection(join.right.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + Some(select.pop_projections()) + }; + let Ok(Some(relation)) = right_relation.build() else { + return internal_err!("Failed to build ASOF right relation"); + }; + let constraint = + self.join_constraint_to_sql(join.join_constraint, &join.on, None)?; + let match_condition = self.expr_to_sql(&Expr::BinaryExpr(BinaryExpr::new( + Box::new(join.match_condition.left.clone()), + join.match_condition.op, + Box::new(join.match_condition.right.clone()), + )))?; + let ast_join = ast::Join { + relation, + global: false, + join_operator: ast::JoinOperator::AsOf { + match_condition, + constraint, + }, + }; + let mut from = select + .pop_from() + .ok_or_else(|| internal_datafusion_err!("ASOF left relation is missing"))?; + from.push_join(ast_join); + select.push_from(from); + + if !already_projected { + let left_projection = left_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF left projection is missing") + })?; + let right_projection = right_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF right projection is missing") + })?; + select.projection( + left_projection + .into_iter() + .chain(right_projection) + .collect(), + ); + } + Ok(()) + } + /// Walk through transparent nodes (SubqueryAlias) to find the inner /// Projection that feeds an Unnest node. /// @@ -2073,6 +2196,74 @@ impl Unparser<'_> { ) } + fn asof_input_requires_derived(plan: &LogicalPlan) -> bool { + let simple_scan = + |scan: &TableScan| scan.filters.is_empty() && scan.fetch.is_none(); + match plan { + LogicalPlan::TableScan(scan) => !simple_scan(scan), + LogicalPlan::SubqueryAlias(alias) => { + !matches!(alias.input.as_ref(), LogicalPlan::TableScan(scan) if simple_scan(scan)) + } + _ => true, + } + } + + fn derive_asof_input( + &self, + plan: &LogicalPlan, + relation: &mut RelationBuilder, + ) -> Result> { + if let LogicalPlan::SubqueryAlias(alias) = plan { + let (inner, columns) = subquery_alias_inner_query_and_columns(alias); + let table_alias = alias.alias.clone(); + if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() + { + let rewritten = + inject_column_aliases_into_subquery(inner.clone(), columns)?; + self.derive( + &rewritten, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), vec![])), + false, + )?; + } else { + self.derive( + inner, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), columns)), + false, + )?; + } + return Ok(Some(table_alias)); + } + + let qualifier = plan + .schema() + .iter() + .find_map(|(qualifier, _)| qualifier.cloned()); + let alias = qualifier + .as_ref() + .map(|qualifier| self.new_table_alias(qualifier.table().to_string(), vec![])); + self.derive(plan, relation, alias, false)?; + Ok(qualifier) + } + + fn derived_input_projection( + &self, + plan: &LogicalPlan, + qualifier: Option<&TableReference>, + ) -> Result> { + plan.schema() + .iter() + .map(|(field_qualifier, field)| { + self.select_item_to_sql(&Expr::Column(Column::new( + qualifier.cloned().or_else(|| field_qualifier.cloned()), + field.name(), + ))) + }) + .collect() + } + fn is_qualified_passthrough_projection(projection: &Projection) -> bool { projection .expr @@ -2084,7 +2275,10 @@ impl Unparser<'_> { plan: Arc, ) -> Arc { if let LogicalPlan::Projection(projection) = plan.as_ref() - && matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + && matches!( + projection.input.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ) && Self::is_qualified_passthrough_projection(projection) { Arc::clone(&projection.input) @@ -2093,26 +2287,30 @@ impl Unparser<'_> { } } - fn qualified_passthrough_join_projection_to_nested_relation( + fn join_input_to_nested_relation( &self, plan: &LogicalPlan, query: &mut Option, ) -> Result> { - let LogicalPlan::Projection(projection) = plan else { - return Ok(None); + let join_plan = match plan { + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) => plan, + LogicalPlan::Projection(projection) + if matches!( + projection.input.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ) && Self::is_qualified_passthrough_projection(projection) => + { + projection.input.as_ref() + } + _ => return Ok(None), }; - if !matches!(projection.input.as_ref(), LogicalPlan::Join(_)) - || !Self::is_qualified_passthrough_projection(projection) - { - return Ok(None); - } let original_query = query.clone(); let mut nested_select = SelectBuilder::default(); nested_select.push_from(TableWithJoinsBuilder::default()); let mut nested_relation = RelationBuilder::default(); self.select_to_sql_recursively( - projection.input.as_ref(), + join_plan, query, &mut nested_select, &mut nested_relation, @@ -2123,11 +2321,11 @@ impl Unparser<'_> { } let Some(mut nested_from) = nested_select.pop_from() else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; nested_from.relation(nested_relation); let Some(table_with_joins) = nested_from.build()? else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; let mut relation = RelationBuilder::default(); diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index d6c31570bf1b0..78c623a8bf453 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -2742,6 +2742,37 @@ fn test_unparse_inner_join_with_table_scan_projection() -> Result<()> { Ok(()) } +#[test] +fn test_unparse_asof_join() -> Result<()> { + let trades_schema = Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("trade_id", DataType::Int32, false), + ]); + let prices_schema = Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ]); + let trades = table_scan(Some("trades"), &trades_schema, None)? + .alias("t")? + .build()?; + let prices = table_scan(Some("prices"), &prices_schema, None)? + .alias("p")? + .build()?; + let plan = LogicalPlanBuilder::from(trades) + .asof_join_on( + prices, + Some(col("t.symbol").eq(col("p.symbol"))), + col("t.ts").gt_eq(col("p.ts")), + )? + .project(vec![col("t.trade_id"), col("p.price")])? + .build()?; + + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT t.trade_id, p.price FROM trades AS t ASOF JOIN prices AS p MATCH_CONDITION ((t.ts >= p.ts)) ON t.symbol = p.symbol"#); + Ok(()) +} + /// Build the three base table scans (`left_table`, `mid_table`, `right_table`) /// shared by the nested passthrough-projection join unparsing tests. fn nested_passthrough_join_tables() -> Result<(LogicalPlan, LogicalPlan, LogicalPlan)> { @@ -2852,6 +2883,65 @@ fn test_unparse_projected_join_unwraps_left_nested_passthrough_projection() -> R Ok(()) } +#[test] +fn test_unparse_nested_asof_join_inputs() -> Result<()> { + let (left, mid, right) = nested_passthrough_join_tables()?; + let nested_left = LogicalPlanBuilder::from(left) + .asof_join_on( + mid, + Some(col("left_table.mid_id").eq(col("mid_table.mid_id"))), + col("left_table.left_id").gt_eq(col("mid_table.mid_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("mid_table.right_id"), + ])? + .build()?; + let plan = LogicalPlanBuilder::from(nested_left) + .asof_join_on( + right, + Some(col("mid_table.right_id").eq(col("right_table.right_id"))), + col("mid_table.right_id").gt_eq(col("right_table.right_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table ASOF JOIN mid_table MATCH_CONDITION ((left_table.left_id >= mid_table.mid_id)) ON left_table.mid_id = mid_table.mid_id ASOF JOIN right_table MATCH_CONDITION ((mid_table.right_id >= right_table.right_id)) ON mid_table.right_id = right_table.right_id"#); + + let (left, mid, right) = nested_passthrough_join_tables()?; + let nested_right = LogicalPlanBuilder::from(mid) + .asof_join_on( + right, + Some(col("mid_table.right_id").eq(col("right_table.right_id"))), + col("mid_table.right_id").gt_eq(col("right_table.right_id")), + )? + .project(vec![ + col("mid_table.mid_id"), + col("mid_table.right_id"), + col("right_table.value"), + ])? + .build()?; + let plan = LogicalPlanBuilder::from(left) + .asof_join_on( + nested_right, + Some(col("left_table.mid_id").eq(col("mid_table.mid_id"))), + col("left_table.left_id").gt_eq(col("mid_table.mid_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table ASOF JOIN (mid_table ASOF JOIN right_table MATCH_CONDITION ((mid_table.right_id >= right_table.right_id)) ON mid_table.right_id = right_table.right_id) MATCH_CONDITION ((left_table.left_id >= mid_table.mid_id)) ON left_table.mid_id = mid_table.mid_id"#); + + Ok(()) +} + #[test] fn test_unparse_left_semi_join_with_table_scan_projection() -> Result<()> { let schema = Schema::new(vec![ diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt new file mode 100644 index 0000000000000..3e84cd037653e --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,287 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement ok +CREATE TABLE asof_left(id INT, grp TEXT, ts TIMESTAMP) AS VALUES + (1, 'A', TIMESTAMP '2024-01-01 09:00:01'), + (2, 'A', TIMESTAMP '2024-01-01 09:00:04'), + (3, 'A', TIMESTAMP '2024-01-01 09:00:07'), + (4, 'B', TIMESTAMP '2024-01-01 09:00:02'), + (5, 'B', TIMESTAMP '2024-01-01 09:00:08'), + (6, NULL, TIMESTAMP '2024-01-01 09:00:03'), + (7, 'A', NULL); + +statement ok +CREATE TABLE asof_right(grp TEXT, ts TIMESTAMP, val TEXT) AS VALUES + ('A', TIMESTAMP '2024-01-01 09:00:02', 'a2'), + ('A', TIMESTAMP '2024-01-01 09:00:04', 'a4'), + ('A', TIMESTAMP '2024-01-01 09:00:06', 'a6'), + ('B', TIMESTAMP '2024-01-01 09:00:01', 'b1'), + ('B', TIMESTAMP '2024-01-01 09:00:06', 'b6'), + (NULL, TIMESTAMP '2024-01-01 09:00:02', 'null-group'), + ('A', NULL, 'null-ts'); + +# Inclusive predecessor per equality group. This also verifies unmatched left +# rows and NULL behavior for equality keys and ordered expressions. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Strict predecessor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts > r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:02 a2 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Inclusive successor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts <= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Strict successor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts < r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:06 a6 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# USING exposes one unqualified equality key. +query TIPT +SELECT grp, l.id, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +A 1 NULL NULL +A 2 2024-01-01T09:00:04 a4 +A 3 2024-01-01T09:00:06 a6 +B 4 2024-01-01T09:00:01 b1 +B 5 2024-01-01T09:00:06 b6 +NULL 6 NULL NULL +A 7 NULL NULL + +# Both qualified equality keys remain addressable. +query ITT +SELECT l.id, l.grp, r.grp +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +1 A NULL +2 A A +3 A A +4 B B +5 B B +6 NULL NULL +7 A NULL + +# Equality keys are optional. +query IT +SELECT l.id, r.label +FROM (VALUES (1, 1), (2, 5), (3, CAST(NULL AS INT))) AS l(id, ts) +ASOF JOIN (VALUES (2, 'r2'), (4, 'r4')) AS r(ts, label) +MATCH_CONDITION (l.ts >= r.ts) +ORDER BY l.id; +---- +1 NULL +2 r4 +3 NULL + +# Multiple equality keys form one candidate group. +query IT +SELECT l.id, r.val +FROM (VALUES + (1, 'X', 'A', TIMESTAMP '2024-01-01 09:00:04'), + (2, 'Y', 'A', TIMESTAMP '2024-01-01 09:00:04') +) AS l(id, venue, grp, ts) +ASOF JOIN (VALUES + ('X', 'A', TIMESTAMP '2024-01-01 09:00:02', 'x-a2'), + ('Y', 'A', TIMESTAMP '2024-01-01 09:00:03', 'y-a3'), + ('X', 'B', TIMESTAMP '2024-01-01 09:00:04', 'x-b4') +) AS r(venue, grp, ts, val) +MATCH_CONDITION (l.ts >= r.ts) +ON l.venue = r.venue AND l.grp = r.grp +ORDER BY l.id; +---- +1 x-a2 +2 y-a3 + +# Candidate selection sees the right input after subquery filtering. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN (SELECT * FROM asof_right WHERE val <> 'a6') r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +WHERE l.id IN (3, 5) +ORDER BY l.id; +---- +3 a4 +5 b6 + +# Equality and match operands use the planner's common coercion types. +query II +SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +1 90 + +# Coercion is applied before the ASOF ordering requirements are planned. +query TT +EXPLAIN SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.payload +02)--AsOf Join: match=[CAST(l.ts AS Int64) >= r.ts], constraint=On, on=[CAST(l.grp AS Int64) = r.grp] +03)----SubqueryAlias: l +04)------Projection: column1 AS id, column2 AS grp, column3 AS ts +05)--------Values: (Int64(1), Int16(5) AS Int64(5), Int32(10) AS Int64(10)) +06)----SubqueryAlias: r +07)------Projection: column1 AS grp, column2 AS ts, column3 AS payload +08)--------Values: (Int64(5), Int64(9), Int64(90)) +physical_plan +01)ProjectionExec: expr=[id@0 as id, payload@5 as payload] +02)--AsOfJoinExec: on=[(CAST(grp AS Int64) = grp)], match=[CAST(ts AS Int64) >= ts] +03)----SortExec: expr=[CAST(grp@1 AS Int64) ASC, CAST(ts@2 AS Int64) ASC], preserve_partitioning=[false] +04)------ProjectionExec: expr=[column1@0 as id, column2@1 as grp, column3@2 as ts] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)----ProjectionExec: expr=[column1@0 as grp, column2@1 as ts, column3@2 as payload] +07)------SortExec: expr=[column1@0 ASC, column2@1 ASC], preserve_partitioning=[false] +08)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Equality operands can name the right input first. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON r.grp = l.grp +WHERE l.id IN (2, 4) +ORDER BY l.id; +---- +2 a4 +4 b1 + +query TT +EXPLAIN SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.val +02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +03)----SubqueryAlias: l +04)------TableScan: asof_left projection=[id, grp, ts] +05)----SubqueryAlias: r +06)------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, val@5 as val] +02)--AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts] +03)----SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts = r.ts) +ON l.grp = r.grp; + +query error ASOF MATCH_CONDITION left operand must reference only the left input +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (r.ts >= l.ts) +ON l.grp = r.grp; + +query error ASOF ON accepts only equality conditions combined with AND +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp > r.grp; + +query error ASOF MATCH_CONDITION must be a single comparison +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts) +ON l.grp = r.grp; diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 75e1fe251ac1b..4a8413718edb9 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::{AsOfMatch, LogicalPlanBuilder, Operator}; + use datafusion::logical_expr::LogicalPlanBuilder; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::serializer; @@ -110,10 +110,10 @@ mod tests { 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( + .asof_join_on( right, - vec![(col("l.b"), col("r.b"))], - AsOfMatch::new(col("l.a"), Operator::GtEq, col("r.a")), + Some(col("l.b").eq(col("r.b"))), + col("l.a").gt_eq(col("r.a")), )? .build()?; let error = to_substrait_plan(&plan, &ctx.state()) diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index af442de6597c1..2cbe26fee3964 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -319,6 +319,7 @@ SELECT a FROM table_name WHERE a > 10; ```text from_item [join_type] JOIN from_item [join_condition] +from_item ASOF JOIN from_item MATCH_CONDITION (condition) [join_condition] from_item CROSS JOIN from_item from_item NATURAL JOIN from_item from_item [join_type] JOIN LATERAL (query) AS alias [join_condition] @@ -400,6 +401,48 @@ SELECT * FROM x LEFT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ ``` +### ASOF JOIN + +DataFusion follows the +[Snowflake `ASOF JOIN` syntax](https://docs.snowflake.com/en/sql-reference/constructs/asof-join). +An `ASOF JOIN` matches each left row with at most one right row according to an +ordered comparison. It preserves every left row and fills the right columns +with `NULL` when no right row matches. + +```sql +SELECT t.*, p.price +FROM trades AS t +ASOF JOIN prices AS p +MATCH_CONDITION (t.ts >= p.ts) +ON t.symbol = p.symbol; +``` + +`MATCH_CONDITION` must compare an expression from the left input with an +expression from the right input using one of the following operators. Operand +order is significant: the left input expression must appear on the left. + +| Condition | Selected right row | +| --------- | ----------------------------------------- | +| `l >= r` | Greatest `r` less than or equal to `l` | +| `l > r` | Greatest `r` strictly less than `l` | +| `l <= r` | Smallest `r` greater than or equal to `l` | +| `l < r` | Smallest `r` strictly greater than `l` | + +An optional `ON` clause containing equality conditions combined with `AND`, or +a `USING` clause, divides rows into equality groups before the ordered match. +An unqualified `USING` key appears once in wildcard output, while both qualified +input keys remain addressable. + +Without equality keys, all rows belong to one group. The initial execution +strategy collects one ordered right partition and shares it across every left +partition, so output partitioning follows the left input. The complete right +input must fit in memory and may be scanned once per left partition; spilling +and repartitioned ASOF execution are not yet supported. + +A `NULL` in either ordered expression or in any equality key does not match. +Both inputs must be bounded. If multiple right rows have the same equality keys +and ordered value, which tied row is selected is nondeterministic. + ### RIGHT OUTER JOIN The keywords `RIGHT JOIN` or `RIGHT OUTER JOIN` define a join that includes all rows from the right table even if there