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/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index bb526895b6b1..dd32830dd5eb 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -17,10 +17,12 @@ use insta::assert_snapshot; -use datafusion::assert_batches_eq; use datafusion::catalog::MemTable; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; +use datafusion::physical_plan::joins::AsOfJoinExec; +use datafusion::physical_plan::{Distribution, ExecutionPlanProperties}; use datafusion::test_util::register_unbounded_file_with_ordering; +use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion_sql::unparser::plan_to_sql; use super::*; @@ -297,3 +299,412 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } + +fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { + let trades_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("trade_id", DataType::Int32, false), + ])); + let trades = vec![ + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(7), Some(2), Some(3)])), + Arc::new(Int32Array::from(vec![3, 4, 6])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&trades_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(1), Some(4), Some(8)])), + Arc::new(Int32Array::from(vec![1, 2, 5])), + ], + )?, + ]; + ctx.register_table( + "trades", + Arc::new(MemTable::try_new( + trades_schema, + trades.into_iter().map(|batch| vec![batch]).collect(), + )?), + )?; + + let prices_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let prices = vec![ + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])), + Arc::new(Int64Array::from(vec![Some(6), Some(1), Some(2)])), + Arc::new(Int32Array::from(vec![60, 101, 999])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&prices_schema), + vec![ + Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])), + Arc::new(Int64Array::from(vec![Some(2), Some(4), Some(6)])), + Arc::new(Int32Array::from(vec![20, 40, 106])), + ], + )?, + ]; + ctx.register_table( + "prices", + Arc::new(MemTable::try_new( + prices_schema, + prices.into_iter().map(|batch| vec![batch]).collect(), + )?), + )?; + Ok(()) +} + +fn find_asof_exec(plan: &Arc) -> Option> { + if plan.downcast_ref::().is_some() { + return Some(Arc::clone(plan)); + } + plan.children().into_iter().find_map(find_asof_exec) +} + +#[tokio::test] +async fn asof_join_all_match_directions_across_batches() -> Result<()> { + let config = SessionConfig::new() + .with_batch_size(2) + .with_target_partitions(2); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx)?; + + for (op, expected) in [ + ( + ">=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + ">", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 20 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<=", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 40 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ( + "<", + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | 20 |", + "| 2 | 60 |", + "| 3 | |", + "| 4 | 106 |", + "| 5 | |", + "| 6 | |", + "+----------+-------+", + ], + ), + ] { + let batches = ctx + .sql(&format!( + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts {op} p.ts) \ + ON t.symbol = p.symbol ORDER BY t.trade_id" + )) + .await? + .collect() + .await?; + assert_batches_eq!(expected, &batches); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_coerces_equality_and_match_types() -> Result<()> { + let ctx = SessionContext::new(); + let batches = ctx + .sql( + "SELECT t.id, p.price \ + FROM (VALUES (CAST(1 AS INT), CAST(4 AS INT), 7)) t(k, ts, id) \ + ASOF JOIN \ + (VALUES (CAST(1 AS BIGINT), CAST(2 AS BIGINT), 20)) p(k, ts, price) \ + MATCH_CONDITION (t.ts >= p.ts) ON t.k = p.k", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----+-------+", + "| id | price |", + "+----+-------+", + "| 7 | 20 |", + "+----+-------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx)?; + let df = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN \ + (SELECT ts, price FROM prices WHERE symbol = 'A') p \ + MATCH_CONDITION (t.ts >= p.ts)", + ) + .await?; + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert_contains!(sql.as_str(), "ASOF JOIN"); + assert!(!sql.contains(" ON "), "unexpected equality clause: {sql}"); + ctx.sql(&sql).await?; + let plan = df.create_physical_plan().await?; + let asof = find_asof_exec(&plan).expect("physical ASOF join must be present"); + let output_partitions = asof.output_partitioning().partition_count(); + assert_eq!( + output_partitions, + asof.children()[0].output_partitioning().partition_count() + ); + assert!( + output_partitions > 1, + "ASOF join did not preserve left-side parallelism" + ); + assert_eq!( + asof.children()[1].output_partitioning().partition_count(), + 1 + ); + let right_plan = displayable(asof.children()[1].as_ref()) + .indent(true) + .to_string(); + assert_contains!(right_plan.as_str(), "SortPreservingMergeExec"); + assert_contains!(right_plan.as_str(), "DataSourceExec: partitions=2"); + assert!(asof.output_ordering().is_some()); + assert!(matches!( + &asof.input_distribution_requirements().into_per_child()[..], + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] + )); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_batches_sorted_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 20 |", + "| 5 | 60 |", + "| 6 | 20 |", + "+----------+-------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + let batches = ctx + .sql( + "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON t.symbol = p.symbol", + ) + .await? + .collect() + .await?; + let explain = arrow::util::pretty::pretty_format_batches(&batches)?.to_string(); + assert_contains!(explain.as_str(), "AsOf Join: match=[t.ts >= p.ts]"); + assert_contains!(explain.as_str(), "on=[t.symbol = p.symbol]"); + assert_contains!(explain.as_str(), "AsOfJoinExec:"); + assert_contains!(explain.as_str(), "on=[(symbol = symbol)]"); + assert_contains!(explain.as_str(), "match=[ts >= ts]"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result<()> { + let ctx = SessionContext::new(); + let tmp_dir = TempDir::new()?; + let schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::UInt32, false), + Field::new("ts", DataType::UInt32, false), + ])); + let ordering = vec![vec![ + col("symbol").sort(true, true), + col("ts").sort(true, true), + ]]; + for table in ["left_stream", "right_stream"] { + let path = tmp_dir.path().join(format!("{table}.csv")); + File::create(&path)?; + register_unbounded_file_with_ordering( + &ctx, + Arc::clone(&schema), + &path, + table, + ordering.clone(), + )?; + } + let error = ctx + .sql( + "SELECT * FROM left_stream l ASOF JOIN right_stream r \ + MATCH_CONDITION (l.ts >= r.ts) ON l.symbol = r.symbol", + ) + .await? + .create_physical_plan() + .await + .expect_err("ASOF physical planning must reject unbounded inputs"); + assert_contains!(error.to_string(), "AsOfJoinExec requires bounded inputs"); + Ok(()) +} + +#[tokio::test] +async fn asof_join_using_preserves_key_access_and_unparser_round_trips() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + let df = ctx + .sql( + "SELECT * FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol)", + ) + .await?; + assert_eq!( + df.schema() + .fields() + .iter() + .map(|field| field.name()) + .collect::>(), + vec!["ts", "trade_id", "symbol", "ts", "price"] + ); + let sql = plan_to_sql(df.logical_plan())?.to_string(); + assert!(sql.contains("ASOF JOIN")); + assert!(sql.contains("MATCH_CONDITION")); + assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}"); + ctx.sql(&sql).await?; + + let batches = ctx + .sql( + "SELECT t.trade_id, t.symbol AS left_symbol, p.symbol AS right_symbol \ + FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + ) + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+----------+-------------+--------------+", + "| trade_id | left_symbol | right_symbol |", + "+----------+-------------+--------------+", + "| 1 | A | |", + "| 2 | A | A |", + "| 3 | A | A |", + "| 4 | B | B |", + "| 5 | B | B |", + "| 6 | | |", + "+----------+-------------+--------------+", + ], + &batches + ); + Ok(()) +} + +#[tokio::test] +async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + for query in [ + "SELECT t.trade_id, p.price FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + "SELECT * FROM trades t \ + ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \ + MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ + ORDER BY t.trade_id", + "SELECT t.trade_id, p.price FROM trades t \ + JOIN prices q ON t.symbol = q.symbol AND t.ts = q.ts \ + ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ + ON q.symbol = p.symbol ORDER BY t.trade_id", + "SELECT t.trade_id, q.trade_id FROM trades t \ + ASOF JOIN (prices p JOIN trades q \ + ON p.symbol = q.symbol AND p.ts = q.ts) \ + MATCH_CONDITION (t.ts >= q.ts) ON t.symbol = p.symbol \ + ORDER BY t.trade_id", + ] { + let expected = ctx.sql(query).await?.collect().await?; + let plan = ctx.sql(query).await?.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + let actual = ctx.sql(&sql).await?.collect().await?; + assert_eq!( + datafusion_common::test_util::batches_to_string(&expected), + datafusion_common::test_util::batches_to_string(&actual), + "unparsed SQL changed ASOF candidate preselection: {sql}" + ); + } + Ok(()) +} + +#[tokio::test] +async fn asof_join_rejects_invalid_contracts() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx)?; + for sql in [ + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts = p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (p.ts >= t.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON t.symbol > p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (1 >= p.ts) ON t.symbol = p.symbol", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON 1 = 1", + "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts AND t.ts > p.ts) ON t.symbol = p.symbol", + ] { + assert!(ctx.sql(sql).await.is_err(), "query should fail: {sql}"); + } + Ok(()) +} diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 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..d7570c080cf7 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -166,6 +166,10 @@ impl AsOfJoinExec { /// floating-point equality keys are not supported because Arrow sorting /// distinguishes signed zero while SQL equality does not. Projection indices /// refer to the full left-then-right join schema. + /// + /// The logical ASOF constructor validates the corresponding pre-coercion + /// contract. Keep the shared operator, side-ownership, and determinism checks + /// aligned across both public entry points. pub fn try_new( left: Arc, right: Arc, diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 73cfc6e710d1..8e49b254d713 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -2224,6 +2224,9 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlan::DescribeTable(_) => Err(proto_error( "LogicalPlan serde is not yet implemented for DescribeTable", )), + LogicalPlan::AsOfJoin(_) => Err(proto_error( + "LogicalPlan serde is not yet implemented for AsOfJoin", + )), LogicalPlan::RecursiveQuery(recursive) => { let static_term = LogicalPlanNode::try_from_logical_plan( recursive.static_term.as_ref(), diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b3809..70c9572598d1 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,8 +16,13 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err}; -use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datafusion_common::{ + Column, DFSchema, Result, not_impl_err, plan_datafusion_err, plan_err, +}; +use datafusion_expr::utils::split_conjunction_owned; +use datafusion_expr::{ + AsOfMatch, BinaryExpr, Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, +}; use sqlparser::ast::{ Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, }; @@ -98,10 +103,129 @@ impl SqlToRel<'_, S> { JoinOperator::CrossJoin(JoinConstraint::None) => { self.parse_cross_join(left, right) } + JoinOperator::AsOf { + match_condition, + constraint, + } => self.parse_asof_join( + left, + right, + match_condition, + constraint, + planner_context, + ), other => not_impl_err!("Unsupported JOIN operator {other:?}"), } } + fn parse_asof_join( + &self, + left: LogicalPlan, + right: LogicalPlan, + sql_match_condition: sqlparser::ast::Expr, + constraint: JoinConstraint, + planner_context: &mut PlannerContext, + ) -> Result { + let join_schema = left.schema().join(right.schema())?; + let match_condition = + self.sql_to_expr(sql_match_condition, &join_schema, planner_context)?; + let Expr::BinaryExpr(BinaryExpr { + left: match_left, + op, + right: match_right, + }) = match_condition + else { + return plan_err!("ASOF MATCH_CONDITION must be a single comparison"); + }; + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {op}" + ); + } + if !expr_owned_by(&match_left, left.schema()) + || !expr_owned_by(&match_right, right.schema()) + { + return plan_err!( + "ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input" + ); + } + let match_condition = AsOfMatch::new(*match_left, op, *match_right); + + match constraint { + JoinConstraint::On(sql_on) => { + let on = self.sql_to_expr(sql_on, &join_schema, planner_context)?; + let on = split_conjunction_owned(on) + .into_iter() + .map(|predicate| { + let Expr::BinaryExpr(BinaryExpr { + left: on_left, + op: Operator::Eq, + right: on_right, + }) = predicate + else { + return plan_err!( + "ASOF ON accepts only equality conditions combined with AND" + ); + }; + if expr_owned_by(&on_left, left.schema()) + && expr_owned_by(&on_right, right.schema()) + { + Ok((*on_left, *on_right)) + } else if expr_owned_by(&on_right, left.schema()) + && expr_owned_by(&on_left, right.schema()) + { + Ok((*on_right, *on_left)) + } else { + plan_err!( + "Each ASOF equality condition must compare one left expression with one right expression" + ) + } + }) + .collect::>()?; + LogicalPlanBuilder::from(left) + .asof_join(right, on, match_condition)? + .build() + } + JoinConstraint::Using(object_names) => { + let keys = object_names + .into_iter() + .map(|object_name| { + let ObjectName(mut object_names) = object_name; + if object_names.len() != 1 { + return not_impl_err!( + "Invalid identifier in ASOF USING clause. Expected single identifier, got {}", + ObjectName(object_names) + ); + } + let id = object_names.swap_remove(0); + id.as_ident() + .ok_or_else(|| { + plan_datafusion_err!( + "Expected identifier in ASOF USING clause" + ) + }) + .map(|ident| { + Column::from_name( + self.ident_normalizer.normalize(ident.clone()), + ) + }) + }) + .collect::>>()?; + LogicalPlanBuilder::from(left) + .asof_join_using(right, keys, match_condition)? + .build() + } + JoinConstraint::None => LogicalPlanBuilder::from(left) + .asof_join(right, vec![], match_condition)? + .build(), + JoinConstraint::Natural => { + not_impl_err!("NATURAL ASOF JOIN is not supported") + } + } + } + fn parse_cross_join( &self, left: LogicalPlan, @@ -180,6 +304,14 @@ impl SqlToRel<'_, S> { } } +fn expr_owned_by(expr: &Expr, schema: &DFSchema) -> bool { + let columns = expr.column_refs(); + !columns.is_empty() + && columns + .iter() + .all(|column| schema.is_column_from_schema(column)) +} + /// Returns `true` if the given [`TableFactor`] is lateral. pub(crate) fn is_lateral(factor: &TableFactor) -> bool { match factor { diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 9922509a0e60..9e2d723d8a45 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -49,7 +49,7 @@ use datafusion_common::{ }; use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX}; use datafusion_expr::{ - Aggregate, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, + Aggregate, AsOfJoin, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, Projection, SkipType, Sort, SortExpr, TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; @@ -178,6 +178,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -1339,11 +1340,8 @@ impl Unparser<'_> { let mut right_relation = RelationBuilder::default(); if already_projected - && let Some(nested_relation) = self - .qualified_passthrough_join_projection_to_nested_relation( - right_plan.as_ref(), - query, - )? + && let Some(nested_relation) = + self.join_input_to_nested_relation(right_plan.as_ref(), query)? { right_relation = nested_relation; } else { @@ -1477,6 +1475,9 @@ impl Unparser<'_> { Ok(()) } + LogicalPlan::AsOfJoin(join) => { + self.asof_join_to_sql(join, query, select, relation) + } LogicalPlan::SubqueryAlias(plan_alias) => { let (plan, mut columns) = subquery_alias_inner_query_and_columns(plan_alias); @@ -1774,6 +1775,126 @@ impl Unparser<'_> { } } + // Keep ASOF-specific locals out of the recursive plan unparser's stack frame. + #[inline(never)] + fn asof_join_to_sql( + &self, + join: &AsOfJoin, + query: &mut Option, + select: &mut SelectBuilder, + relation: &mut RelationBuilder, + ) -> Result<()> { + let already_projected = select.already_projected(); + let left_plan = + Self::unwrap_qualified_passthrough_join_projection(Arc::clone(&join.left)); + let inline_left_join = matches!(left_plan.as_ref(), LogicalPlan::Join(_)); + let left_projection = if already_projected { + None + } else if inline_left_join { + self.select_to_sql_recursively(left_plan.as_ref(), query, select, relation)?; + select.pop_projections(); + Some(self.derived_input_projection(join.left.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + let qualifier = self.derive_asof_input(join.left.as_ref(), relation)?; + Some(self.derived_input_projection(join.left.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively(join.left.as_ref(), query, select, relation)?; + Some(select.pop_projections()) + }; + if already_projected { + if inline_left_join { + self.select_to_sql_recursively( + left_plan.as_ref(), + query, + select, + relation, + )?; + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + self.derive_asof_input(join.left.as_ref(), relation)?; + } else { + self.select_to_sql_recursively( + join.left.as_ref(), + query, + select, + relation, + )?; + } + } + + let mut right_relation = RelationBuilder::default(); + let nested_right = + self.join_input_to_nested_relation(join.right.as_ref(), query)?; + let right_projection = if already_projected { + if let Some(nested_right) = nested_right { + right_relation = nested_right; + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + } + None + } else if let Some(nested_right) = nested_right { + right_relation = nested_right; + Some(self.derived_input_projection(join.right.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + let qualifier = + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + Some(self.derived_input_projection(join.right.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + Some(select.pop_projections()) + }; + let Ok(Some(relation)) = right_relation.build() else { + return internal_err!("Failed to build ASOF right relation"); + }; + let constraint = + self.join_constraint_to_sql(join.join_constraint, &join.on, None)?; + let match_condition = self.expr_to_sql(&Expr::BinaryExpr(BinaryExpr::new( + Box::new(join.match_condition.left.clone()), + join.match_condition.op, + Box::new(join.match_condition.right.clone()), + )))?; + let ast_join = ast::Join { + relation, + global: false, + join_operator: ast::JoinOperator::AsOf { + match_condition, + constraint, + }, + }; + let mut from = select + .pop_from() + .ok_or_else(|| internal_datafusion_err!("ASOF left relation is missing"))?; + from.push_join(ast_join); + select.push_from(from); + + if !already_projected { + let left_projection = left_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF left projection is missing") + })?; + let right_projection = right_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF right projection is missing") + })?; + select.projection( + left_projection + .into_iter() + .chain(right_projection) + .collect(), + ); + } + Ok(()) + } + /// Walk through transparent nodes (SubqueryAlias) to find the inner /// Projection that feeds an Unnest node. /// @@ -2072,6 +2193,74 @@ impl Unparser<'_> { ) } + fn asof_input_requires_derived(plan: &LogicalPlan) -> bool { + let simple_scan = + |scan: &TableScan| scan.filters.is_empty() && scan.fetch.is_none(); + match plan { + LogicalPlan::TableScan(scan) => !simple_scan(scan), + LogicalPlan::SubqueryAlias(alias) => { + !matches!(alias.input.as_ref(), LogicalPlan::TableScan(scan) if simple_scan(scan)) + } + _ => true, + } + } + + fn derive_asof_input( + &self, + plan: &LogicalPlan, + relation: &mut RelationBuilder, + ) -> Result> { + if let LogicalPlan::SubqueryAlias(alias) = plan { + let (inner, columns) = subquery_alias_inner_query_and_columns(alias); + let table_alias = alias.alias.clone(); + if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() + { + let rewritten = + inject_column_aliases_into_subquery(inner.clone(), columns)?; + self.derive( + &rewritten, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), vec![])), + false, + )?; + } else { + self.derive( + inner, + relation, + Some(self.new_table_alias(table_alias.table().to_string(), columns)), + false, + )?; + } + return Ok(Some(table_alias)); + } + + let qualifier = plan + .schema() + .iter() + .find_map(|(qualifier, _)| qualifier.cloned()); + let alias = qualifier + .as_ref() + .map(|qualifier| self.new_table_alias(qualifier.table().to_string(), vec![])); + self.derive(plan, relation, alias, false)?; + Ok(qualifier) + } + + fn derived_input_projection( + &self, + plan: &LogicalPlan, + qualifier: Option<&TableReference>, + ) -> Result> { + plan.schema() + .iter() + .map(|(field_qualifier, field)| { + self.select_item_to_sql(&Expr::Column(Column::new( + qualifier.cloned().or_else(|| field_qualifier.cloned()), + field.name(), + ))) + }) + .collect() + } + fn is_qualified_passthrough_projection(projection: &Projection) -> bool { projection .expr @@ -2092,26 +2281,28 @@ impl Unparser<'_> { } } - fn qualified_passthrough_join_projection_to_nested_relation( + fn join_input_to_nested_relation( &self, plan: &LogicalPlan, query: &mut Option, ) -> Result> { - let LogicalPlan::Projection(projection) = plan else { - return Ok(None); + let join_plan = match plan { + LogicalPlan::Join(_) => plan, + LogicalPlan::Projection(projection) + if matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + && Self::is_qualified_passthrough_projection(projection) => + { + projection.input.as_ref() + } + _ => return Ok(None), }; - if !matches!(projection.input.as_ref(), LogicalPlan::Join(_)) - || !Self::is_qualified_passthrough_projection(projection) - { - return Ok(None); - } let original_query = query.clone(); let mut nested_select = SelectBuilder::default(); nested_select.push_from(TableWithJoinsBuilder::default()); let mut nested_relation = RelationBuilder::default(); self.select_to_sql_recursively( - projection.input.as_ref(), + join_plan, query, &mut nested_select, &mut nested_relation, @@ -2122,11 +2313,11 @@ impl Unparser<'_> { } let Some(mut nested_from) = nested_select.pop_from() else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; nested_from.relation(nested_relation); let Some(table_with_joins) = nested_from.build()? else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; let mut relation = RelationBuilder::default(); diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt new file mode 100644 index 000000000000..daca2d13c4d8 --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,245 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement ok +CREATE TABLE asof_left(id INT, grp TEXT, ts TIMESTAMP) AS VALUES + (1, 'A', TIMESTAMP '2024-01-01 09:00:01'), + (2, 'A', TIMESTAMP '2024-01-01 09:00:04'), + (3, 'A', TIMESTAMP '2024-01-01 09:00:07'), + (4, 'B', TIMESTAMP '2024-01-01 09:00:02'), + (5, 'B', TIMESTAMP '2024-01-01 09:00:08'), + (6, NULL, TIMESTAMP '2024-01-01 09:00:03'), + (7, 'A', NULL); + +statement ok +CREATE TABLE asof_right(grp TEXT, ts TIMESTAMP, val TEXT) AS VALUES + ('A', TIMESTAMP '2024-01-01 09:00:02', 'a2'), + ('A', TIMESTAMP '2024-01-01 09:00:04', 'a4'), + ('A', TIMESTAMP '2024-01-01 09:00:06', 'a6'), + ('B', TIMESTAMP '2024-01-01 09:00:01', 'b1'), + ('B', TIMESTAMP '2024-01-01 09:00:06', 'b6'), + (NULL, TIMESTAMP '2024-01-01 09:00:02', 'null-group'), + ('A', NULL, 'null-ts'); + +# Inclusive predecessor per equality group. This also verifies unmatched left +# rows and NULL behavior for equality keys and ordered expressions. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Strict predecessor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts > r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 NULL NULL +2 2024-01-01T09:00:04 2024-01-01T09:00:02 a2 +3 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Inclusive successor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts <= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# Strict successor per equality group. +query IPPT +SELECT l.id, l.ts, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts < r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 2024-01-01T09:00:01 2024-01-01T09:00:02 a2 +2 2024-01-01T09:00:04 2024-01-01T09:00:06 a6 +3 2024-01-01T09:00:07 NULL NULL +4 2024-01-01T09:00:02 2024-01-01T09:00:06 b6 +5 2024-01-01T09:00:08 NULL NULL +6 2024-01-01T09:00:03 NULL NULL +7 NULL NULL NULL + +# USING exposes one unqualified equality key. +query TIPT +SELECT grp, l.id, r.ts, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +A 1 NULL NULL +A 2 2024-01-01T09:00:04 a4 +A 3 2024-01-01T09:00:06 a6 +B 4 2024-01-01T09:00:01 b1 +B 5 2024-01-01T09:00:06 b6 +NULL 6 NULL NULL +A 7 NULL NULL + +# Both qualified equality keys remain addressable. +query ITT +SELECT l.id, l.grp, r.grp +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +1 A NULL +2 A A +3 A A +4 B B +5 B B +6 NULL NULL +7 A NULL + +# Equality keys are optional. +query IT +SELECT l.id, r.label +FROM (VALUES (1, 1), (2, 5), (3, CAST(NULL AS INT))) AS l(id, ts) +ASOF JOIN (VALUES (2, 'r2'), (4, 'r4')) AS r(ts, label) +MATCH_CONDITION (l.ts >= r.ts) +ORDER BY l.id; +---- +1 NULL +2 r4 +3 NULL + +# Multiple equality keys form one candidate group. +query IT +SELECT l.id, r.val +FROM (VALUES + (1, 'X', 'A', TIMESTAMP '2024-01-01 09:00:04'), + (2, 'Y', 'A', TIMESTAMP '2024-01-01 09:00:04') +) AS l(id, venue, grp, ts) +ASOF JOIN (VALUES + ('X', 'A', TIMESTAMP '2024-01-01 09:00:02', 'x-a2'), + ('Y', 'A', TIMESTAMP '2024-01-01 09:00:03', 'y-a3'), + ('X', 'B', TIMESTAMP '2024-01-01 09:00:04', 'x-b4') +) AS r(venue, grp, ts, val) +MATCH_CONDITION (l.ts >= r.ts) +ON l.venue = r.venue AND l.grp = r.grp +ORDER BY l.id; +---- +1 x-a2 +2 y-a3 + +# Candidate selection sees the right input after subquery filtering. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN (SELECT * FROM asof_right WHERE val <> 'a6') r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +WHERE l.id IN (3, 5) +ORDER BY l.id; +---- +3 a4 +5 b6 + +# Equality and match operands use the planner's common coercion types. +query II +SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +1 90 + +query TT +EXPLAIN SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.val +02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +03)----SubqueryAlias: l +04)------TableScan: asof_left projection=[id, grp, ts] +05)----SubqueryAlias: r +06)------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, val@5 as val] +02)--AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts] +03)----SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts = r.ts) +ON l.grp = r.grp; + +query error ASOF MATCH_CONDITION left operand must reference only the left input +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (r.ts >= l.ts) +ON l.grp = r.grp; + +query error ASOF ON accepts only equality conditions combined with AND +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp > r.grp; + +query error ASOF MATCH_CONDITION must be a single comparison +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts) +ON l.grp = r.grp; diff --git a/datafusion/substrait/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?; diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index af442de6597c..73c256d062f1 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -319,6 +319,7 @@ SELECT a FROM table_name WHERE a > 10; ```text from_item [join_type] JOIN from_item [join_condition] +from_item ASOF JOIN from_item MATCH_CONDITION (condition) [join_condition] from_item CROSS JOIN from_item from_item NATURAL JOIN from_item from_item [join_type] JOIN LATERAL (query) AS alias [join_condition] @@ -400,6 +401,45 @@ SELECT * FROM x LEFT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ ``` +### ASOF JOIN + +An `ASOF JOIN` matches each left row with at most one right row according to an +ordered comparison. It preserves every left row and fills the right columns +with `NULL` when no right row matches. + +```sql +SELECT t.*, p.price +FROM trades AS t +ASOF JOIN prices AS p +MATCH_CONDITION (t.ts >= p.ts) +ON t.symbol = p.symbol; +``` + +`MATCH_CONDITION` must compare an expression from the left input with an +expression from the right input using one of the following operators: + +| Condition | Selected right row | +| --------- | ----------------------------------------- | +| `l >= r` | Greatest `r` less than or equal to `l` | +| `l > r` | Greatest `r` strictly less than `l` | +| `l <= r` | Smallest `r` greater than or equal to `l` | +| `l < r` | Smallest `r` strictly greater than `l` | + +An optional `ON` clause containing equality conditions combined with `AND`, or +a `USING` clause, divides rows into equality groups before the ordered match. +An unqualified `USING` key appears once in wildcard output, while both qualified +input keys remain addressable. + +Without equality keys, all rows belong to one group. The initial execution +strategy collects one ordered right partition and shares it across every left +partition, so output partitioning follows the left input. The complete right +input must fit in memory and may be scanned once per left partition; spilling +and repartitioned ASOF execution are not yet supported. + +A `NULL` in either ordered expression or in any equality key does not match. +Both inputs must be bounded. If multiple right rows have the same equality keys +and ordered value, which tied row is selected is nondeterministic. + ### RIGHT OUTER JOIN The keywords `RIGHT JOIN` or `RIGHT OUTER JOIN` define a join that includes all rows from the right table even if there