diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aad392693723..0122df6fb7d5 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -41,7 +41,8 @@ use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::filter::FilterExecBuilder; use crate::physical_plan::joins::utils as join_utils; use crate::physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -1790,6 +1791,51 @@ impl DefaultPhysicalPlanner { join } } + LogicalPlan::AsOfJoin(join) => { + let [physical_left, physical_right] = children.two()?; + let join_on = join + .on + .iter() + .map(|(left, right)| { + Ok(( + create_physical_expr( + left, + join.left.schema(), + execution_props, + planning_ctx, + )?, + create_physical_expr( + right, + join.right.schema(), + execution_props, + planning_ctx, + )?, + )) + }) + .collect::>()?; + let match_condition = AsOfMatchExpr::new( + create_physical_expr( + &join.match_condition.left, + join.left.schema(), + execution_props, + planning_ctx, + )?, + join.match_condition.op, + create_physical_expr( + &join.match_condition.right, + join.right.schema(), + execution_props, + planning_ctx, + )?, + ); + Arc::new(AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + match_condition, + None, + )?) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, @@ -2291,6 +2337,7 @@ fn extract_dml_filters( | LogicalPlan::Sort(_) | LogicalPlan::Union(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Window(_) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index ef5e496b0d7d..7f090eeb833f 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -31,10 +31,10 @@ use crate::expr_rewriter::{ rewrite_sort_cols_by_aggs, }; use crate::logical_plan::{ - Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare, - Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest, - Values, Window, + Aggregate, Analyze, AsOfJoin, AsOfMatch, Distinct, DistinctOn, EmptyRelation, + Explain, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, + PlanType, Prepare, Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, + Union, Unnest, Values, Window, }; use crate::select_expr::SelectExpr; use crate::utils::{ @@ -1007,6 +1007,68 @@ impl LogicalPlanBuilder { ) } + /// Apply a left-preserving ASOF join using equality expressions and one + /// ordered match condition. + pub fn asof_join( + self, + right: LogicalPlan, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + ) -> Result { + self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On) + } + + /// Apply a left-preserving ASOF join using `USING` equality keys. + pub fn asof_join_using( + self, + right: LogicalPlan, + using_keys: Vec, + match_condition: AsOfMatch, + ) -> Result { + let on = using_keys + .into_iter() + .map(|key| { + let left = Self::normalize(&self.plan, key.clone())?; + let right = Self::normalize(&right, key)?; + Ok((Expr::Column(left), Expr::Column(right))) + }) + .collect::>()?; + self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using) + } + + fn asof_join_with_constraint( + self, + right: LogicalPlan, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + join_constraint: JoinConstraint, + ) -> Result { + let normalize = |expr, schema: &DFSchema| { + normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[]) + }; + let on = on + .into_iter() + .map(|(left, right_expr)| { + Ok(( + normalize(left, self.plan.schema())?, + normalize(right_expr, right.schema())?, + )) + }) + .collect::>()?; + let match_condition = AsOfMatch { + left: normalize(match_condition.left, self.plan.schema())?, + op: match_condition.op, + right: normalize(match_condition.right, right.schema())?, + }; + Ok(Self::new(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + self.plan, + Arc::new(right), + on, + match_condition, + join_constraint, + )?))) + } + pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result { if column.relation.is_some() { // column is already normalized @@ -1776,6 +1838,14 @@ pub fn build_join_schema( dfschema.with_functional_dependencies(func_dependencies) } +/// Creates the schema for a left-preserving ASOF join. +/// +/// Both `ON` and `USING` preserve all qualified input fields. SQL wildcard +/// expansion handles the unqualified `USING` key as a single column. +pub fn build_asof_join_schema(left: &DFSchema, right: &DFSchema) -> Result { + build_join_schema(left, right, &JoinType::Left) +} + /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise /// conflict with the columns from the other. /// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 09f41c94f64f..c5cd003d1f34 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -21,10 +21,10 @@ use std::collections::HashMap; use std::fmt; use crate::{ - Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, - Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, - Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values, - Window, expr_vec_fmt, + Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, + Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, + Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, + Values, Window, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -493,6 +493,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Filter": format!("{}", filter_expr) }) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let join_expr: Vec = + on.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + json!({ + "Node Type": "AsOf Join", + "Join Constraint": format!("{join_constraint:?}"), + "Join Keys": join_expr.join(", "), + "Match Condition": match_condition.to_string(), + }) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 4766c3f33379..98113d12c1b4 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -28,8 +28,8 @@ pub mod tree_node; pub use builder::{ LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE, - build_join_schema, requalify_sides_if_needed, table_scan, union, - wrap_projection_for_join_if_necessary, + build_asof_join_schema, build_join_schema, requalify_sides_if_needed, table_scan, + union, wrap_projection_for_join_if_necessary, }; pub use ddl::{ CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction, @@ -41,12 +41,12 @@ pub use dml::{ WriteOp, }; pub use plan::{ - Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, - EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, - Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, - Unnest, Values, Window, projection_schema, + Aggregate, Analyze, AsOfJoin, AsOfMatch, ColumnUnnestList, DescribeTable, Distinct, + DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, + Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, + Projection, RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, + StringifiedPlan, Subquery, SubqueryAlias, TableScan, TableScanBuilder, + ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 1a8cd81aa74b..1aa1fd561cbd 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,14 +41,15 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement, WriteOp}; use crate::utils::{ - check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, - find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, - merge_schema, split_conjunction, + check_aggregate_and_window_nesting, enumerate_grouping_sets, expr_to_columns, + exprlist_to_fields, find_out_reference_exprs, grouping_set_expr_count, + grouping_set_to_exprlist, merge_schema, split_conjunction, }; use crate::{ BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet, LogicalPlanBuilder, Operator, Prepare, TableProviderFilterPushDown, TableSource, - WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed, + WindowFunctionDefinition, build_asof_join_schema, build_join_schema, expr_vec_fmt, + requalify_sides_if_needed, }; use crate::statistics::StatisticsRequest; @@ -295,6 +296,9 @@ pub enum LogicalPlan { Unnest(Unnest), /// A variadic query (e.g. "Recursive CTEs") RecursiveQuery(RecursiveQuery), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), } impl Default for LogicalPlan { @@ -342,6 +346,7 @@ impl LogicalPlan { LogicalPlan::Aggregate(Aggregate { schema, .. }) => schema, LogicalPlan::Sort(Sort { input, .. }) => input.schema(), LogicalPlan::Join(Join { schema, .. }) => schema, + LogicalPlan::AsOfJoin(AsOfJoin { schema, .. }) => schema, LogicalPlan::Repartition(Repartition { input, .. }) => input.schema(), LogicalPlan::Limit(Limit { input, .. }) => input.schema(), LogicalPlan::Statement(statement) => statement.schema(), @@ -370,7 +375,8 @@ impl LogicalPlan { | LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Unnest(_) - | LogicalPlan::Join(_) => self + | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) => self .inputs() .iter() .map(|input| input.schema().as_ref()) @@ -460,6 +466,9 @@ impl LogicalPlan { LogicalPlan::Aggregate(Aggregate { input, .. }) => vec![input], LogicalPlan::Sort(Sort { input, .. }) => vec![input], LogicalPlan::Join(Join { left, right, .. }) => vec![left, right], + LogicalPlan::AsOfJoin(AsOfJoin { left, right, .. }) => { + vec![left, right] + } LogicalPlan::Limit(Limit { input, .. }) => vec![input], LogicalPlan::Subquery(Subquery { subquery, .. }) => vec![subquery], LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => vec![input], @@ -495,12 +504,20 @@ impl LogicalPlan { let mut using_columns: Vec> = vec![]; self.apply_with_subqueries(|plan| { - if let LogicalPlan::Join(Join { - join_constraint: JoinConstraint::Using, - on, - .. - }) = plan - { + let on = match plan { + LogicalPlan::Join(Join { + join_constraint: JoinConstraint::Using, + on, + .. + }) + | LogicalPlan::AsOfJoin(AsOfJoin { + join_constraint: JoinConstraint::Using, + on, + .. + }) => Some(on), + _ => None, + }; + if let Some(on) = on { // The join keys in using-join must be columns. let columns = on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| { @@ -568,6 +585,7 @@ impl LogicalPlan { right.head_output_expr() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => left.head_output_expr(), LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => { static_term.head_output_expr() } @@ -691,6 +709,26 @@ impl LogicalPlan { null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema: _, + }) => Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + left, + right, + on.into_iter() + .map(|(left, right)| (left.unalias(), right.unalias())) + .collect(), + AsOfMatch { + left: match_condition.left.unalias(), + op: match_condition.op, + right: match_condition.right.unalias(), + }, + join_constraint, + )?)), LogicalPlan::Subquery(_) => Ok(self), LogicalPlan::SubqueryAlias(SubqueryAlias { input, @@ -994,6 +1032,45 @@ impl LogicalPlan { null_aware: *null_aware, })) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let (left, right) = self.only_two_inputs(inputs)?; + let expected = on.len() * 2 + 2; + assert_eq_or_internal_err!( + expected, + expr.len(), + "Invalid number of new ASOF join expressions: expected {}, got {}", + expected, + expr.len() + ); + + let mut iter = expr.into_iter(); + let mut new_on = Vec::with_capacity(on.len()); + for _ in 0..on.len() { + let left = iter.next().expect("expression count checked").unalias(); + let right = iter.next().expect("expression count checked").unalias(); + new_on.push((left, right)); + } + let match_left = iter.next().expect("expression count checked").unalias(); + let match_right = + iter.next().expect("expression count checked").unalias(); + + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + new_on, + AsOfMatch { + left: match_left, + op: match_condition.op, + right: match_right, + }, + *join_constraint, + )?)) + } LogicalPlan::Subquery(Subquery { outer_ref_columns, spans, @@ -1419,6 +1496,7 @@ impl LogicalPlan { right.max_rows() } }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => left.max_rows(), LogicalPlan::Repartition(Repartition { input, .. }) => input.max_rows(), LogicalPlan::Union(Union { inputs, .. }) => { inputs.iter().try_fold(0usize, |mut acc, plan| { @@ -1494,6 +1572,7 @@ impl LogicalPlan { | JoinType::LeftAnti | JoinType::RightAnti => 0, }, + LogicalPlan::AsOfJoin(AsOfJoin { left, .. }) => left.min_rows(), LogicalPlan::Union(Union { inputs, .. }) => inputs .iter() .fold(0, |rows, input| rows.saturating_add(input.min_rows())), @@ -1547,6 +1626,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), @@ -1585,6 +1665,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), @@ -2213,6 +2294,25 @@ impl LogicalPlan { } } } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let equality = on + .iter() + .map(|(left, right)| format!("{left} = {right}")) + .join(", "); + write!( + f, + "AsOf Join: match=[{match_condition}], constraint={join_constraint:?}" + )?; + if !equality.is_empty() { + write!(f, ", on=[{equality}]")?; + } + Ok(()) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. @@ -4343,6 +4443,169 @@ pub struct Join { pub null_aware: bool, } +/// The ordered comparison used by an [`AsOfJoin`]. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct AsOfMatch { + /// Expression evaluated against the left input. + pub left: Expr, + /// One of [`Operator::Lt`], [`Operator::LtEq`], [`Operator::Gt`], or + /// [`Operator::GtEq`]. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: Expr, +} + +impl AsOfMatch { + /// Creates an ordered ASOF match condition. + pub fn new(left: Expr, op: Operator, right: Expr) -> Self { + Self { left, op, right } + } +} + +impl Display for AsOfMatch { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{} {} {}", self.left, self.op, self.right) + } +} + +/// Match each left row with at most one ordered row from the right input. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AsOfJoin { + /// Left input. Every left row is preserved exactly once. + pub left: Arc, + /// Right input. + pub right: Arc, + /// Equality clauses expressed as pairs of left and right expressions. + pub on: Vec<(Expr, Expr)>, + /// Ordered match condition. + pub match_condition: Box, + /// Whether equality keys came from `ON` or `USING`. + pub join_constraint: JoinConstraint, + /// Output schema. + pub schema: DFSchemaRef, +} + +impl AsOfJoin { + /// Creates an ASOF join and validates its logical contract. + /// + /// This is the pre-coercion boundary. The physical ASOF constructor repeats + /// the shared operator, side-ownership, and determinism checks for direct + /// physical-plan callers and adds execution-only constraints. Keep the + /// shared checks aligned across both entry points. + pub fn try_new( + left: Arc, + right: Arc, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + join_constraint: JoinConstraint, + ) -> Result { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + + Self::validate_side(&match_condition.left, left.schema(), "left match")?; + Self::validate_side(&match_condition.right, right.schema(), "right match")?; + if match_condition.left.is_volatile() || match_condition.right.is_volatile() { + return plan_err!("ASOF MATCH_CONDITION must be deterministic"); + } + + let left_type = match_condition.left.get_type(left.schema())?; + let right_type = match_condition.right.get_type(right.schema())?; + if crate::type_coercion::binary::comparison_coercion(&left_type, &right_type) + .is_none() + { + return plan_err!( + "ASOF match expressions have incompatible types {left_type} and {right_type}" + ); + } + + for (left_expr, right_expr) in &on { + Self::validate_side(left_expr, left.schema(), "left equality")?; + Self::validate_side(right_expr, right.schema(), "right equality")?; + if left_expr.is_volatile() || right_expr.is_volatile() { + return plan_err!("ASOF equality expressions must be deterministic"); + } + let left_type = left_expr.get_type(left.schema())?; + let right_type = right_expr.get_type(right.schema())?; + let Some(common_type) = crate::type_coercion::binary::comparison_coercion( + &left_type, + &right_type, + ) else { + return plan_err!( + "ASOF equality expressions have incompatible types {left_type} and {right_type}" + ); + }; + if !crate::utils::can_hash(&common_type) { + return plan_err!( + "ASOF equality expressions have unsupported hash type {common_type}" + ); + } + } + + if join_constraint == JoinConstraint::Using + && on.iter().any(|(left, right)| { + left.get_as_join_column().is_none() + || right.get_as_join_column().is_none() + }) + { + return plan_err!("ASOF USING keys must be columns"); + } + + let schema = build_asof_join_schema(left.schema(), right.schema())?; + Ok(Self { + left, + right, + on, + match_condition: Box::new(match_condition), + join_constraint, + schema: Arc::new(schema), + }) + } + + fn validate_side(expr: &Expr, schema: &DFSchema, name: &str) -> Result<()> { + let mut columns = HashSet::new(); + expr_to_columns(expr, &mut columns)?; + if columns.is_empty() { + return plan_err!("ASOF {name} expression must reference its input"); + } + if let Some(column) = columns + .iter() + .find(|column| !schema.is_column_from_schema(column)) + { + return plan_err!( + "ASOF {name} expression references column {column} outside its input" + ); + } + Ok(()) + } +} + +impl PartialOrd for AsOfJoin { + fn partial_cmp(&self, other: &Self) -> Option { + ( + &self.left, + &self.right, + &self.on, + &self.match_condition, + &self.join_constraint, + ) + .partial_cmp(&( + &other.left, + &other.right, + &other.on, + &other.match_condition, + &other.join_constraint, + )) + .filter(|cmp| *cmp != Ordering::Equal || self == other) + } +} + impl Join { /// Creates a new Join operator with automatically computed schema. /// @@ -5915,6 +6178,15 @@ mod tests { .build()?; assert_eq!(cross_join.min_rows(), 2); + let asof_join = LogicalPlanBuilder::from(two_rows.clone()) + .asof_join( + one_row.clone(), + vec![], + AsOfMatch::new(col("l.column1"), Operator::GtEq, col("r.column1")), + )? + .build()?; + assert_eq!(asof_join.min_rows(), 2); + for (join_type, expected_min_rows) in [ // An inner join with a join condition may filter out every row, // while outer joins preserve the rows of the outer side(s). @@ -6838,6 +7110,32 @@ mod tests { Ok(()) } + #[test] + fn test_asof_using_preserves_qualified_keys() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("ts", DataType::Int64, false), + ]); + let left = Arc::new(table_scan(Some("t1"), &schema, None)?.build()?); + let right = Arc::new(table_scan(Some("t2"), &schema, None)?.build()?); + let join = AsOfJoin::try_new( + left, + right, + vec![(col("t1.id"), col("t2.id"))], + AsOfMatch::new(col("t1.ts"), Operator::GtEq, col("t2.ts")), + JoinConstraint::Using, + )?; + + assert_eq!(join.schema.fields().len(), 4); + assert_eq!( + join.schema + .index_of_column(&Column::from_qualified_name("t2.id"))?, + 2 + ); + assert!(join.schema.field(2).is_nullable()); + Ok(()) + } + #[test] fn test_join_try_new_schema_validation() -> Result<()> { let left_schema = Schema::new(vec![ diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c4c1d743b58b..ee43666736fe 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -41,11 +41,12 @@ use std::sync::Arc; use crate::logical_plan::plan::RangePartitioning; use crate::{ - Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, - DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, - LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort, - Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode, - Values, Window, WriteOp, builder::unnest_with_options, dml::CopyTo, + Aggregate, Analyze, AsOfJoin, AsOfMatch, CreateMemoryTable, CreateView, DdlStatement, + Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, + Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, + Sort, Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, + UserDefinedLogicalNode, Values, Window, WriteOp, builder::unnest_with_options, + dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -150,6 +151,23 @@ impl TreeNode for LogicalPlan { null_aware, }) }), + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) => (left, right).map_elements(f)?.update_data(|(left, right)| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) + }), LogicalPlan::Limit(Limit { skip, fetch, input }) => input .map_elements(f)? .update_data(|input| LogicalPlan::Limit(Limit { skip, fetch, input })), @@ -447,6 +465,13 @@ impl LogicalPlan { LogicalPlan::Join(Join { on, filter, .. }) => { (on, filter).apply_ref_elements(f) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + .. + }) => on.apply_elements(&mut f)?.visit_sibling(|| { + (&match_condition.left, &match_condition.right).apply_ref_elements(&mut f) + }), LogicalPlan::Sort(Sort { expr, .. }) => expr.apply_elements(f), LogicalPlan::Extension(extension) => { // would be nice to avoid this copy -- maybe can @@ -614,6 +639,29 @@ impl LogicalPlan { null_aware, }) }), + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + schema, + }) => (on, (match_condition.left, match_condition.right)) + .map_elements(f)? + .update_data(|(on, (left_match, right_match))| { + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition: Box::new(AsOfMatch { + left: left_match, + op: match_condition.op, + right: right_match, + }), + join_constraint, + schema, + }) + }), LogicalPlan::Sort(Sort { expr, input, fetch }) => expr .map_elements(f)? .update_data(|expr| LogicalPlan::Sort(Sort { expr, input, fetch })), diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 0f82f0b0df76..29d3017b338d 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,10 +57,10 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, - Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, - WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, - lit, not, + AsOfJoin, AsOfMatch, Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, + LogicalPlan, Operator, Projection, Union, ValueOrLambda, WindowFrame, + WindowFrameBound, WindowFrameUnits, WriteOp, is_false, is_not_false, is_not_true, + is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -191,6 +191,7 @@ impl<'a> TypeCoercionRewriter<'a> { pub fn coerce_plan(&mut self, plan: LogicalPlan) -> Result { match plan { LogicalPlan::Join(join) => self.coerce_join(join), + LogicalPlan::AsOfJoin(join) => self.coerce_asof_join(join), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), LogicalPlan::Dml(dml) => self.coerce_dml(dml), @@ -284,6 +285,36 @@ impl<'a> TypeCoercionRewriter<'a> { Ok(LogicalPlan::Join(join)) } + /// Coerce ASOF equality and ordered match expressions across input schemas. + pub fn coerce_asof_join(&mut self, mut join: AsOfJoin) -> Result { + join.on = join + .on + .into_iter() + .map(|(left, right)| { + self.coerce_binary_op( + left, + join.left.schema(), + Operator::Eq, + right, + join.right.schema(), + ) + }) + .collect::>()?; + let (left, right) = self.coerce_binary_op( + join.match_condition.left, + join.left.schema(), + join.match_condition.op, + join.match_condition.right, + join.right.schema(), + )?; + join.match_condition = Box::new(AsOfMatch { + left, + op: join.match_condition.op, + right, + }); + Ok(LogicalPlan::AsOfJoin(join)) + } + /// Coerce the union’s inputs to a common schema compatible with all inputs. /// This occurs after wildcard expansion and the coercion of the input expressions. pub fn coerce_union(union_plan: Union) -> Result { diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 94d61b74879c..6004ccb86e15 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -566,6 +566,7 @@ impl OptimizerRule for CommonSubexprEliminate { LogicalPlan::Window(window) => self.try_optimize_window(window, config)?, LogicalPlan::Aggregate(agg) => self.try_optimize_aggregate(agg, config)?, LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 413efd95588d..33fe1aac0b53 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -407,6 +407,26 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + let mut left_required = Vec::new(); + let mut right_required = Vec::new(); + for index in indices.indices() { + if *index < left_len { + left_required.push(*index); + } else { + right_required.push(*index - left_len); + } + } + let left_indices = RequiredIndices::new_from_indices(left_required) + .with_plan_exprs(&plan, join.left.schema())?; + let right_indices = RequiredIndices::new_from_indices(right_required) + .with_plan_exprs(&plan, join.right.schema())?; + vec![ + left_indices.with_projection_beneficial(), + right_indices.with_projection_beneficial(), + ] + } // these nodes are explicitly rewritten in the match statement above LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index 46c01180958c..ca4a6688b50c 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -411,6 +411,11 @@ fn map_children_mut Result>( let r = f(Arc::make_mut(right))?; l || r } + LogicalPlan::AsOfJoin(join) => { + let l = f(Arc::make_mut(&mut join.left))?; + let r = f(Arc::make_mut(&mut join.right))?; + l || r + } LogicalPlan::Union(Union { inputs, .. }) => { let mut changed = false; for input in inputs { diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 22625afca1f4..e0912f06e815 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -166,6 +166,10 @@ impl AsOfJoinExec { /// floating-point equality keys are not supported because Arrow sorting /// distinguishes signed zero while SQL equality does not. Projection indices /// refer to the full left-then-right join schema. + /// + /// The logical ASOF constructor validates the corresponding pre-coercion + /// contract. Keep the shared operator, side-ownership, and determinism checks + /// aligned across both public entry points. pub fn try_new( left: Arc, right: Arc, @@ -546,6 +550,156 @@ impl ExecutionPlan for AsOfJoinExec { column_statistics, })) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(&self.left)?; + let right = ctx.encode_child(&self.right)?; + let on = self + .on + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::>>()?; + let match_operator = match self.match_condition.op { + Operator::Lt => protobuf::AsOfMatchOperator::Lt, + Operator::LtEq => protobuf::AsOfMatchOperator::LtEq, + Operator::Gt => protobuf::AsOfMatchOperator::Gt, + Operator::GtEq => protobuf::AsOfMatchOperator::GtEq, + op => { + return internal_err!( + "AsOfJoinExec cannot serialize unsupported match operator {op}" + ); + } + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + left_match_expr: Some( + ctx.encode_expr(&self.match_condition.left)?, + ), + right_match_expr: Some( + ctx.encode_expr(&self.match_condition.right)?, + ), + match_operator: match_operator.into(), + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`; preserve the empty projection with + // the invalid column-index sentinel used by hash join. + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(projection) if projection.is_empty() => vec![u32::MAX], + Some(projection) => { + projection.iter().map(|index| *index as u32).collect() + } + }, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AsOfJoinExec { + /// Reconstruct an [`AsOfJoinExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let asof_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::AsOfJoin, + "AsOfJoinExec", + ); + let left = + ctx.decode_required_child(asof_join.left.as_deref(), "AsOfJoinExec", "left")?; + let right = ctx.decode_required_child( + asof_join.right.as_deref(), + "AsOfJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = asof_join + .on + .iter() + .map(|pair| { + let left = ctx.decode_required_expr( + pair.left.as_ref(), + left_schema.as_ref(), + "AsOfJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + pair.right.as_ref(), + right_schema.as_ref(), + "AsOfJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::>()?; + let left_match = ctx.decode_required_expr( + asof_join.left_match_expr.as_ref(), + left_schema.as_ref(), + "AsOfJoinExec", + "left_match_expr", + )?; + let right_match = ctx.decode_required_expr( + asof_join.right_match_expr.as_ref(), + right_schema.as_ref(), + "AsOfJoinExec", + "right_match_expr", + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from( + asof_join.match_operator, + ) + .map_err(|_| { + datafusion_common::internal_datafusion_err!( + "AsOfJoinExec: unknown AsOfMatchOperator {}", + asof_join.match_operator + ) + })?; + let op = match match_operator { + protobuf::AsOfMatchOperator::Lt => Operator::Lt, + protobuf::AsOfMatchOperator::LtEq => Operator::LtEq, + protobuf::AsOfMatchOperator::Gt => Operator::Gt, + protobuf::AsOfMatchOperator::GtEq => Operator::GtEq, + protobuf::AsOfMatchOperator::Unspecified => { + return internal_err!("AsOfJoinExec match operator must be specified"); + } + }; + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match asof_join.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|index| *index as usize).collect()), + }; + + Ok(Arc::new(Self::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, op, right_match), + projection, + )?)) + } } /// Materialized right input shared by every left output partition. diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 4b63631613ae..7860757386c3 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -63,6 +63,7 @@ message LogicalPlanNode { CteWorkTableScanNode cte_work_table_scan = 32; DmlNode dml = 33; EmptyTableScanNode empty_table_scan = 34; + AsOfJoinNode as_of_join = 35; } } @@ -276,6 +277,25 @@ message JoinNode { bool null_aware = 9; } +enum AsOfMatchOperator { + AS_OF_MATCH_OPERATOR_UNSPECIFIED = 0; + AS_OF_MATCH_OPERATOR_LT = 1; + AS_OF_MATCH_OPERATOR_LT_EQ = 2; + AS_OF_MATCH_OPERATOR_GT = 3; + AS_OF_MATCH_OPERATOR_GT_EQ = 4; +} + +message AsOfJoinNode { + LogicalPlanNode left = 1; + LogicalPlanNode right = 2; + repeated LogicalExprNode left_join_key = 3; + repeated LogicalExprNode right_join_key = 4; + LogicalExprNode left_match_expr = 5; + LogicalExprNode right_match_expr = 6; + AsOfMatchOperator match_operator = 7; + datafusion_common.JoinConstraint join_constraint = 8; +} + message DistinctNode { LogicalPlanNode input = 1; } @@ -900,6 +920,7 @@ message PhysicalPlanNode { ArrowScanExecNode arrow_scan = 38; ScalarSubqueryExecNode scalar_subquery = 39; PiecewiseMergeJoinExecNode piecewise_merge_join = 40; + AsOfJoinExecNode as_of_join = 41; } } @@ -1714,6 +1735,16 @@ message PiecewiseMergeJoinExecNode { uint64 num_partitions = 7; } +message AsOfJoinExecNode { + PhysicalPlanNode left = 1; + PhysicalPlanNode right = 2; + repeated JoinOn on = 3; + PhysicalExprNode left_match_expr = 4; + PhysicalExprNode right_match_expr = 5; + AsOfMatchOperator match_operator = 6; + repeated uint32 projection = 7; +} + message AsyncFuncExecNode { PhysicalPlanNode input = 1; repeated PhysicalExprNode async_exprs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 9811357f1dd5..2f4c2c20fe68 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1627,6 +1627,507 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode { deserializer.deserialize_struct("datafusion.ArrowScanExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for AsOfJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.on.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if !self.projection.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.on.is_empty() { + struct_ser.serialize_field("on", &self.on)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if !self.projection.is_empty() { + struct_ser.serialize_field("projection", &self.projection)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "on", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "projection", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + On, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + Projection, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "on" => Ok(GeneratedField::On), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "projection" => Ok(GeneratedField::Projection), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut on__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut projection__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::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()) + ; + } + } + } + Ok(AsOfJoinExecNode { + left: left__, + right: right__, + on: on__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + projection: projection__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfJoinNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.left_join_key.is_empty() { + len += 1; + } + if !self.right_join_key.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if self.join_constraint != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.left_join_key.is_empty() { + struct_ser.serialize_field("leftJoinKey", &self.left_join_key)?; + } + if !self.right_join_key.is_empty() { + struct_ser.serialize_field("rightJoinKey", &self.right_join_key)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if self.join_constraint != 0 { + let v = super::datafusion_common::JoinConstraint::try_from(self.join_constraint) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_constraint)))?; + struct_ser.serialize_field("joinConstraint", &v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "left_join_key", + "leftJoinKey", + "right_join_key", + "rightJoinKey", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "join_constraint", + "joinConstraint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + LeftJoinKey, + RightJoinKey, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + JoinConstraint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "leftJoinKey" | "left_join_key" => Ok(GeneratedField::LeftJoinKey), + "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "joinConstraint" | "join_constraint" => Ok(GeneratedField::JoinConstraint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut left_join_key__ = None; + let mut right_join_key__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut join_constraint__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::LeftJoinKey => { + if left_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("leftJoinKey")); + } + left_join_key__ = Some(map_.next_value()?); + } + GeneratedField::RightJoinKey => { + if right_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("rightJoinKey")); + } + right_join_key__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::JoinConstraint => { + if join_constraint__.is_some() { + return Err(serde::de::Error::duplicate_field("joinConstraint")); + } + join_constraint__ = Some(map_.next_value::()? as i32); + } + } + } + Ok(AsOfJoinNode { + left: left__, + right: right__, + left_join_key: left_join_key__.unwrap_or_default(), + right_join_key: right_join_key__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + join_constraint: join_constraint__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfMatchOperator { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for AsOfMatchOperator { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + "AS_OF_MATCH_OPERATOR_LT", + "AS_OF_MATCH_OPERATOR_LT_EQ", + "AS_OF_MATCH_OPERATOR_GT", + "AS_OF_MATCH_OPERATOR_GT_EQ", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = AsOfMatchOperator; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Ok(AsOfMatchOperator::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Ok(AsOfMatchOperator::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Ok(AsOfMatchOperator::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Ok(AsOfMatchOperator::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Ok(AsOfMatchOperator::GtEq), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for AsyncFuncExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -13879,6 +14380,9 @@ impl serde::Serialize for LogicalPlanNode { logical_plan_node::LogicalPlanType::EmptyTableScan(v) => { struct_ser.serialize_field("emptyTableScan", v)?; } + logical_plan_node::LogicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -13940,6 +14444,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "dml", "empty_table_scan", "emptyTableScan", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -13977,6 +14483,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { CteWorkTableScan, Dml, EmptyTableScan, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -14031,6 +14538,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "cteWorkTableScan" | "cte_work_table_scan" => Ok(GeneratedField::CteWorkTableScan), "dml" => Ok(GeneratedField::Dml), "emptyTableScan" | "empty_table_scan" => Ok(GeneratedField::EmptyTableScan), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14282,6 +14790,13 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { return Err(serde::de::Error::duplicate_field("emptyTableScan")); } logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::EmptyTableScan) +; + } + GeneratedField::AsOfJoin => { + if logical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::AsOfJoin) ; } } @@ -20585,6 +21100,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin(v) => { struct_ser.serialize_field("piecewiseMergeJoin", v)?; } + physical_plan_node::PhysicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -20659,6 +21177,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "scalarSubquery", "piecewise_merge_join", "piecewiseMergeJoin", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -20702,6 +21222,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { ArrowScan, ScalarSubquery, PiecewiseMergeJoin, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20762,6 +21283,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), "piecewiseMergeJoin" | "piecewise_merge_join" => Ok(GeneratedField::PiecewiseMergeJoin), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -21055,6 +21577,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("piecewiseMergeJoin")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin) +; + } + GeneratedField::AsOfJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::AsOfJoin) ; } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index a20632860a4c..ca1ab5542a6a 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -5,7 +5,7 @@ pub struct LogicalPlanNode { #[prost( oneof = "logical_plan_node::LogicalPlanType", - tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" + tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35" )] pub logical_plan_type: ::core::option::Option, } @@ -79,6 +79,8 @@ pub mod logical_plan_node { Dml(::prost::alloc::boxed::Box), #[prost(message, tag = "34")] EmptyTableScan(super::EmptyTableScanNode), + #[prost(message, tag = "35")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -421,6 +423,29 @@ pub struct JoinNode { pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub left_join_key: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "4")] + pub right_join_key: ::prost::alloc::vec::Vec, + #[prost(message, optional, boxed, tag = "5")] + pub left_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(message, optional, boxed, tag = "6")] + pub right_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(enumeration = "AsOfMatchOperator", tag = "7")] + pub match_operator: i32, + #[prost(enumeration = "super::datafusion_common::JoinConstraint", tag = "8")] + pub join_constraint: i32, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1342,7 +1367,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41" )] pub physical_plan_type: ::core::option::Option, } @@ -1432,6 +1457,8 @@ pub mod physical_plan_node { PiecewiseMergeJoin( ::prost::alloc::boxed::Box, ), + #[prost(message, tag = "41")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2600,6 +2627,23 @@ pub struct PiecewiseMergeJoinExecNode { pub num_partitions: u64, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub on: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "4")] + pub left_match_expr: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub right_match_expr: ::core::option::Option, + #[prost(enumeration = "AsOfMatchOperator", tag = "6")] + pub match_operator: i32, + #[prost(uint32, repeated, tag = "7")] + pub projection: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -2631,6 +2675,41 @@ pub struct PhysicalScalarSubqueryExprNode { #[prost(uint32, tag = "3")] pub index: u32, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum AsOfMatchOperator { + Unspecified = 0, + Lt = 1, + LtEq = 2, + Gt = 3, + GtEq = 4, +} +impl AsOfMatchOperator { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Some(Self::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Some(Self::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Some(Self::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Some(Self::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Some(Self::GtEq), + _ => None, + } + } +} /// Identifies a built-in file format supported by DataFusion. /// Used by DefaultLogicalExtensionCodec to serialize/deserialize /// FileFormatFactory instances (e.g. in CopyTo plans). diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 73cfc6e710d1..23607bbe91cb 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -59,17 +59,17 @@ use datafusion_datasource_json::file_format::{ use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, Operator, RangePartitioning, RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, }; use datafusion_expr::{ DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ - Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection, - Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window, - builder::project, + Aggregate, AsOfJoin, AsOfMatch, CreateCatalog, CreateCatalogSchema, + CreateExternalTable, CreateView, DdlStatement, Distinct, EmptyRelation, + Extension, Join, Prepare, Projection, Repartition, Sort, SubqueryAlias, + TableScan, TableScanBuilder, Values, Window, builder::project, }, }; use datafusion_proto_common::protobuf_common; @@ -1061,6 +1061,69 @@ impl AsLogicalPlan for LogicalPlanNode { join.null_aware, )?)) } + LogicalPlanType::AsOfJoin(join) => { + let left_keys = + from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; + let right_keys = + from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + if left_keys.len() != right_keys.len() { + return Err(proto_error(format!( + "Received an AsOfJoinNode with left_join_key and right_join_key of different lengths: {} and {}", + left_keys.len(), + right_keys.len() + ))); + } + let left_match = from_proto::parse_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode left_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let right_match = from_proto::parse_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode right_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from( + join.match_operator, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = match match_operator { + protobuf::AsOfMatchOperator::Lt => Operator::Lt, + protobuf::AsOfMatchOperator::LtEq => Operator::LtEq, + protobuf::AsOfMatchOperator::Gt => Operator::Gt, + protobuf::AsOfMatchOperator::GtEq => Operator::GtEq, + protobuf::AsOfMatchOperator::Unspecified => { + return Err(proto_error("ASOF match operator must be specified")); + } + }; + let join_constraint = protobuf::JoinConstraint::try_from( + join.join_constraint, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF JoinConstraint {}", + join.join_constraint + )) + })?; + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + left_keys.into_iter().zip(right_keys).collect(), + AsOfMatch::new(left_match, op, right_match), + JoinConstraint::from(join_constraint), + )?)) + } LogicalPlanType::Union(union) => { assert_or_internal_err!( union.inputs.len() >= 2, @@ -1717,6 +1780,68 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + .. + }) => { + let left = LogicalPlanNode::try_from_logical_plan( + left.as_ref(), + extension_codec, + )?; + let right = LogicalPlanNode::try_from_logical_plan( + right.as_ref(), + extension_codec, + )?; + let (left_join_key, right_join_key) = on + .iter() + .map(|(left, right)| { + Ok(( + serialize_expr(left, extension_codec)?, + serialize_expr(right, extension_codec)?, + )) + }) + .collect::, ToProtoError>>()? + .into_iter() + .unzip(); + let match_operator = match match_condition.op { + Operator::Lt => protobuf::AsOfMatchOperator::Lt, + Operator::LtEq => protobuf::AsOfMatchOperator::LtEq, + Operator::Gt => protobuf::AsOfMatchOperator::Gt, + Operator::GtEq => protobuf::AsOfMatchOperator::GtEq, + op => { + return Err(proto_error(format!( + "Unsupported ASOF match operator {op}" + ))); + } + }; + Ok(LogicalPlanNode { + logical_plan_type: Some(LogicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + left_join_key, + right_join_key, + left_match_expr: Some(Box::new(serialize_expr( + &match_condition.left, + extension_codec, + )?)), + right_match_expr: Some(Box::new(serialize_expr( + &match_condition.right, + extension_codec, + )?)), + match_operator: match_operator.into(), + join_constraint: protobuf::JoinConstraint::from( + *join_constraint, + ) + .into(), + }, + ))), + }) + } LogicalPlan::Subquery(subquery) => { // Serialize the inner subquery plan directly — the // LogicalPlan::Subquery wrapper is reconstructed during diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 26d1a8ed83d3..d1ddcccdac2e 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -62,8 +62,8 @@ use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::explain::ExplainExec; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PiecewiseMergeJoinExec, - SortMergeJoinExec, SymmetricHashJoinExec, + AsOfJoinExec, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PiecewiseMergeJoinExec, SortMergeJoinExec, SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -1231,6 +1231,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::SortMergeJoin(_) => { SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } + PhysicalPlanType::AsOfJoin(_) => { + AsOfJoinExec::try_from_proto(self.node(), &decode_ctx) + } PhysicalPlanType::AsyncFunc(_) => { AsyncFuncExec::try_from_proto(self.node(), &decode_ctx) } diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 84b7b46b0f17..7dadeace010b 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -29,8 +29,9 @@ use datafusion::physical_plan::expressions::{ }; use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, PiecewiseMergeJoinExec, - SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, HashJoinExec, NestedLoopJoinExec, PartitionMode, + PiecewiseMergeJoinExec, SortMergeJoinExec, StreamJoinPartitionMode, + SymmetricHashJoinExec, }; use datafusion::prelude::SessionContext; use datafusion_common::ScalarValue; @@ -89,6 +90,41 @@ fn roundtrip_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_asof_join() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let on = vec![( + Arc::new(Column::new("symbol", 0)) as _, + Arc::new(Column::new("symbol", 0)) as _, + )]; + + for projection in [None, Some(vec![]), Some(vec![0, 5])] { + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + roundtrip_test(Arc::new(AsOfJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&left_schema))), + Arc::new(EmptyExec::new(Arc::clone(&right_schema))), + on.clone(), + AsOfMatchExpr::new( + Arc::new(Column::new("ts", 1)), + op, + Arc::new(Column::new("ts", 1)), + ), + projection.clone(), + )?))?; + } + } + Ok(()) +} + #[test] fn roundtrip_nested_loop_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 7c916db9c3cd..12f099cc2b1c 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -74,8 +74,9 @@ use datafusion_common::format::{ }; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, - TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, + Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, + SplitPoint, TableReference, internal_datafusion_err, internal_err, not_impl_err, + plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -90,7 +91,7 @@ use datafusion_expr::logical_plan::{ ExplainOption, Extension, UserDefinedLogicalNodeCore, }; use datafusion_expr::{ - Accumulator, AggregateUDF, ColumnarValue, DmlStatement, ExprFunctionExt, + Accumulator, AggregateUDF, AsOfMatch, ColumnarValue, DmlStatement, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, @@ -3827,6 +3828,39 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_asof_join() -> Result<()> { + let ctx = SessionContext::new(); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + ctx.register_table("trades", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("prices", Arc::new(EmptyTable::new(right_schema)))?; + + let left = ctx.table("trades").await?.into_optimized_plan()?; + let right = ctx.table("prices").await?.into_optimized_plan()?; + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + let plan = LogicalPlanBuilder::from(left.clone()) + .asof_join_using( + right.clone(), + vec![Column::from_name("symbol")], + AsOfMatch::new(col("trades.ts"), op, col("prices.ts")), + )? + .build()?; + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{round_trip:?}")); + } + Ok(()) +} + // Single column, single split point range partitioning #[tokio::test] async fn roundtrip_range_partitioning_single_col() -> Result<()> { diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 9922509a0e60..155fd12fa797 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -198,6 +198,7 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs index c3599a2635ff..15f59919a2a9 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs @@ -51,6 +51,9 @@ pub fn to_substrait_rel( LogicalPlan::Aggregate(plan) => producer.handle_aggregate(plan), LogicalPlan::Sort(plan) => producer.handle_sort(plan), LogicalPlan::Join(plan) => producer.handle_join(plan), + LogicalPlan::AsOfJoin(plan) => { + not_impl_err!("Substrait ASOF join is not supported: {plan:?}")? + } LogicalPlan::Repartition(plan) => producer.handle_repartition(plan), LogicalPlan::Union(plan) => producer.handle_union(plan), LogicalPlan::TableScan(plan) => producer.handle_table_scan(plan), diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 1981ef66db37..75e1fe251ac1 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -18,7 +18,7 @@ #[cfg(test)] mod tests { use datafusion::datasource::provider_as_source; - use datafusion::logical_expr::LogicalPlanBuilder; + use datafusion::logical_expr::{AsOfMatch, LogicalPlanBuilder, Operator}; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::serializer; @@ -27,7 +27,7 @@ mod tests { use datafusion::prelude::*; use insta::assert_snapshot; - use std::fs; + use std::{fs, sync::Arc}; use substrait::proto::expression::field_reference::{ReferenceType, RootType}; use substrait::proto::expression::reference_segment; use substrait::proto::expression::{ReferenceSegment, RexType}; @@ -103,6 +103,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn asof_join_fails_closed_until_substrait_has_an_extension() -> Result<()> { + let ctx = create_context().await?; + let table = provider_as_source(ctx.table_provider("data").await?); + let left = LogicalPlanBuilder::scan("l", Arc::clone(&table), None)?.build()?; + let right = LogicalPlanBuilder::scan("r", table, None)?.build()?; + let plan = LogicalPlanBuilder::from(left) + .asof_join( + right, + vec![(col("l.b"), col("r.b"))], + AsOfMatch::new(col("l.a"), Operator::GtEq, col("r.a")), + )? + .build()?; + let error = to_substrait_plan(&plan, &ctx.state()) + .expect_err("ASOF must not be lowered to a generic Substrait join"); + assert!( + error + .to_string() + .contains("Substrait ASOF join is not supported") + ); + Ok(()) + } + #[tokio::test] async fn include_remaps_for_projects() -> Result<()> { let ctx = create_context().await?;