diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index bb526895b6b12..873fadd67a5ec 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,312 @@ 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_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_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_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?; + 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?; + 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(()) +} diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index a6a254ea25bac..b205738d72697 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -38,13 +38,15 @@ use crate::logical_plan::{ }; use crate::select_expr::SelectExpr; use crate::utils::{ - can_hash, columnize_expr, compare_sort_expr, expand_qualified_wildcard, - expand_wildcard, expr_to_columns, find_valid_equijoin_key_pair, - group_window_expr_by_sort_keys, + can_hash, check_all_columns_from_schema, columnize_expr, compare_sort_expr, + expand_qualified_wildcard, expand_wildcard, expr_to_columns, + find_valid_equijoin_key_pair, group_window_expr_by_sort_keys, + split_conjunction_owned, }; use crate::{ - DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, RecursiveQuery, - Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, binary_expr, lit, + BinaryExpr, DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, + RecursiveQuery, Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, + binary_expr, lit, }; use super::dml::InsertOp; @@ -1007,23 +1009,67 @@ impl LogicalPlanBuilder { ) } - /// Apply a left-preserving ASOF join using equality expressions and one - /// ordered match condition. - pub fn asof_join( + /// Apply a left-preserving ASOF join using an optional equality condition + /// and one ordered match condition. + /// + /// When present, `on_expr` must contain equality comparisons combined with + /// `AND`. Each comparison must have one operand that references only the + /// left input and one that references only `right`; their order does not + /// matter. `match_condition` must be a single `<`, `<=`, `>`, or `>=` + /// comparison whose left operand references only the left input and whose + /// right operand references only `right`. + pub fn asof_join_on( self, right: LogicalPlan, - on: Vec<(Expr, Expr)>, - match_condition: AsOfMatch, + on_expr: Option, + match_condition: Expr, ) -> Result { - self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On) + let on = on_expr + .into_iter() + .flat_map(split_conjunction_owned) + .map(|predicate| { + let Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::Eq, + right: right_expr, + }) = predicate + else { + return plan_err!( + "ASOF ON accepts only equality conditions combined with AND" + ); + }; + find_valid_equijoin_key_pair( + &left, + &right_expr, + self.plan.schema(), + right.schema(), + )? + .ok_or_else(|| { + plan_datafusion_err!( + "Each ASOF equality condition must compare one left expression with one right expression" + ) + }) + }) + .collect::>()?; + self.asof_join_with_constraint( + right, + on, + AsOfMatch::try_from(match_condition)?, + JoinConstraint::On, + ) } - /// Apply a left-preserving ASOF join using `USING` equality keys. + /// Apply a left-preserving ASOF join using `USING` equality keys and one + /// ordered match condition. + /// + /// Every key in `using_keys` must resolve in both inputs. + /// `match_condition` follows the same operand and operator requirements as + /// [`asof_join_on`](Self::asof_join_on). pub fn asof_join_using( self, right: LogicalPlan, using_keys: Vec, - match_condition: AsOfMatch, + match_condition: Expr, ) -> Result { let on = using_keys .into_iter() @@ -1033,7 +1079,12 @@ impl LogicalPlanBuilder { Ok((Expr::Column(left), Expr::Column(right))) }) .collect::>()?; - self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using) + self.asof_join_with_constraint( + right, + on, + AsOfMatch::try_from(match_condition)?, + JoinConstraint::Using, + ) } fn asof_join_with_constraint( @@ -1043,6 +1094,17 @@ impl LogicalPlanBuilder { match_condition: AsOfMatch, join_constraint: JoinConstraint, ) -> Result { + let left_columns = match_condition.left.column_refs(); + let right_columns = match_condition.right.column_refs(); + if left_columns.is_empty() + || right_columns.is_empty() + || !check_all_columns_from_schema(&left_columns, self.plan.schema())? + || !check_all_columns_from_schema(&right_columns, right.schema())? + { + return plan_err!( + "ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input" + ); + } let normalize = |expr, schema: &DFSchema| { normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[]) }; @@ -2914,6 +2976,70 @@ mod tests { Ok(()) } + #[test] + fn asof_join_on_extracts_and_validates_conditions() -> Result<()> { + let values = vec![vec![lit(1), lit(2)]]; + let left = LogicalPlanBuilder::values(values.clone())? + .alias("l")? + .build()?; + let right = LogicalPlanBuilder::values(values)?.alias("r")?.build()?; + + let plan = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some( + col("r.column1") + .eq(col("l.column1")) + .and(col("l.column2").eq(col("r.column2"))), + ), + col("l.column2").gt_eq(col("r.column2")), + )? + .build()?; + let LogicalPlan::AsOfJoin(join) = plan else { + panic!("expected ASOF join") + }; + assert_eq!( + join.on, + vec![ + (col("l.column1"), col("r.column1")), + (col("l.column2"), col("r.column2")), + ] + ); + assert_eq!( + join.match_condition.as_ref(), + &AsOfMatch::new(col("l.column2"), Operator::GtEq, col("r.column2")) + ); + + let invalid_on = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some(col("l.column1").gt(col("r.column1"))), + col("l.column2").gt_eq(col("r.column2")), + ) + .expect_err("non-equality ASOF ON should fail"); + assert_snapshot!(invalid_on.strip_backtrace(), @r#"Error during planning: ASOF ON accepts only equality conditions combined with AND"#); + + let invalid_match = LogicalPlanBuilder::from(left.clone()) + .asof_join_on( + right.clone(), + Some(col("l.column1").eq(col("r.column1"))), + col("l.column2").eq(col("r.column2")), + ) + .expect_err("equality ASOF MATCH_CONDITION should fail"); + assert_snapshot!(invalid_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION requires <, <=, >, or >=, found ="#); + + let reversed_match = LogicalPlanBuilder::from(left) + .asof_join_on( + right, + Some(col("l.column1").eq(col("r.column1"))), + col("r.column2").gt_eq(col("l.column2")), + ) + .expect_err("reversed ASOF MATCH_CONDITION should fail"); + assert_snapshot!(reversed_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input"#); + + Ok(()) + } + #[test] fn plan_builder_from_logical_plan() -> Result<()> { let plan = diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 74b5c42353ad8..edd182a00e854 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -4462,6 +4462,25 @@ impl AsOfMatch { } } +impl TryFrom for AsOfMatch { + type Error = DataFusionError; + + fn try_from(condition: Expr) -> Result { + let Expr::BinaryExpr(BinaryExpr { left, op, right }) = condition else { + return plan_err!("ASOF MATCH_CONDITION must be a single comparison"); + }; + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {op}" + ); + } + Ok(Self::new(*left, op, *right)) + } +} + impl Display for AsOfMatch { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { write!(f, "{} {} {}", self.left, self.op, self.right) @@ -6179,10 +6198,10 @@ mod tests { assert_eq!(cross_join.min_rows(), 2); let asof_join = LogicalPlanBuilder::from(two_rows.clone()) - .asof_join( + .asof_join_on( one_row.clone(), - vec![], - AsOfMatch::new(col("l.column1"), Operator::GtEq, col("r.column1")), + None, + col("l.column1").gt_eq(col("r.column1")), )? .build()?; assert_eq!(asof_join.min_rows(), 2); diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b38099..73c2f05504269 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,7 +16,7 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err}; +use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err, plan_err}; use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{ Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, @@ -98,10 +98,75 @@ impl SqlToRel<'_, S> { JoinOperator::CrossJoin(JoinConstraint::None) => { self.parse_cross_join(left, right) } + JoinOperator::AsOf { + match_condition, + constraint, + } => self.parse_asof_join( + left, + right, + match_condition, + constraint, + planner_context, + ), other => not_impl_err!("Unsupported JOIN operator {other:?}"), } } + fn parse_asof_join( + &self, + left: LogicalPlan, + right: LogicalPlan, + sql_match_condition: sqlparser::ast::Expr, + constraint: JoinConstraint, + planner_context: &mut PlannerContext, + ) -> Result { + let join_schema = left.schema().join(right.schema())?; + let match_condition = + self.sql_to_expr(sql_match_condition, &join_schema, planner_context)?; + + match constraint { + JoinConstraint::On(sql_on) => { + let on = self.sql_to_expr(sql_on, &join_schema, planner_context)?; + LogicalPlanBuilder::from(left) + .asof_join_on(right, Some(on), match_condition)? + .build() + } + JoinConstraint::Using(object_names) => { + let keys = object_names + .into_iter() + .map(|object_name| { + let ObjectName(mut object_names) = object_name; + if object_names.len() != 1 { + return not_impl_err!( + "Invalid identifier in ASOF USING clause. Expected single identifier, got {}", + ObjectName(object_names) + ); + } + let id = object_names.swap_remove(0); + id.as_ident() + .ok_or_else(|| { + plan_datafusion_err!( + "Expected identifier in ASOF USING clause" + ) + }) + .map(|ident| { + Column::from_name( + self.ident_normalizer.normalize(ident.clone()), + ) + }) + }) + .collect::>>()?; + LogicalPlanBuilder::from(left) + .asof_join_using(right, keys, match_condition)? + .build() + } + JoinConstraint::None => LogicalPlanBuilder::from(left) + .asof_join_on(right, None, match_condition)? + .build(), + JoinConstraint::Natural => plan_err!("NATURAL ASOF JOIN is not supported"), + } + } + fn parse_cross_join( &self, left: LogicalPlan, diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 29d4d1b5b8b4f..f15e3eb887247 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, }; @@ -184,6 +184,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -204,7 +205,6 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) - | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } @@ -1479,11 +1479,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 { @@ -1617,6 +1614,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); @@ -1917,6 +1917,129 @@ impl Unparser<'_> { } } + // Keep ASOF-specific locals out of the recursive plan unparser's stack frame. + #[inline(never)] + fn asof_join_to_sql( + &self, + join: &AsOfJoin, + query: &mut Option, + select: &mut SelectBuilder, + relation: &mut RelationBuilder, + ) -> Result<()> { + let already_projected = select.already_projected(); + let left_plan = + Self::unwrap_qualified_passthrough_join_projection(Arc::clone(&join.left)); + let inline_left_join = matches!( + left_plan.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ); + let left_projection = if already_projected { + None + } else if inline_left_join { + self.select_to_sql_recursively(left_plan.as_ref(), query, select, relation)?; + select.pop_projections(); + Some(self.derived_input_projection(join.left.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + let qualifier = self.derive_asof_input(join.left.as_ref(), relation)?; + Some(self.derived_input_projection(join.left.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively(join.left.as_ref(), query, select, relation)?; + Some(select.pop_projections()) + }; + if already_projected { + if inline_left_join { + self.select_to_sql_recursively( + left_plan.as_ref(), + query, + select, + relation, + )?; + } else if Self::asof_input_requires_derived(join.left.as_ref()) { + self.derive_asof_input(join.left.as_ref(), relation)?; + } else { + self.select_to_sql_recursively( + join.left.as_ref(), + query, + select, + relation, + )?; + } + } + + let mut right_relation = RelationBuilder::default(); + let nested_right = + self.join_input_to_nested_relation(join.right.as_ref(), query)?; + let right_projection = if already_projected { + if let Some(nested_right) = nested_right { + right_relation = nested_right; + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + } + None + } else if let Some(nested_right) = nested_right { + right_relation = nested_right; + Some(self.derived_input_projection(join.right.as_ref(), None)?) + } else if Self::asof_input_requires_derived(join.right.as_ref()) { + let qualifier = + self.derive_asof_input(join.right.as_ref(), &mut right_relation)?; + Some(self.derived_input_projection(join.right.as_ref(), qualifier.as_ref())?) + } else { + self.select_to_sql_recursively( + join.right.as_ref(), + query, + select, + &mut right_relation, + )?; + Some(select.pop_projections()) + }; + let Ok(Some(relation)) = right_relation.build() else { + return internal_err!("Failed to build ASOF right relation"); + }; + let constraint = + self.join_constraint_to_sql(join.join_constraint, &join.on, None)?; + let match_condition = self.expr_to_sql(&Expr::BinaryExpr(BinaryExpr::new( + Box::new(join.match_condition.left.clone()), + join.match_condition.op, + Box::new(join.match_condition.right.clone()), + )))?; + let ast_join = ast::Join { + relation, + global: false, + join_operator: ast::JoinOperator::AsOf { + match_condition, + constraint, + }, + }; + let mut from = select + .pop_from() + .ok_or_else(|| internal_datafusion_err!("ASOF left relation is missing"))?; + from.push_join(ast_join); + select.push_from(from); + + if !already_projected { + let left_projection = left_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF left projection is missing") + })?; + let right_projection = right_projection.ok_or_else(|| { + internal_datafusion_err!("ASOF right projection is missing") + })?; + select.projection( + left_projection + .into_iter() + .chain(right_projection) + .collect(), + ); + } + Ok(()) + } + /// Walk through transparent nodes (SubqueryAlias) to find the inner /// Projection that feeds an Unnest node. /// @@ -2215,6 +2338,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 @@ -2226,7 +2417,10 @@ impl Unparser<'_> { plan: Arc, ) -> Arc { if let LogicalPlan::Projection(projection) = plan.as_ref() - && matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + && matches!( + projection.input.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ) && Self::is_qualified_passthrough_projection(projection) { Arc::clone(&projection.input) @@ -2235,26 +2429,30 @@ impl Unparser<'_> { } } - fn qualified_passthrough_join_projection_to_nested_relation( + fn join_input_to_nested_relation( &self, plan: &LogicalPlan, query: &mut Option, ) -> Result> { - let LogicalPlan::Projection(projection) = plan else { - return Ok(None); + let join_plan = match plan { + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) => plan, + LogicalPlan::Projection(projection) + if matches!( + projection.input.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ) && Self::is_qualified_passthrough_projection(projection) => + { + projection.input.as_ref() + } + _ => return Ok(None), }; - if !matches!(projection.input.as_ref(), LogicalPlan::Join(_)) - || !Self::is_qualified_passthrough_projection(projection) - { - return Ok(None); - } let original_query = query.clone(); let mut nested_select = SelectBuilder::default(); nested_select.push_from(TableWithJoinsBuilder::default()); let mut nested_relation = RelationBuilder::default(); self.select_to_sql_recursively( - projection.input.as_ref(), + join_plan, query, &mut nested_select, &mut nested_relation, @@ -2265,11 +2463,11 @@ impl Unparser<'_> { } let Some(mut nested_from) = nested_select.pop_from() else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; nested_from.relation(nested_relation); let Some(table_with_joins) = nested_from.build()? else { - return internal_err!("Failed to build nested join relation"); + return internal_err!("Failed to build nested join input relation"); }; let mut relation = RelationBuilder::default(); diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 32dfade056dd2..191425416c119 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -2823,6 +2823,37 @@ fn test_unparse_inner_join_with_table_scan_projection() -> Result<()> { Ok(()) } +#[test] +fn test_unparse_asof_join() -> Result<()> { + let trades_schema = Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("trade_id", DataType::Int32, false), + ]); + let prices_schema = Schema::new(vec![ + Field::new("symbol", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ]); + let trades = table_scan(Some("trades"), &trades_schema, None)? + .alias("t")? + .build()?; + let prices = table_scan(Some("prices"), &prices_schema, None)? + .alias("p")? + .build()?; + let plan = LogicalPlanBuilder::from(trades) + .asof_join_on( + prices, + Some(col("t.symbol").eq(col("p.symbol"))), + col("t.ts").gt_eq(col("p.ts")), + )? + .project(vec![col("t.trade_id"), col("p.price")])? + .build()?; + + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT t.trade_id, p.price FROM trades AS t ASOF JOIN prices AS p MATCH_CONDITION ((t.ts >= p.ts)) ON t.symbol = p.symbol"#); + Ok(()) +} + /// Build the three base table scans (`left_table`, `mid_table`, `right_table`) /// shared by the nested passthrough-projection join unparsing tests. fn nested_passthrough_join_tables() -> Result<(LogicalPlan, LogicalPlan, LogicalPlan)> { @@ -2933,6 +2964,65 @@ fn test_unparse_projected_join_unwraps_left_nested_passthrough_projection() -> R Ok(()) } +#[test] +fn test_unparse_nested_asof_join_inputs() -> Result<()> { + let (left, mid, right) = nested_passthrough_join_tables()?; + let nested_left = LogicalPlanBuilder::from(left) + .asof_join_on( + mid, + Some(col("left_table.mid_id").eq(col("mid_table.mid_id"))), + col("left_table.left_id").gt_eq(col("mid_table.mid_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("mid_table.right_id"), + ])? + .build()?; + let plan = LogicalPlanBuilder::from(nested_left) + .asof_join_on( + right, + Some(col("mid_table.right_id").eq(col("right_table.right_id"))), + col("mid_table.right_id").gt_eq(col("right_table.right_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table ASOF JOIN mid_table MATCH_CONDITION ((left_table.left_id >= mid_table.mid_id)) ON left_table.mid_id = mid_table.mid_id ASOF JOIN right_table MATCH_CONDITION ((mid_table.right_id >= right_table.right_id)) ON mid_table.right_id = right_table.right_id"#); + + let (left, mid, right) = nested_passthrough_join_tables()?; + let nested_right = LogicalPlanBuilder::from(mid) + .asof_join_on( + right, + Some(col("mid_table.right_id").eq(col("right_table.right_id"))), + col("mid_table.right_id").gt_eq(col("right_table.right_id")), + )? + .project(vec![ + col("mid_table.mid_id"), + col("mid_table.right_id"), + col("right_table.value"), + ])? + .build()?; + let plan = LogicalPlanBuilder::from(left) + .asof_join_on( + nested_right, + Some(col("left_table.mid_id").eq(col("mid_table.mid_id"))), + col("left_table.left_id").gt_eq(col("mid_table.mid_id")), + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + assert_snapshot!(plan_to_sql(&plan)?, @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table ASOF JOIN (mid_table ASOF JOIN right_table MATCH_CONDITION ((mid_table.right_id >= right_table.right_id)) ON mid_table.right_id = right_table.right_id) MATCH_CONDITION ((left_table.left_id >= mid_table.mid_id)) ON left_table.mid_id = mid_table.mid_id"#); + + Ok(()) +} + #[test] fn test_unparse_left_semi_join_with_table_scan_projection() -> Result<()> { let schema = Schema::new(vec![ diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt new file mode 100644 index 0000000000000..77933d0a1247c --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,511 @@ +# 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 + +# SELECT * exposes the USING key once. +query ITPPT +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +1 A 2024-01-01T09:00:01 NULL NULL +2 A 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 A 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 B 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 B 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 NULL 2024-01-01T09:00:03 NULL NULL +7 A NULL NULL 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 + +# USING accepts multiple equality keys. +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) +USING (venue, 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 + +# A filter on a right output column is applied after candidate selection. +query IT +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 +WHERE r.val <> 'a6' +ORDER BY l.id; +---- +2 a4 +4 b1 +5 b6 + +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 +WHERE r.val <> 'a6'; +---- +logical_plan +01)Filter: r.val != Utf8View("a6") +02)--Projection: l.id, r.val +03)----AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +04)------SubqueryAlias: l +05)--------TableScan: asof_left projection=[id, grp, ts] +06)------SubqueryAlias: r +07)--------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)FilterExec: val@1 != a6 +02)--ProjectionExec: expr=[id@0 as id, val@5 as val] +03)----AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts] +04)------SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)------SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +07)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Projection pruning keeps only columns required by the ASOF contract. +query TT +EXPLAIN SELECT l.id +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 +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] +physical_plan +01)ProjectionExec: expr=[id@0 as id] +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] + +# Equality and ordered operands may be expressions. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts + INTERVAL '1 second') +ON lower(l.grp) = lower(r.grp) +WHERE l.id IN (2, 3, 4) +ORDER BY l.id; +---- +2 a2 +3 a6 +4 b1 + +# Empty inputs preserve left-join semantics. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN (SELECT * FROM asof_right WHERE false) r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 NULL +2 NULL +3 NULL +4 NULL +5 NULL +6 NULL +7 NULL + +query I +SELECT l.id +FROM (SELECT * FROM asof_left WHERE false) l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- + +# Duplicate left rows are emitted independently. +query IT +SELECT l.id, r.val +FROM ( + SELECT * FROM asof_left + UNION ALL + SELECT * FROM asof_left +) l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id, r.val; +---- +1 NULL +1 NULL +2 a4 +2 a4 +3 a6 +3 a6 +4 b1 +4 b1 +5 b6 +5 b6 +6 NULL +6 NULL +7 NULL +7 NULL + +# ASOF JOIN supports independently aliased self joins. +query II +SELECT a.id, b.id +FROM asof_left a +ASOF JOIN asof_left b +MATCH_CONDITION (a.ts > b.ts) +ON a.grp = b.grp +ORDER BY a.id; +---- +1 NULL +2 1 +3 2 +4 NULL +5 4 +6 NULL +7 NULL + +# Equality and match operands use the planner's common coercion types. +query II +SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +1 90 + +# Coercion is applied before the ASOF ordering requirements are planned. +query TT +EXPLAIN SELECT l.id, r.payload +FROM (VALUES (1, CAST(5 AS SMALLINT), CAST(10 AS INT))) AS l(id, grp, ts) +ASOF JOIN ( + VALUES (CAST(5 AS BIGINT), CAST(9 AS BIGINT), 90) +) AS r(grp, ts, payload) +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.payload +02)--AsOf Join: match=[CAST(l.ts AS Int64) >= r.ts], constraint=On, on=[CAST(l.grp AS Int64) = r.grp] +03)----SubqueryAlias: l +04)------Projection: column1 AS id, column2 AS grp, column3 AS ts +05)--------Values: (Int64(1), Int16(5) AS Int64(5), Int32(10) AS Int64(10)) +06)----SubqueryAlias: r +07)------Projection: column1 AS grp, column2 AS ts, column3 AS payload +08)--------Values: (Int64(5), Int64(9), Int64(90)) +physical_plan +01)ProjectionExec: expr=[id@0 as id, payload@5 as payload] +02)--AsOfJoinExec: on=[(CAST(grp AS Int64) = grp)], match=[CAST(ts AS Int64) >= ts] +03)----SortExec: expr=[CAST(grp@1 AS Int64) ASC, CAST(ts@2 AS Int64) ASC], preserve_partitioning=[false] +04)------ProjectionExec: expr=[column1@0 as id, column2@1 as grp, column3@2 as ts] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)----ProjectionExec: expr=[column1@0 as grp, column2@1 as ts, column3@2 as payload] +07)------SortExec: expr=[column1@0 ASC, column2@1 ASC], preserve_partitioning=[false] +08)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Equality operands can name the right input first. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON r.grp = l.grp +WHERE l.id IN (2, 4) +ORDER BY l.id; +---- +2 a4 +4 b1 + +query TT +EXPLAIN SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id, r.val +02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +03)----SubqueryAlias: l +04)------TableScan: asof_left projection=[id, grp, ts] +05)----SubqueryAlias: r +06)------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, val@5 as val] +02)--AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts] +03)----SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts = r.ts) +ON l.grp = r.grp; + +query error ASOF MATCH_CONDITION left operand must reference only the left input +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (r.ts >= l.ts) +ON l.grp = r.grp; + +query error ASOF ON accepts only equality conditions combined with AND +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp > r.grp; + +query error ASOF MATCH_CONDITION must be a single comparison +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts) +ON l.grp = r.grp; + +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 (1 >= r.ts) +ON l.grp = r.grp; + +query error Each ASOF equality condition must compare one left expression with one right expression +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON 1 = 1; + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts AND l.ts > r.ts) +ON l.grp = r.grp; + +query error Each ASOF equality condition must compare one left expression with one right expression +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = l.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 AND l.ts >= r.ts; + +query error Ambiguous reference to unqualified field ts +SELECT ts +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp); diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 75e1fe251ac1b..4a8413718edb9 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -18,7 +18,7 @@ #[cfg(test)] mod tests { use datafusion::datasource::provider_as_source; - use datafusion::logical_expr::{AsOfMatch, LogicalPlanBuilder, Operator}; + use datafusion::logical_expr::LogicalPlanBuilder; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::serializer; @@ -110,10 +110,10 @@ mod tests { let left = LogicalPlanBuilder::scan("l", Arc::clone(&table), None)?.build()?; let right = LogicalPlanBuilder::scan("r", table, None)?.build()?; let plan = LogicalPlanBuilder::from(left) - .asof_join( + .asof_join_on( right, - vec![(col("l.b"), col("r.b"))], - AsOfMatch::new(col("l.a"), Operator::GtEq, col("r.a")), + Some(col("l.b").eq(col("r.b"))), + col("l.a").gt_eq(col("r.a")), )? .build()?; let error = to_substrait_plan(&plan, &ctx.state()) diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index af442de6597c1..2cbe26fee3964 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -319,6 +319,7 @@ SELECT a FROM table_name WHERE a > 10; ```text from_item [join_type] JOIN from_item [join_condition] +from_item ASOF JOIN from_item MATCH_CONDITION (condition) [join_condition] from_item CROSS JOIN from_item from_item NATURAL JOIN from_item from_item [join_type] JOIN LATERAL (query) AS alias [join_condition] @@ -400,6 +401,48 @@ SELECT * FROM x LEFT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ ``` +### ASOF JOIN + +DataFusion follows the +[Snowflake `ASOF JOIN` syntax](https://docs.snowflake.com/en/sql-reference/constructs/asof-join). +An `ASOF JOIN` matches each left row with at most one right row according to an +ordered comparison. It preserves every left row and fills the right columns +with `NULL` when no right row matches. + +```sql +SELECT t.*, p.price +FROM trades AS t +ASOF JOIN prices AS p +MATCH_CONDITION (t.ts >= p.ts) +ON t.symbol = p.symbol; +``` + +`MATCH_CONDITION` must compare an expression from the left input with an +expression from the right input using one of the following operators. Operand +order is significant: the left input expression must appear on the left. + +| Condition | Selected right row | +| --------- | ----------------------------------------- | +| `l >= r` | Greatest `r` less than or equal to `l` | +| `l > r` | Greatest `r` strictly less than `l` | +| `l <= r` | Smallest `r` greater than or equal to `l` | +| `l < r` | Smallest `r` strictly greater than `l` | + +An optional `ON` clause containing equality conditions combined with `AND`, or +a `USING` clause, divides rows into equality groups before the ordered match. +An unqualified `USING` key appears once in wildcard output, while both qualified +input keys remain addressable. + +Without equality keys, all rows belong to one group. The initial execution +strategy collects one ordered right partition and shares it across every left +partition, so output partitioning follows the left input. The complete right +input must fit in memory and may be scanned once per left partition; spilling +and repartitioned ASOF execution are not yet supported. + +A `NULL` in either ordered expression or in any equality key does not match. +Both inputs must be bounded. If multiple right rows have the same equality keys +and ordered value, which tied row is selected is nondeterministic. + ### RIGHT OUTER JOIN The keywords `RIGHT JOIN` or `RIGHT OUTER JOIN` define a join that includes all rows from the right table even if there