From 17eadf107c29c0cfaeab6779f6842e3549f02c85 Mon Sep 17 00:00:00 2001 From: Jonathan Date: Mon, 6 Jul 2026 10:21:03 -0400 Subject: [PATCH 1/3] asof join attempt --- datafusion-testing | 2 +- .../physical-optimizer/src/join_selection.rs | 325 ++++- .../physical-plan/src/joins/asof_join.rs | 1134 +++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 2 + testing | 2 +- 5 files changed, 1412 insertions(+), 53 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/asof_join.rs diff --git a/datafusion-testing b/datafusion-testing index eccb0e4a42634..f72ac4075ada5 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit eccb0e4a426344ef3faf534cd60e02e9c3afd3ac +Subproject commit f72ac4075ada5ea9810551bc0c3e3161c61204a2 diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 29bbc8e10888b..86d04cfd5e38c 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -28,14 +28,16 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{JoinSide, JoinType, internal_err}; +use datafusion_expr::Operator; use datafusion_expr_common::sort_properties::SortProperties; use datafusion_physical_expr::LexOrdering; -use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{BinaryExpr, Column}; use datafusion_physical_plan::execution_plan::EmissionType; -use datafusion_physical_plan::joins::utils::ColumnIndex; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinCondition, AsOfJoinExec, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; @@ -260,61 +262,62 @@ fn statistical_join_selection_subrule( collect_threshold_byte_size: usize, collect_threshold_num_rows: usize, ) -> Result>> { - let transformed = - if let Some(hash_join) = plan.as_any().downcast_ref::() { - match hash_join.partition_mode() { - PartitionMode::Auto => try_collect_left( - hash_join, - false, - collect_threshold_byte_size, - collect_threshold_num_rows, - )? + let transformed = if let Some(asof_join) = try_asof_join(&plan)? { + Some(asof_join) + } else if let Some(hash_join) = plan.as_any().downcast_ref::() { + match hash_join.partition_mode() { + PartitionMode::Auto => try_collect_left( + hash_join, + false, + collect_threshold_byte_size, + collect_threshold_num_rows, + )? + .map_or_else( + || partitioned_hash_join(hash_join).map(Some), + |v| Ok(Some(v)), + )?, + PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? .map_or_else( || partitioned_hash_join(hash_join).map(Some), |v| Ok(Some(v)), )?, - PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? - .map_or_else( - || partitioned_hash_join(hash_join).map(Some), - |v| Ok(Some(v)), - )?, - PartitionMode::Partitioned => { - let left = hash_join.left(); - let right = hash_join.right(); - // Don't swap null-aware anti joins as they have specific side requirements - if hash_join.join_type().supports_swap() - && !hash_join.null_aware - && should_swap_join_order(&**left, &**right)? - { - hash_join - .swap_inputs(PartitionMode::Partitioned) - .map(Some)? - } else { - None - } + PartitionMode::Partitioned => { + let left = hash_join.left(); + let right = hash_join.right(); + // Don't swap null-aware anti joins as they have specific side requirements + if hash_join.join_type().supports_swap() + && !hash_join.null_aware + && should_swap_join_order(&**left, &**right)? + { + hash_join + .swap_inputs(PartitionMode::Partitioned) + .map(Some)? + } else { + None } } - } else if let Some(cross_join) = plan.as_any().downcast_ref::() { - let left = cross_join.left(); - let right = cross_join.right(); - if should_swap_join_order(&**left, &**right)? { - cross_join.swap_inputs().map(Some)? - } else { - None - } - } else if let Some(nl_join) = plan.as_any().downcast_ref::() { - let left = nl_join.left(); - let right = nl_join.right(); - if nl_join.join_type().supports_swap() - && should_swap_join_order(&**left, &**right)? - { - nl_join.swap_inputs().map(Some)? - } else { - None - } + } + } else if let Some(cross_join) = plan.as_any().downcast_ref::() { + let left = cross_join.left(); + let right = cross_join.right(); + if should_swap_join_order(&**left, &**right)? { + cross_join.swap_inputs().map(Some)? } else { None - }; + } + } else if let Some(nl_join) = plan.as_any().downcast_ref::() { + let left = nl_join.left(); + let right = nl_join.right(); + if nl_join.join_type().supports_swap() + && should_swap_join_order(&**left, &**right)? + { + nl_join.swap_inputs().map(Some)? + } else { + None + } + } else { + None + }; Ok(if let Some(transformed) = transformed { Transformed::yes(transformed) @@ -323,6 +326,226 @@ fn statistical_join_selection_subrule( }) } +fn try_asof_join( + plan: &Arc, +) -> Result>> { + if let Some(hash_join) = plan.as_any().downcast_ref::() { + if !matches!(hash_join.join_type(), JoinType::Inner | JoinType::Left) { + return Ok(None); + } + + let Some(filter) = hash_join.filter() else { + return Ok(None); + }; + let Some((filter_on, asof_condition)) = + extract_asof_join_predicates(filter, hash_join.left(), hash_join.right())? + else { + return Ok(None); + }; + + let mut on = hash_join.on().to_vec(); + on.extend(filter_on); + + return Ok(Some(Arc::new(AsOfJoinExec::try_new( + Arc::clone(hash_join.left()), + Arc::clone(hash_join.right()), + on, + asof_condition, + *hash_join.join_type(), + hash_join.projection.as_ref().map(|p| p.to_vec()), + hash_join.null_equality(), + )?))); + } + + if let Some(nl_join) = plan.as_any().downcast_ref::() { + if !matches!(nl_join.join_type(), JoinType::Inner | JoinType::Left) { + return Ok(None); + } + + let Some(filter) = nl_join.filter() else { + return Ok(None); + }; + let Some((on, asof_condition)) = + extract_asof_join_predicates(filter, nl_join.left(), nl_join.right())? + else { + return Ok(None); + }; + + return Ok(Some(Arc::new(AsOfJoinExec::try_new( + Arc::clone(nl_join.left()), + Arc::clone(nl_join.right()), + on, + asof_condition, + *nl_join.join_type(), + nl_join.projection().as_ref().map(|p| p.to_vec()), + datafusion_common::NullEquality::NullEqualsNothing, + )?))); + } + + Ok(None) +} + +fn extract_asof_join_predicates( + filter: &JoinFilter, + left: &Arc, + right: &Arc, +) -> Result> { + let mut visitor = AsOfPredicateVisitor { + filter, + left, + right, + on: vec![], + asof_condition: None, + inequality_count: 0, + unsupported: false, + }; + visitor.visit(filter.expression())?; + + if visitor.unsupported || visitor.inequality_count != 1 { + return Ok(None); + } + + Ok(visitor + .asof_condition + .map(|asof_condition| (visitor.on, asof_condition))) +} + +struct AsOfPredicateVisitor<'a> { + filter: &'a JoinFilter, + left: &'a Arc, + right: &'a Arc, + on: JoinOn, + asof_condition: Option, + inequality_count: usize, + unsupported: bool, +} + +impl AsOfPredicateVisitor<'_> { + fn visit(&mut self, expr: &PhysicalExprRef) -> Result<()> { + let Some(binary) = expr.as_any().downcast_ref::() else { + self.unsupported = true; + return Ok(()); + }; + + if binary.op() == &Operator::And { + self.visit(binary.left())?; + self.visit(binary.right())?; + return Ok(()); + } + + match binary.op() { + Operator::Eq => { + if let Some((left, right)) = + self.extract_side_pair(binary.left(), binary.right(), *binary.op())? + { + self.on.push((left, right)); + } else { + self.unsupported = true; + } + } + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => { + self.inequality_count += 1; + if self.inequality_count > 1 { + self.unsupported = true; + return Ok(()); + } + + if let Some((left, right, op)) = + self.extract_inequality(binary.left(), binary.right(), *binary.op())? + { + self.asof_condition = + Some(AsOfJoinCondition::try_new(left, op, right)?); + } else { + self.unsupported = true; + } + } + _ => self.unsupported = true, + } + + Ok(()) + } + + fn extract_side_pair( + &self, + left_expr: &PhysicalExprRef, + right_expr: &PhysicalExprRef, + op: Operator, + ) -> Result> { + let Some((left_side, left_expr)) = self.filter_column(left_expr)? else { + return Ok(None); + }; + let Some((right_side, right_expr)) = self.filter_column(right_expr)? else { + return Ok(None); + }; + + match (left_side, right_side, op) { + (JoinSide::Left, JoinSide::Right, Operator::Eq) => { + Ok(Some((left_expr, right_expr))) + } + (JoinSide::Right, JoinSide::Left, Operator::Eq) => { + Ok(Some((right_expr, left_expr))) + } + _ => Ok(None), + } + } + + fn extract_inequality( + &self, + left_expr: &PhysicalExprRef, + right_expr: &PhysicalExprRef, + op: Operator, + ) -> Result> { + let Some((left_side, left_expr)) = self.filter_column(left_expr)? else { + return Ok(None); + }; + let Some((right_side, right_expr)) = self.filter_column(right_expr)? else { + return Ok(None); + }; + + match (left_side, right_side) { + (JoinSide::Left, JoinSide::Right) => Ok(Some((left_expr, right_expr, op))), + (JoinSide::Right, JoinSide::Left) => { + Ok(Some((right_expr, left_expr, flip_inequality(op)?))) + } + _ => Ok(None), + } + } + + fn filter_column( + &self, + expr: &PhysicalExprRef, + ) -> Result> { + let Some(column) = expr.as_any().downcast_ref::() else { + return Ok(None); + }; + let Some(column_index) = self.filter.column_indices().get(column.index()) else { + return Ok(None); + }; + + let (schema, side) = match column_index.side { + JoinSide::Left => (self.left.schema(), JoinSide::Left), + JoinSide::Right => (self.right.schema(), JoinSide::Right), + JoinSide::None => return Ok(None), + }; + + let field = schema.field(column_index.index); + Ok(Some(( + side, + Arc::new(Column::new(field.name(), column_index.index)) as _, + ))) + } +} + +fn flip_inequality(op: Operator) -> Result { + match op { + Operator::Lt => Ok(Operator::Gt), + Operator::LtEq => Ok(Operator::GtEq), + Operator::Gt => Ok(Operator::Lt), + Operator::GtEq => Ok(Operator::LtEq), + _ => internal_err!("Can not flip non-inequality operator {op}"), + } +} + /// Pipeline-fixing join selection subrule. pub type PipelineFixerSubrule = dyn Fn(Arc, &ConfigOptions) -> Result>; 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..95eaa87344498 --- /dev/null +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -0,0 +1,1134 @@ +// 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. + +//! nearest-match inequality join + +use std::any::Any; +use std::cmp::Ordering; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::utils::{ + build_batch_from_indices, build_join_schema, check_join_is_valid, + estimate_join_statistics, ColumnIndex, +}; +use super::JoinOn; +use crate::common::can_project; +use crate::execution_plan::{boundedness_from_children, EmissionType}; +use crate::joins::JoinOnRef; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::projection::{EmbeddedProjection, ProjectionExec}; +use crate::sorts::sort::sort_batch; +use crate::{ + common, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, PhysicalSortExpr, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, Statistics, +}; + +use arrow::array::{ArrayRef, UInt32Builder, UInt64Builder}; +use arrow::compute::{concat_batches, SortOptions}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion_common::{ + exec_err, internal_err, plan_err, project_schema, JoinSide, JoinType, NullEquality, + Result, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::{ColumnarValue, Operator}; +use datafusion_physical_expr::equivalence::{ + join_equivalence_properties, ProjectionMapping, +}; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr_common::physical_expr::fmt_sql; +use datafusion_physical_expr_common::sort_expr::LexOrdering; + +use futures::future::{BoxFuture, FutureExt}; +use futures::{ready, Stream}; + +/// whether an as-of join searches backward or forward from each left row +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AsOfJoinDirection { + /// looks for the greatest right-side as-of key that is less than, or equal + /// to, the left-side key + Backward, + /// looks for the smallest right-side as-of key that is greater than, or + /// equal to, the left-side key + Forward, +} + +/// the inequality predicate that drives an as-of join +#[derive(Debug, Clone)] +pub struct AsOfJoinCondition { + /// as-of key expression evaluated against the left input + pub left: PhysicalExprRef, + /// as-of key expression evaluated against the right input + pub right: PhysicalExprRef, + /// search direction derived from the inequality + pub direction: AsOfJoinDirection, + /// whether the inequality is strict + pub strict: bool, +} + +impl AsOfJoinCondition { + /// creates an as-of condition from a normalized `left op right` + /// inequality + pub fn try_new( + left: PhysicalExprRef, + op: Operator, + right: PhysicalExprRef, + ) -> Result { + let (direction, strict) = match op { + Operator::Gt => (AsOfJoinDirection::Backward, true), + Operator::GtEq => (AsOfJoinDirection::Backward, false), + Operator::Lt => (AsOfJoinDirection::Forward, true), + Operator::LtEq => (AsOfJoinDirection::Forward, false), + _ => { + return plan_err!( + "AsOfJoinCondition requires an inequality operator, got {op}" + ) + } + }; + + Ok(Self { + left, + right, + direction, + strict, + }) + } + + fn sort_options(&self) -> SortOptions { + // for a forward asof join we flip the asof sort order so that lets the + // same merge code keep the "last row that still passes" and still get + // the nearest match instead of the farthest one + let descending = matches!(self.direction, AsOfJoinDirection::Forward); + SortOptions::new(descending, false) + } +} + +/// join execution plan that partitions rows by equality keys and emits at most +/// one nearest right-side row for each left-side row that satisfies the as-of +/// inequality +/// +/// this initial implementation is a bounded, in-memory operator and it requests +/// single-partition inputs, collects both sides, sorts them by equality keys and +/// the as-of key, then does a linear merge inside each equality partition +#[derive(Debug, Clone)] +pub struct AsOfJoinExec { + /// left input + pub(crate) left: Arc, + /// right input + pub(crate) right: Arc, + /// equality predicates used as partition keys + pub(crate) on: JoinOn, + /// inequality predicate used as the as-of key + pub(crate) asof_condition: AsOfJoinCondition, + /// how the join is performed + pub(crate) join_type: JoinType, + /// unprojected join schema + join_schema: SchemaRef, + /// information of index and left / right placement of columns + column_indices: Vec, + /// projection to apply to the output of the join + pub projection: Option>, + /// defines null equality for the equality partition keys + pub(crate) null_equality: NullEquality, + /// execution metrics + metrics: ExecutionPlanMetricsSet, + /// cache holding plan properties + cache: Arc, +} + +impl AsOfJoinExec { + /// tries to create a new as-of join + #[allow(clippy::too_many_arguments)] + pub fn try_new( + left: Arc, + right: Arc, + on: JoinOn, + asof_condition: AsOfJoinCondition, + join_type: JoinType, + projection: Option>, + null_equality: NullEquality, + ) -> Result { + if !matches!(join_type, JoinType::Inner | JoinType::Left) { + return plan_err!( + "AsOfJoinExec supports Inner and Left joins, got {join_type:?}" + ); + } + + let left_schema = left.schema(); + let right_schema = right.schema(); + check_join_is_valid(&left_schema, &right_schema, &on)?; + + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + let join_schema = Arc::new(join_schema); + can_project(&join_schema, projection.as_deref())?; + + let cache = Arc::new(Self::compute_properties( + &left, + &right, + Arc::clone(&join_schema), + join_type, + &on, + projection.as_ref(), + )?); + + Ok(Self { + left, + right, + on, + asof_condition, + join_type, + join_schema, + column_indices, + projection, + null_equality, + metrics: ExecutionPlanMetricsSet::new(), + cache, + }) + } + + /// left input + pub fn left(&self) -> &Arc { + &self.left + } + + /// right input + pub fn right(&self) -> &Arc { + &self.right + } + + /// equality predicates used as partition keys + pub fn on(&self) -> JoinOnRef<'_> { + &self.on + } + + /// inequality predicate used as the as-of key + pub fn asof_condition(&self) -> &AsOfJoinCondition { + &self.asof_condition + } + + /// how the join is performed + pub fn join_type(&self) -> &JoinType { + &self.join_type + } + + /// defines null equality for equality partition keys + pub fn null_equality(&self) -> NullEquality { + self.null_equality + } + + /// returns whether the join contains a projection + pub fn contains_projection(&self) -> bool { + self.projection.is_some() + } + + /// returns a new as-of join with the given projection + pub fn with_projection(&self, projection: Option>) -> Result { + can_project(&self.schema(), projection.as_deref())?; + let projection = match projection { + Some(projection) => match &self.projection { + Some(p) => Some(projection.iter().map(|i| p[*i]).collect()), + None => Some(projection), + }, + None => None, + }; + Self::try_new( + Arc::clone(&self.left), + Arc::clone(&self.right), + self.on.clone(), + self.asof_condition.clone(), + self.join_type, + projection, + self.null_equality, + ) + } + + fn maintains_input_order(_join_type: JoinType) -> Vec { + vec![false, false] + } + + fn compute_properties( + left: &Arc, + right: &Arc, + schema: SchemaRef, + join_type: JoinType, + on: JoinOnRef<'_>, + projection: Option<&Vec>, + ) -> Result { + let mut eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + Arc::clone(&schema), + &Self::maintains_input_order(join_type), + None, + on, + )?; + + let mut output_partitioning = + datafusion_physical_expr::Partitioning::UnknownPartitioning(1); + + if let Some(projection) = projection { + let projection_mapping = + ProjectionMapping::from_indices(projection, &schema)?; + let out_schema = project_schema(&schema, Some(projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, out_schema); + } + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Final, + boundedness_from_children([left, right]), + )) + } +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let display_on = self + .on + .iter() + .map(|(l, r)| { + format!( + "({l}, {r})", + l = fmt_sql(l.as_ref()), + r = fmt_sql(r.as_ref()) + ) + }) + .collect::>() + .join(", "); + let op = match (self.asof_condition.direction, self.asof_condition.strict) + { + (AsOfJoinDirection::Backward, true) => ">", + (AsOfJoinDirection::Backward, false) => ">=", + (AsOfJoinDirection::Forward, true) => "<", + (AsOfJoinDirection::Forward, false) => "<=", + }; + let display_projection = if self.contains_projection() { + format!( + ", projection=[{}]", + self.projection + .as_ref() + .unwrap() + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.fields().get(*index).unwrap().name(), + index + )) + .collect::>() + .join(", ") + ) + } else { + "".to_string() + }; + write!( + f, + "AsOfJoinExec: join_type={:?}, on=[{}], asof={} {} {}{}", + self.join_type, + display_on, + fmt_sql(self.asof_condition.left.as_ref()), + op, + fmt_sql(self.asof_condition.right.as_ref()), + display_projection + ) + } + DisplayFormatType::TreeRender => { + if self.join_type != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type) + } else { + Ok(()) + } + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition, Distribution::SinglePartition] + } + + fn maintains_input_order(&self) -> Vec { + Self::maintains_input_order(self.join_type) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.asof_condition.clone(), + self.join_type, + self.projection.clone(), + self.null_equality, + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return internal_err!( + "Invalid AsOfJoinExec partition {partition}; this operator produces one partition" + ); + } + + let left = self.left.execute(0, Arc::clone(&context))?; + let right = self.right.execute(0, Arc::clone(&context))?; + + let column_indices_after_projection = match &self.projection { + Some(projection) => projection + .iter() + .map(|i| self.column_indices[*i].clone()) + .collect(), + None => self.column_indices.clone(), + }; + + Ok(Box::pin(AsOfJoinStream::new( + self.schema(), + left, + right, + Arc::clone(&self.join_schema), + column_indices_after_projection, + self.on.clone(), + self.asof_condition.clone(), + self.join_type, + self.null_equality, + context.session_config().batch_size(), + BaselineMetrics::new(&self.metrics, partition), + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn partition_statistics(&self, partition: Option) -> Result { + if partition.is_some() { + return Ok(Statistics::new_unknown(&self.schema())); + } + estimate_join_statistics( + self.left.partition_statistics(None)?, + self.right.partition_statistics(None)?, + &self.on, + &self.join_type, + &self.join_schema, + ) + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + if self.contains_projection() { + return Ok(None); + } + + crate::projection::try_embed_projection(projection, self) + } +} + +impl EmbeddedProjection for AsOfJoinExec { + fn with_projection(&self, projection: Option>) -> Result { + self.with_projection(projection) + } +} + +/// stream returned by the as-of join +pub struct AsOfJoinStream { + schema: SchemaRef, + state: AsOfJoinStreamState, + batch_size: usize, + baseline_metrics: BaselineMetrics, +} + +enum AsOfJoinStreamState { + Loading(BoxFuture<'static, Result>), + Draining { batch: RecordBatch, offset: usize }, + Done, +} + +impl AsOfJoinStream { + #[allow(clippy::too_many_arguments)] + fn new( + schema: SchemaRef, + left: SendableRecordBatchStream, + right: SendableRecordBatchStream, + join_schema: SchemaRef, + column_indices: Vec, + on: JoinOn, + asof_condition: AsOfJoinCondition, + join_type: JoinType, + null_equality: NullEquality, + batch_size: usize, + baseline_metrics: BaselineMetrics, + ) -> Self { + let state = AsOfJoinStreamState::Loading( + async move { + let left_schema = left.schema(); + let right_schema = right.schema(); + let left = concat_or_empty(left_schema, common::collect(left).await?)?; + let right = concat_or_empty(right_schema, common::collect(right).await?)?; + + join_asof( + left, + right, + join_schema, + column_indices, + on, + asof_condition, + join_type, + null_equality, + ) + } + .boxed(), + ); + + Self { + schema, + state, + batch_size, + baseline_metrics, + } + } +} + +impl Stream for AsOfJoinStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + loop { + match &mut this.state { + AsOfJoinStreamState::Loading(fut) => { + match ready!(fut.as_mut().poll(cx)) { + Ok(batch) if batch.num_rows() == 0 => { + this.state = AsOfJoinStreamState::Done; + return Poll::Ready(None); + } + Ok(batch) => { + this.state = + AsOfJoinStreamState::Draining { batch, offset: 0 }; + } + Err(e) => { + this.state = AsOfJoinStreamState::Done; + return Poll::Ready(Some(Err(e))); + } + } + } + AsOfJoinStreamState::Draining { batch, offset } => { + if *offset >= batch.num_rows() { + this.state = AsOfJoinStreamState::Done; + return Poll::Ready(None); + } + let len = (batch.num_rows() - *offset).min(this.batch_size); + let output = batch.slice(*offset, len); + *offset += len; + this.baseline_metrics.record_output(len); + return Poll::Ready(Some(Ok(output))); + } + AsOfJoinStreamState::Done => return Poll::Ready(None), + } + } + } +} + +impl RecordBatchStream for AsOfJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +fn concat_or_empty(schema: SchemaRef, batches: Vec) -> Result { + if batches.is_empty() { + Ok(RecordBatch::new_empty(schema)) + } else { + Ok(concat_batches(&schema, &batches)?) + } +} + +#[allow(clippy::too_many_arguments)] +fn join_asof( + left: RecordBatch, + right: RecordBatch, + schema: SchemaRef, + column_indices: Vec, + on: JoinOn, + asof_condition: AsOfJoinCondition, + join_type: JoinType, + null_equality: NullEquality, +) -> Result { + if right.num_rows() > u32::MAX as usize { + return exec_err!( + "AsOfJoinExec does not support right inputs with more than {} rows", + u32::MAX + ); + } + + let left = sort_join_input(&left, &on, &asof_condition, JoinSide::Left)?; + let right = sort_join_input(&right, &on, &asof_condition, JoinSide::Right)?; + + let left_keys = SortedKeyValues::try_new( + &left, + &on, + &asof_condition, + JoinSide::Left, + null_equality, + )?; + let right_keys = SortedKeyValues::try_new( + &right, + &on, + &asof_condition, + JoinSide::Right, + null_equality, + )?; + + let mut left_indices = UInt64Builder::new(); + let mut right_indices = UInt32Builder::new(); + + let mut left_pos = 0; + let mut right_pos = 0; + + // both sides are sorted by the equality keys first, so the outer loop can + // move partition by partition, when a left partition has no matching right + // partition, left joins emit nulls and inner joins simply skip it + while left_pos < left.num_rows() { + if right_pos >= right.num_rows() { + append_unmatched_left_partition( + left_pos, + left.num_rows(), + join_type, + &mut left_indices, + &mut right_indices, + ); + break; + } + + match left_keys.partition_cmp(left_pos, &right_keys, right_pos) { + Ordering::Less => { + let left_end = left_keys.partition_end(left_pos); + append_unmatched_left_partition( + left_pos, + left_end, + join_type, + &mut left_indices, + &mut right_indices, + ); + left_pos = left_end; + } + Ordering::Greater => { + right_pos = right_keys.partition_end(right_pos); + } + Ordering::Equal => { + let left_end = left_keys.partition_end(left_pos); + let right_end = right_keys.partition_end(right_pos); + append_asof_partition( + &left_keys, + &right_keys, + left_pos, + left_end, + right_pos, + right_end, + join_type, + asof_condition.strict, + &mut left_indices, + &mut right_indices, + ); + left_pos = left_end; + right_pos = right_end; + } + } + } + + build_batch_from_indices( + &schema, + &left, + &right, + &left_indices.finish(), + &right_indices.finish(), + &column_indices, + JoinSide::Left, + join_type, + ) +} + +fn append_unmatched_left_partition( + left_start: usize, + left_end: usize, + join_type: JoinType, + left_indices: &mut UInt64Builder, + right_indices: &mut UInt32Builder, +) { + if join_type != JoinType::Left { + return; + } + + for left_idx in left_start..left_end { + left_indices.append_value(left_idx as u64); + right_indices.append_null(); + } +} + +#[allow(clippy::too_many_arguments)] +fn append_asof_partition( + left_keys: &SortedKeyValues, + right_keys: &SortedKeyValues, + left_start: usize, + left_end: usize, + right_start: usize, + right_end: usize, + join_type: JoinType, + strict: bool, + left_indices: &mut UInt64Builder, + right_indices: &mut UInt32Builder, +) { + let mut right_pos = right_start; + let mut last_match = None; + + // within one equality partition, the asof keys are sorted in the direction + // that makes a single right-side cursor enough and every right row we pass is + // a better candidate than the previous one, so `last_match` is the nearest + // right row seen so far for the current and later left rows + for left_idx in left_start..left_end { + if !left_keys.is_joinable(left_idx) { + append_unmatched_left(left_idx, join_type, left_indices, right_indices); + continue; + } + + while right_pos < right_end { + if !right_keys.is_joinable(right_pos) { + right_pos += 1; + continue; + } + + if right_keys.passes_asof(right_pos, left_keys, left_idx, strict) { + last_match = Some(right_pos); + right_pos += 1; + } else { + break; + } + } + + if let Some(right_idx) = last_match { + left_indices.append_value(left_idx as u64); + right_indices.append_value(right_idx as u32); + } else { + append_unmatched_left(left_idx, join_type, left_indices, right_indices); + } + } +} + +fn append_unmatched_left( + left_idx: usize, + join_type: JoinType, + left_indices: &mut UInt64Builder, + right_indices: &mut UInt32Builder, +) { + if join_type == JoinType::Left { + left_indices.append_value(left_idx as u64); + right_indices.append_null(); + } +} + +fn sort_join_input( + batch: &RecordBatch, + on: JoinOnRef<'_>, + asof_condition: &AsOfJoinCondition, + side: JoinSide, +) -> Result { + let mut sort_exprs = on + .iter() + .map(|(left, right)| PhysicalSortExpr { + expr: Arc::clone(if side == JoinSide::Left { left } else { right }), + options: SortOptions::new(false, false), + }) + .collect::>(); + + sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(if side == JoinSide::Left { + &asof_condition.left + } else { + &asof_condition.right + }), + options: asof_condition.sort_options(), + }); + + let Some(sort_exprs) = LexOrdering::new(sort_exprs) else { + return plan_err!("AsOfJoinExec requires at least one as-of sort expression"); + }; + + sort_batch(batch, &sort_exprs, None) +} + +struct SortedKeyValues { + partition_rows: Option, + asof_rows: Rows, + partition_valid: Vec, + asof_valid: Vec, + row_count: usize, +} + +impl SortedKeyValues { + fn try_new( + batch: &RecordBatch, + on: JoinOnRef<'_>, + asof_condition: &AsOfJoinCondition, + side: JoinSide, + null_equality: NullEquality, + ) -> Result { + let partition_arrays = on + .iter() + .map(|(left, right)| { + evaluate_expr(if side == JoinSide::Left { left } else { right }, batch) + }) + .collect::>>()?; + let asof_array = evaluate_expr( + if side == JoinSide::Left { + &asof_condition.left + } else { + &asof_condition.right + }, + batch, + )?; + + let partition_valid = + build_partition_validity(&partition_arrays, batch.num_rows(), null_equality); + let asof_valid = build_validity(&asof_array, batch.num_rows()); + let partition_rows = if partition_arrays.is_empty() { + None + } else { + Some(convert_to_rows( + &partition_arrays, + SortOptions::new(false, false), + )?) + }; + let asof_rows = convert_to_rows(&[asof_array], asof_condition.sort_options())?; + + Ok(Self { + partition_rows, + asof_rows, + partition_valid, + asof_valid, + row_count: batch.num_rows(), + }) + } + + fn partition_cmp( + &self, + left_idx: usize, + right: &SortedKeyValues, + right_idx: usize, + ) -> Ordering { + match (&self.partition_rows, &right.partition_rows) { + (Some(left_rows), Some(right_rows)) => { + left_rows.row(left_idx).cmp(&right_rows.row(right_idx)) + } + (None, None) => Ordering::Equal, + _ => unreachable!("AsOfJoinExec left/right partition key counts differ"), + } + } + + fn partition_end(&self, start: usize) -> usize { + let Some(partition_rows) = &self.partition_rows else { + return self.row_count; + }; + + let mut end = start + 1; + while end < self.row_count && partition_rows.row(start) == partition_rows.row(end) + { + end += 1; + } + end + } + + fn is_joinable(&self, row: usize) -> bool { + self.partition_valid[row] && self.asof_valid[row] + } + + fn passes_asof( + &self, + row: usize, + left: &SortedKeyValues, + left_row: usize, + strict: bool, + ) -> bool { + // `self` is the right side here because forward joins sort the asof key + // descending, the comparison still reads as "has this right row moved + // far enough toward the left row to be a valid candidate?" + let ord = self.asof_rows.row(row).cmp(&left.asof_rows.row(left_row)); + if strict { + ord == Ordering::Less + } else { + ord != Ordering::Greater + } + } +} + +fn evaluate_expr(expr: &PhysicalExprRef, batch: &RecordBatch) -> Result { + match expr.evaluate(batch)? { + ColumnarValue::Array(array) => Ok(array), + ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(batch.num_rows()), + } +} + +fn convert_to_rows(arrays: &[ArrayRef], options: SortOptions) -> Result { + let sort_fields = arrays + .iter() + .map(|array| SortField::new_with_options(array.data_type().clone(), options)) + .collect(); + Ok(RowConverter::new(sort_fields)?.convert_columns(arrays)?) +} + +fn build_partition_validity( + arrays: &[ArrayRef], + row_count: usize, + null_equality: NullEquality, +) -> Vec { + if null_equality == NullEquality::NullEqualsNull { + return vec![true; row_count]; + } + + let mut valid = vec![true; row_count]; + for array in arrays { + if array.null_count() == 0 { + continue; + } + for (row, is_valid) in valid.iter_mut().enumerate() { + *is_valid &= !array.is_null(row); + } + } + valid +} + +fn build_validity(array: &ArrayRef, row_count: usize) -> Vec { + if array.null_count() == 0 { + return vec![true; row_count]; + } + + (0..row_count).map(|row| !array.is_null(row)).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::assert_batches_eq; + use datafusion_execution::TaskContext; + use datafusion_physical_expr::expressions::Column; + + use crate::test::TestMemoryExec; + + fn col(name: &str, index: usize) -> PhysicalExprRef { + Arc::new(Column::new(name, index)) + } + + fn table( + names: [&str; 3], + cols: [Vec>; 3], + ) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new(names[0], DataType::Int32, true), + Field::new(names[1], DataType::Int32, true), + Field::new(names[2], DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(cols[0].clone())), + Arc::new(Int32Array::from(cols[1].clone())), + Arc::new(Int32Array::from(cols[2].clone())), + ], + )?; + + Ok(TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?) + } + + async fn collect_join(join: AsOfJoinExec) -> Result> { + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + common::collect(stream).await + } + + #[tokio::test] + async fn backward_left_asof_join_partitions_and_null_extends() -> Result<()> { + let left = table( + ["l_sym", "l_t", "l_v"], + [ + vec![Some(1), Some(1), Some(2), Some(2), Some(3)], + vec![Some(10), Some(12), Some(5), Some(2), Some(7)], + vec![Some(100), Some(120), Some(200), Some(220), Some(300)], + ], + )?; + let right = table( + ["r_sym", "r_t", "r_v"], + [ + vec![Some(1), Some(1), Some(2), Some(2), Some(4)], + vec![Some(9), Some(11), Some(3), Some(6), Some(1)], + vec![Some(90), Some(110), Some(230), Some(260), Some(400)], + ], + )?; + + let join = AsOfJoinExec::try_new( + left, + right, + vec![(col("l_sym", 0), col("r_sym", 0))], + AsOfJoinCondition::try_new(col("l_t", 1), Operator::GtEq, col("r_t", 1))?, + JoinType::Left, + None, + NullEquality::NullEqualsNothing, + )?; + + let batches = collect_join(join).await?; + let expected = [ + "+-------+-----+-----+-------+-----+-----+", + "| l_sym | l_t | l_v | r_sym | r_t | r_v |", + "+-------+-----+-----+-------+-----+-----+", + "| 1 | 10 | 100 | 1 | 9 | 90 |", + "| 1 | 12 | 120 | 1 | 11 | 110 |", + "| 2 | 2 | 220 | | | |", + "| 2 | 5 | 200 | 2 | 3 | 230 |", + "| 3 | 7 | 300 | | | |", + "+-------+-----+-----+-------+-----+-----+", + ]; + assert_batches_eq!(expected, &batches); + + Ok(()) + } + + #[tokio::test] + async fn forward_inner_asof_join_flips_sort_direction() -> Result<()> { + let left = table( + ["l_sym", "l_t", "l_v"], + [ + vec![Some(1), Some(1), Some(1)], + vec![Some(10), Some(12), Some(15)], + vec![Some(100), Some(120), Some(150)], + ], + )?; + let right = table( + ["r_sym", "r_t", "r_v"], + [ + vec![Some(1), Some(1), Some(1)], + vec![Some(9), Some(11), Some(20)], + vec![Some(90), Some(110), Some(200)], + ], + )?; + + let join = AsOfJoinExec::try_new( + left, + right, + vec![(col("l_sym", 0), col("r_sym", 0))], + AsOfJoinCondition::try_new(col("l_t", 1), Operator::LtEq, col("r_t", 1))?, + JoinType::Inner, + None, + NullEquality::NullEqualsNothing, + )?; + + let batches = collect_join(join).await?; + let expected = [ + "+-------+-----+-----+-------+-----+-----+", + "| l_sym | l_t | l_v | r_sym | r_t | r_v |", + "+-------+-----+-----+-------+-----+-----+", + "| 1 | 15 | 150 | 1 | 20 | 200 |", + "| 1 | 12 | 120 | 1 | 20 | 200 |", + "| 1 | 10 | 100 | 1 | 11 | 110 |", + "+-------+-----+-----+-------+-----+-----+", + ]; + assert_batches_eq!(expected, &batches); + + Ok(()) + } + + #[tokio::test] + async fn null_keys_do_not_match_for_default_null_equality() -> Result<()> { + let left = table( + ["l_sym", "l_t", "l_v"], + [ + vec![None, Some(1), Some(1)], + vec![Some(5), None, Some(10)], + vec![Some(50), Some(100), Some(110)], + ], + )?; + let right = table( + ["r_sym", "r_t", "r_v"], + [ + vec![None, Some(1), Some(1)], + vec![Some(4), Some(9), None], + vec![Some(40), Some(90), Some(999)], + ], + )?; + + let join = AsOfJoinExec::try_new( + left, + right, + vec![(col("l_sym", 0), col("r_sym", 0))], + AsOfJoinCondition::try_new(col("l_t", 1), Operator::GtEq, col("r_t", 1))?, + JoinType::Left, + None, + NullEquality::NullEqualsNothing, + )?; + + let batches = collect_join(join).await?; + let expected = [ + "+-------+-----+-----+-------+-----+-----+", + "| l_sym | l_t | l_v | r_sym | r_t | r_v |", + "+-------+-----+-----+-------+-----+-----+", + "| 1 | 10 | 110 | 1 | 9 | 90 |", + "| 1 | | 100 | | | |", + "| | 5 | 50 | | | |", + "+-------+-----+-----+-------+-----+-----+", + ]; + assert_batches_eq!(expected, &batches); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 2cdfa1e6ac020..b73a871337014 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::{AsOfJoinCondition, AsOfJoinDirection, AsOfJoinExec}; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ @@ -30,6 +31,7 @@ pub use piecewise_merge_join::PiecewiseMergeJoinExec; pub use sort_merge_join::SortMergeJoinExec; pub use symmetric_hash_join::SymmetricHashJoinExec; pub mod chain; +mod asof_join; mod cross_join; mod hash_join; mod nested_loop_join; diff --git a/testing b/testing index 7df2b70baf4f0..d2a1371230349 160000 --- a/testing +++ b/testing @@ -1 +1 @@ -Subproject commit 7df2b70baf4f081ebf8e0c6bd22745cf3cbfd824 +Subproject commit d2a13712303498963395318a4eb42872e66aead7 From 1605743b063b446cf7a873f2b6c78922b6eac2ce Mon Sep 17 00:00:00 2001 From: Jonathan Date: Mon, 6 Jul 2026 11:14:31 -0400 Subject: [PATCH 2/3] benchmark --- datafusion/core/Cargo.toml | 4 + datafusion/core/benches/asof_join_sql.rs | 152 ++++++++++ datafusion/physical-plan/Cargo.toml | 5 + datafusion/physical-plan/benches/asof_join.rs | 275 ++++++++++++++++++ 4 files changed, 436 insertions(+) create mode 100644 datafusion/core/benches/asof_join_sql.rs create mode 100644 datafusion/physical-plan/benches/asof_join.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 8965948a0f4e2..3ee2123310721 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -204,6 +204,10 @@ name = "csv_load" harness = false name = "distinct_query_sql" +[[bench]] +harness = false +name = "asof_join_sql" + [[bench]] harness = false name = "push_down_filter" diff --git a/datafusion/core/benches/asof_join_sql.rs b/datafusion/core/benches/asof_join_sql.rs new file mode 100644 index 0000000000000..f9b3704294c7c --- /dev/null +++ b/datafusion/core/benches/asof_join_sql.rs @@ -0,0 +1,152 @@ +// 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. + +//! Criterion benchmark comparing an as-of join expressed as SQL across builds. +//! +//! The query +//! +//! ```sql +//! SELECT ... FROM left l JOIN right r ON l.sym = r.sym AND l.t >= r.t +//! ``` +//! +//! is run through a full `SessionContext`. On a build that carries the AsOf +//! join optimizer rule (`try_asof_join` inside `JoinSelection`), the physical +//! plan is rewritten to `AsOfJoinExec`, which emits a single nearest match per +//! left row. On a build without the rule (e.g. `branch-53`), the planner falls +//! back to a hash join plus an inequality filter, which materializes *every* +//! right row with `r.t <= l.t`. +//! +//! The two strategies are therefore "similar SQL" rather than identical +//! semantics, but running this same binary on both branches gives an +//! apples-to-apples comparison of how each executes the query. On startup the +//! benchmark prints the physical plan so the active strategy is visible in the +//! output. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::Int64Array; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::datasource::MemTable; +use datafusion::error::Result; +use datafusion::execution::context::SessionContext; +use tokio::runtime::Runtime; + +/// Left/right table sizes and partition count. Kept modest because the +/// no-asof (hash join + filter) plan materializes an inequality join whose +/// output grows with the per-`sym` partition size. +const LEFT_ROWS: usize = 50_000; +const RIGHT_ROWS: usize = 50_000; +const NUM_SYMS: usize = 2_000; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("sym", DataType::Int64, false), + Field::new("t", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])) +} + +fn batches(num_rows: usize, num_syms: usize, schema: &SchemaRef) -> Vec { + let syms: Vec = (0..num_rows).map(|i| (i % num_syms) as i64).collect(); + let ts: Vec = (0..num_rows).map(|i| i as i64).collect(); + let vals: Vec = (0..num_rows).map(|i| i as i64).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(syms)), + Arc::new(Int64Array::from(ts)), + Arc::new(Int64Array::from(vals)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut out = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + out.push(batch.slice(offset, len)); + offset += len; + } + out +} + +fn make_ctx() -> Result { + let ctx = SessionContext::new(); + let s = schema(); + let left = MemTable::try_new(Arc::clone(&s), vec![batches(LEFT_ROWS, NUM_SYMS, &s)])?; + let right = + MemTable::try_new(Arc::clone(&s), vec![batches(RIGHT_ROWS, NUM_SYMS, &s)])?; + ctx.register_table("left_tbl", Arc::new(left))?; + ctx.register_table("right_tbl", Arc::new(right))?; + Ok(ctx) +} + +/// Backward as-of: for each left row, the right row(s) with `r.t <= l.t`. +/// +/// All six columns are selected so no projection is pushed into the join; this +/// keeps the comparison on the raw join/match cost and sidesteps the operator's +/// (currently broken) projection path. +const SQL_INNER: &str = "SELECT l.sym, l.t, l.v, r.sym AS r_sym, r.t AS r_t, r.v AS r_v \ + FROM left_tbl l JOIN right_tbl r ON l.sym = r.sym AND l.t >= r.t"; +const SQL_LEFT: &str = "SELECT l.sym, l.t, l.v, r.sym AS r_sym, r.t AS r_t, r.v AS r_v \ + FROM left_tbl l LEFT JOIN right_tbl r ON l.sym = r.sym AND l.t >= r.t"; + +fn run_sql(ctx: &SessionContext, rt: &Runtime, sql: &str) { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + black_box(rt.block_on(df.collect()).unwrap()); +} + +/// Print the physical plan once so the output records which execution strategy +/// (AsOfJoinExec vs hash join + filter) this build selected. +fn report_plan(ctx: &SessionContext, rt: &Runtime, label: &str, sql: &str) { + let plan = rt + .block_on(async { + let df = ctx.sql(sql).await?; + df.create_physical_plan().await + }) + .unwrap(); + let displayable = + datafusion::physical_plan::displayable(plan.as_ref()).indent(true); + eprintln!("\n===== physical plan [{label}] =====\n{displayable}"); +} + +fn bench_asof_sql(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let ctx = make_ctx().unwrap(); + + report_plan(&ctx, &rt, "inner_backward", SQL_INNER); + report_plan(&ctx, &rt, "left_backward", SQL_LEFT); + + let mut group = c.benchmark_group("asof_join_sql"); + + group.bench_function("inner_backward", |b| { + b.iter(|| run_sql(&ctx, &rt, SQL_INNER)) + }); + group.bench_function("left_backward", |b| { + b.iter(|| run_sql(&ctx, &rt, SQL_LEFT)) + }); + + group.finish(); +} + +criterion_group!(benches, bench_asof_sql); +criterion_main!(benches); diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 6a28486cca5dc..75400454a77dc 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -107,3 +107,8 @@ required-features = ["test_utils"] harness = false name = "aggregate_vectorized" required-features = ["test_utils"] + +[[bench]] +harness = false +name = "asof_join" +required-features = ["test_utils"] diff --git a/datafusion/physical-plan/benches/asof_join.rs b/datafusion/physical-plan/benches/asof_join.rs new file mode 100644 index 0000000000000..7edbf167db724 --- /dev/null +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -0,0 +1,275 @@ +// 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. + +//! Criterion benchmarks for AsOf Join +//! +//! These benchmarks measure `AsOfJoinExec` end-to-end: each side is fed as +//! in-memory `RecordBatch`es and the operator collects, sorts by the equality +//! keys and the time column, and performs the as-of match. The internal sort +//! is therefore included in the measurement, which reflects how the operator +//! behaves in a plan. +//! +//! Data model: `(sym: Int64, t: Int64, v: Int64)` where `sym` is the equality +//! ("partition") key and `t` is the ordering / time column. `num_syms` +//! controls how many distinct partitions the rows are spread across, which +//! drives the partition sizes the match loop walks over. + +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion_common::{JoinType, NullEquality}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::collect; +use datafusion_physical_plan::joins::utils::JoinOn; +use datafusion_physical_plan::joins::{AsOfJoinCondition, AsOfJoinExec}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::ExecutionPlan; +use tokio::runtime::Runtime; + +/// Build in-memory batches (split into ~8192-row chunks). +/// +/// Schema: `(sym: Int64, t: Int64, v: Int64)`. +/// +/// `sym = row_index % num_syms` and `t = row_index`, so each `sym` partition +/// carries a strictly increasing sequence of timestamps once sorted. +/// +/// The operator always sorts its inputs internally. When `sorted` is true the +/// rows are pre-sorted by `(sym, t)`, so that internal sort sees already-ordered +/// input and is effectively free — isolating the join/match kernel. When false +/// the rows are in `t` order (syms interleaved), so the internal sort does real +/// work and is included in the measurement. +fn build_batches( + num_rows: usize, + num_syms: usize, + sorted: bool, + schema: &SchemaRef, +) -> Vec { + // Row order: by default `t == row_index` ascending with syms interleaved. + // When `sorted`, stably reorder by sym so rows come out in (sym, t) order. + let mut order: Vec = (0..num_rows).collect(); + if sorted { + order.sort_by_key(|&i| (i % num_syms) as i64); + } + let syms: Vec = order.iter().map(|&i| (i % num_syms) as i64).collect(); + let ts: Vec = order.iter().map(|&i| i as i64).collect(); + let vals: Vec = order.iter().map(|&i| i as i64).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(syms)), + Arc::new(Int64Array::from(ts)), + Arc::new(Int64Array::from(vals)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec(batches: &[RecordBatch], schema: &SchemaRef) -> Arc { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("sym", DataType::Int64, false), + Field::new("t", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])) +} + +/// Run one as-of join to completion and return the produced row count. +/// +/// `op` selects the as-of direction: `GtEq`/`Gt` are backward (nearest prior +/// right row), `LtEq`/`Lt` are forward (nearest following right row). +fn do_asof_join( + left: Arc, + right: Arc, + op: Operator, + join_type: JoinType, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("sym", &left.schema()).unwrap(), + col("sym", &right.schema()).unwrap(), + )]; + let condition = AsOfJoinCondition::try_new( + col("t", &left.schema()).unwrap(), + op, + col("t", &right.schema()).unwrap(), + ) + .unwrap(); + let join = AsOfJoinExec::try_new( + left, + right, + on, + condition, + join_type, + None, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let task_ctx = Arc::new(TaskContext::default()); + rt.block_on(async { + let batches = collect(Arc::new(join), task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +/// A single benchmark configuration. +struct Case { + name: &'static str, + left_rows: usize, + right_rows: usize, + num_syms: usize, + op: Operator, + join_type: JoinType, +} + +/// Register one case: build both sides once, then re-create the leaf execs +/// each iteration (cheap `Arc` clones of the shared batches). +/// +/// `sorted` selects whether the inputs are pre-sorted by `(sym, t)`: the +/// `presorted` variant makes the operator's internal sort a no-op so the +/// measurement isolates the join/match kernel; `sort_included` measures the +/// operator as it runs on unordered input. +fn run_case( + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, + rt: &Runtime, + s: &SchemaRef, + case: &Case, + sorted: bool, +) { + let left_batches = build_batches(case.left_rows, case.num_syms, sorted, s); + let right_batches = build_batches(case.right_rows, case.num_syms, sorted, s); + let variant = if sorted { "presorted" } else { "sort_included" }; + let id = BenchmarkId::new(format!("{}/{variant}", case.name), case.left_rows); + group.bench_function(id, |b| { + b.iter(|| { + let left = make_exec(&left_batches, s); + let right = make_exec(&right_batches, s); + do_asof_join(left, right, case.op, case.join_type, rt) + }) + }); +} + +fn bench_asof(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + let n = 100_000; + + let cases = [ + // Backward Inner — nearest prior right row per (sym), 1K partitions. + Case { + name: "inner_backward", + left_rows: n, + right_rows: n, + num_syms: 1_000, + op: Operator::GtEq, + join_type: JoinType::Inner, + }, + // Backward Left — like above but unmatched left rows are null-extended. + Case { + name: "left_backward", + left_rows: n, + right_rows: n, + num_syms: 1_000, + op: Operator::GtEq, + join_type: JoinType::Left, + }, + // Backward Left, strict inequality (Gt) — excludes equal-timestamp matches. + Case { + name: "left_backward_strict", + left_rows: n, + right_rows: n, + num_syms: 1_000, + op: Operator::Gt, + join_type: JoinType::Left, + }, + // Forward Inner — nearest following right row (flips the sort direction). + Case { + name: "inner_forward", + left_rows: n, + right_rows: n, + num_syms: 1_000, + op: Operator::LtEq, + join_type: JoinType::Inner, + }, + // Few partitions — 10 syms, so the match loop walks long per-sym runs. + Case { + name: "left_backward_few_partitions", + left_rows: n, + right_rows: n, + num_syms: 10, + op: Operator::GtEq, + join_type: JoinType::Left, + }, + // Many partitions — 50K syms, so partitions are tiny (~2 rows each). + Case { + name: "left_backward_many_partitions", + left_rows: n, + right_rows: n, + num_syms: 50_000, + op: Operator::GtEq, + join_type: JoinType::Left, + }, + // Asymmetric — large left probed against a small right side. + Case { + name: "left_backward_asymmetric", + left_rows: n, + right_rows: n / 10, + num_syms: 1_000, + op: Operator::GtEq, + join_type: JoinType::Left, + }, + ]; + + let mut group = c.benchmark_group("asof_join"); + + // Every case measured on unordered input (internal sort included). + for case in &cases { + run_case(&mut group, &rt, &s, case, false); + } + + // Pre-sorted counterparts for the two primary cases, to isolate the + // join/match kernel from the internal sort cost. + for case in cases + .iter() + .filter(|c| matches!(c.name, "inner_backward" | "left_backward")) + { + run_case(&mut group, &rt, &s, case, true); + } + + group.finish(); +} + +criterion_group!(benches, bench_asof); +criterion_main!(benches); From c39b77bd31a6667e8934c577a06e5034228eb910 Mon Sep 17 00:00:00 2001 From: Jonathan Date: Thu, 23 Jul 2026 09:52:15 -0400 Subject: [PATCH 3/3] feat: ETE integration into DF + partitioned fast path --- datafusion/common/src/config.rs | 6 + datafusion/core/benches/asof_join_sql.rs | 9 +- datafusion/core/src/physical_planner.rs | 169 +++++- datafusion/core/tests/asof_join.rs | 225 ++++++++ datafusion/expr/src/expr_rewriter/mod.rs | 1 + datafusion/expr/src/logical_plan/display.rs | 22 +- datafusion/expr/src/logical_plan/mod.rs | 2 +- datafusion/expr/src/logical_plan/plan.rs | 196 ++++++- datafusion/expr/src/logical_plan/tree_node.rs | 52 +- .../optimizer/src/analyzer/type_coercion.rs | 27 +- .../optimizer/src/common_subexpr_eliminate.rs | 1 + .../optimizer/src/optimize_projections/mod.rs | 13 + .../physical-optimizer/src/join_selection.rs | 325 ++---------- datafusion/physical-plan/benches/asof_join.rs | 4 +- .../physical-plan/src/joins/asof_join.rs | 181 +++++-- datafusion/physical-plan/src/joins/mod.rs | 2 +- datafusion/proto/proto/datafusion.proto | 25 + datafusion/proto/src/generated/pbjson.rs | 487 ++++++++++++++++++ datafusion/proto/src/generated/prost.rs | 48 +- datafusion/proto/src/logical_plan/mod.rs | 92 +++- datafusion/proto/src/physical_plan/mod.rs | 173 ++++++- .../tests/cases/roundtrip_logical_plan.rs | 44 +- .../tests/cases/roundtrip_physical_plan.rs | 39 +- datafusion/sql/src/relation/join.rs | 106 +++- datafusion/sql/src/unparser/plan.rs | 1 + datafusion/sql/tests/sql_integration.rs | 64 +++ .../src/logical_plan/producer/rel/mod.rs | 1 + docs/source/user-guide/configs.md | 1 + 28 files changed, 1965 insertions(+), 351 deletions(-) create mode 100644 datafusion/core/tests/asof_join.rs diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index d71af206c78d5..eb51b4a237547 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1077,6 +1077,12 @@ config_namespace! { /// one range filter. pub enable_piecewise_merge_join: bool, default = false + /// When set to true, an as-of join (`ASOF JOIN`) requires its inputs to be sorted by + /// the equality keys and then the as-of key. The planner adds a sort only when an input + /// is not already so ordered, letting pre-sorted inputs skip sorting entirely. When + /// false, the operator instead collects and sorts each input in memory. + pub asof_join_use_sorted_input: bool, default = true + /// The maximum estimated size in bytes for one input side of a HashJoin /// will be collected into a single partition pub hash_join_single_partition_threshold: usize, default = 1024 * 1024 diff --git a/datafusion/core/benches/asof_join_sql.rs b/datafusion/core/benches/asof_join_sql.rs index f9b3704294c7c..230444052b7c0 100644 --- a/datafusion/core/benches/asof_join_sql.rs +++ b/datafusion/core/benches/asof_join_sql.rs @@ -42,7 +42,7 @@ use std::sync::Arc; use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; -use criterion::{criterion_group, criterion_main, Criterion}; +use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::datasource::MemTable; use datafusion::error::Result; use datafusion::execution::context::SessionContext; @@ -124,8 +124,7 @@ fn report_plan(ctx: &SessionContext, rt: &Runtime, label: &str, sql: &str) { df.create_physical_plan().await }) .unwrap(); - let displayable = - datafusion::physical_plan::displayable(plan.as_ref()).indent(true); + let displayable = datafusion::physical_plan::displayable(plan.as_ref()).indent(true); eprintln!("\n===== physical plan [{label}] =====\n{displayable}"); } @@ -141,9 +140,7 @@ fn bench_asof_sql(c: &mut Criterion) { group.bench_function("inner_backward", |b| { b.iter(|| run_sql(&ctx, &rt, SQL_INNER)) }); - group.bench_function("left_backward", |b| { - b.iter(|| run_sql(&ctx, &rt, SQL_LEFT)) - }); + group.bench_function("left_backward", |b| b.iter(|| run_sql(&ctx, &rt, SQL_LEFT))); group.finish(); } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d20b80ffa9f40..7e474550d79dd 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -29,7 +29,8 @@ use crate::error::{DataFusionError, Result}; use crate::execution::context::{ExecutionProps, SessionState}; use crate::logical_expr::utils::generate_sort_key; use crate::logical_expr::{ - Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, + Aggregate, AsOfJoin, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, + Values, Window, }; use crate::logical_expr::{ Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType, Repartition, @@ -42,7 +43,8 @@ use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::filter::FilterExecBuilder; use crate::physical_plan::joins::utils as join_utils; use crate::physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + AsOfJoinCondition, AsOfJoinExec, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -1709,6 +1711,137 @@ impl DefaultPhysicalPlanner { join } } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + null_equality, + .. + }) => { + let [physical_left, physical_right] = children.two()?; + + let left_df_schema = left.schema(); + let right_df_schema = right.schema(); + let execution_props = session_state.execution_props(); + + // Equijoin keys: `.0` references the LEFT input, `.1` the RIGHT. + let join_on = on + .iter() + .map(|(l, r)| { + let l = create_physical_expr(l, left_df_schema, execution_props)?; + let r = + create_physical_expr(r, right_df_schema, execution_props)?; + Ok((l, r)) + }) + .collect::>()?; + + // The match condition must be a single inequality comparison. + let Expr::BinaryExpr(BinaryExpr { + left: match_left, + op, + right: match_right, + }) = match_condition + else { + return plan_err!( + "AsOfJoin match_condition must be a binary inequality expression, got {match_condition}" + ); + }; + + let mut op = *op; + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "AsOfJoin match_condition requires an inequality operator (<, <=, >, >=), got {op:?}" + ); + } + + fn reverse_ineq(op: Operator) -> Operator { + match op { + Operator::Lt => Operator::Gt, + Operator::LtEq => Operator::GtEq, + Operator::Gt => Operator::Lt, + Operator::GtEq => Operator::LtEq, + _ => op, + } + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Side { + Left, + Right, + Both, + } + + let side_of = |e: &Expr| -> Result { + let cols = e.column_refs(); + let any_left = cols + .iter() + .any(|c| left_df_schema.index_of_column(c).is_ok()); + let any_right = cols + .iter() + .any(|c| right_df_schema.index_of_column(c).is_ok()); + Ok(match (any_left, any_right) { + (true, false) => Side::Left, + (false, true) => Side::Right, + (true, true) => Side::Both, + (false, false) => { + return plan_err!( + "AsOfJoin match_condition operand must reference an input column" + ); + } + }) + }; + + // Orient the condition so it reads `left_input op right_input`, + // flipping the operator if it was written the other way around. + let mut lhs_logical = match_left.as_ref(); + let mut rhs_logical = match_right.as_ref(); + let lhs_side = side_of(lhs_logical)?; + let rhs_side = side_of(rhs_logical)?; + + if lhs_side == Side::Both || rhs_side == Side::Both { + return plan_err!( + "AsOfJoin match_condition operands must each reference a single input side" + ); + } + + if lhs_side == Side::Right && rhs_side == Side::Left { + std::mem::swap(&mut lhs_logical, &mut rhs_logical); + op = reverse_ineq(op); + } else if !(lhs_side == Side::Left && rhs_side == Side::Right) { + return plan_err!( + "AsOfJoin match_condition must compare the left input against the right input" + ); + } + + let condition_left = + create_physical_expr(lhs_logical, left_df_schema, execution_props)?; + let condition_right = + create_physical_expr(rhs_logical, right_df_schema, execution_props)?; + let asof_condition = + AsOfJoinCondition::try_new(condition_left, op, condition_right)?; + + let use_sorted_input = session_state + .config_options() + .optimizer + .asof_join_use_sorted_input; + Arc::new( + AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + asof_condition, + *join_type, + None, + *null_equality, + )? + .with_required_input_ordering(use_sorted_input), + ) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, .. }) => { @@ -2167,6 +2300,7 @@ fn extract_dml_filters( | LogicalPlan::Sort(_) | LogicalPlan::Union(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Window(_) @@ -3784,6 +3918,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_asof_join_planning() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("t", DataType::Int64, false), + ]); + + let left = scan_empty(Some("l"), &schema, None)?.build()?; + let right = scan_empty(Some("r"), &schema, None)?.build()?; + + let asof = AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + vec![(col("l.k"), col("r.k"))], + col("l.t").gt_eq(col("r.t")), + JoinType::Inner, + datafusion_common::NullEquality::NullEqualsNothing, + )?; + let logical_plan = LogicalPlan::AsOfJoin(asof); + + let session_state = make_session_state(); + let planner = DefaultPhysicalPlanner::default(); + let physical_plan = planner + .create_physical_plan(&logical_plan, &session_state) + .await?; + + let displayed = displayable(physical_plan.as_ref()).indent(true).to_string(); + assert_contains!(displayed, "AsOfJoinExec"); + Ok(()) + } + #[tokio::test] async fn test_explain() { let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]); diff --git a/datafusion/core/tests/asof_join.rs b/datafusion/core/tests/asof_join.rs new file mode 100644 index 0000000000000..b4061e0eb61d2 --- /dev/null +++ b/datafusion/core/tests/asof_join.rs @@ -0,0 +1,225 @@ +// 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. + +use std::sync::Arc; + +use arrow::array::{Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion::physical_plan::displayable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::{Result, assert_batches_eq, assert_contains}; + +/// trades(symbol, t, price): rows the quotes align against. +fn trades_batch() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("t", DataType::Int64, false), + Field::new("price", DataType::Int64, false), + ])); + let symbol = Arc::new(StringArray::from(vec!["AAPL", "AAPL", "MSFT"])); + let t = Arc::new(Int64Array::from(vec![10, 20, 10])); + let price = Arc::new(Int64Array::from(vec![100, 200, 50])); + RecordBatch::try_new(schema, vec![symbol, t, price]).map_err(Into::into) +} + +/// quotes(symbol, t, bid): the left/driving side of the ASOF join. +fn quotes_batch() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("t", DataType::Int64, false), + Field::new("bid", DataType::Int64, false), + ])); + let symbol = Arc::new(StringArray::from(vec!["AAPL", "AAPL", "AAPL", "MSFT"])); + let t = Arc::new(Int64Array::from(vec![5, 15, 25, 5])); + let bid = Arc::new(Int64Array::from(vec![1, 2, 3, 9])); + RecordBatch::try_new(schema, vec![symbol, t, bid]).map_err(Into::into) +} + +const ASOF_QUERY: &str = "SELECT q.symbol, q.t, tr.t AS trade_t, tr.price \ + FROM quotes q \ + ASOF JOIN trades tr \ + MATCH_CONDITION (q.t >= tr.t) \ + ON q.symbol = tr.symbol \ + ORDER BY q.symbol, q.t"; + +/// For each quote, the join must return the single nearest earlier trade of the +/// same symbol (largest `tr.t` with `tr.t <= q.t`), and NULLs when none exists +/// (LEFT-join semantics: every quote row survives). +#[tokio::test] +async fn asof_join_nearest_earlier_trade() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_batch("trades", trades_batch()?)?; + ctx.register_batch("quotes", quotes_batch()?)?; + + let results = ctx.sql(ASOF_QUERY).await?.collect().await?; + + assert_batches_eq!( + &[ + "+--------+----+---------+-------+", + "| symbol | t | trade_t | price |", + "+--------+----+---------+-------+", + "| AAPL | 5 | | |", + "| AAPL | 15 | 10 | 100 |", + "| AAPL | 25 | 20 | 200 |", + "| MSFT | 5 | | |", + "+--------+----+---------+-------+", + ], + &results + ); + Ok(()) +} + +/// The query must lower to the dedicated physical operator, not a rewrite into +/// a regular join + filter. +#[tokio::test] +async fn asof_join_uses_physical_operator() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_batch("trades", trades_batch()?)?; + ctx.register_batch("quotes", quotes_batch()?)?; + + let physical_plan = ctx.sql(ASOF_QUERY).await?.create_physical_plan().await?; + let displayed = displayable(physical_plan.as_ref()).indent(true).to_string(); + + assert_contains!(displayed, "AsOfJoinExec"); + Ok(()) +} + +/// With multi-partition inputs and `target_partitions > 1`, the planner must +/// hash-partition both inputs on the equality keys (`RepartitionExec`) so that +/// same-key rows split across input partitions are brought together; each +/// partition is then joined independently and the results are unchanged. +#[tokio::test] +async fn asof_join_runs_partitioned_by_equality_keys() -> Result<()> { + use datafusion::datasource::MemTable; + + fn batch( + schema: &Arc, + symbols: Vec<&str>, + ts: Vec, + vals: Vec, + ) -> Result { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from(symbols)), + Arc::new(Int64Array::from(ts)), + Arc::new(Int64Array::from(vals)), + ], + ) + .map_err(Into::into) + } + + let trades_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("t", DataType::Int64, false), + Field::new("price", DataType::Int64, false), + ])); + let quotes_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("t", DataType::Int64, false), + Field::new("bid", DataType::Int64, false), + ])); + + // Same rows as the single-partition test, but split so AAPL straddles two + // input partitions on each side (exercises cross-partition co-location). + let trades = MemTable::try_new( + Arc::clone(&trades_schema), + vec![ + vec![batch(&trades_schema, vec!["AAPL"], vec![10], vec![100])?], + vec![batch( + &trades_schema, + vec!["AAPL", "MSFT"], + vec![20, 10], + vec![200, 50], + )?], + ], + )?; + let quotes = MemTable::try_new( + Arc::clone("es_schema), + vec![ + vec![batch( + "es_schema, + vec!["AAPL", "AAPL"], + vec![5, 25], + vec![1, 3], + )?], + vec![batch( + "es_schema, + vec!["AAPL", "MSFT"], + vec![15, 5], + vec![2, 9], + )?], + ], + )?; + + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + ctx.register_table("trades", Arc::new(trades))?; + ctx.register_table("quotes", Arc::new(quotes))?; + + // The plan must hash-partition both inputs on the equality keys. + let physical_plan = ctx.sql(ASOF_QUERY).await?.create_physical_plan().await?; + let displayed = displayable(physical_plan.as_ref()).indent(true).to_string(); + assert_contains!(displayed.clone(), "AsOfJoinExec"); + assert_contains!(displayed, "RepartitionExec: partitioning=Hash"); + + // Results must be identical to single-partition execution. + let results = ctx.sql(ASOF_QUERY).await?.collect().await?; + assert_batches_eq!( + &[ + "+--------+----+---------+-------+", + "| symbol | t | trade_t | price |", + "+--------+----+---------+-------+", + "| AAPL | 5 | | |", + "| AAPL | 15 | 10 | 100 |", + "| AAPL | 25 | 20 | 200 |", + "| MSFT | 5 | | |", + "+--------+----+---------+-------+", + ], + &results + ); + Ok(()) +} + +/// With the ordering-aware path disabled via config, the operator sorts its +/// inputs internally (no required input ordering) and still returns the +/// correct nearest-match results. +#[tokio::test] +async fn asof_join_sorts_internally_when_disabled() -> Result<()> { + let mut config = SessionConfig::new(); + config.options_mut().optimizer.asof_join_use_sorted_input = false; + let ctx = SessionContext::new_with_config(config); + ctx.register_batch("trades", trades_batch()?)?; + ctx.register_batch("quotes", quotes_batch()?)?; + + let results = ctx.sql(ASOF_QUERY).await?.collect().await?; + assert_batches_eq!( + &[ + "+--------+----+---------+-------+", + "| symbol | t | trade_t | price |", + "+--------+----+---------+-------+", + "| AAPL | 5 | | |", + "| AAPL | 15 | 10 | 100 |", + "| AAPL | 25 | 20 | 200 |", + "| MSFT | 5 | | |", + "+--------+----+---------+-------+", + ], + &results + ); + Ok(()) +} diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 32a88ab8cf310..d125f83ca39cd 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -324,6 +324,7 @@ impl NamePreserver { plan, LogicalPlan::Filter(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::TableScan(_) | LogicalPlan::Limit(_) | LogicalPlan::Statement(_) diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 58c7feb616179..f1167b23a2266 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,20 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Filter": format!("{}", filter_expr) }) } + LogicalPlan::AsOfJoin(AsOfJoin { + on: keys, + match_condition, + join_type, + .. + }) => { + let join_expr: Vec = + keys.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + json!({ + "Node Type": format!("{} AsOf Join", join_type), + "Join Keys": join_expr.join(", "), + "Match Condition": format!("{}", match_condition) + }) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index c2b01868c97f3..7e1884e39c53f 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -38,7 +38,7 @@ pub use ddl::{ }; pub use dml::{DmlStatement, WriteOp}; pub use plan::{ - Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, + Aggregate, Analyze, AsOfJoin, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 99688a52a75c1..44a7ca38298fd 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -234,6 +234,10 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), + /// AsOf join: match each left row to the single nearest right row per an + /// explicit `MATCH_CONDITION` inequality. Produced only by `ASOF JOIN` + /// syntax, never by an optimizer rewrite, so semantics are always opt-in. + AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -334,6 +338,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(), @@ -365,7 +370,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()) @@ -455,6 +461,7 @@ 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], @@ -563,6 +570,13 @@ impl LogicalPlan { right.head_output_expr() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, right, .. }) => { + if left.schema().fields().is_empty() { + right.head_output_expr() + } else { + left.head_output_expr() + } + } LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => { static_term.head_output_expr() } @@ -686,6 +700,36 @@ impl LogicalPlan { null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema: _, + null_equality, + }) => { + let schema = + build_join_schema(left.schema(), right.schema(), &join_type)?; + + let new_on: Vec<_> = on + .into_iter() + .map(|equi_expr| { + // SimplifyExpression rule may add alias to the equi_expr. + (equi_expr.0.unalias(), equi_expr.1.unalias()) + }) + .collect(); + + Ok(LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on: new_on, + match_condition, + join_type, + schema: DFSchemaRef::new(schema), + null_equality, + })) + } LogicalPlan::Subquery(_) => Ok(self), LogicalPlan::SubqueryAlias(SubqueryAlias { input, @@ -948,6 +992,52 @@ impl LogicalPlan { null_aware: *null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + join_type, + on, + null_equality, + .. + }) => { + let (left, right) = self.only_two_inputs(inputs)?; + let schema = build_join_schema(left.schema(), right.schema(), join_type)?; + + let equi_expr_count = on.len() * 2; + // The equijoin expressions are followed by the single required + // match_condition expression. + assert_eq!(expr.len(), equi_expr_count + 1); + + let Some(match_condition) = expr.pop() else { + internal_err!( + "Expected a match_condition expression to construct the AsOf join" + )? + }; + + // The first part of expr is equi-exprs, + // and the struct of each equi-expr is like `left-expr = right-expr`. + assert_eq!(expr.len(), equi_expr_count); + let mut new_on = Vec::with_capacity(on.len()); + let mut iter = expr.into_iter(); + while let Some(left) = iter.next() { + let Some(right) = iter.next() else { + internal_err!( + "Expected a pair of expressions to construct the join on expression" + )? + }; + + // SimplifyExpression rule may add alias to the equi_expr. + new_on.push((left.unalias(), right.unalias())); + } + + Ok(LogicalPlan::AsOfJoin(AsOfJoin { + left: Arc::new(left), + right: Arc::new(right), + on: new_on, + match_condition, + join_type: *join_type, + schema: DFSchemaRef::new(schema), + null_equality: *null_equality, + })) + } LogicalPlan::Subquery(Subquery { outer_ref_columns, spans, @@ -1360,6 +1450,11 @@ impl LogicalPlan { right.max_rows() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => { + // AsOf joins (Inner/Left) emit at most one right match per + // left row, so the output is bounded by the left input. + left.max_rows() + } LogicalPlan::Repartition(Repartition { input, .. }) => input.max_rows(), LogicalPlan::Union(Union { inputs, .. }) => { inputs.iter().try_fold(0usize, |mut acc, plan| { @@ -1410,6 +1505,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), @@ -1448,6 +1544,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), @@ -2071,6 +2168,20 @@ impl LogicalPlan { } } } + LogicalPlan::AsOfJoin(AsOfJoin { + on: keys, + match_condition, + join_type, + .. + }) => { + let join_expr: Vec = + keys.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + write!(f, "{join_type} AsOf Join:")?; + if !join_expr.is_empty() { + write!(f, " {}", join_expr.join(", "))?; + } + write!(f, " Match: {match_condition}") + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. @@ -3845,6 +3956,89 @@ pub struct Sort { pub fetch: Option, } +/// AsOf join: for each left row, emit the single nearest right row that +/// satisfies the `match_condition` inequality, within groups defined by the +/// `on` equijoin keys. +/// +/// Unlike a regular [`Join`] with an inequality filter (which returns *all* +/// matching right rows), an AsOf join returns only the closest match. It is +/// produced only by explicit `ASOF JOIN ... MATCH_CONDITION (...)` syntax, +/// never by an optimizer rewrite, so query semantics are always opt-in. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AsOfJoin { + /// Left input + pub left: Arc, + /// Right input + pub right: Arc, + /// Equijoin keys from the `ON` clause, as (left, right) expression pairs. + pub on: Vec<(Expr, Expr)>, + /// The single `MATCH_CONDITION` inequality: a binary comparison + /// (`<`, `<=`, `>`, `>=`) oriented as `left_expr op right_expr`. + pub match_condition: Expr, + /// Join type. Only `Inner` and `Left` are supported. + pub join_type: JoinType, + /// The output schema: left fields followed by right fields. + pub schema: DFSchemaRef, + /// Null equality for the equijoin keys. + pub null_equality: NullEquality, +} + +impl AsOfJoin { + /// Creates a new AsOf join, computing the output schema from the inputs. + pub fn try_new( + left: Arc, + right: Arc, + on: Vec<(Expr, Expr)>, + match_condition: Expr, + join_type: JoinType, + null_equality: NullEquality, + ) -> Result { + let join_schema = build_join_schema(left.schema(), right.schema(), &join_type)?; + Ok(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema: Arc::new(join_schema), + null_equality, + }) + } +} + +impl PartialOrd for AsOfJoin { + fn partial_cmp(&self, other: &Self) -> Option { + #[derive(PartialEq, PartialOrd)] + struct ComparableAsOfJoin<'a> { + pub left: &'a Arc, + pub right: &'a Arc, + pub on: &'a Vec<(Expr, Expr)>, + pub match_condition: &'a Expr, + pub join_type: &'a JoinType, + pub null_equality: &'a NullEquality, + } + let comparable_self = ComparableAsOfJoin { + left: &self.left, + right: &self.right, + on: &self.on, + match_condition: &self.match_condition, + join_type: &self.join_type, + null_equality: &self.null_equality, + }; + let comparable_other = ComparableAsOfJoin { + left: &other.left, + right: &other.right, + on: &other.on, + match_condition: &other.match_condition, + join_type: &other.join_type, + null_equality: &other.null_equality, + }; + comparable_self + .partial_cmp(&comparable_other) + .filter(|cmp| *cmp != Ordering::Equal || self == other) + } +} + /// Join two logical plans on one or more join columns #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Join { diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index 90b18dd48fbd4..0c8ab3d8fc0ce 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -38,7 +38,7 @@ //! * [`LogicalPlan::expressions`]: Return a copy of the plan's expressions use crate::{ - Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, + Aggregate, Analyze, AsOfJoin, 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, @@ -147,6 +147,25 @@ impl TreeNode for LogicalPlan { null_aware, }) }), + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema, + null_equality, + }) => (left, right).map_elements(f)?.update_data(|(left, right)| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema, + null_equality, + }) + }), LogicalPlan::Limit(Limit { skip, fetch, input }) => input .map_elements(f)? .update_data(|input| LogicalPlan::Limit(Limit { skip, fetch, input })), @@ -430,6 +449,13 @@ impl LogicalPlan { LogicalPlan::Join(Join { on, filter, .. }) => { (on, filter).apply_ref_elements(f) } + // The rewritable expressions of an AsOf join are the equijoin (on) + // pairs followed by the single required match_condition. + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + .. + }) => (on, match_condition).apply_ref_elements(f), LogicalPlan::Sort(Sort { expr, .. }) => expr.apply_elements(f), LogicalPlan::Extension(extension) => { // would be nice to avoid this copy -- maybe can @@ -580,6 +606,30 @@ impl LogicalPlan { null_aware, }) }), + // The rewritable expressions of an AsOf join are the equijoin (on) + // pairs followed by the single required match_condition, in the same + // order as observed by apply_expressions. + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema, + null_equality, + }) => (on, match_condition).map_elements(f)?.update_data( + |(on, match_condition)| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + schema, + null_equality, + }) + }, + ), 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 ed04aa4285d1a..ebad19dd07fae 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -49,9 +49,9 @@ use datafusion_expr::type_coercion::other::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, is_not_false, is_not_true, - is_not_unknown, is_true, is_unknown, lit, not, + AsOfJoin, Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, + Union, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, is_not_false, + is_not_true, is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -167,6 +167,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(asof) => self.coerce_asof_join(asof), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), _ => Ok(plan), @@ -210,6 +211,26 @@ impl<'a> TypeCoercionRewriter<'a> { Ok(LogicalPlan::Join(join)) } + /// Coerce the equijoin key pairs of an AsOf join. + /// + /// Like [`Self::coerce_join`], the equality keys are stored as a parallel + /// list of left/right expressions, so each pair must be coerced together. + /// The `match_condition` is a single binary comparison and is coerced by the + /// generic expression pass. + pub fn coerce_asof_join(&mut self, mut asof: AsOfJoin) -> Result { + asof.on = asof + .on + .into_iter() + .map(|(lhs, rhs)| { + let left_schema = asof.left.schema(); + let right_schema = asof.right.schema(); + self.coerce_binary_op(lhs, left_schema, Operator::Eq, rhs, right_schema) + }) + .collect::>>()?; + + Ok(LogicalPlan::AsOfJoin(asof)) + } + /// 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 2096c42770315..5a70460b837ef 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -570,6 +570,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 93df300bb50b4..d04894f06fd59 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -398,6 +398,19 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + let (left_req_indices, right_req_indices) = + split_join_requirements(left_len, indices, &join.join_type); + let left_indices = + left_req_indices.with_plan_exprs(&plan, join.left.schema())?; + let right_indices = + right_req_indices.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/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 86d04cfd5e38c..29bbc8e10888b 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -28,16 +28,14 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{JoinSide, JoinType, internal_err}; -use datafusion_expr::Operator; use datafusion_expr_common::sort_properties::SortProperties; use datafusion_physical_expr::LexOrdering; -use datafusion_physical_expr::PhysicalExprRef; -use datafusion_physical_expr::expressions::{BinaryExpr, Column}; +use datafusion_physical_expr::expressions::Column; use datafusion_physical_plan::execution_plan::EmissionType; -use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; +use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::{ - AsOfJoinCondition, AsOfJoinExec, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, - PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, + StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; @@ -262,62 +260,61 @@ fn statistical_join_selection_subrule( collect_threshold_byte_size: usize, collect_threshold_num_rows: usize, ) -> Result>> { - let transformed = if let Some(asof_join) = try_asof_join(&plan)? { - Some(asof_join) - } else if let Some(hash_join) = plan.as_any().downcast_ref::() { - match hash_join.partition_mode() { - PartitionMode::Auto => try_collect_left( - hash_join, - false, - collect_threshold_byte_size, - collect_threshold_num_rows, - )? - .map_or_else( - || partitioned_hash_join(hash_join).map(Some), - |v| Ok(Some(v)), - )?, - PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? + let transformed = + if let Some(hash_join) = plan.as_any().downcast_ref::() { + match hash_join.partition_mode() { + PartitionMode::Auto => try_collect_left( + hash_join, + false, + collect_threshold_byte_size, + collect_threshold_num_rows, + )? .map_or_else( || partitioned_hash_join(hash_join).map(Some), |v| Ok(Some(v)), )?, - PartitionMode::Partitioned => { - let left = hash_join.left(); - let right = hash_join.right(); - // Don't swap null-aware anti joins as they have specific side requirements - if hash_join.join_type().supports_swap() - && !hash_join.null_aware - && should_swap_join_order(&**left, &**right)? - { - hash_join - .swap_inputs(PartitionMode::Partitioned) - .map(Some)? - } else { - None + PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? + .map_or_else( + || partitioned_hash_join(hash_join).map(Some), + |v| Ok(Some(v)), + )?, + PartitionMode::Partitioned => { + let left = hash_join.left(); + let right = hash_join.right(); + // Don't swap null-aware anti joins as they have specific side requirements + if hash_join.join_type().supports_swap() + && !hash_join.null_aware + && should_swap_join_order(&**left, &**right)? + { + hash_join + .swap_inputs(PartitionMode::Partitioned) + .map(Some)? + } else { + None + } } } - } - } else if let Some(cross_join) = plan.as_any().downcast_ref::() { - let left = cross_join.left(); - let right = cross_join.right(); - if should_swap_join_order(&**left, &**right)? { - cross_join.swap_inputs().map(Some)? - } else { - None - } - } else if let Some(nl_join) = plan.as_any().downcast_ref::() { - let left = nl_join.left(); - let right = nl_join.right(); - if nl_join.join_type().supports_swap() - && should_swap_join_order(&**left, &**right)? - { - nl_join.swap_inputs().map(Some)? + } else if let Some(cross_join) = plan.as_any().downcast_ref::() { + let left = cross_join.left(); + let right = cross_join.right(); + if should_swap_join_order(&**left, &**right)? { + cross_join.swap_inputs().map(Some)? + } else { + None + } + } else if let Some(nl_join) = plan.as_any().downcast_ref::() { + let left = nl_join.left(); + let right = nl_join.right(); + if nl_join.join_type().supports_swap() + && should_swap_join_order(&**left, &**right)? + { + nl_join.swap_inputs().map(Some)? + } else { + None + } } else { None - } - } else { - None - }; + }; Ok(if let Some(transformed) = transformed { Transformed::yes(transformed) @@ -326,226 +323,6 @@ fn statistical_join_selection_subrule( }) } -fn try_asof_join( - plan: &Arc, -) -> Result>> { - if let Some(hash_join) = plan.as_any().downcast_ref::() { - if !matches!(hash_join.join_type(), JoinType::Inner | JoinType::Left) { - return Ok(None); - } - - let Some(filter) = hash_join.filter() else { - return Ok(None); - }; - let Some((filter_on, asof_condition)) = - extract_asof_join_predicates(filter, hash_join.left(), hash_join.right())? - else { - return Ok(None); - }; - - let mut on = hash_join.on().to_vec(); - on.extend(filter_on); - - return Ok(Some(Arc::new(AsOfJoinExec::try_new( - Arc::clone(hash_join.left()), - Arc::clone(hash_join.right()), - on, - asof_condition, - *hash_join.join_type(), - hash_join.projection.as_ref().map(|p| p.to_vec()), - hash_join.null_equality(), - )?))); - } - - if let Some(nl_join) = plan.as_any().downcast_ref::() { - if !matches!(nl_join.join_type(), JoinType::Inner | JoinType::Left) { - return Ok(None); - } - - let Some(filter) = nl_join.filter() else { - return Ok(None); - }; - let Some((on, asof_condition)) = - extract_asof_join_predicates(filter, nl_join.left(), nl_join.right())? - else { - return Ok(None); - }; - - return Ok(Some(Arc::new(AsOfJoinExec::try_new( - Arc::clone(nl_join.left()), - Arc::clone(nl_join.right()), - on, - asof_condition, - *nl_join.join_type(), - nl_join.projection().as_ref().map(|p| p.to_vec()), - datafusion_common::NullEquality::NullEqualsNothing, - )?))); - } - - Ok(None) -} - -fn extract_asof_join_predicates( - filter: &JoinFilter, - left: &Arc, - right: &Arc, -) -> Result> { - let mut visitor = AsOfPredicateVisitor { - filter, - left, - right, - on: vec![], - asof_condition: None, - inequality_count: 0, - unsupported: false, - }; - visitor.visit(filter.expression())?; - - if visitor.unsupported || visitor.inequality_count != 1 { - return Ok(None); - } - - Ok(visitor - .asof_condition - .map(|asof_condition| (visitor.on, asof_condition))) -} - -struct AsOfPredicateVisitor<'a> { - filter: &'a JoinFilter, - left: &'a Arc, - right: &'a Arc, - on: JoinOn, - asof_condition: Option, - inequality_count: usize, - unsupported: bool, -} - -impl AsOfPredicateVisitor<'_> { - fn visit(&mut self, expr: &PhysicalExprRef) -> Result<()> { - let Some(binary) = expr.as_any().downcast_ref::() else { - self.unsupported = true; - return Ok(()); - }; - - if binary.op() == &Operator::And { - self.visit(binary.left())?; - self.visit(binary.right())?; - return Ok(()); - } - - match binary.op() { - Operator::Eq => { - if let Some((left, right)) = - self.extract_side_pair(binary.left(), binary.right(), *binary.op())? - { - self.on.push((left, right)); - } else { - self.unsupported = true; - } - } - Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => { - self.inequality_count += 1; - if self.inequality_count > 1 { - self.unsupported = true; - return Ok(()); - } - - if let Some((left, right, op)) = - self.extract_inequality(binary.left(), binary.right(), *binary.op())? - { - self.asof_condition = - Some(AsOfJoinCondition::try_new(left, op, right)?); - } else { - self.unsupported = true; - } - } - _ => self.unsupported = true, - } - - Ok(()) - } - - fn extract_side_pair( - &self, - left_expr: &PhysicalExprRef, - right_expr: &PhysicalExprRef, - op: Operator, - ) -> Result> { - let Some((left_side, left_expr)) = self.filter_column(left_expr)? else { - return Ok(None); - }; - let Some((right_side, right_expr)) = self.filter_column(right_expr)? else { - return Ok(None); - }; - - match (left_side, right_side, op) { - (JoinSide::Left, JoinSide::Right, Operator::Eq) => { - Ok(Some((left_expr, right_expr))) - } - (JoinSide::Right, JoinSide::Left, Operator::Eq) => { - Ok(Some((right_expr, left_expr))) - } - _ => Ok(None), - } - } - - fn extract_inequality( - &self, - left_expr: &PhysicalExprRef, - right_expr: &PhysicalExprRef, - op: Operator, - ) -> Result> { - let Some((left_side, left_expr)) = self.filter_column(left_expr)? else { - return Ok(None); - }; - let Some((right_side, right_expr)) = self.filter_column(right_expr)? else { - return Ok(None); - }; - - match (left_side, right_side) { - (JoinSide::Left, JoinSide::Right) => Ok(Some((left_expr, right_expr, op))), - (JoinSide::Right, JoinSide::Left) => { - Ok(Some((right_expr, left_expr, flip_inequality(op)?))) - } - _ => Ok(None), - } - } - - fn filter_column( - &self, - expr: &PhysicalExprRef, - ) -> Result> { - let Some(column) = expr.as_any().downcast_ref::() else { - return Ok(None); - }; - let Some(column_index) = self.filter.column_indices().get(column.index()) else { - return Ok(None); - }; - - let (schema, side) = match column_index.side { - JoinSide::Left => (self.left.schema(), JoinSide::Left), - JoinSide::Right => (self.right.schema(), JoinSide::Right), - JoinSide::None => return Ok(None), - }; - - let field = schema.field(column_index.index); - Ok(Some(( - side, - Arc::new(Column::new(field.name(), column_index.index)) as _, - ))) - } -} - -fn flip_inequality(op: Operator) -> Result { - match op { - Operator::Lt => Ok(Operator::Gt), - Operator::LtEq => Ok(Operator::GtEq), - Operator::Gt => Ok(Operator::Lt), - Operator::GtEq => Ok(Operator::LtEq), - _ => internal_err!("Can not flip non-inequality operator {op}"), - } -} - /// Pipeline-fixing join selection subrule. pub type PipelineFixerSubrule = dyn Fn(Arc, &ConfigOptions) -> Result>; diff --git a/datafusion/physical-plan/benches/asof_join.rs b/datafusion/physical-plan/benches/asof_join.rs index 7edbf167db724..36a4f07b24c17 100644 --- a/datafusion/physical-plan/benches/asof_join.rs +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -32,16 +32,16 @@ use std::sync::Arc; use arrow::array::{Int64Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_common::{JoinType, NullEquality}; use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::collect; use datafusion_physical_plan::joins::utils::JoinOn; use datafusion_physical_plan::joins::{AsOfJoinCondition, AsOfJoinExec}; use datafusion_physical_plan::test::TestMemoryExec; -use datafusion_physical_plan::ExecutionPlan; use tokio::runtime::Runtime; /// Build in-memory batches (split into ~8192-row chunks). diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 95eaa87344498..deff0f3f34615 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -24,43 +24,43 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use super::JoinOn; use super::utils::{ - build_batch_from_indices, build_join_schema, check_join_is_valid, - estimate_join_statistics, ColumnIndex, + ColumnIndex, build_batch_from_indices, build_join_schema, check_join_is_valid, + estimate_join_statistics, symmetric_join_output_partitioning, }; -use super::JoinOn; use crate::common::can_project; -use crate::execution_plan::{boundedness_from_children, EmissionType}; +use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::joins::JoinOnRef; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::projection::{EmbeddedProjection, ProjectionExec}; use crate::sorts::sort::sort_batch; use crate::{ - common, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PhysicalSortExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + PhysicalSortExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, + Statistics, common, }; use arrow::array::{ArrayRef, UInt32Builder, UInt64Builder}; -use arrow::compute::{concat_batches, SortOptions}; +use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::{ - exec_err, internal_err, plan_err, project_schema, JoinSide, JoinType, NullEquality, - Result, + JoinSide, JoinType, NullEquality, Result, exec_err, internal_err, plan_err, + project_schema, }; use datafusion_execution::TaskContext; use datafusion_expr::{ColumnarValue, Operator}; +use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::equivalence::{ - join_equivalence_properties, ProjectionMapping, + ProjectionMapping, join_equivalence_properties, }; -use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr_common::physical_expr::fmt_sql; -use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; use futures::future::{BoxFuture, FutureExt}; -use futures::{ready, Stream}; +use futures::{Stream, ready}; /// whether an as-of join searches backward or forward from each left row #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -102,7 +102,7 @@ impl AsOfJoinCondition { _ => { return plan_err!( "AsOfJoinCondition requires an inequality operator, got {op}" - ) + ); } }; @@ -127,9 +127,11 @@ impl AsOfJoinCondition { /// one nearest right-side row for each left-side row that satisfies the as-of /// inequality /// -/// this initial implementation is a bounded, in-memory operator and it requests -/// single-partition inputs, collects both sides, sorts them by equality keys and -/// the as-of key, then does a linear merge inside each equality partition +/// this implementation hash-partitions both inputs by the equality keys so +/// partitions can be joined independently in parallel; within each partition it +/// is a bounded, in-memory operator that collects both sides, sorts them by the +/// equality keys and the as-of key, then does a linear merge inside each +/// equality group. with no equality keys it falls back to a single partition #[derive(Debug, Clone)] pub struct AsOfJoinExec { /// left input @@ -150,6 +152,11 @@ pub struct AsOfJoinExec { pub projection: Option>, /// defines null equality for the equality partition keys pub(crate) null_equality: NullEquality, + /// when true the operator requires its inputs to already be sorted by the + /// equality keys and then the as-of key; `required_input_ordering` + /// advertises that order (so pre-sorted inputs skip sorting) and the kernel + /// does not sort internally. when false it collects and sorts internally. + pub(crate) require_input_ordering: bool, /// execution metrics metrics: ExecutionPlanMetricsSet, /// cache holding plan properties @@ -202,11 +209,21 @@ impl AsOfJoinExec { column_indices, projection, null_equality, + require_input_ordering: false, metrics: ExecutionPlanMetricsSet::new(), cache, }) } + /// Returns a copy that requires its inputs to already be ordered by the + /// equality keys and then the as-of key. When set, [`Self::required_input_ordering`] + /// advertises that order — so a pre-sorted input avoids a sort entirely and + /// otherwise the planner inserts one — and the kernel skips its own sort. + pub fn with_required_input_ordering(mut self, require: bool) -> Self { + self.require_input_ordering = require; + self + } + /// left input pub fn left(&self) -> &Arc { &self.left @@ -237,6 +254,12 @@ impl AsOfJoinExec { self.null_equality } + /// whether the operator requires pre-sorted inputs and skips its internal + /// sort (see [`Self::with_required_input_ordering`]) + pub fn require_input_ordering(&self) -> bool { + self.require_input_ordering + } + /// returns whether the join contains a projection pub fn contains_projection(&self) -> bool { self.projection.is_some() @@ -261,6 +284,7 @@ impl AsOfJoinExec { projection, self.null_equality, ) + .map(|exec| exec.with_required_input_ordering(self.require_input_ordering)) } fn maintains_input_order(_join_type: JoinType) -> Vec { @@ -285,8 +309,14 @@ impl AsOfJoinExec { on, )?; - let mut output_partitioning = - datafusion_physical_expr::Partitioning::UnknownPartitioning(1); + // With equality keys, both inputs are hash-partitioned on those keys + // (see `required_input_distribution`), so the output carries the same + // partitioning. Without keys we fall back to a single partition. + let mut output_partitioning = if on.is_empty() { + datafusion_physical_expr::Partitioning::UnknownPartitioning(1) + } else { + symmetric_join_output_partitioning(left, right, &join_type)? + }; if let Some(projection) = projection { let projection_mapping = @@ -383,7 +413,38 @@ impl ExecutionPlan for AsOfJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + if self.on.is_empty() { + // Without equality keys there is nothing to hash-partition on, so + // the nearest match must be computed over a single partition. + vec![Distribution::SinglePartition, Distribution::SinglePartition] + } else { + // Hash-partition both sides on the equality keys so rows sharing a + // key co-locate in the same partition on both sides; each partition + // can then be joined independently in parallel. + let (left_keys, right_keys): (Vec<_>, Vec<_>) = self + .on + .iter() + .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) + .unzip(); + vec![ + Distribution::HashPartitioned(left_keys), + Distribution::HashPartitioned(right_keys), + ] + } + } + + fn required_input_ordering(&self) -> Vec> { + if !self.require_input_ordering { + return vec![None, None]; + } + // This order must match `sort_join_input` exactly so the kernel can + // safely skip its internal sort; when the input already provides it the + // planner inserts no sort. + let left = input_sort_ordering(&self.on, &self.asof_condition, JoinSide::Left) + .map(OrderingRequirements::from); + let right = input_sort_ordering(&self.on, &self.asof_condition, JoinSide::Right) + .map(OrderingRequirements::from); + vec![left, right] } fn maintains_input_order(&self) -> Vec { @@ -398,15 +459,18 @@ impl ExecutionPlan for AsOfJoinExec { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self::try_new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.on.clone(), - self.asof_condition.clone(), - self.join_type, - self.projection.clone(), - self.null_equality, - )?)) + Ok(Arc::new( + Self::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.asof_condition.clone(), + self.join_type, + self.projection.clone(), + self.null_equality, + )? + .with_required_input_ordering(self.require_input_ordering), + )) } fn execute( @@ -414,14 +478,22 @@ impl ExecutionPlan for AsOfJoinExec { partition: usize, context: Arc, ) -> Result { - if partition != 0 { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + if left_partitions != right_partitions { return internal_err!( - "Invalid AsOfJoinExec partition {partition}; this operator produces one partition" + "Invalid AsOfJoinExec: partition count mismatch \ + {left_partitions}!={right_partitions}, consider using RepartitionExec" ); } - let left = self.left.execute(0, Arc::clone(&context))?; - let right = self.right.execute(0, Arc::clone(&context))?; + // Each output partition joins the matching input partition from both + // sides. Because the inputs are hash-partitioned on the equality keys + // (see `required_input_distribution`), all rows sharing a key land in + // the same partition on both sides, so joining per-partition is + // complete and correct. + let left = self.left.execute(partition, Arc::clone(&context))?; + let right = self.right.execute(partition, Arc::clone(&context))?; let column_indices_after_projection = match &self.projection { Some(projection) => projection @@ -435,12 +507,16 @@ impl ExecutionPlan for AsOfJoinExec { self.schema(), left, right, - Arc::clone(&self.join_schema), + // The output batch is assembled directly from the (possibly + // projected) `column_indices`, so its schema must match the + // projected output schema rather than the full join schema. + self.schema(), column_indices_after_projection, self.on.clone(), self.asof_condition.clone(), self.join_type, self.null_equality, + self.require_input_ordering, context.session_config().batch_size(), BaselineMetrics::new(&self.metrics, partition), ))) @@ -507,6 +583,7 @@ impl AsOfJoinStream { asof_condition: AsOfJoinCondition, join_type: JoinType, null_equality: NullEquality, + sorted_input: bool, batch_size: usize, baseline_metrics: BaselineMetrics, ) -> Self { @@ -526,6 +603,7 @@ impl AsOfJoinStream { asof_condition, join_type, null_equality, + sorted_input, ) } .boxed(), @@ -605,6 +683,7 @@ fn join_asof( asof_condition: AsOfJoinCondition, join_type: JoinType, null_equality: NullEquality, + sorted_input: bool, ) -> Result { if right.num_rows() > u32::MAX as usize { return exec_err!( @@ -613,8 +692,16 @@ fn join_asof( ); } - let left = sort_join_input(&left, &on, &asof_condition, JoinSide::Left)?; - let right = sort_join_input(&right, &on, &asof_condition, JoinSide::Right)?; + let (left, right) = if sorted_input { + // Inputs are already ordered by the equality keys then the as-of key + // (guaranteed by `required_input_ordering`), so skip the internal sort. + (left, right) + } else { + ( + sort_join_input(&left, &on, &asof_condition, JoinSide::Left)?, + sort_join_input(&right, &on, &asof_condition, JoinSide::Right)?, + ) + }; let left_keys = SortedKeyValues::try_new( &left, @@ -778,12 +865,15 @@ fn append_unmatched_left( } } -fn sort_join_input( - batch: &RecordBatch, +/// The ordering a join input must have for the merge: the equality keys +/// ascending (nulls last), then the as-of key in the direction the merge +/// expects. Single source of truth shared by the internal sort and +/// [`AsOfJoinExec::required_input_ordering`], so the two can never diverge. +fn input_sort_ordering( on: JoinOnRef<'_>, asof_condition: &AsOfJoinCondition, side: JoinSide, -) -> Result { +) -> Option { let mut sort_exprs = on .iter() .map(|(left, right)| PhysicalSortExpr { @@ -801,7 +891,16 @@ fn sort_join_input( options: asof_condition.sort_options(), }); - let Some(sort_exprs) = LexOrdering::new(sort_exprs) else { + LexOrdering::new(sort_exprs) +} + +fn sort_join_input( + batch: &RecordBatch, + on: JoinOnRef<'_>, + asof_condition: &AsOfJoinCondition, + side: JoinSide, +) -> Result { + let Some(sort_exprs) = input_sort_ordering(on, asof_condition, side) else { return plan_err!("AsOfJoinExec requires at least one as-of sort expression"); }; diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index b73a871337014..f572fbe7e75d1 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -30,8 +30,8 @@ use parking_lot::Mutex; pub use piecewise_merge_join::PiecewiseMergeJoinExec; pub use sort_merge_join::SortMergeJoinExec; pub use symmetric_hash_join::SymmetricHashJoinExec; -pub mod chain; mod asof_join; +pub mod chain; mod cross_join; mod hash_join; mod nested_loop_join; diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 99c4205375468..cc33d7e09fc04 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -62,6 +62,7 @@ message LogicalPlanNode { RecursiveQueryNode recursive_query = 31; CteWorkTableScanNode cte_work_table_scan = 32; DmlNode dml = 33; + AsOfJoinNode as_of_join = 34; } } @@ -252,6 +253,16 @@ message JoinNode { LogicalExprNode filter = 8; } +message AsOfJoinNode { + LogicalPlanNode left = 1; + LogicalPlanNode right = 2; + datafusion_common.JoinType join_type = 3; + repeated LogicalExprNode left_join_key = 4; + repeated LogicalExprNode right_join_key = 5; + datafusion_common.NullEquality null_equality = 6; + LogicalExprNode match_condition = 7; +} + message DistinctNode { LogicalPlanNode input = 1; } @@ -752,6 +763,7 @@ message PhysicalPlanNode { AsyncFuncExecNode async_func = 36; BufferExecNode buffer = 37; ArrowScanExecNode arrow_scan = 38; + AsOfJoinExecNode as_of_join = 39; } } @@ -1167,6 +1179,19 @@ message HashJoinExecNode { bool null_aware = 10; } +message AsOfJoinExecNode { + PhysicalPlanNode left = 1; + PhysicalPlanNode right = 2; + repeated JoinOn on = 3; + datafusion_common.JoinType join_type = 4; + datafusion_common.NullEquality null_equality = 5; + PhysicalExprNode match_condition_left = 6; + PhysicalExprNode match_condition_right = 7; + string match_condition_op = 8; + repeated uint32 projection = 9; + bool require_input_ordering = 10; +} + enum StreamPartitionMode { SINGLE_PARTITION = 0; PARTITIONED_EXEC = 1; diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 2f2958d98487a..51d72efeb1a7f 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -1390,6 +1390,465 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode { deserializer.deserialize_struct("datafusion.ArrowScanExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for AsOfJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.on.is_empty() { + len += 1; + } + if self.join_type != 0 { + len += 1; + } + if self.null_equality != 0 { + len += 1; + } + if self.match_condition_left.is_some() { + len += 1; + } + if self.match_condition_right.is_some() { + len += 1; + } + if !self.match_condition_op.is_empty() { + len += 1; + } + if !self.projection.is_empty() { + len += 1; + } + if self.require_input_ordering { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.on.is_empty() { + struct_ser.serialize_field("on", &self.on)?; + } + if self.join_type != 0 { + let v = super::datafusion_common::JoinType::try_from(self.join_type) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_type)))?; + struct_ser.serialize_field("joinType", &v)?; + } + if self.null_equality != 0 { + let v = super::datafusion_common::NullEquality::try_from(self.null_equality) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.null_equality)))?; + struct_ser.serialize_field("nullEquality", &v)?; + } + if let Some(v) = self.match_condition_left.as_ref() { + struct_ser.serialize_field("matchConditionLeft", v)?; + } + if let Some(v) = self.match_condition_right.as_ref() { + struct_ser.serialize_field("matchConditionRight", v)?; + } + if !self.match_condition_op.is_empty() { + struct_ser.serialize_field("matchConditionOp", &self.match_condition_op)?; + } + if !self.projection.is_empty() { + struct_ser.serialize_field("projection", &self.projection)?; + } + if self.require_input_ordering { + struct_ser.serialize_field("requireInputOrdering", &self.require_input_ordering)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "on", + "join_type", + "joinType", + "null_equality", + "nullEquality", + "match_condition_left", + "matchConditionLeft", + "match_condition_right", + "matchConditionRight", + "match_condition_op", + "matchConditionOp", + "projection", + "require_input_ordering", + "requireInputOrdering", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + On, + JoinType, + NullEquality, + MatchConditionLeft, + MatchConditionRight, + MatchConditionOp, + Projection, + RequireInputOrdering, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "on" => Ok(GeneratedField::On), + "joinType" | "join_type" => Ok(GeneratedField::JoinType), + "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), + "matchConditionLeft" | "match_condition_left" => Ok(GeneratedField::MatchConditionLeft), + "matchConditionRight" | "match_condition_right" => Ok(GeneratedField::MatchConditionRight), + "matchConditionOp" | "match_condition_op" => Ok(GeneratedField::MatchConditionOp), + "projection" => Ok(GeneratedField::Projection), + "requireInputOrdering" | "require_input_ordering" => Ok(GeneratedField::RequireInputOrdering), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut on__ = None; + let mut join_type__ = None; + let mut null_equality__ = None; + let mut match_condition_left__ = None; + let mut match_condition_right__ = None; + let mut match_condition_op__ = None; + let mut projection__ = None; + let mut require_input_ordering__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = Some(map_.next_value()?); + } + GeneratedField::JoinType => { + if join_type__.is_some() { + return Err(serde::de::Error::duplicate_field("joinType")); + } + join_type__ = Some(map_.next_value::()? as i32); + } + GeneratedField::NullEquality => { + if null_equality__.is_some() { + return Err(serde::de::Error::duplicate_field("nullEquality")); + } + null_equality__ = Some(map_.next_value::()? as i32); + } + GeneratedField::MatchConditionLeft => { + if match_condition_left__.is_some() { + return Err(serde::de::Error::duplicate_field("matchConditionLeft")); + } + match_condition_left__ = map_.next_value()?; + } + GeneratedField::MatchConditionRight => { + if match_condition_right__.is_some() { + return Err(serde::de::Error::duplicate_field("matchConditionRight")); + } + match_condition_right__ = map_.next_value()?; + } + GeneratedField::MatchConditionOp => { + if match_condition_op__.is_some() { + return Err(serde::de::Error::duplicate_field("matchConditionOp")); + } + match_condition_op__ = Some(map_.next_value()?); + } + GeneratedField::Projection => { + if projection__.is_some() { + return Err(serde::de::Error::duplicate_field("projection")); + } + projection__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + GeneratedField::RequireInputOrdering => { + if require_input_ordering__.is_some() { + return Err(serde::de::Error::duplicate_field("requireInputOrdering")); + } + require_input_ordering__ = Some(map_.next_value()?); + } + } + } + Ok(AsOfJoinExecNode { + left: left__, + right: right__, + on: on__.unwrap_or_default(), + join_type: join_type__.unwrap_or_default(), + null_equality: null_equality__.unwrap_or_default(), + match_condition_left: match_condition_left__, + match_condition_right: match_condition_right__, + match_condition_op: match_condition_op__.unwrap_or_default(), + projection: projection__.unwrap_or_default(), + require_input_ordering: require_input_ordering__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfJoinNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if self.join_type != 0 { + len += 1; + } + if !self.left_join_key.is_empty() { + len += 1; + } + if !self.right_join_key.is_empty() { + len += 1; + } + if self.null_equality != 0 { + len += 1; + } + if self.match_condition.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if self.join_type != 0 { + let v = super::datafusion_common::JoinType::try_from(self.join_type) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_type)))?; + struct_ser.serialize_field("joinType", &v)?; + } + if !self.left_join_key.is_empty() { + struct_ser.serialize_field("leftJoinKey", &self.left_join_key)?; + } + if !self.right_join_key.is_empty() { + struct_ser.serialize_field("rightJoinKey", &self.right_join_key)?; + } + if self.null_equality != 0 { + let v = super::datafusion_common::NullEquality::try_from(self.null_equality) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.null_equality)))?; + struct_ser.serialize_field("nullEquality", &v)?; + } + if let Some(v) = self.match_condition.as_ref() { + struct_ser.serialize_field("matchCondition", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "join_type", + "joinType", + "left_join_key", + "leftJoinKey", + "right_join_key", + "rightJoinKey", + "null_equality", + "nullEquality", + "match_condition", + "matchCondition", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + JoinType, + LeftJoinKey, + RightJoinKey, + NullEquality, + MatchCondition, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "joinType" | "join_type" => Ok(GeneratedField::JoinType), + "leftJoinKey" | "left_join_key" => Ok(GeneratedField::LeftJoinKey), + "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), + "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), + "matchCondition" | "match_condition" => Ok(GeneratedField::MatchCondition), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut join_type__ = None; + let mut left_join_key__ = None; + let mut right_join_key__ = None; + let mut null_equality__ = None; + let mut match_condition__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::JoinType => { + if join_type__.is_some() { + return Err(serde::de::Error::duplicate_field("joinType")); + } + join_type__ = Some(map_.next_value::()? as i32); + } + GeneratedField::LeftJoinKey => { + if left_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("leftJoinKey")); + } + left_join_key__ = Some(map_.next_value()?); + } + GeneratedField::RightJoinKey => { + if right_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("rightJoinKey")); + } + right_join_key__ = Some(map_.next_value()?); + } + GeneratedField::NullEquality => { + if null_equality__.is_some() { + return Err(serde::de::Error::duplicate_field("nullEquality")); + } + null_equality__ = Some(map_.next_value::()? as i32); + } + GeneratedField::MatchCondition => { + if match_condition__.is_some() { + return Err(serde::de::Error::duplicate_field("matchCondition")); + } + match_condition__ = map_.next_value()?; + } + } + } + Ok(AsOfJoinNode { + left: left__, + right: right__, + join_type: join_type__.unwrap_or_default(), + left_join_key: left_join_key__.unwrap_or_default(), + right_join_key: right_join_key__.unwrap_or_default(), + null_equality: null_equality__.unwrap_or_default(), + match_condition: match_condition__, + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for AsyncFuncExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -12634,6 +13093,9 @@ impl serde::Serialize for LogicalPlanNode { logical_plan_node::LogicalPlanType::Dml(v) => { struct_ser.serialize_field("dml", v)?; } + logical_plan_node::LogicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -12693,6 +13155,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "cte_work_table_scan", "cteWorkTableScan", "dml", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -12729,6 +13193,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { RecursiveQuery, CteWorkTableScan, Dml, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12782,6 +13247,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "recursiveQuery" | "recursive_query" => Ok(GeneratedField::RecursiveQuery), "cteWorkTableScan" | "cte_work_table_scan" => Ok(GeneratedField::CteWorkTableScan), "dml" => Ok(GeneratedField::Dml), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -13026,6 +13492,13 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { return Err(serde::de::Error::duplicate_field("dml")); } logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::Dml) +; + } + GeneratedField::AsOfJoin => { + if logical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::AsOfJoin) ; } } @@ -18099,6 +18572,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::ArrowScan(v) => { struct_ser.serialize_field("arrowScan", v)?; } + physical_plan_node::PhysicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -18169,6 +18645,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "buffer", "arrow_scan", "arrowScan", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -18210,6 +18688,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { AsyncFunc, Buffer, ArrowScan, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -18268,6 +18747,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "asyncFunc" | "async_func" => Ok(GeneratedField::AsyncFunc), "buffer" => Ok(GeneratedField::Buffer), "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18547,6 +19027,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("arrowScan")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::ArrowScan) +; + } + GeneratedField::AsOfJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::AsOfJoin) ; } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 3ff96dc8dcdcf..5f11dd1b3d103 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -5,7 +5,7 @@ pub struct LogicalPlanNode { #[prost( oneof = "logical_plan_node::LogicalPlanType", - tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33" + tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" )] pub logical_plan_type: ::core::option::Option, } @@ -77,6 +77,8 @@ pub mod logical_plan_node { CteWorkTableScan(super::CteWorkTableScanNode), #[prost(message, tag = "33")] Dml(::prost::alloc::boxed::Box), + #[prost(message, tag = "34")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -386,6 +388,23 @@ pub struct JoinNode { pub filter: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(enumeration = "super::datafusion_common::JoinType", tag = "3")] + pub join_type: i32, + #[prost(message, repeated, tag = "4")] + pub left_join_key: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "5")] + pub right_join_key: ::prost::alloc::vec::Vec, + #[prost(enumeration = "super::datafusion_common::NullEquality", tag = "6")] + pub null_equality: i32, + #[prost(message, optional, tag = "7")] + pub match_condition: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1079,7 +1098,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" )] pub physical_plan_type: ::core::option::Option, } @@ -1163,6 +1182,8 @@ pub mod physical_plan_node { Buffer(::prost::alloc::boxed::Box), #[prost(message, tag = "38")] ArrowScan(super::ArrowScanExecNode), + #[prost(message, tag = "39")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1749,6 +1770,29 @@ pub struct HashJoinExecNode { pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub on: ::prost::alloc::vec::Vec, + #[prost(enumeration = "super::datafusion_common::JoinType", tag = "4")] + pub join_type: i32, + #[prost(enumeration = "super::datafusion_common::NullEquality", tag = "5")] + pub null_equality: i32, + #[prost(message, optional, tag = "6")] + pub match_condition_left: ::core::option::Option, + #[prost(message, optional, tag = "7")] + pub match_condition_right: ::core::option::Option, + #[prost(string, tag = "8")] + pub match_condition_op: ::prost::alloc::string::String, + #[prost(uint32, repeated, tag = "9")] + pub projection: ::prost::alloc::vec::Vec, + #[prost(bool, tag = "10")] + pub require_input_ordering: bool, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct SymmetricHashJoinExecNode { #[prost(message, optional, boxed, tag = "1")] pub left: ::core::option::Option<::prost::alloc::boxed::Box>, diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 218c2e4e47d04..1923dbc5df03a 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -58,10 +58,10 @@ use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ - Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, JoinConstraint, Prepare, - Projection, Repartition, Sort, SubqueryAlias, TableScan, Values, Window, - builder::project, + Aggregate, AsOfJoin, CreateCatalog, CreateCatalogSchema, CreateExternalTable, + CreateView, DdlStatement, Distinct, EmptyRelation, Extension, Join, + JoinConstraint, Prepare, Projection, Repartition, Sort, SubqueryAlias, TableScan, + Values, Window, builder::project, }, }; @@ -776,6 +776,44 @@ impl AsLogicalPlan for LogicalPlanNode { builder.build() } + LogicalPlanType::AsOfJoin(join) => { + let left_keys: Vec = + from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; + let right_keys: Vec = + from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + let join_type = + protobuf::JoinType::try_from(join.join_type).map_err(|_| { + proto_error(format!( + "Received an AsOfJoinNode message with unknown JoinType {}", + join.join_type + )) + })?; + let null_equality = + protobuf::NullEquality::try_from(join.null_equality).map_err(|_| { + proto_error(format!( + "Received an AsOfJoinNode message with unknown NullEquality {}", + join.null_equality + )) + })?; + let match_condition = from_proto::parse_expr( + join.match_condition.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode is missing its match_condition") + })?, + ctx, + extension_codec, + )?; + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + let on = left_keys.into_iter().zip(right_keys).collect::>(); + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + on, + match_condition, + join_type.into(), + null_equality.into(), + )?)) + } LogicalPlanType::Union(union) => { assert_or_internal_err!( union.inputs.len() >= 2, @@ -1366,6 +1404,52 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_type, + null_equality, + .. + }) => { + let left: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan( + left.as_ref(), + extension_codec, + )?; + let right: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan( + right.as_ref(), + extension_codec, + )?; + let (left_join_key, right_join_key) = on + .iter() + .map(|(l, r)| { + Ok(( + serialize_expr(l, extension_codec)?, + serialize_expr(r, extension_codec)?, + )) + }) + .collect::, ToProtoError>>()? + .into_iter() + .unzip(); + let join_type: protobuf::JoinType = join_type.to_owned().into(); + let null_equality: protobuf::NullEquality = + null_equality.to_owned().into(); + let match_condition = serialize_expr(match_condition, extension_codec)?; + Ok(LogicalPlanNode { + logical_plan_type: Some(LogicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + join_type: join_type.into(), + left_join_key, + right_join_key, + null_equality: null_equality.into(), + match_condition: Some(match_condition), + }, + ))), + }) + } LogicalPlan::Subquery(_) => { not_impl_err!("LogicalPlan serde is not yet implemented for subqueries") } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index d8435e293249a..c060d1b02101a 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -51,7 +51,7 @@ use datafusion_datasource_parquet::source::ParquetSource; #[cfg(feature = "parquet")] use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; -use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, Operator, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; @@ -76,8 +76,9 @@ use datafusion_physical_plan::expressions::PhysicalSortExpr; use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinCondition, AsOfJoinDirection, AsOfJoinExec, CrossJoinExec, HashJoinExec, + NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, + SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -313,6 +314,13 @@ impl protobuf::PhysicalPlanNode { PhysicalPlanType::SortMergeJoin(sort_join) => { self.try_into_sort_join(sort_join, ctx, codec, proto_converter) } + PhysicalPlanType::AsOfJoin(asof_join) => self + .try_into_as_of_join_physical_plan( + asof_join, + ctx, + codec, + proto_converter, + ), PhysicalPlanType::AsyncFunc(async_func) => self .try_into_async_func_physical_plan( async_func, @@ -405,6 +413,14 @@ impl protobuf::PhysicalPlanNode { ); } + if let Some(exec) = plan.downcast_ref::() { + return protobuf::PhysicalPlanNode::try_from_as_of_join_exec( + exec, + codec, + proto_converter, + ); + } + if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_cross_join_exec( exec, @@ -1422,6 +1438,97 @@ impl protobuf::PhysicalPlanNode { )?)) } + fn try_into_as_of_join_physical_plan( + &self, + asof_join: &protobuf::AsOfJoinExecNode, + ctx: &TaskContext, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let left: Arc = + into_physical_plan(&asof_join.left, ctx, codec, proto_converter)?; + let right: Arc = + into_physical_plan(&asof_join.right, ctx, codec, proto_converter)?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = asof_join + .on + .iter() + .map(|col| { + let left = proto_converter.proto_to_physical_expr( + &col.left.clone().unwrap(), + ctx, + left_schema.as_ref(), + codec, + )?; + let right = proto_converter.proto_to_physical_expr( + &col.right.clone().unwrap(), + ctx, + right_schema.as_ref(), + codec, + )?; + Ok((left, right)) + }) + .collect::>()?; + let join_type = + protobuf::JoinType::try_from(asof_join.join_type).map_err(|_| { + proto_error(format!( + "Received an AsOfJoinExecNode message with unknown JoinType {}", + asof_join.join_type + )) + })?; + let null_equality = protobuf::NullEquality::try_from(asof_join.null_equality) + .map_err(|_| { + proto_error(format!( + "Received an AsOfJoinExecNode message with unknown NullEquality {}", + asof_join.null_equality + )) + })?; + let match_left = proto_converter.proto_to_physical_expr( + asof_join.match_condition_left.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode is missing match_condition_left") + })?, + ctx, + left_schema.as_ref(), + codec, + )?; + let match_right = proto_converter.proto_to_physical_expr( + asof_join.match_condition_right.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode is missing match_condition_right") + })?, + ctx, + right_schema.as_ref(), + codec, + )?; + let op = crate::logical_plan::from_proto::from_proto_binary_op( + &asof_join.match_condition_op, + )?; + let asof_condition = AsOfJoinCondition::try_new(match_left, op, match_right)?; + let projection = if !asof_join.projection.is_empty() { + Some( + asof_join + .projection + .iter() + .map(|i| *i as usize) + .collect::>(), + ) + } else { + None + }; + Ok(Arc::new( + AsOfJoinExec::try_new( + left, + right, + on, + asof_condition, + join_type.into(), + projection, + null_equality.into(), + )? + .with_required_input_ordering(asof_join.require_input_ordering), + )) + } + fn try_into_symmetric_hash_join_physical_plan( &self, sym_join: &protobuf::SymmetricHashJoinExecNode, @@ -2516,6 +2623,66 @@ impl protobuf::PhysicalPlanNode { }) } + fn try_from_as_of_join_exec( + exec: &AsOfJoinExec, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + exec.left().to_owned(), + codec, + proto_converter, + )?; + let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + exec.right().to_owned(), + codec, + proto_converter, + )?; + let on: Vec = exec + .on() + .iter() + .map(|tuple| { + let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; + let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; + Ok::<_, DataFusionError>(protobuf::JoinOn { + left: Some(l), + right: Some(r), + }) + }) + .collect::>()?; + let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); + let null_equality: protobuf::NullEquality = exec.null_equality().into(); + let condition = exec.asof_condition(); + let match_condition_left = + proto_converter.physical_expr_to_proto(&condition.left, codec)?; + let match_condition_right = + proto_converter.physical_expr_to_proto(&condition.right, codec)?; + let op = match (condition.direction, condition.strict) { + (AsOfJoinDirection::Backward, true) => Operator::Gt, + (AsOfJoinDirection::Backward, false) => Operator::GtEq, + (AsOfJoinDirection::Forward, true) => Operator::Lt, + (AsOfJoinDirection::Forward, false) => Operator::LtEq, + }; + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + null_equality: null_equality.into(), + match_condition_left: Some(match_condition_left), + match_condition_right: Some(match_condition_right), + match_condition_op: format!("{op:?}"), + projection: exec.projection.as_ref().map_or_else(Vec::new, |v| { + v.iter().map(|x| *x as u32).collect::>() + }), + require_input_ordering: exec.require_input_ordering(), + }, + ))), + }) + } + fn try_from_symmetric_hash_join_exec( exec: &SymmetricHashJoinExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 9407cbf9a0749..035895b7f28a5 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -78,7 +78,7 @@ use datafusion_expr::expr::{ self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, Like, NullTreatment, ScalarFunction, Unnest, WildcardOptions, }; -use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; +use datafusion_expr::logical_plan::{AsOfJoin, Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, PartitionEvaluator, @@ -777,6 +777,48 @@ async fn create_parquet_scan( Ok(input) } +#[tokio::test] +async fn roundtrip_logical_plan_as_of_join() -> Result<()> { + use datafusion_common::{JoinType, NullEquality}; + + let ctx = SessionContext::new(); + ctx.register_csv("t1", "tests/testdata/test.csv", CsvReadOptions::default()) + .await?; + ctx.register_csv("t2", "tests/testdata/test.csv", CsvReadOptions::default()) + .await?; + + let left = ctx.table("t1").await?.into_optimized_plan()?; + let right = ctx.table("t2").await?.into_optimized_plan()?; + + for (join_type, op) in [ + (JoinType::Inner, Operator::GtEq), + (JoinType::Left, Operator::Gt), + (JoinType::Left, Operator::Lt), + (JoinType::Inner, Operator::LtEq), + ] { + let match_condition = Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("t1.b")), + op, + Box::new(col("t2.b")), + )); + + let plan = LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left.clone()), + Arc::new(right.clone()), + vec![(col("t1.a"), col("t2.a"))], + match_condition, + join_type, + NullEquality::NullEqualsNothing, + )?); + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan}"), format!("{logical_round_trip}")); + } + + Ok(()) +} + #[tokio::test] async fn roundtrip_logical_plan_distinct_on() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 0f3634bdb6c3e..325dfb483157d 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -67,8 +67,8 @@ use datafusion::physical_plan::expressions::{ }; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinCondition, AsOfJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::metrics::MetricType; @@ -309,6 +309,41 @@ fn roundtrip_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_as_of_join() -> Result<()> { + let field_sym = Field::new("sym", DataType::Int64, false); + let field_t = Field::new("t", DataType::Int64, false); + let schema_left = Arc::new(Schema::new(vec![field_sym.clone(), field_t.clone()])); + let schema_right = Arc::new(Schema::new(vec![field_sym, field_t])); + + let on = vec![( + Arc::new(Column::new("sym", 0)) as _, + Arc::new(Column::new("sym", 0)) as _, + )]; + + for join_type in &[JoinType::Inner, JoinType::Left] { + for op in &[Operator::Gt, Operator::GtEq, Operator::Lt, Operator::LtEq] { + for projection in [None, Some(vec![0usize, 2, 3])] { + let condition = AsOfJoinCondition::try_new( + Arc::new(Column::new("t", 1)), + *op, + Arc::new(Column::new("t", 1)), + )?; + roundtrip_test(Arc::new(AsOfJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + condition, + *join_type, + projection, + NullEquality::NullEqualsNothing, + )?))?; + } + } + } + Ok(()) +} + #[test] fn roundtrip_nested_loop_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 8e1a8817309f0..ca05f290d4b7d 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,12 +16,19 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err}; -use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datafusion_common::{ + Column, NullEquality, Result, not_impl_err, plan_datafusion_err, plan_err, +}; +use datafusion_expr::utils::{find_valid_equijoin_key_pair, split_conjunction}; +use datafusion_expr::{ + AsOfJoin, BinaryExpr, Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, +}; use sqlparser::ast::{ - Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, + Expr as SqlExpr, Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, + TableWithJoins, }; use std::collections::HashSet; +use std::sync::Arc; impl SqlToRel<'_, S> { pub(crate) fn plan_table_with_joins( @@ -98,6 +105,16 @@ 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:?}"), } } @@ -110,6 +127,89 @@ impl SqlToRel<'_, S> { LogicalPlanBuilder::from(left).cross_join(right)?.build() } + fn parse_asof_join( + &self, + left: LogicalPlan, + right: LogicalPlan, + match_condition: SqlExpr, + constraint: JoinConstraint, + planner_context: &mut PlannerContext, + ) -> Result { + let join_schema = left.schema().join(right.schema())?; + + // The MATCH_CONDITION must be a single inequality comparison. The + // physical planner auto-orients its operands, so we only validate here. + let match_condition = + self.sql_to_expr(match_condition, &join_schema, planner_context)?; + match &match_condition { + Expr::BinaryExpr(BinaryExpr { op, .. }) + if matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) => {} + other => { + return plan_err!( + "ASOF JOIN MATCH_CONDITION must be a single inequality \ + (<, <=, >, >=), got {other}" + ); + } + } + + // Build the equijoin key pairs from the ON clause, oriented so that + // `.0` references the left input and `.1` the right input. + let on = match constraint { + JoinConstraint::On(sql_expr) => { + let on_expr = + self.sql_to_expr(sql_expr, &join_schema, planner_context)?; + let mut keys: Vec<(Expr, Expr)> = vec![]; + for conjunct in split_conjunction(&on_expr) { + let Expr::BinaryExpr(BinaryExpr { + left: lhs, + op: Operator::Eq, + right: rhs, + }) = conjunct + else { + return plan_err!( + "ASOF JOIN ON clause must be equality conditions \ + between the two inputs" + ); + }; + match find_valid_equijoin_key_pair( + lhs, + rhs, + left.schema(), + right.schema(), + )? { + Some(pair) => keys.push(pair), + None => { + return plan_err!( + "ASOF JOIN ON clause must be equality conditions \ + between the two inputs" + ); + } + } + } + keys + } + // ASOF with no partition keys is allowed. + JoinConstraint::None => vec![], + JoinConstraint::Using(_) | JoinConstraint::Natural => { + return not_impl_err!("USING/NATURAL not supported for ASOF JOIN"); + } + }; + + // Snowflake ASOF keeps every left row, emitting NULLs when there is no + // matching right row, so default to a LEFT join. + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + on, + match_condition, + JoinType::Left, + NullEquality::NullEqualsNothing, + )?)) + } + fn parse_join( &self, left: LogicalPlan, diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index ca8dfa431b4f5..f4150e9b812c5 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -131,6 +131,7 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 444bdae73ac26..1e17c127aafb1 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2513,6 +2513,70 @@ fn join_with_using() { ); } +#[test] +fn asof_join_with_on() { + let sql = "SELECT j1_id, j2_id \ + FROM j1 \ + ASOF JOIN j2 \ + MATCH_CONDITION (j1.j1_id >= j2.j2_id) \ + ON j1.j1_string = j2.j2_string"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r" + Projection: j1.j1_id, j2.j2_id + Left AsOf Join: j1.j1_string = j2.j2_string Match: j1.j1_id >= j2.j2_id + TableScan: j1 + TableScan: j2 + " + ); +} + +#[test] +fn asof_join_without_on() { + let sql = "SELECT j1_id, j2_id \ + FROM j1 \ + ASOF JOIN j2 \ + MATCH_CONDITION (j1.j1_id >= j2.j2_id)"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r" + Projection: j1.j1_id, j2.j2_id + Left AsOf Join: Match: j1.j1_id >= j2.j2_id + TableScan: j1 + TableScan: j2 + " + ); +} + +#[test] +fn asof_join_match_condition_not_inequality() { + let sql = "SELECT j1_id, j2_id \ + FROM j1 \ + ASOF JOIN j2 \ + MATCH_CONDITION (j1.j1_id = j2.j2_id) \ + ON j1.j1_string = j2.j2_string"; + let err = logical_plan(sql).unwrap_err().strip_backtrace(); + assert_contains!( + err, + "ASOF JOIN MATCH_CONDITION must be a single inequality (<, <=, >, >=)" + ); +} + +#[test] +fn asof_join_match_condition_conjunction() { + let sql = "SELECT j1_id, j2_id \ + FROM j1 \ + ASOF JOIN j2 \ + MATCH_CONDITION (j1.j1_id >= j2.j2_id AND j1.j1_id <= j2.j2_id)"; + let err = logical_plan(sql).unwrap_err().strip_backtrace(); + assert_contains!( + err, + "ASOF JOIN MATCH_CONDITION must be a single inequality (<, <=, >, >=)" + ); +} + #[test] fn equijoin_explicit_syntax_3_tables() { let sql = "SELECT id, order_id, l_description \ diff --git a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs index c3599a2635ffa..c340f42d8adc2 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs @@ -51,6 +51,7 @@ 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!("Unsupported plan type: {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/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 1245e59d477a7..ab90e04562f74 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -157,6 +157,7 @@ The following configuration settings are available: | datafusion.optimizer.top_down_join_key_reordering | true | When set to true, the physical plan optimizer will run a top down process to reorder the join keys | | datafusion.optimizer.prefer_hash_join | true | When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. HashJoin can work more efficiently than SortMergeJoin but consumes more memory | | datafusion.optimizer.enable_piecewise_merge_join | false | When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently experimental. Physical planner will opt for PiecewiseMergeJoin when there is only one range filter. | +| datafusion.optimizer.asof_join_use_sorted_input | true | When set to true, an as-of join (`ASOF JOIN`) requires its inputs to be sorted by the equality keys and then the as-of key. The planner adds a sort only when an input is not already so ordered, letting pre-sorted inputs skip sorting entirely. When false, the operator instead collects and sorts each input in memory. | | datafusion.optimizer.hash_join_single_partition_threshold | 1048576 | The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition | | datafusion.optimizer.hash_join_single_partition_threshold_rows | 131072 | The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition | | datafusion.optimizer.hash_join_inlist_pushdown_max_size | 131072 | Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides larger than this will use hash table lookups instead. Set to 0 to always use hash table lookups. InList pushdown can be more efficient for small build sides because it can result in better statistics pruning as well as use any bloom filters present on the scan side. InList expressions are also more transparent and easier to serialize over the network in distributed uses of DataFusion. On the other hand InList pushdown requires making a copy of the data and thus adds some overhead to the build side and uses more memory. This setting is per-partition, so we may end up using `hash_join_inlist_pushdown_max_size` \* `target_partitions` memory. The default is 128kB per partition. This should allow point lookup joins (e.g. joining on a unique primary key) to use InList pushdown in most cases but avoids excessive memory usage or overhead for larger joins. |