From 13ef0651144fda009e5bf5e4e268c753381949e5 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 15:05:10 +0800 Subject: [PATCH 01/18] feat: support ASOF joins --- benchmarks/src/asof.rs | 194 +++ benchmarks/src/bin/dfbench.rs | 4 +- benchmarks/src/lib.rs | 1 + datafusion/core/src/dataframe/mod.rs | 43 +- datafusion/core/src/physical_planner.rs | 69 +- datafusion/core/src/prelude.rs | 2 +- .../physical_optimizer/filter_pushdown.rs | 42 + datafusion/core/tests/sql/joins.rs | 421 +++++ datafusion/expr/src/logical_plan/builder.rs | 107 +- datafusion/expr/src/logical_plan/display.rs | 23 +- datafusion/expr/src/logical_plan/mod.rs | 16 +- datafusion/expr/src/logical_plan/plan.rs | 279 +++- datafusion/expr/src/logical_plan/tree_node.rs | 57 +- .../optimizer/src/analyzer/type_coercion.rs | 37 +- .../optimizer/src/common_subexpr_eliminate.rs | 1 + .../optimizer/src/optimize_projections/mod.rs | 43 +- datafusion/optimizer/src/optimizer.rs | 5 + datafusion/optimizer/src/push_down_filter.rs | 26 + datafusion/physical-plan/Cargo.toml | 5 + datafusion/physical-plan/benches/asof_join.rs | 196 +++ .../physical-plan/src/joins/asof_join.rs | 1419 +++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 2 + .../proto-models/proto/datafusion.proto | 31 + .../proto-models/src/generated/pbjson.rs | 530 ++++++ .../proto-models/src/generated/prost.rs | 83 +- .../proto/src/logical_plan/from_proto.rs | 16 + datafusion/proto/src/logical_plan/mod.rs | 118 +- datafusion/proto/src/logical_plan/to_proto.rs | 20 +- datafusion/proto/src/physical_plan/mod.rs | 140 +- .../tests/cases/roundtrip_logical_plan.rs | 31 + .../tests/cases/roundtrip_physical_plan.rs | 37 +- datafusion/sql/src/relation/join.rs | 136 +- datafusion/sql/src/unparser/plan.rs | 250 ++- .../src/logical_plan/producer/rel/mod.rs | 3 + datafusion/substrait/tests/cases/serialize.rs | 20 + 35 files changed, 4332 insertions(+), 75 deletions(-) create mode 100644 benchmarks/src/asof.rs create mode 100644 datafusion/physical-plan/benches/asof_join.rs create mode 100644 datafusion/physical-plan/src/joins/asof_join.rs diff --git a/benchmarks/src/asof.rs b/benchmarks/src/asof.rs new file mode 100644 index 0000000000000..d74607550f4ed --- /dev/null +++ b/benchmarks/src/asof.rs @@ -0,0 +1,194 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; +use clap::Args; +use datafusion::physical_plan::execute_stream; +use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; +use futures::StreamExt; + +/// Run end-to-end ASOF join benchmarks. +/// +/// The cases cover ordered-input reuse, optimizer-inserted sort/repartition, +/// wide payload materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 4). If not specified, runs all queries + #[arg(short, long)] + query: Option, + + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option, +} + +const ASOF_QUERIES: &[&str] = &[ + // Q1: predecessor without equality keys, ordered range input can be reused + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + "#, + // Q2: grouped predecessor, optimizer supplies repartitioning and ordering + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + value AS payload + FROM range(1000000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + value AS payload + FROM range(1000000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q3: grouped predecessor with a wide payload + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + repeat('x', 256) AS payload + FROM range(250000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + repeat('y', 256) AS payload + FROM range(250000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q4: successor matching requires descending input order + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(500000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(500000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts <= r.ts) + "#, +]; + +impl RunOpt { + pub async fn run(self) -> Result<()> { + println!("Running ASOF benchmarks with the following options: {self:#?}\n"); + + let query_range = match self.query { + Some(query_id) if (1..=ASOF_QUERIES.len()).contains(&query_id) => { + query_id..=query_id + } + Some(query_id) => { + return exec_err!( + "Query {query_id} not found. Available queries: 1 to {}", + ASOF_QUERIES.len() + ); + } + None => 1..=ASOF_QUERIES.len(), + }; + + let config = self.common.config()?; + let runtime = self.common.build_runtime()?; + let ctx = SessionContext::new_with_config_rt(config, runtime); + let mut benchmark_run = BenchmarkRun::new(); + + for query_id in query_range { + let sql = ASOF_QUERIES[query_id - 1]; + benchmark_run.start_new_case(&format!("Query {query_id}")); + match self.benchmark_query(sql, &query_id.to_string(), &ctx).await { + Ok(results) => { + for result in results { + benchmark_run.write_iter(result.elapsed, result.row_count); + } + } + Err(error) => { + return Err(DataFusionError::Context( + format!("ASOF benchmark Q{query_id} failed with error:"), + Box::new(error), + )); + } + } + } + + benchmark_run.maybe_write_json(self.output_path.as_ref())?; + Ok(()) + } + + async fn benchmark_query( + &self, + sql: &str, + query_name: &str, + ctx: &SessionContext, + ) -> Result> { + let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?; + let plan_string = format!("{physical_plan:#?}"); + if !plan_string.contains("AsOfJoinExec") { + return Err(exec_datafusion_err!( + "Query {query_name} does not use AsOfJoinExec. Physical plan: {plan_string}" + )); + } + + let mut query_results = Vec::with_capacity(self.common.iterations); + for iteration in 0..self.common.iterations { + let start = Instant::now(); + let row_count = Self::execute_sql_without_result_buffering(sql, ctx).await?; + let elapsed = start.elapsed(); + println!( + "Query {query_name} iteration {iteration} returned {row_count} rows in {elapsed:?}" + ); + query_results.push(QueryResult { elapsed, row_count }); + } + Ok(query_results) + } + + async fn execute_sql_without_result_buffering( + sql: &str, + ctx: &SessionContext, + ) -> Result { + let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?; + let mut stream = execute_stream(physical_plan, ctx.task_ctx())?; + let mut row_count = 0; + while let Some(batch) = stream.next().await { + row_count += batch?.num_rows(); + } + Ok(row_count) + } +} diff --git a/benchmarks/src/bin/dfbench.rs b/benchmarks/src/bin/dfbench.rs index 50dd99368b7f0..1e478747ceea1 100644 --- a/benchmarks/src/bin/dfbench.rs +++ b/benchmarks/src/bin/dfbench.rs @@ -32,7 +32,7 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; use datafusion_benchmarks::{ - cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch, + asof, cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch, }; #[derive(Debug, Parser)] @@ -44,6 +44,7 @@ struct Cli { #[derive(Debug, Subcommand)] enum Options { + Asof(asof::RunOpt), Cancellation(cancellation::RunOpt), Clickbench(clickbench::RunOpt), Dict(dict::RunOpt), @@ -65,6 +66,7 @@ pub async fn main() -> Result<()> { let cli = Cli::parse(); match cli.command { + Options::Asof(opt) => opt.run().await, Options::Cancellation(opt) => opt.run().await, Options::Clickbench(opt) => opt.run().await, Options::Dict(opt) => opt.run().await, diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 8d24d44a174e3..74b10ae86d78e 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -16,6 +16,7 @@ // under the License. //! DataFusion benchmark runner +pub mod asof; pub mod cancellation; pub mod clickbench; pub mod dict; diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 325ae91d27bbf..faf7cd94fce14 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -58,8 +58,8 @@ use datafusion_common::{ }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ - ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case, - dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, + AsOfMatch, ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, + UNNAMED_TABLE, case, dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, }; use datafusion_functions::core::coalesce; use datafusion_functions::math::nanvl; @@ -1379,6 +1379,45 @@ impl DataFrame { }) } + /// Join this `DataFrame` to the closest eligible row in `right`. + /// + /// Every left row is emitted exactly once. `on` contains optional equality + /// expressions and `match_condition` selects the ordered predecessor or + /// successor from the matching right group. + pub fn join_asof( + self, + right: DataFrame, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + ) -> Result { + let plan = LogicalPlanBuilder::from(self.plan) + .asof_join(right.plan, on, match_condition)? + .build()?; + Ok(DataFrame { + session_state: self.session_state, + plan, + projection_requires_validation: true, + }) + } + + /// Join this `DataFrame` to the closest eligible row in `right` using + /// same-named equality keys. + pub fn join_asof_using( + self, + right: DataFrame, + using_keys: Vec, + match_condition: AsOfMatch, + ) -> Result { + let plan = LogicalPlanBuilder::from(self.plan) + .asof_join_using(right.plan, using_keys, match_condition)? + .build()?; + Ok(DataFrame { + session_state: self.session_state, + plan, + projection_requires_validation: true, + }) + } + /// Repartition a DataFrame based on a logical partitioning scheme. /// /// # Example diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aef8036c749a8..58b4da0f91a94 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -43,7 +43,8 @@ use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::filter::FilterExecBuilder; use crate::physical_plan::joins::utils as join_utils; use crate::physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -93,8 +94,8 @@ use datafusion_expr::physical_planning_context::{ use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, - FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, - WindowFrame, WindowFrameBound, WriteOp, + FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType, + StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, @@ -1860,6 +1861,67 @@ 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, + )?, + ); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect(); + Arc::new(AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + match_condition, + right_output_indices, + )?) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, @@ -2359,6 +2421,7 @@ fn extract_dml_filters( | LogicalPlan::Sort(_) | LogicalPlan::Union(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Aggregate(_) | LogicalPlan::Window(_) diff --git a/datafusion/core/src/prelude.rs b/datafusion/core/src/prelude.rs index 31d9d7eb471f0..12a4f393ca4ac 100644 --- a/datafusion/core/src/prelude.rs +++ b/datafusion/core/src/prelude.rs @@ -34,7 +34,7 @@ pub use crate::execution::options::{ pub use datafusion_common::Column; pub use datafusion_expr::{ - Expr, + AsOfMatch, Expr, Operator, expr_fn::*, lit, lit_timestamp_nano, logical_plan::{JoinType, Partitioning}, diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 909b80cadaae3..2f613264f45d9 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -55,6 +55,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{AsOfJoinExec, AsOfMatchExpr}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -123,6 +124,47 @@ fn test_pushdown_volatile_functions_not_allowed() { ); } +#[test] +fn test_asof_join_pushes_only_left_filters() { + let left = TestScanBuilder::new(schema()).with_support(true).build(); + let right = TestScanBuilder::new(schema()).with_support(true).build(); + let join = Arc::new( + AsOfJoinExec::try_new( + left, + right, + vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())], + AsOfMatchExpr::new( + col("b", &schema()).unwrap(), + Operator::GtEq, + col("b", &schema()).unwrap(), + ), + vec![0, 1, 2], + ) + .unwrap(), + ); + let left_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let right_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 5)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter)); + let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap()); + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); + let formatted = format_plan_for_test(&optimized); + + assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}"); + assert!(formatted.contains("AsOfJoinExec:"), "{formatted}"); + assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}"); + assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}"); +} + /// Show that we can use config options to determine how to do pushdown. #[test] fn test_pushdown_into_scan_with_config_options() { diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 7c0e89ee96418..392a40323c226 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -20,6 +20,9 @@ use insta::assert_snapshot; use datafusion::assert_batches_eq; use datafusion::catalog::MemTable; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; +use datafusion::logical_expr::AsOfJoin; +use datafusion::physical_plan::joins::AsOfJoinExec; +use datafusion::physical_plan::{Distribution, ExecutionPlanProperties}; use datafusion::test_util::register_unbounded_file_with_ordering; use datafusion_sql::unparser::plan_to_sql; @@ -299,3 +302,421 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } + +async 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, vec![trades])?), + )?; + + 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, vec![prices])?), + )?; + 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) +} + +fn find_asof_plan(plan: &LogicalPlan) -> Option<&AsOfJoin> { + if let LogicalPlan::AsOfJoin(join) = plan { + return Some(join); + } + plan.inputs().into_iter().find_map(find_asof_plan) +} + +fn has_filter_predicate(plan: &LogicalPlan, needle: &str) -> bool { + let found_here = match plan { + LogicalPlan::Filter(filter) => filter.predicate.to_string().contains(needle), + LogicalPlan::TableScan(scan) => scan + .filters + .iter() + .any(|filter| filter.to_string().contains(needle)), + _ => false, + }; + found_here + || plan + .inputs() + .into_iter() + .any(|input| has_filter_predicate(input, needle)) +} + +#[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).await?; + + 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_dataframe_api() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + let left = ctx.table("trades").await?; + let right = ctx.table("prices").await?; + let batches = left + .join_asof( + right, + vec![(col("symbol"), col("symbol"))], + AsOfMatch::new(col("ts"), Operator::GtEq, col("ts")), + )? + .select(vec![col("trade_id"), col("price")])? + .sort(vec![col("trade_id").sort(true, true)])? + .collect() + .await?; + assert_batches_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 101 |", + "| 5 | 106 |", + "| 6 | |", + "+----------+-------+", + ], + &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_without_equality_keys_is_single_partition() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx).await?; + let df = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices 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"); + assert_eq!(asof.output_partitioning().partition_count(), 1); + assert!(asof.output_ordering().is_some()); + assert!(matches!( + &asof.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); + Ok(()) +} + +#[tokio::test] +async fn asof_join_filter_pushdown_respects_selection_boundary() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + let plan = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts >= p.ts) ON t.symbol = p.symbol \ + WHERE t.trade_id > 1 AND p.price > 20", + ) + .await? + .into_optimized_plan()?; + let join = find_asof_plan(&plan).expect("logical ASOF join must be present"); + assert!(has_filter_predicate(join.left.as_ref(), "trade_id")); + assert!(!has_filter_predicate(join.right.as_ref(), "price")); + assert!(has_filter_predicate(&plan, "price")); + Ok(()) +} + +#[tokio::test] +async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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_merges_key_and_unparser_round_trips() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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!["symbol", "ts", "trade_id", "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?; + Ok(()) +} + +#[tokio::test] +async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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).await?; + 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 2ecb12c30afad..d5dda08bf6a45 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,43 @@ pub fn build_join_schema( dfschema.with_functional_dependencies(func_dependencies) } +/// Creates the schema for a left-preserving ASOF join. +/// +/// `ON` emits all left fields followed by nullable right fields. `USING` emits +/// each equality key once by omitting the corresponding right field. +pub fn build_asof_join_schema( + left: &DFSchema, + right: &DFSchema, + on: &[(Expr, Expr)], + join_constraint: JoinConstraint, +) -> Result { + let omitted_right_indices = if join_constraint == JoinConstraint::Using { + on.iter() + .map(|(_, right_expr)| { + let column = right_expr.get_as_join_column().ok_or_else(|| { + plan_datafusion_err!("ASOF USING keys must be columns") + })?; + right.index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + + let full_schema = build_join_schema(left, right, &JoinType::Left)?; + let left_len = left.fields().len(); + let fields = full_schema + .iter() + .enumerate() + .filter(|(index, _)| { + *index < left_len || !omitted_right_indices.contains(&(*index - left_len)) + }) + .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field))) + .collect(); + DFSchema::new_with_metadata(fields, full_schema.metadata().clone())? + .with_functional_dependencies(left.functional_dependencies().clone()) +} + /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise /// conflict with the columns from the other. /// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 09f41c94f64fa..c5cd003d1f34e 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -21,10 +21,10 @@ use std::collections::HashMap; use std::fmt; use crate::{ - Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, - Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, - Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values, - Window, expr_vec_fmt, + Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, + Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, + Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, + Values, Window, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -493,6 +493,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Filter": format!("{}", filter_expr) }) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let join_expr: Vec = + on.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + json!({ + "Node Type": "AsOf Join", + "Join Constraint": format!("{join_constraint:?}"), + "Join Keys": join_expr.join(", "), + "Match Condition": match_condition.to_string(), + }) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 4766c3f33379f..98113d12c1b4a 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -28,8 +28,8 @@ pub mod tree_node; pub use builder::{ LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE, - build_join_schema, requalify_sides_if_needed, table_scan, union, - wrap_projection_for_join_if_necessary, + build_asof_join_schema, build_join_schema, requalify_sides_if_needed, table_scan, + union, wrap_projection_for_join_if_necessary, }; pub use ddl::{ CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction, @@ -41,12 +41,12 @@ pub use dml::{ WriteOp, }; pub use plan::{ - Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, - EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, - Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, - Unnest, Values, Window, projection_schema, + Aggregate, Analyze, AsOfJoin, AsOfMatch, ColumnUnnestList, DescribeTable, Distinct, + DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, + Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, + Projection, RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, + StringifiedPlan, Subquery, SubqueryAlias, TableScan, TableScanBuilder, + ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9cfab21a0395e..cac596b8bd58a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,13 +41,15 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement}; use crate::utils::{ - enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, - grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, + 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; @@ -238,6 +240,9 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -341,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(), @@ -369,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()) @@ -459,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], @@ -494,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)| { @@ -567,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() } @@ -690,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, @@ -985,6 +1024,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, @@ -1410,6 +1488,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| { @@ -1460,6 +1539,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), @@ -1498,6 +1578,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), @@ -2126,6 +2207,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, .. @@ -4237,6 +4337,165 @@ 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. + 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(), &on, join_constraint)?; + 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. /// diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c10ac92eef4f5..daf31489b08c7 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -41,11 +41,11 @@ 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, 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, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -150,6 +150,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 +464,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 @@ -610,6 +634,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 afd4e980b5424..5fb234ace4ac7 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,9 +57,9 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + AsOfJoin, AsOfMatch, Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -175,6 +175,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), _ => Ok(plan), @@ -218,6 +219,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 2775d62144c56..1f1bcbe60c53f 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -566,6 +566,7 @@ impl OptimizerRule for CommonSubexprEliminate { LogicalPlan::Window(window) => self.try_optimize_window(window, config)?, LogicalPlan::Aggregate(agg) => self.try_optimize_aggregate(agg, config)?, LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 80aceb8cad44c..da5c715d280cb 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,8 +24,9 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, - get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, + Column, DFSchema, HashMap, JoinConstraint, JoinType, Result, + assert_eq_or_internal_err, get_required_group_by_exprs_indices, + internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; use datafusion_expr::{ @@ -407,6 +408,44 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + std::collections::HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect::>(); + 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 if let Some(right_index) = + right_output_indices.get(*index - left_len) + { + right_required.push(*right_index); + } + } + let left_indices = RequiredIndices::new_from_indices(left_required) + .with_plan_exprs(&plan, join.left.schema())?; + let right_indices = RequiredIndices::new_from_indices(right_required) + .with_plan_exprs(&plan, join.right.schema())?; + vec![ + left_indices.with_projection_beneficial(), + right_indices.with_projection_beneficial(), + ] + } // these nodes are explicitly rewritten in the match statement above LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index db7ad8475273a..0abb09bea8768 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -411,6 +411,11 @@ fn map_children_mut Result>( let r = f(Arc::make_mut(right))?; l || r } + LogicalPlan::AsOfJoin(join) => { + let l = f(Arc::make_mut(&mut join.left))?; + let r = f(Arc::make_mut(&mut join.right))?; + l || r + } LogicalPlan::Union(Union { inputs, .. }) => { let mut changed = false; for input in inputs { diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index f30b1187b7bca..b2ca187e05315 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1104,6 +1104,32 @@ impl OptimizerRule for PushDownFilter { result.map_data(|plan| Ok(with_filters(keep_predicates, plan))) } LogicalPlan::Join(join) => push_down_join(join, Some(filter.predicate)), + LogicalPlan::AsOfJoin(mut join) => { + let (push, keep): (Vec<_>, Vec<_>) = + split_conjunction_owned(filter.predicate) + .into_iter() + .partition(|predicate| { + !predicate.is_volatile() + && predicate.column_refs().iter().all(|column| { + join.left.schema().is_column_from_schema(column) + }) + }); + if push.is_empty() { + filter.predicate = + conjunction(keep).expect("filter predicates are not empty"); + filter.input = Arc::new(LogicalPlan::AsOfJoin(join)); + Ok(Transformed::no(LogicalPlan::Filter(filter))) + } else { + join.left = Arc::new(LogicalPlan::Filter(Filter::new( + conjunction(push).expect("push predicates are not empty"), + join.left, + ))); + Ok(Transformed::yes(with_filters( + keep, + LogicalPlan::AsOfJoin(join), + ))) + } + } LogicalPlan::TableScan(mut scan) => { let filter_predicates = split_conjunction(&filter.predicate); // Filters containing scalar subqueries cannot be pushed to diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 58c2f0d7da537..d243915c205cf 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -121,6 +121,11 @@ harness = false name = "sort_merge_join" required-features = ["test_utils"] +[[bench]] +harness = false +name = "asof_join" +required-features = ["test_utils"] + [[bench]] harness = false name = "aggregate_vectorized" diff --git a/datafusion/physical-plan/benches/asof_join.rs b/datafusion/physical-plan/benches/asof_join.rs new file mode 100644 index 0000000000000..da39080c94f90 --- /dev/null +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -0,0 +1,196 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmarks for the pre-sorted ASOF join kernel. + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, Int64Array, RecordBatch, StringArray, StringDictionaryBuilder, +}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::{TaskContext, config::SessionConfig}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::joins::{AsOfJoinExec, AsOfMatchExpr, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; +use tokio::runtime::Runtime; + +#[derive(Clone, Copy)] +enum Payload { + Int64, + WideUtf8, + Dictionary, +} + +impl Payload { + fn name(self) -> &'static str { + match self { + Self::Int64 => "int64", + Self::WideUtf8 => "wide_utf8", + Self::Dictionary => "dictionary", + } + } + + fn data_type(self) -> DataType { + match self { + Self::Int64 => DataType::Int64, + Self::WideUtf8 => DataType::Utf8, + Self::Dictionary => { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + } + } + + fn array(self, rows: &[(i64, i64, usize)]) -> ArrayRef { + match self { + Self::Int64 => Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, _, row)| *row as i64), + )), + Self::WideUtf8 => Arc::new(StringArray::from_iter_values( + rows.iter() + .map(|(_, _, row)| format!("row_{row:08}_{}", "x".repeat(244))), + )), + Self::Dictionary => { + let mut builder = StringDictionaryBuilder::::new(); + for (_, _, row) in rows { + builder.append_value(format!("category_{}", row % 64)); + } + Arc::new(builder.finish()) + } + } + } +} + +fn schema(payload: Payload) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", payload.data_type(), false), + ])) +} + +fn build_sorted_batches( + num_rows: usize, + num_groups: usize, + time_offset: i64, + payload: Payload, + schema: &SchemaRef, +) -> Vec { + let mut rows = (0..num_rows) + .map(|row| { + ( + (row % num_groups) as i64, + (row / num_groups) as i64 + time_offset, + row, + ) + }) + .collect::>(); + rows.sort_unstable_by_key(|(key, ts, _)| (*key, *ts)); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(key, _, _)| *key), + )), + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, ts, _)| *ts), + )), + payload.array(&rows), + ], + ) + .unwrap(); + + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(8192); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec(batches: &[RecordBatch], schema: &SchemaRef) -> Arc { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn do_join( + left: Arc, + right: Arc, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let left_match = col("ts", &left.schema()).unwrap(); + let right_match = col("ts", &right.schema()).unwrap(); + let join = AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, Operator::GtEq, right_match), + vec![2], + ) + .unwrap(); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(8192)), + ); + rt.block_on(async { + collect(Arc::new(join), task_ctx) + .await + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum() + }) +} + +fn bench_asof_join(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let num_rows = 100_000; + let num_groups = 10_000; + let mut group = c.benchmark_group("asof_join"); + + for payload in [Payload::Int64, Payload::WideUtf8, Payload::Dictionary] { + let schema = schema(payload); + let left_batches = + build_sorted_batches(num_rows, num_groups, 1, payload, &schema); + let right_batches = + build_sorted_batches(num_rows, num_groups, 0, payload, &schema); + group.bench_function( + BenchmarkId::new(payload.name(), format!("{num_rows}_rows_10_per_key")), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &schema); + let right = make_exec(&right_batches, &schema); + do_join(left, right, &rt) + }) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_asof_join); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs new file mode 100644 index 0000000000000..5945c8d750716 --- /dev/null +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -0,0 +1,1419 @@ +// 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. + +//! Ordered, left-preserving ASOF join execution. + +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; +use datafusion_common::utils::{ + compare_rows, get_row_at_idx, normalize_float_zero_scalar, +}; +use datafusion_common::{ + ColumnStatistics, JoinType, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{Partitioning, PhysicalSortExpr}; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{StreamExt, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::joins::utils::{JoinOn, build_join_schema}; +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricsSet, RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, + check_if_same_properties, +}; + +/// Physical ordered comparison for an ASOF join. +#[derive(Debug, Clone)] +pub struct AsOfMatchExpr { + /// Expression evaluated against the left input. + pub left: PhysicalExprRef, + /// Ordered comparison operator. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: PhysicalExprRef, +} + +impl AsOfMatchExpr { + /// Creates a physical ASOF match expression. + pub fn new(left: PhysicalExprRef, op: Operator, right: PhysicalExprRef) -> Self { + Self { left, op, right } + } +} + +/// A sort-merge ASOF join that emits exactly one row for every left row. +#[derive(Debug, Clone)] +pub struct AsOfJoinExec { + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + schema: SchemaRef, + metrics: ExecutionPlanMetricsSet, + left_ordering: LexOrdering, + right_ordering: LexOrdering, + cache: Arc, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + pub fn try_new( + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + ) -> Result { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "AsOfJoinExec requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + if left.boundedness().is_unbounded() || right.boundedness().is_unbounded() { + return plan_err!("AsOfJoinExec requires bounded inputs"); + } + if is_volatile(&match_condition.left) || is_volatile(&match_condition.right) { + return plan_err!("AsOfJoinExec match expression must be deterministic"); + } + if on + .iter() + .any(|(left, right)| is_volatile(left) || is_volatile(right)) + { + return plan_err!("AsOfJoinExec equality expressions must be deterministic"); + } + + let left_schema = left.schema(); + let right_schema = right.schema(); + validate_expr_side(&match_condition.left, &left_schema, "left match")?; + validate_expr_side(&match_condition.right, &right_schema, "right match")?; + for (left_expr, right_expr) in &on { + validate_expr_side(left_expr, &left_schema, "left equality")?; + validate_expr_side(right_expr, &right_schema, "right equality")?; + let left_type = left_expr.data_type(&left_schema)?; + let right_type = right_expr.data_type(&right_schema)?; + if left_type != right_type { + return plan_err!( + "AsOfJoinExec equality expression types differ: {left_type} and {right_type}" + ); + } + if !datafusion_expr::utils::can_hash(&left_type) { + return plan_err!( + "AsOfJoinExec equality expressions have unsupported hash type {left_type}" + ); + } + } + let left_match_type = match_condition.left.data_type(&left_schema)?; + let right_match_type = match_condition.right.data_type(&right_schema)?; + if left_match_type != right_match_type { + return plan_err!( + "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" + ); + } + if let Some(index) = right_output_indices + .iter() + .find(|index| **index >= right_schema.fields().len()) + { + return plan_err!( + "AsOfJoinExec right output index {index} is outside schema with {} fields", + right_schema.fields().len() + ); + } + if !right_output_indices + .windows(2) + .all(|pair| pair[0] < pair[1]) + { + return plan_err!( + "AsOfJoinExec right output indices must be strictly increasing" + ); + } + + let schema = + build_output_schema(&left_schema, &right_schema, &right_output_indices); + let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); + let equality_options = SortOptions { + descending: false, + nulls_first: true, + }; + let match_options = SortOptions { + descending, + nulls_first: true, + }; + let mut left_sort_exprs = on + .iter() + .map(|(left, _)| PhysicalSortExpr { + expr: Arc::clone(left), + options: equality_options, + }) + .collect::>(); + left_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.left), + options: match_options, + }); + let mut right_sort_exprs = on + .iter() + .map(|(_, right)| PhysicalSortExpr { + expr: Arc::clone(right), + options: equality_options, + }) + .collect::>(); + right_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.right), + options: match_options, + }); + let left_ordering = LexOrdering::new(left_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF left ordering must not be empty" + ) + })?; + let right_ordering = LexOrdering::new(right_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF right ordering must not be empty" + ) + })?; + let cache = Arc::new(Self::compute_properties(&left, &schema, on.is_empty())?); + + Ok(Self { + left, + right, + on, + match_condition, + right_output_indices, + schema, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + cache, + }) + } + + fn compute_properties( + left: &Arc, + schema: &SchemaRef, + single_partition: bool, + ) -> Result { + let left_schema = left.schema(); + let mapping = ProjectionMapping::try_new( + left_schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + ( + Arc::new(PhysicalColumn::new(field.name(), index)) + as PhysicalExprRef, + field.name().to_string(), + ) + }), + &left_schema, + )?; + let input_eq_properties = left.equivalence_properties(); + let eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); + let output_partitioning = if single_partition { + Partitioning::UnknownPartitioning(1) + } else { + left.output_partitioning() + .project(&mapping, input_eq_properties) + }; + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + + /// Equality expressions. + pub fn on(&self) -> &JoinOn { + &self.on + } + + /// Ordered match expression. + pub fn match_condition(&self) -> &AsOfMatchExpr { + &self.match_condition + } + + /// Indices of right input columns emitted after the left columns. + pub fn right_output_indices(&self) -> &[usize] { + &self.right_output_indices + } + + /// Left input. + pub fn left(&self) -> &Arc { + &self.left + } + + /// Right input. + pub fn right(&self) -> &Arc { + &self.right + } +} + +fn build_output_schema( + left: &SchemaRef, + right: &SchemaRef, + right_output_indices: &[usize], +) -> SchemaRef { + let full_schema = build_join_schema(left, right, &JoinType::Left).0; + let left_len = left.fields().len(); + let fields = full_schema + .fields() + .iter() + .take(left_len) + .cloned() + .chain( + right_output_indices + .iter() + .map(|index| Arc::clone(&full_schema.fields()[left_len + *index])), + ) + .collect::>(); + Arc::new(Schema::new_with_metadata( + fields, + full_schema.metadata().clone(), + )) +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + let on = self + .on + .iter() + .map(|(left, right)| { + format!("({} = {})", fmt_sql(left.as_ref()), fmt_sql(right.as_ref())) + }) + .collect::>() + .join(", "); + let match_condition = format!( + "{} {} {}", + fmt_sql(self.match_condition.left.as_ref()), + self.match_condition.op, + fmt_sql(self.match_condition.right.as_ref()) + ); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]", + Self::static_name(), + on, + match_condition + ), + DisplayFormatType::TreeRender => { + writeln!(f, "on={on}")?; + writeln!(f, "match={match_condition}") + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.on.is_empty() { + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) + } else { + let (left, right) = self + .on + .iter() + .map(|(left, right)| (Arc::clone(left), Arc::clone(right))) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left), + Distribution::KeyPartitioned(right), + ]) + } + } + + fn required_input_ordering(&self) -> Vec> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec { + vec![true, false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + check_if_same_properties!(self, children); + match &children[..] { + [left, right] => Ok(Arc::new(Self::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.match_condition.clone(), + self.right_output_indices.clone(), + )?)), + _ => internal_err!("AsOfJoinExec requires two children"), + } + } + + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + assert_eq_or_internal_err!( + children.len(), + 2, + "AsOfJoinExec requires two children" + ); + let left = children.remove(0); + let right = children.remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&self) + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + left_partitions, + right_partitions, + "AsOfJoinExec partition count mismatch: {left_partitions} != {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let right_stream = self.right.execute(partition, Arc::clone(&context))?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let state = AsOfJoinStreamState::new( + Arc::clone(&self.schema), + InputCursor::new( + left_stream, + left_keys, + Arc::clone(&self.match_condition.left), + ), + InputCursor::new( + right_stream, + right_keys, + Arc::clone(&self.match_condition.right), + ), + self.match_condition.op, + self.right_output_indices.clone(), + context.session_config().batch_size(), + AsOfJoinMetrics::new(partition, &self.metrics), + ); + let stream = stream::try_unfold(state, |mut state| async move { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, state))), + None => Ok(None), + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition), ChildStats::Skip] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let left = &input_stats[0]; + let mut column_statistics = left.column_statistics.clone(); + column_statistics.truncate(self.left.schema().fields().len()); + column_statistics.resize_with( + self.left.schema().fields().len(), + ColumnStatistics::new_unknown, + ); + column_statistics.extend( + self.right_output_indices + .iter() + .map(|_| ColumnStatistics::new_unknown()), + ); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec, + _config: &ConfigOptions, + ) -> Result { + let left_indices = (0..self.left.schema().fields().len()).collect::>(); + let left = ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + left_indices, + &self.left, + )?; + let right = ChildFilterDescription::all_unsupported(&parent_filters); + Ok(FilterDescription::new().with_child(left).with_child(right)) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result>> { + Ok(FilterPushdownPropagation::if_any(child_pushdown_result)) + } +} + +#[derive(Clone)] +struct Candidate { + batch: Arc, + row: usize, + group: Vec, +} + +struct InputCursor { + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + batch: Option>, + key_arrays: Vec, + match_array: Option, + row: usize, + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: vec![], + match_array: None, + row: 0, + eof: false, + } + } + + async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result { + loop { + if let Some(batch) = &self.batch + && self.row < batch.num_rows() + { + return Ok(true); + } + self.batch = None; + self.key_arrays.clear(); + self.match_array = None; + self.row = 0; + if self.eof { + return Ok(false); + } + let Some(batch) = self.stream.next().await.transpose()? else { + self.eof = true; + return Ok(false); + }; + if batch.num_rows() == 0 { + continue; + } + let batch = Arc::new(batch); + let _timer = elapsed_compute.timer(); + self.key_arrays = self + .key_exprs + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::>()?; + self.match_array = Some( + self.match_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?, + ); + self.batch = Some(batch); + } + } + + fn group(&self) -> Result> { + get_row_at_idx(&self.key_arrays, self.row) + .map(|row| row.into_iter().map(normalize_float_zero_scalar).collect()) + } + + fn match_value(&self) -> Result { + let array = self.match_array.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF match array is missing") + })?; + ScalarValue::try_from_array(array, self.row).map(normalize_float_zero_scalar) + } + + fn batch_row(&self) -> Result<(Arc, usize)> { + let batch = self.batch.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF input batch is missing") + })?; + Ok((Arc::clone(batch), self.row)) + } + + fn advance(&mut self) { + self.row += 1; + } +} + +struct AsOfJoinMetrics { + baseline: BaselineMetrics, + matched_rows: Count, + unmatched_left_rows: Count, +} + +impl AsOfJoinMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + matched_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("matched_rows", partition), + unmatched_left_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("unmatched_left_rows", partition), + } + } +} + +#[derive(Default)] +struct PendingRows { + sources: Vec>, + source_by_ptr: HashMap, + indices: Vec>, +} + +impl PendingRows { + fn len(&self) -> usize { + self.indices.len() + } + + fn is_empty(&self) -> bool { + self.indices.is_empty() + } + + fn push(&mut self, batch: Arc, row: usize) { + let ptr = Arc::as_ptr(&batch) as usize; + let source = *self.source_by_ptr.entry(ptr).or_insert_with(|| { + let source = self.sources.len(); + self.sources.push(batch); + source + }); + self.indices.push(Some((source, row))); + } + + fn push_null(&mut self) { + self.indices.push(None); + } + + fn materialize_column( + &self, + source_column: usize, + data_type: &arrow::datatypes::DataType, + ) -> Result { + if self.indices.is_empty() { + return internal_err!("ASOF output materialization has no pending rows"); + } + + if self.sources.len() == 1 + && self.indices.iter().all(Option::is_some) + && let Some((0, first_row)) = self.indices[0] + && self + .indices + .iter() + .enumerate() + .all(|(offset, index)| *index == Some((0, first_row + offset))) + { + return Ok(self.sources[0] + .column(source_column) + .slice(first_row, self.indices.len())); + } + + let has_null = self.indices.iter().any(Option::is_none); + let null_array = has_null.then(|| new_null_array(data_type, 1)); + let mut source_arrays: Vec<&dyn Array> = + Vec::with_capacity(self.sources.len() + usize::from(has_null)); + if let Some(null_array) = &null_array { + source_arrays.push(null_array.as_ref()); + } + source_arrays.extend( + self.sources + .iter() + .map(|batch| batch.column(source_column).as_ref()), + ); + let source_offset = usize::from(has_null); + let interleave_indices = self + .indices + .iter() + .map(|index| match index { + Some((source, row)) => (source + source_offset, *row), + None => (0, 0), + }) + .collect::>(); + interleave(&source_arrays, &interleave_indices).map_err(Into::into) + } + + fn clear(&mut self) { + self.sources.clear(); + self.source_by_ptr.clear(); + self.indices.clear(); + } +} + +struct AsOfJoinStreamState { + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + candidate: Option, + group_sort_options: Vec, + pending_left: PendingRows, + pending_right: PendingRows, + batch_size: usize, + metrics: AsOfJoinMetrics, +} + +impl AsOfJoinStreamState { + fn new( + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + batch_size: usize, + metrics: AsOfJoinMetrics, + ) -> Self { + let group_sort_options = vec![ + SortOptions { + descending: false, + nulls_first: true, + }; + left.key_exprs.len() + ]; + Self { + pending_left: PendingRows::default(), + pending_right: PendingRows::default(), + schema, + left, + right, + op, + right_output_indices, + candidate: None, + group_sort_options, + batch_size: batch_size.max(1), + metrics, + } + } + + async fn next_batch(&mut self) -> Result> { + loop { + if self.pending_left.len() >= self.batch_size { + return self.flush().map(Some); + } + if !self + .left + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + if !self.pending_left.is_empty() { + return self.flush().map(Some); + } + self.metrics.baseline.done(); + return Ok(None); + } + + let (left_group, left_match) = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + (self.left.group()?, self.left.match_value()?) + }; + if left_match.is_null() || left_group.iter().any(ScalarValue::is_null) { + self.candidate = None; + self.push_current_left(None)?; + self.left.advance(); + continue; + } + let candidate_is_other_group = if let Some(candidate) = &self.candidate { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + compare_rows(&candidate.group, &left_group, &self.group_sort_options)? + != Ordering::Equal + } else { + false + }; + if candidate_is_other_group { + self.candidate = None; + } + + loop { + if !self + .right + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + break; + } + let action = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_group = self.right.group()?; + if right_group.iter().any(ScalarValue::is_null) { + RightAction::Advance + } else { + match compare_rows( + &right_group, + &left_group, + &self.group_sort_options, + )? { + Ordering::Less => RightAction::Advance, + Ordering::Greater => RightAction::Stop, + Ordering::Equal => { + let right_match = self.right.match_value()?; + if right_match.is_null() { + RightAction::Advance + } else if is_eligible(self.op, &left_match, &right_match)? + { + let (batch, row) = self.right.batch_row()?; + RightAction::Candidate(Candidate { + batch, + row, + group: right_group, + }) + } else { + RightAction::Stop + } + } + } + } + }; + match action { + RightAction::Advance => self.right.advance(), + RightAction::Candidate(candidate) => { + self.candidate = Some(candidate); + self.right.advance(); + } + RightAction::Stop => break, + } + } + + self.push_current_left(self.candidate.clone())?; + self.left.advance(); + } + } + + fn push_current_left(&mut self, candidate: Option) -> Result<()> { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let (left_batch, left_row) = self.left.batch_row()?; + self.pending_left.push(left_batch, left_row); + match candidate { + Some(candidate) => { + if !self.right_output_indices.is_empty() { + self.pending_right.push(candidate.batch, candidate.row); + } + self.metrics.matched_rows.add(1); + } + None => { + if !self.right_output_indices.is_empty() { + self.pending_right.push_null(); + } + self.metrics.unmatched_left_rows.add(1); + } + } + Ok(()) + } + + fn flush(&mut self) -> Result { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let left_len = self.schema.fields().len() - self.right_output_indices.len(); + let mut arrays = Vec::with_capacity(self.schema.fields().len()); + for index in 0..left_len { + arrays.push( + self.pending_left + .materialize_column(index, self.schema.field(index).data_type())?, + ); + } + for (offset, source_index) in self.right_output_indices.iter().enumerate() { + arrays.push(self.pending_right.materialize_column( + *source_index, + self.schema.field(left_len + offset).data_type(), + )?); + } + self.pending_left.clear(); + self.pending_right.clear(); + let batch = RecordBatch::try_new(Arc::clone(&self.schema), arrays)?; + (&batch).record_output(&self.metrics.baseline); + Ok(batch) + } +} + +fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Result<()> { + let columns = collect_columns(expr); + if columns.is_empty() { + return plan_err!("AsOfJoinExec {name} expression must reference its input"); + } + if let Some(column) = columns.iter().find(|column| { + schema + .fields() + .get(column.index()) + .is_none_or(|field| field.name() != column.name()) + }) { + return plan_err!( + "AsOfJoinExec {name} expression references column {column} outside its input" + ); + } + Ok(()) +} + +enum RightAction { + Advance, + Candidate(Candidate), + Stop, +} + +fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result { + let ordering = right.try_cmp(left)?; + Ok(match op { + Operator::Gt => ordering == Ordering::Less, + Operator::GtEq => ordering != Ordering::Greater, + Operator::Lt => ordering == Ordering::Greater, + Operator::LtEq => ordering != Ordering::Less, + _ => false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::test::TestMemoryExec; + use arrow::array::{ + DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Int8Type}; + use datafusion_execution::config::SessionConfig; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr_common::metrics::MetricValue; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + struct VolatileExpr; + + impl std::fmt::Display for VolatileExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile") + } + } + + impl PhysicalExpr for VolatileExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Int64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, _batch: &RecordBatch) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn is_volatile_node(&self) -> bool { + true + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile()") + } + } + + fn make_batch( + schema: &SchemaRef, + keys: Vec>, + times: Vec>, + values: Vec, + ) -> Result { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from(keys)), + Arc::new(Int64Array::from(times)), + Arc::new(Int32Array::from(values)), + ], + ) + .map_err(Into::into) + } + + fn test_exec() -> Result> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let left_batches = vec![ + RecordBatch::new_empty(Arc::clone(&left_schema)), + make_batch(&left_schema, vec![None], vec![Some(3)], vec![0])?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![None, Some(1)], + vec![1, 2], + )?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(4), Some(7)], + vec![3, 4], + )?, + make_batch( + &left_schema, + vec![Some("B"), Some("C")], + vec![Some(2), Some(3)], + vec![5, 6], + )?, + ]; + let left = TestMemoryExec::try_new_exec( + &[left_batches], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let right_batches = vec![ + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![None, Some("A")], + vec![Some(2), None], + vec![999, 777], + )?, + make_batch(&right_schema, vec![Some("A")], vec![Some(2)], vec![20])?, + make_batch(&right_schema, vec![Some("A")], vec![Some(4)], vec![40])?, + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![Some("A"), Some("B")], + vec![Some(6), Some(1)], + vec![60, 101], + )?, + ]; + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let on: JoinOn = vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )]; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?)) + } + + #[test] + fn eligibility_matches_public_semantics() -> Result<()> { + let left = ScalarValue::Int64(Some(10)); + let lower = ScalarValue::Int64(Some(9)); + let equal = ScalarValue::Int64(Some(10)); + let higher = ScalarValue::Int64(Some(11)); + assert!(is_eligible(Operator::Gt, &left, &lower)?); + assert!(!is_eligible(Operator::Gt, &left, &equal)?); + assert!(is_eligible(Operator::GtEq, &left, &equal)?); + assert!(is_eligible(Operator::Lt, &left, &higher)?); + assert!(!is_eligible(Operator::Lt, &left, &equal)?); + assert!(is_eligible(Operator::LtEq, &left, &equal)?); + Ok(()) + } + + #[tokio::test] + async fn state_survives_empty_input_batches_and_output_flushes() -> Result<()> { + let exec = test_exec()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let batches = collect(Arc::clone(&exec) as _, context).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 2, 2, 1] + ); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!( + ids, + vec![ + Some(0), + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + ] + ); + assert_eq!( + prices, + vec![None, None, None, Some(40), Some(60), Some(101), None] + ); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!(metrics.output_rows(), Some(7)); + assert_eq!( + metrics + .sum_by_name("matched_rows") + .map(|value| value.as_usize()), + Some(3) + ); + assert_eq!( + metrics + .sum_by_name("unmatched_left_rows") + .map(|value| value.as_usize()), + Some(4) + ); + assert!(metrics.elapsed_compute().is_some()); + assert!( + metrics.iter().any(|metric| { + matches!(metric.value(), MetricValue::ElapsedCompute(_)) + }) + ); + Ok(()) + } + + #[tokio::test] + async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut left_payload = StringDictionaryBuilder::::new(); + for _ in 0..129 { + left_payload.append_value("left"); + } + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(StringArray::from(vec!["A"; 129])), + Arc::new(Int64Array::from_iter_values(-1..128)), + Arc::new(left_payload.finish()), + ], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut right_payload = StringDictionaryBuilder::::new(); + right_payload.append_value("right"); + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(Int64Array::from(vec![0])), + Arc::new(right_payload.finish()), + ], + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(256)), + ); + let batches = collect(exec, context).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 129); + assert_eq!(batches[0].column(2).data_type(), &dictionary_type); + assert_eq!(batches[0].column(3).data_type(), &dictionary_type); + + let right_output = batches[0] + .column(3) + .as_any() + .downcast_ref::>() + .expect("right output must remain Dictionary(Int8, Utf8)"); + assert!(right_output.is_null(0)); + assert_eq!(right_output.null_count(), 1); + let values = right_output + .values() + .as_any() + .downcast_ref::() + .expect("dictionary values must be Utf8"); + for row in 1..129 { + assert_eq!( + values.value(right_output.keys().value(row) as usize), + "right" + ); + } + Ok(()) + } + + #[test] + fn rejects_volatile_physical_expressions() -> Result<()> { + let exec = test_exec()?; + let volatile = Arc::new(VolatileExpr) as PhysicalExprRef; + let match_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + exec.on().clone(), + AsOfMatchExpr::new( + Arc::clone(&volatile), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + ) + .expect_err("volatile match expression must be rejected"); + assert!(match_error.to_string().contains("must be deterministic")); + + let equality_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![(volatile, Arc::new(PhysicalColumn::new("key", 0)))], + exec.match_condition().clone(), + vec![2], + ) + .expect_err("volatile equality expression must be rejected"); + assert!(equality_error.to_string().contains("must be deterministic")); + Ok(()) + } + + #[test] + fn properties_and_statistics_follow_left_preserving_contract() -> Result<()> { + let exec = test_exec()?; + let exec_plan: Arc = Arc::clone(&exec) as _; + assert_eq!(exec.maintains_input_order(), vec![true, false]); + assert_eq!(exec_plan.pipeline_behavior(), EmissionType::Incremental); + assert_eq!(exec_plan.boundedness(), Boundedness::Bounded); + assert!(matches!( + &exec.input_distribution_requirements().into_per_child()[..], + [ + Distribution::KeyPartitioned(_), + Distribution::KeyPartitioned(_) + ] + )); + for ordering in exec.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 2); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + assert_eq!( + requirement[1].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + } + + let no_keys: Arc = Arc::new(AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::Lt, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + assert_eq!(no_keys.output_partitioning().partition_count(), 1); + assert!(matches!( + &no_keys.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + for ordering in no_keys.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 1); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: true, + nulls_first: true, + }) + ); + } + + let mut key_stats = ColumnStatistics::new_unknown(); + key_stats.null_count = Precision::Exact(1); + key_stats.distinct_count = Precision::Exact(4); + let mut ts_stats = ColumnStatistics::new_unknown(); + ts_stats.min_value = Precision::Exact(ScalarValue::Int64(Some(1))); + ts_stats.max_value = Precision::Exact(ScalarValue::Int64(Some(7))); + let mut id_stats = ColumnStatistics::new_unknown(); + id_stats.null_count = Precision::Exact(0); + id_stats.distinct_count = Precision::Exact(7); + let left_column_statistics = vec![key_stats, ts_stats, id_stats]; + let left_stats = Arc::new(Statistics { + num_rows: Precision::Exact(7), + total_byte_size: Precision::Exact(128), + column_statistics: left_column_statistics.clone(), + }); + let right_stats = Arc::new(Statistics::new_unknown(&exec.right().schema())); + let stats = exec + .statistics_from_inputs(&[left_stats, right_stats], &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(7)); + assert_eq!(stats.total_byte_size, Precision::Absent); + assert_eq!(stats.column_statistics.len(), 4); + assert_eq!( + &stats.column_statistics[..3], + left_column_statistics.as_slice() + ); + assert_eq!(stats.column_statistics[3], ColumnStatistics::new_unknown()); + assert_eq!( + exec.child_stats_requests(None), + vec![ChildStats::At(None), ChildStats::Skip] + ); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index bbb25dda65165..820f60b09b3ba 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -18,6 +18,7 @@ //! DataFusion Join implementations use arrow::array::BooleanBufferBuilder; +pub use asof_join::{AsOfJoinExec, AsOfMatchExpr}; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ @@ -29,6 +30,7 @@ use parking_lot::Mutex; pub use piecewise_merge_join::PiecewiseMergeJoinExec; pub use sort_merge_join::SortMergeJoinExec; pub use symmetric_hash_join::SymmetricHashJoinExec; +mod asof_join; pub mod chain; mod cross_join; mod hash_join; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 205cf89abed1b..f61a59f058cbd 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -63,6 +63,7 @@ message LogicalPlanNode { CteWorkTableScanNode cte_work_table_scan = 32; DmlNode dml = 33; EmptyTableScanNode empty_table_scan = 34; + AsOfJoinNode as_of_join = 35; } } @@ -276,6 +277,25 @@ message JoinNode { bool null_aware = 9; } +enum AsOfMatchOperator { + AS_OF_MATCH_OPERATOR_UNSPECIFIED = 0; + AS_OF_MATCH_OPERATOR_LT = 1; + AS_OF_MATCH_OPERATOR_LT_EQ = 2; + AS_OF_MATCH_OPERATOR_GT = 3; + AS_OF_MATCH_OPERATOR_GT_EQ = 4; +} + +message AsOfJoinNode { + LogicalPlanNode left = 1; + LogicalPlanNode right = 2; + repeated LogicalExprNode left_join_key = 3; + repeated LogicalExprNode right_join_key = 4; + LogicalExprNode left_match_expr = 5; + LogicalExprNode right_match_expr = 6; + AsOfMatchOperator match_operator = 7; + datafusion_common.JoinConstraint join_constraint = 8; +} + message DistinctNode { LogicalPlanNode input = 1; } @@ -881,6 +901,7 @@ message PhysicalPlanNode { BufferExecNode buffer = 37; ArrowScanExecNode arrow_scan = 38; ScalarSubqueryExecNode scalar_subquery = 39; + AsOfJoinExecNode as_of_join = 40; } } @@ -1637,6 +1658,16 @@ message SortMergeJoinExecNode { datafusion_common.NullEquality null_equality = 7; } +message AsOfJoinExecNode { + PhysicalPlanNode left = 1; + PhysicalPlanNode right = 2; + repeated JoinOn on = 3; + PhysicalExprNode left_match_expr = 4; + PhysicalExprNode right_match_expr = 5; + AsOfMatchOperator match_operator = 6; + repeated uint32 right_output_indices = 7; +} + message AsyncFuncExecNode { PhysicalPlanNode input = 1; repeated PhysicalExprNode async_exprs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d23f8eee5fd2c..d1f5c8d6985fa 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1520,6 +1520,508 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode { deserializer.deserialize_struct("datafusion.ArrowScanExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for AsOfJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.on.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if !self.right_output_indices.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.on.is_empty() { + struct_ser.serialize_field("on", &self.on)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if !self.right_output_indices.is_empty() { + struct_ser.serialize_field("rightOutputIndices", &self.right_output_indices)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "on", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "right_output_indices", + "rightOutputIndices", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + On, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + RightOutputIndices, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "on" => Ok(GeneratedField::On), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "rightOutputIndices" | "right_output_indices" => Ok(GeneratedField::RightOutputIndices), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut on__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut right_output_indices__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::RightOutputIndices => { + if right_output_indices__.is_some() { + return Err(serde::de::Error::duplicate_field("rightOutputIndices")); + } + right_output_indices__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(AsOfJoinExecNode { + left: left__, + right: right__, + on: on__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + right_output_indices: right_output_indices__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfJoinNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.left_join_key.is_empty() { + len += 1; + } + if !self.right_join_key.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if self.join_constraint != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.left_join_key.is_empty() { + struct_ser.serialize_field("leftJoinKey", &self.left_join_key)?; + } + if !self.right_join_key.is_empty() { + struct_ser.serialize_field("rightJoinKey", &self.right_join_key)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if self.join_constraint != 0 { + let v = super::datafusion_common::JoinConstraint::try_from(self.join_constraint) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_constraint)))?; + struct_ser.serialize_field("joinConstraint", &v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "left_join_key", + "leftJoinKey", + "right_join_key", + "rightJoinKey", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "join_constraint", + "joinConstraint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + LeftJoinKey, + RightJoinKey, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + JoinConstraint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "leftJoinKey" | "left_join_key" => Ok(GeneratedField::LeftJoinKey), + "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "joinConstraint" | "join_constraint" => Ok(GeneratedField::JoinConstraint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut left_join_key__ = None; + let mut right_join_key__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut join_constraint__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::LeftJoinKey => { + if left_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("leftJoinKey")); + } + left_join_key__ = Some(map_.next_value()?); + } + GeneratedField::RightJoinKey => { + if right_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("rightJoinKey")); + } + right_join_key__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::JoinConstraint => { + if join_constraint__.is_some() { + return Err(serde::de::Error::duplicate_field("joinConstraint")); + } + join_constraint__ = Some(map_.next_value::()? as i32); + } + } + } + Ok(AsOfJoinNode { + left: left__, + right: right__, + left_join_key: left_join_key__.unwrap_or_default(), + right_join_key: right_join_key__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + join_constraint: join_constraint__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfMatchOperator { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for AsOfMatchOperator { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + "AS_OF_MATCH_OPERATOR_LT", + "AS_OF_MATCH_OPERATOR_LT_EQ", + "AS_OF_MATCH_OPERATOR_GT", + "AS_OF_MATCH_OPERATOR_GT_EQ", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = AsOfMatchOperator; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Ok(AsOfMatchOperator::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Ok(AsOfMatchOperator::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Ok(AsOfMatchOperator::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Ok(AsOfMatchOperator::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Ok(AsOfMatchOperator::GtEq), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for AsyncFuncExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -13674,6 +14176,9 @@ impl serde::Serialize for LogicalPlanNode { logical_plan_node::LogicalPlanType::EmptyTableScan(v) => { struct_ser.serialize_field("emptyTableScan", v)?; } + logical_plan_node::LogicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -13735,6 +14240,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "dml", "empty_table_scan", "emptyTableScan", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -13772,6 +14279,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { CteWorkTableScan, Dml, EmptyTableScan, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -13826,6 +14334,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "cteWorkTableScan" | "cte_work_table_scan" => Ok(GeneratedField::CteWorkTableScan), "dml" => Ok(GeneratedField::Dml), "emptyTableScan" | "empty_table_scan" => Ok(GeneratedField::EmptyTableScan), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14077,6 +14586,13 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { return Err(serde::de::Error::duplicate_field("emptyTableScan")); } logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::EmptyTableScan) +; + } + GeneratedField::AsOfJoin => { + if logical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::AsOfJoin) ; } } @@ -20331,6 +20847,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::ScalarSubquery(v) => { struct_ser.serialize_field("scalarSubquery", v)?; } + physical_plan_node::PhysicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -20403,6 +20922,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "arrowScan", "scalar_subquery", "scalarSubquery", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -20445,6 +20966,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { Buffer, ArrowScan, ScalarSubquery, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20504,6 +21026,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "buffer" => Ok(GeneratedField::Buffer), "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -20790,6 +21313,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("scalarSubquery")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::ScalarSubquery) +; + } + GeneratedField::AsOfJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::AsOfJoin) ; } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 6baabbf37a41c..0ece21da6dcb8 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -5,7 +5,7 @@ pub struct LogicalPlanNode { #[prost( oneof = "logical_plan_node::LogicalPlanType", - tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" + tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35" )] pub logical_plan_type: ::core::option::Option, } @@ -79,6 +79,8 @@ pub mod logical_plan_node { Dml(::prost::alloc::boxed::Box), #[prost(message, tag = "34")] EmptyTableScan(super::EmptyTableScanNode), + #[prost(message, tag = "35")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -421,6 +423,29 @@ pub struct JoinNode { pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub left_join_key: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "4")] + pub right_join_key: ::prost::alloc::vec::Vec, + #[prost(message, optional, boxed, tag = "5")] + pub left_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(message, optional, boxed, tag = "6")] + pub right_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(enumeration = "AsOfMatchOperator", tag = "7")] + pub match_operator: i32, + #[prost(enumeration = "super::datafusion_common::JoinConstraint", tag = "8")] + pub join_constraint: i32, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1292,7 +1317,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40" )] pub physical_plan_type: ::core::option::Option, } @@ -1378,6 +1403,8 @@ pub mod physical_plan_node { ArrowScan(super::ArrowScanExecNode), #[prost(message, tag = "39")] ScalarSubquery(::prost::alloc::boxed::Box), + #[prost(message, tag = "40")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2473,6 +2500,23 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub on: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "4")] + pub left_match_expr: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub right_match_expr: ::core::option::Option, + #[prost(enumeration = "AsOfMatchOperator", tag = "6")] + pub match_operator: i32, + #[prost(uint32, repeated, tag = "7")] + pub right_output_indices: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -2504,6 +2548,41 @@ pub struct PhysicalScalarSubqueryExprNode { #[prost(uint32, tag = "3")] pub index: u32, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum AsOfMatchOperator { + Unspecified = 0, + Lt = 1, + LtEq = 2, + Gt = 3, + GtEq = 4, +} +impl AsOfMatchOperator { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Some(Self::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Some(Self::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Some(Self::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Some(Self::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Some(Self::GtEq), + _ => None, + } + } +} /// Identifies a built-in file format supported by DataFusion. /// Used by DefaultLogicalExtensionCodec to serialize/deserialize /// FileFormatFactory instances (e.g. in CopyTo plans). diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 6d9a73e06ff45..8e6fe2f3d9c29 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -234,6 +234,22 @@ impl FromProto for JoinConstraint { } } +impl TryFromProto for Operator { + type Error = Error; + + fn try_from_proto(value: protobuf::AsOfMatchOperator) -> Result { + match value { + protobuf::AsOfMatchOperator::Lt => Ok(Self::Lt), + protobuf::AsOfMatchOperator::LtEq => Ok(Self::LtEq), + protobuf::AsOfMatchOperator::Gt => Ok(Self::Gt), + protobuf::AsOfMatchOperator::GtEq => Ok(Self::GtEq), + protobuf::AsOfMatchOperator::Unspecified => Err(Error::General( + "ASOF match operator must be specified".to_string(), + )), + } + } +} + impl FromProto for NullEquality { fn from_proto(t: protobuf::NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 732676a3c0a0f..0fa9d501cd804 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -60,17 +60,17 @@ use datafusion_datasource_json::file_format::{ use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, Operator, RangePartitioning, RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, }; use datafusion_expr::{ DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ - Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection, - Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window, - builder::project, + Aggregate, AsOfJoin, AsOfMatch, CreateCatalog, CreateCatalogSchema, + CreateExternalTable, CreateView, DdlStatement, Distinct, EmptyRelation, + Extension, Join, Prepare, Projection, Repartition, Sort, SubqueryAlias, + TableScan, TableScanBuilder, Values, Window, builder::project, }, }; use datafusion_proto_common::protobuf_common; @@ -1033,6 +1033,61 @@ impl AsLogicalPlan for LogicalPlanNode { join.null_aware, )?)) } + LogicalPlanType::AsOfJoin(join) => { + let left_keys = + from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; + let right_keys = + from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + if left_keys.len() != right_keys.len() { + return Err(proto_error(format!( + "Received an AsOfJoinNode with left_join_key and right_join_key of different lengths: {} and {}", + left_keys.len(), + right_keys.len() + ))); + } + let left_match = from_proto::parse_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode left_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let right_match = from_proto::parse_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode right_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from( + join.match_operator, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + let join_constraint = protobuf::JoinConstraint::try_from( + join.join_constraint, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF JoinConstraint {}", + join.join_constraint + )) + })?; + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + left_keys.into_iter().zip(right_keys).collect(), + AsOfMatch::new(left_match, op, right_match), + JoinConstraint::from_proto(join_constraint), + )?)) + } LogicalPlanType::Union(union) => { assert_or_internal_err!( union.inputs.len() >= 2, @@ -1685,6 +1740,59 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + .. + }) => { + let left = LogicalPlanNode::try_from_logical_plan( + left.as_ref(), + extension_codec, + )?; + let right = LogicalPlanNode::try_from_logical_plan( + right.as_ref(), + extension_codec, + )?; + let (left_join_key, right_join_key) = on + .iter() + .map(|(left, right)| { + Ok(( + serialize_expr(left, extension_codec)?, + serialize_expr(right, extension_codec)?, + )) + }) + .collect::, ToProtoError>>()? + .into_iter() + .unzip(); + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(match_condition.op)?; + Ok(LogicalPlanNode { + logical_plan_type: Some(LogicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + left_join_key, + right_join_key, + left_match_expr: Some(Box::new(serialize_expr( + &match_condition.left, + extension_codec, + )?)), + right_match_expr: Some(Box::new(serialize_expr( + &match_condition.right, + extension_codec, + )?)), + match_operator: match_operator.into(), + join_constraint: protobuf::JoinConstraint::from_proto( + *join_constraint, + ) + .into(), + }, + ))), + }) + } LogicalPlan::Subquery(subquery) => { // Serialize the inner subquery plan directly — the // LogicalPlan::Subquery wrapper is reconstructed during diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 23ce254e99a40..1fc0a03944c71 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -32,8 +32,8 @@ use datafusion_expr::expr::{ }; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ - Expr, JoinConstraint, JoinType, SortExpr, TryCast, WindowFrame, WindowFrameBound, - WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, + Expr, JoinConstraint, JoinType, Operator, SortExpr, TryCast, WindowFrame, + WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, logical_plan::StringifiedPlan, }; @@ -782,6 +782,22 @@ impl FromProto for protobuf::JoinConstraint { } } +impl TryFromProto for protobuf::AsOfMatchOperator { + type Error = Error; + + fn try_from_proto(value: Operator) -> Result { + match value { + Operator::Lt => Ok(Self::Lt), + Operator::LtEq => Ok(Self::LtEq), + Operator::Gt => Ok(Self::Gt), + Operator::GtEq => Ok(Self::GtEq), + op => Err(Error::General(format!( + "Unsupported ASOF match operator {op}" + ))), + } + } +} + impl FromProto for protobuf::NullEquality { fn from_proto(t: NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index fa18b1ffc6684..a4fd7808ce4ae 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -54,7 +54,7 @@ use datafusion_datasource_parquet::source::ParquetSource; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, Operator, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; @@ -83,8 +83,8 @@ use datafusion_physical_plan::expressions::PhysicalSortExpr; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -865,6 +865,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::SortMergeJoin(sort_join) => { self.try_into_sort_join(sort_join, ctx, proto_converter) } + PhysicalPlanType::AsOfJoin(asof_join) => { + self.try_into_asof_join(asof_join, ctx, proto_converter) + } PhysicalPlanType::AsyncFunc(async_func) => { self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) } @@ -957,6 +960,14 @@ pub trait PhysicalPlanNodeExt: Sized { ); } + if let Some(exec) = plan.downcast_ref::() { + return protobuf::PhysicalPlanNode::try_from_asof_join_exec( + exec, + codec, + proto_converter, + ); + } + if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_cross_join_exec( exec, @@ -2654,6 +2665,73 @@ pub trait PhysicalPlanNodeExt: Sized { )?)) } + fn try_into_asof_join( + &self, + join: &protobuf::AsOfJoinExecNode, + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let left = into_physical_plan(&join.left, ctx, proto_converter)?; + let right = into_physical_plan(&join.right, ctx, proto_converter)?; + let on = join + .on + .iter() + .map(|pair| { + let left_expr = proto_converter.proto_to_physical_expr( + pair.left.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode left equality expression is missing", + ) + })?, + left.schema().as_ref(), + ctx, + )?; + let right_expr = proto_converter.proto_to_physical_expr( + pair.right.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode right equality expression is missing", + ) + })?, + right.schema().as_ref(), + ctx, + )?; + Ok((left_expr, right_expr)) + }) + .collect::>()?; + let left_match = proto_converter.proto_to_physical_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode left_match_expr is missing") + })?, + left.schema().as_ref(), + ctx, + )?; + let right_match = proto_converter.proto_to_physical_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode right_match_expr is missing") + })?, + right.schema().as_ref(), + ctx, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from(join.match_operator) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, op, right_match), + join.right_output_indices + .iter() + .map(|index| *index as usize) + .collect(), + )?)) + } + fn try_into_generate_series_physical_plan( &self, generate_series: &protobuf::GenerateSeriesNode, @@ -3299,6 +3377,62 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + fn try_from_asof_join_exec( + exec: &AsOfJoinExec, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.left()), + codec, + proto_converter, + )?; + let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.right()), + codec, + proto_converter, + )?; + let on = exec + .on() + .iter() + .map(|(left, right)| { + Ok::<_, DataFusionError>(protobuf::JoinOn { + left: Some(proto_converter.physical_expr_to_proto(left, codec)?), + right: Some(proto_converter.physical_expr_to_proto(right, codec)?), + }) + }) + .collect::>()?; + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(exec.match_condition().op)?; + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + left_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().left, + codec, + )?, + ), + right_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().right, + codec, + )?, + ), + match_operator: match_operator.into(), + right_output_indices: exec + .right_output_indices() + .iter() + .map(|index| *index as u32) + .collect(), + }, + ))), + }) + } + fn try_from_cross_join_exec( exec: &CrossJoinExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 74f7253386764..c10ba089adc37 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -3773,6 +3773,37 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_asof_join() -> Result<()> { + let ctx = SessionContext::new(); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + ctx.register_table("trades", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("prices", Arc::new(EmptyTable::new(right_schema)))?; + + for op in ["<", "<=", ">", ">="] { + let plan = ctx + .sql(&format!( + "SELECT * FROM trades t ASOF JOIN prices p \ + MATCH_CONDITION (t.ts {op} p.ts) USING (symbol)" + )) + .await? + .into_optimized_plan()?; + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{round_trip:?}")); + } + Ok(()) +} + // Single column, single split point range partitioning #[tokio::test] async fn roundtrip_range_partitioning_single_col() -> Result<()> { diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 8027c0fa7899a..24b7b904bbd35 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -72,8 +72,8 @@ use datafusion::physical_plan::expressions::{ }; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, HashJoinExec, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; @@ -478,6 +478,39 @@ fn roundtrip_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_asof_join() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let on = vec![( + Arc::new(Column::new("symbol", 0)) as _, + Arc::new(Column::new("symbol", 0)) as _, + )]; + + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + roundtrip_test(Arc::new(AsOfJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&left_schema))), + Arc::new(EmptyExec::new(Arc::clone(&right_schema))), + on.clone(), + AsOfMatchExpr::new( + Arc::new(Column::new("ts", 1)), + op, + Arc::new(Column::new("ts", 1)), + ), + vec![2], + )?))?; + } + Ok(()) +} + #[test] fn roundtrip_nested_loop_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b38099..70c9572598d1b 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 5eef9b82d975e..79733eee6ce84 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -115,6 +115,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -1298,11 +1299,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 { @@ -1436,6 +1434,156 @@ impl Unparser<'_> { Ok(()) } + LogicalPlan::AsOfJoin(join) => { + 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") + })?; + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .filter_map(|(_, right)| right.get_as_join_column()) + .collect::>() + } else { + vec![] + }; + let right_projection = right_projection.into_iter().filter(|item| { + let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier( + ids, + )) = item + else { + return true; + }; + let Some(name) = ids.last() else { + return true; + }; + !omitted_right.iter().any(|column| column.name == name.value) + }); + select.projection( + left_projection + .into_iter() + .chain(right_projection) + .collect(), + ); + } + Ok(()) + } LogicalPlan::SubqueryAlias(plan_alias) => { let (plan, mut columns) = subquery_alias_inner_query_and_columns(plan_alias); @@ -2031,6 +2179,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 @@ -2051,26 +2267,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, @@ -2081,11 +2299,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/substrait/src/logical_plan/producer/rel/mod.rs b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs index c3599a2635ffa..15f59919a2a95 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs @@ -51,6 +51,9 @@ pub fn to_substrait_rel( LogicalPlan::Aggregate(plan) => producer.handle_aggregate(plan), LogicalPlan::Sort(plan) => producer.handle_sort(plan), LogicalPlan::Join(plan) => producer.handle_join(plan), + LogicalPlan::AsOfJoin(plan) => { + not_impl_err!("Substrait ASOF join is not supported: {plan:?}")? + } LogicalPlan::Repartition(plan) => producer.handle_repartition(plan), LogicalPlan::Union(plan) => producer.handle_union(plan), LogicalPlan::TableScan(plan) => producer.handle_table_scan(plan), diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 1981ef66db377..10a29c1e9df1a 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -103,6 +103,26 @@ mod tests { Ok(()) } + #[tokio::test] + async fn asof_join_fails_closed_until_substrait_has_an_extension() -> Result<()> { + let ctx = create_context().await?; + let plan = ctx + .sql( + "SELECT * FROM data l ASOF JOIN data r \ + MATCH_CONDITION (l.a >= r.a) ON l.b = r.b", + ) + .await? + .into_optimized_plan()?; + let error = to_substrait_plan(&plan, &ctx.state()) + .expect_err("ASOF must not be lowered to a generic Substrait join"); + assert!( + error + .to_string() + .contains("Substrait ASOF join is not supported") + ); + Ok(()) + } + #[tokio::test] async fn include_remaps_for_projects() -> Result<()> { let ctx = create_context().await?; From 082268d7e5eedb5d99b343954102085c38b5b2b2 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 17:32:04 +0800 Subject: [PATCH 02/18] fix: reduce ASOF unparser stack usage --- datafusion/optimizer/src/push_down_filter.rs | 15 +- .../proto-models/src/generated/pbjson.rs | 2 +- datafusion/sql/src/unparser/plan.rs | 290 +++++++++--------- .../sqllogictest/test_files/asof_join.slt | 184 +++++++++++ docs/source/user-guide/sql/select.md | 34 ++ 5 files changed, 369 insertions(+), 156 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/asof_join.slt diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index b2ca187e05315..44170d6702eb2 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1115,15 +1115,18 @@ impl OptimizerRule for PushDownFilter { }) }); if push.is_empty() { - filter.predicate = - conjunction(keep).expect("filter predicates are not empty"); + let Some(predicate) = conjunction(keep) else { + return internal_err!("ASOF join filter predicates are empty"); + }; + filter.predicate = predicate; filter.input = Arc::new(LogicalPlan::AsOfJoin(join)); Ok(Transformed::no(LogicalPlan::Filter(filter))) } else { - join.left = Arc::new(LogicalPlan::Filter(Filter::new( - conjunction(push).expect("push predicates are not empty"), - join.left, - ))); + let Some(predicate) = conjunction(push) else { + return internal_err!("ASOF join push-down predicates are empty"); + }; + join.left = + Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.left))); Ok(Transformed::yes(with_filters( keep, LogicalPlan::AsOfJoin(join), diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d1f5c8d6985fa..eb2f554095b2d 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1701,7 +1701,7 @@ impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { if right_output_indices__.is_some() { return Err(serde::de::Error::duplicate_field("rightOutputIndices")); } - right_output_indices__ = + right_output_indices__ = Some(map_.next_value::>>()? .into_iter().map(|x| x.0).collect()) ; diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 79733eee6ce84..d81cc735011d2 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, }; @@ -1435,154 +1435,7 @@ impl Unparser<'_> { Ok(()) } LogicalPlan::AsOfJoin(join) => { - 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") - })?; - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .filter_map(|(_, right)| right.get_as_join_column()) - .collect::>() - } else { - vec![] - }; - let right_projection = right_projection.into_iter().filter(|item| { - let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier( - ids, - )) = item - else { - return true; - }; - let Some(name) = ids.last() else { - return true; - }; - !omitted_right.iter().any(|column| column.name == name.value) - }); - select.projection( - left_projection - .into_iter() - .chain(right_projection) - .collect(), - ); - } - Ok(()) + self.asof_join_to_sql(join, query, select, relation) } LogicalPlan::SubqueryAlias(plan_alias) => { let (plan, mut columns) = @@ -1881,6 +1734,145 @@ 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") + })?; + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .filter_map(|(_, right)| right.get_as_join_column()) + .collect::>() + } else { + vec![] + }; + let right_projection = right_projection.into_iter().filter(|item| { + let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(ids)) = + item + else { + return true; + }; + let Some(name) = ids.last() else { + return true; + }; + !omitted_right.iter().any(|column| column.name == name.value) + }); + 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. /// diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt new file mode 100644 index 0000000000000..ec300fca74f24 --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,184 @@ +# 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 INT) AS VALUES + (1, 'A', 1), + (2, 'A', 4), + (3, 'A', 7), + (4, 'B', 2), + (5, 'B', 8), + (6, NULL, 3), + (7, 'A', NULL); + +statement ok +CREATE TABLE asof_right(grp TEXT, ts INT, val TEXT) AS VALUES + ('A', 2, 'a2'), + ('A', 4, 'a4'), + ('A', 6, 'a6'), + ('B', 1, 'b1'), + ('B', 6, 'b6'), + (NULL, 2, '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 IIIT +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 1 NULL NULL +2 4 4 a4 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Strict predecessor per equality group. +query IIIT +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 1 NULL NULL +2 4 2 a2 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Inclusive successor per equality group. +query IIIT +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 1 2 a2 +2 4 4 a4 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# Strict successor per equality group. +query IIIT +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 1 2 a2 +2 4 6 a6 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# USING merges the equality key into one output column. +query TIIT +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 4 a4 +A 3 6 a6 +B 4 1 b1 +B 5 6 b6 +NULL 6 NULL NULL +A 7 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 + +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/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index ea96f6ae4528d..50235c169fb20 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -296,6 +296,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] @@ -377,6 +378,39 @@ 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. +Without equality keys, all rows belong to one group and DataFusion executes the +join in a single partition. + +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 From e546de5510de599c91f756bef4bf45243d93011a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:05:01 +0800 Subject: [PATCH 03/18] feat: add ASOF join physical operator --- .../physical-plan/src/joins/asof_join.rs | 1419 +++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 2 + 2 files changed, 1421 insertions(+) create mode 100644 datafusion/physical-plan/src/joins/asof_join.rs diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs new file mode 100644 index 0000000000000..5945c8d750716 --- /dev/null +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -0,0 +1,1419 @@ +// 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. + +//! Ordered, left-preserving ASOF join execution. + +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; +use datafusion_common::utils::{ + compare_rows, get_row_at_idx, normalize_float_zero_scalar, +}; +use datafusion_common::{ + ColumnStatistics, JoinType, Result, ScalarValue, Statistics, + assert_eq_or_internal_err, internal_err, plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::Column as PhysicalColumn; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{Partitioning, PhysicalSortExpr}; +use datafusion_physical_expr_common::physical_expr::{ + PhysicalExprRef, fmt_sql, is_volatile, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; +use futures::{StreamExt, stream}; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::joins::utils::{JoinOn, build_join_schema}; +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricsSet, RecordOutput, Time, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, + check_if_same_properties, +}; + +/// Physical ordered comparison for an ASOF join. +#[derive(Debug, Clone)] +pub struct AsOfMatchExpr { + /// Expression evaluated against the left input. + pub left: PhysicalExprRef, + /// Ordered comparison operator. + pub op: Operator, + /// Expression evaluated against the right input. + pub right: PhysicalExprRef, +} + +impl AsOfMatchExpr { + /// Creates a physical ASOF match expression. + pub fn new(left: PhysicalExprRef, op: Operator, right: PhysicalExprRef) -> Self { + Self { left, op, right } + } +} + +/// A sort-merge ASOF join that emits exactly one row for every left row. +#[derive(Debug, Clone)] +pub struct AsOfJoinExec { + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + schema: SchemaRef, + metrics: ExecutionPlanMetricsSet, + left_ordering: LexOrdering, + right_ordering: LexOrdering, + cache: Arc, +} + +impl AsOfJoinExec { + /// Creates a bounded ASOF join over sorted inputs. + pub fn try_new( + left: Arc, + right: Arc, + on: JoinOn, + match_condition: AsOfMatchExpr, + right_output_indices: Vec, + ) -> Result { + if !matches!( + match_condition.op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return plan_err!( + "AsOfJoinExec requires <, <=, >, or >=, found {}", + match_condition.op + ); + } + if left.boundedness().is_unbounded() || right.boundedness().is_unbounded() { + return plan_err!("AsOfJoinExec requires bounded inputs"); + } + if is_volatile(&match_condition.left) || is_volatile(&match_condition.right) { + return plan_err!("AsOfJoinExec match expression must be deterministic"); + } + if on + .iter() + .any(|(left, right)| is_volatile(left) || is_volatile(right)) + { + return plan_err!("AsOfJoinExec equality expressions must be deterministic"); + } + + let left_schema = left.schema(); + let right_schema = right.schema(); + validate_expr_side(&match_condition.left, &left_schema, "left match")?; + validate_expr_side(&match_condition.right, &right_schema, "right match")?; + for (left_expr, right_expr) in &on { + validate_expr_side(left_expr, &left_schema, "left equality")?; + validate_expr_side(right_expr, &right_schema, "right equality")?; + let left_type = left_expr.data_type(&left_schema)?; + let right_type = right_expr.data_type(&right_schema)?; + if left_type != right_type { + return plan_err!( + "AsOfJoinExec equality expression types differ: {left_type} and {right_type}" + ); + } + if !datafusion_expr::utils::can_hash(&left_type) { + return plan_err!( + "AsOfJoinExec equality expressions have unsupported hash type {left_type}" + ); + } + } + let left_match_type = match_condition.left.data_type(&left_schema)?; + let right_match_type = match_condition.right.data_type(&right_schema)?; + if left_match_type != right_match_type { + return plan_err!( + "AsOfJoinExec match expression types differ: {left_match_type} and {right_match_type}" + ); + } + if let Some(index) = right_output_indices + .iter() + .find(|index| **index >= right_schema.fields().len()) + { + return plan_err!( + "AsOfJoinExec right output index {index} is outside schema with {} fields", + right_schema.fields().len() + ); + } + if !right_output_indices + .windows(2) + .all(|pair| pair[0] < pair[1]) + { + return plan_err!( + "AsOfJoinExec right output indices must be strictly increasing" + ); + } + + let schema = + build_output_schema(&left_schema, &right_schema, &right_output_indices); + let descending = matches!(match_condition.op, Operator::Lt | Operator::LtEq); + let equality_options = SortOptions { + descending: false, + nulls_first: true, + }; + let match_options = SortOptions { + descending, + nulls_first: true, + }; + let mut left_sort_exprs = on + .iter() + .map(|(left, _)| PhysicalSortExpr { + expr: Arc::clone(left), + options: equality_options, + }) + .collect::>(); + left_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.left), + options: match_options, + }); + let mut right_sort_exprs = on + .iter() + .map(|(_, right)| PhysicalSortExpr { + expr: Arc::clone(right), + options: equality_options, + }) + .collect::>(); + right_sort_exprs.push(PhysicalSortExpr { + expr: Arc::clone(&match_condition.right), + options: match_options, + }); + let left_ordering = LexOrdering::new(left_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF left ordering must not be empty" + ) + })?; + let right_ordering = LexOrdering::new(right_sort_exprs).ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ASOF right ordering must not be empty" + ) + })?; + let cache = Arc::new(Self::compute_properties(&left, &schema, on.is_empty())?); + + Ok(Self { + left, + right, + on, + match_condition, + right_output_indices, + schema, + metrics: ExecutionPlanMetricsSet::new(), + left_ordering, + right_ordering, + cache, + }) + } + + fn compute_properties( + left: &Arc, + schema: &SchemaRef, + single_partition: bool, + ) -> Result { + let left_schema = left.schema(); + let mapping = ProjectionMapping::try_new( + left_schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + ( + Arc::new(PhysicalColumn::new(field.name(), index)) + as PhysicalExprRef, + field.name().to_string(), + ) + }), + &left_schema, + )?; + let input_eq_properties = left.equivalence_properties(); + let eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); + let output_partitioning = if single_partition { + Partitioning::UnknownPartitioning(1) + } else { + left.output_partitioning() + .project(&mapping, input_eq_properties) + }; + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + + /// Equality expressions. + pub fn on(&self) -> &JoinOn { + &self.on + } + + /// Ordered match expression. + pub fn match_condition(&self) -> &AsOfMatchExpr { + &self.match_condition + } + + /// Indices of right input columns emitted after the left columns. + pub fn right_output_indices(&self) -> &[usize] { + &self.right_output_indices + } + + /// Left input. + pub fn left(&self) -> &Arc { + &self.left + } + + /// Right input. + pub fn right(&self) -> &Arc { + &self.right + } +} + +fn build_output_schema( + left: &SchemaRef, + right: &SchemaRef, + right_output_indices: &[usize], +) -> SchemaRef { + let full_schema = build_join_schema(left, right, &JoinType::Left).0; + let left_len = left.fields().len(); + let fields = full_schema + .fields() + .iter() + .take(left_len) + .cloned() + .chain( + right_output_indices + .iter() + .map(|index| Arc::clone(&full_schema.fields()[left_len + *index])), + ) + .collect::>(); + Arc::new(Schema::new_with_metadata( + fields, + full_schema.metadata().clone(), + )) +} + +impl DisplayAs for AsOfJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + let on = self + .on + .iter() + .map(|(left, right)| { + format!("({} = {})", fmt_sql(left.as_ref()), fmt_sql(right.as_ref())) + }) + .collect::>() + .join(", "); + let match_condition = format!( + "{} {} {}", + fmt_sql(self.match_condition.left.as_ref()), + self.match_condition.op, + fmt_sql(self.match_condition.right.as_ref()) + ); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "{}: on=[{}], match=[{}]", + Self::static_name(), + on, + match_condition + ), + DisplayFormatType::TreeRender => { + writeln!(f, "on={on}")?; + writeln!(f, "match={match_condition}") + } + } + } +} + +impl ExecutionPlan for AsOfJoinExec { + fn name(&self) -> &'static str { + "AsOfJoinExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.on.is_empty() { + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) + } else { + let (left, right) = self + .on + .iter() + .map(|(left, right)| (Arc::clone(left), Arc::clone(right))) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left), + Distribution::KeyPartitioned(right), + ]) + } + } + + fn required_input_ordering(&self) -> Vec> { + vec![ + Some(OrderingRequirements::from(self.left_ordering.clone())), + Some(OrderingRequirements::from(self.right_ordering.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec { + vec![true, false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + check_if_same_properties!(self, children); + match &children[..] { + [left, right] => Ok(Arc::new(Self::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.match_condition.clone(), + self.right_output_indices.clone(), + )?)), + _ => internal_err!("AsOfJoinExec requires two children"), + } + } + + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + assert_eq_or_internal_err!( + children.len(), + 2, + "AsOfJoinExec requires two children" + ); + let left = children.remove(0); + let right = children.remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&self) + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + left_partitions, + right_partitions, + "AsOfJoinExec partition count mismatch: {left_partitions} != {right_partitions}" + ); + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + let right_stream = self.right.execute(partition, Arc::clone(&context))?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let state = AsOfJoinStreamState::new( + Arc::clone(&self.schema), + InputCursor::new( + left_stream, + left_keys, + Arc::clone(&self.match_condition.left), + ), + InputCursor::new( + right_stream, + right_keys, + Arc::clone(&self.match_condition.right), + ), + self.match_condition.op, + self.right_output_indices.clone(), + context.session_config().batch_size(), + AsOfJoinMetrics::new(partition, &self.metrics), + ); + let stream = stream::try_unfold(state, |mut state| async move { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, state))), + None => Ok(None), + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition), ChildStats::Skip] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let left = &input_stats[0]; + let mut column_statistics = left.column_statistics.clone(); + column_statistics.truncate(self.left.schema().fields().len()); + column_statistics.resize_with( + self.left.schema().fields().len(), + ColumnStatistics::new_unknown, + ); + column_statistics.extend( + self.right_output_indices + .iter() + .map(|_| ColumnStatistics::new_unknown()), + ); + Ok(Arc::new(Statistics { + num_rows: left.num_rows, + total_byte_size: Precision::Absent, + column_statistics, + })) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec, + _config: &ConfigOptions, + ) -> Result { + let left_indices = (0..self.left.schema().fields().len()).collect::>(); + let left = ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + left_indices, + &self.left, + )?; + let right = ChildFilterDescription::all_unsupported(&parent_filters); + Ok(FilterDescription::new().with_child(left).with_child(right)) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result>> { + Ok(FilterPushdownPropagation::if_any(child_pushdown_result)) + } +} + +#[derive(Clone)] +struct Candidate { + batch: Arc, + row: usize, + group: Vec, +} + +struct InputCursor { + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + batch: Option>, + key_arrays: Vec, + match_array: Option, + row: usize, + eof: bool, +} + +impl InputCursor { + fn new( + stream: SendableRecordBatchStream, + key_exprs: Vec, + match_expr: PhysicalExprRef, + ) -> Self { + Self { + stream, + key_exprs, + match_expr, + batch: None, + key_arrays: vec![], + match_array: None, + row: 0, + eof: false, + } + } + + async fn ensure_row(&mut self, elapsed_compute: &Time) -> Result { + loop { + if let Some(batch) = &self.batch + && self.row < batch.num_rows() + { + return Ok(true); + } + self.batch = None; + self.key_arrays.clear(); + self.match_array = None; + self.row = 0; + if self.eof { + return Ok(false); + } + let Some(batch) = self.stream.next().await.transpose()? else { + self.eof = true; + return Ok(false); + }; + if batch.num_rows() == 0 { + continue; + } + let batch = Arc::new(batch); + let _timer = elapsed_compute.timer(); + self.key_arrays = self + .key_exprs + .iter() + .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) + .collect::>()?; + self.match_array = Some( + self.match_expr + .evaluate(&batch)? + .into_array(batch.num_rows())?, + ); + self.batch = Some(batch); + } + } + + fn group(&self) -> Result> { + get_row_at_idx(&self.key_arrays, self.row) + .map(|row| row.into_iter().map(normalize_float_zero_scalar).collect()) + } + + fn match_value(&self) -> Result { + let array = self.match_array.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF match array is missing") + })?; + ScalarValue::try_from_array(array, self.row).map(normalize_float_zero_scalar) + } + + fn batch_row(&self) -> Result<(Arc, usize)> { + let batch = self.batch.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!("ASOF input batch is missing") + })?; + Ok((Arc::clone(batch), self.row)) + } + + fn advance(&mut self) { + self.row += 1; + } +} + +struct AsOfJoinMetrics { + baseline: BaselineMetrics, + matched_rows: Count, + unmatched_left_rows: Count, +} + +impl AsOfJoinMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + matched_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("matched_rows", partition), + unmatched_left_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("unmatched_left_rows", partition), + } + } +} + +#[derive(Default)] +struct PendingRows { + sources: Vec>, + source_by_ptr: HashMap, + indices: Vec>, +} + +impl PendingRows { + fn len(&self) -> usize { + self.indices.len() + } + + fn is_empty(&self) -> bool { + self.indices.is_empty() + } + + fn push(&mut self, batch: Arc, row: usize) { + let ptr = Arc::as_ptr(&batch) as usize; + let source = *self.source_by_ptr.entry(ptr).or_insert_with(|| { + let source = self.sources.len(); + self.sources.push(batch); + source + }); + self.indices.push(Some((source, row))); + } + + fn push_null(&mut self) { + self.indices.push(None); + } + + fn materialize_column( + &self, + source_column: usize, + data_type: &arrow::datatypes::DataType, + ) -> Result { + if self.indices.is_empty() { + return internal_err!("ASOF output materialization has no pending rows"); + } + + if self.sources.len() == 1 + && self.indices.iter().all(Option::is_some) + && let Some((0, first_row)) = self.indices[0] + && self + .indices + .iter() + .enumerate() + .all(|(offset, index)| *index == Some((0, first_row + offset))) + { + return Ok(self.sources[0] + .column(source_column) + .slice(first_row, self.indices.len())); + } + + let has_null = self.indices.iter().any(Option::is_none); + let null_array = has_null.then(|| new_null_array(data_type, 1)); + let mut source_arrays: Vec<&dyn Array> = + Vec::with_capacity(self.sources.len() + usize::from(has_null)); + if let Some(null_array) = &null_array { + source_arrays.push(null_array.as_ref()); + } + source_arrays.extend( + self.sources + .iter() + .map(|batch| batch.column(source_column).as_ref()), + ); + let source_offset = usize::from(has_null); + let interleave_indices = self + .indices + .iter() + .map(|index| match index { + Some((source, row)) => (source + source_offset, *row), + None => (0, 0), + }) + .collect::>(); + interleave(&source_arrays, &interleave_indices).map_err(Into::into) + } + + fn clear(&mut self) { + self.sources.clear(); + self.source_by_ptr.clear(); + self.indices.clear(); + } +} + +struct AsOfJoinStreamState { + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + candidate: Option, + group_sort_options: Vec, + pending_left: PendingRows, + pending_right: PendingRows, + batch_size: usize, + metrics: AsOfJoinMetrics, +} + +impl AsOfJoinStreamState { + fn new( + schema: SchemaRef, + left: InputCursor, + right: InputCursor, + op: Operator, + right_output_indices: Vec, + batch_size: usize, + metrics: AsOfJoinMetrics, + ) -> Self { + let group_sort_options = vec![ + SortOptions { + descending: false, + nulls_first: true, + }; + left.key_exprs.len() + ]; + Self { + pending_left: PendingRows::default(), + pending_right: PendingRows::default(), + schema, + left, + right, + op, + right_output_indices, + candidate: None, + group_sort_options, + batch_size: batch_size.max(1), + metrics, + } + } + + async fn next_batch(&mut self) -> Result> { + loop { + if self.pending_left.len() >= self.batch_size { + return self.flush().map(Some); + } + if !self + .left + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + if !self.pending_left.is_empty() { + return self.flush().map(Some); + } + self.metrics.baseline.done(); + return Ok(None); + } + + let (left_group, left_match) = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + (self.left.group()?, self.left.match_value()?) + }; + if left_match.is_null() || left_group.iter().any(ScalarValue::is_null) { + self.candidate = None; + self.push_current_left(None)?; + self.left.advance(); + continue; + } + let candidate_is_other_group = if let Some(candidate) = &self.candidate { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + compare_rows(&candidate.group, &left_group, &self.group_sort_options)? + != Ordering::Equal + } else { + false + }; + if candidate_is_other_group { + self.candidate = None; + } + + loop { + if !self + .right + .ensure_row(self.metrics.baseline.elapsed_compute()) + .await? + { + break; + } + let action = { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_group = self.right.group()?; + if right_group.iter().any(ScalarValue::is_null) { + RightAction::Advance + } else { + match compare_rows( + &right_group, + &left_group, + &self.group_sort_options, + )? { + Ordering::Less => RightAction::Advance, + Ordering::Greater => RightAction::Stop, + Ordering::Equal => { + let right_match = self.right.match_value()?; + if right_match.is_null() { + RightAction::Advance + } else if is_eligible(self.op, &left_match, &right_match)? + { + let (batch, row) = self.right.batch_row()?; + RightAction::Candidate(Candidate { + batch, + row, + group: right_group, + }) + } else { + RightAction::Stop + } + } + } + } + }; + match action { + RightAction::Advance => self.right.advance(), + RightAction::Candidate(candidate) => { + self.candidate = Some(candidate); + self.right.advance(); + } + RightAction::Stop => break, + } + } + + self.push_current_left(self.candidate.clone())?; + self.left.advance(); + } + } + + fn push_current_left(&mut self, candidate: Option) -> Result<()> { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let (left_batch, left_row) = self.left.batch_row()?; + self.pending_left.push(left_batch, left_row); + match candidate { + Some(candidate) => { + if !self.right_output_indices.is_empty() { + self.pending_right.push(candidate.batch, candidate.row); + } + self.metrics.matched_rows.add(1); + } + None => { + if !self.right_output_indices.is_empty() { + self.pending_right.push_null(); + } + self.metrics.unmatched_left_rows.add(1); + } + } + Ok(()) + } + + fn flush(&mut self) -> Result { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let left_len = self.schema.fields().len() - self.right_output_indices.len(); + let mut arrays = Vec::with_capacity(self.schema.fields().len()); + for index in 0..left_len { + arrays.push( + self.pending_left + .materialize_column(index, self.schema.field(index).data_type())?, + ); + } + for (offset, source_index) in self.right_output_indices.iter().enumerate() { + arrays.push(self.pending_right.materialize_column( + *source_index, + self.schema.field(left_len + offset).data_type(), + )?); + } + self.pending_left.clear(); + self.pending_right.clear(); + let batch = RecordBatch::try_new(Arc::clone(&self.schema), arrays)?; + (&batch).record_output(&self.metrics.baseline); + Ok(batch) + } +} + +fn validate_expr_side(expr: &PhysicalExprRef, schema: &Schema, name: &str) -> Result<()> { + let columns = collect_columns(expr); + if columns.is_empty() { + return plan_err!("AsOfJoinExec {name} expression must reference its input"); + } + if let Some(column) = columns.iter().find(|column| { + schema + .fields() + .get(column.index()) + .is_none_or(|field| field.name() != column.name()) + }) { + return plan_err!( + "AsOfJoinExec {name} expression references column {column} outside its input" + ); + } + Ok(()) +} + +enum RightAction { + Advance, + Candidate(Candidate), + Stop, +} + +fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result { + let ordering = right.try_cmp(left)?; + Ok(match op { + Operator::Gt => ordering == Ordering::Less, + Operator::GtEq => ordering != Ordering::Greater, + Operator::Lt => ordering == Ordering::Greater, + Operator::LtEq => ordering != Ordering::Less, + _ => false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::test::TestMemoryExec; + use arrow::array::{ + DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Int8Type}; + use datafusion_execution::config::SessionConfig; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr_common::metrics::MetricValue; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + struct VolatileExpr; + + impl std::fmt::Display for VolatileExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile") + } + } + + impl PhysicalExpr for VolatileExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Int64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, _batch: &RecordBatch) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn is_volatile_node(&self) -> bool { + true + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "volatile()") + } + } + + fn make_batch( + schema: &SchemaRef, + keys: Vec>, + times: Vec>, + values: Vec, + ) -> Result { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from(keys)), + Arc::new(Int64Array::from(times)), + Arc::new(Int32Array::from(values)), + ], + ) + .map_err(Into::into) + } + + fn test_exec() -> Result> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let left_batches = vec![ + RecordBatch::new_empty(Arc::clone(&left_schema)), + make_batch(&left_schema, vec![None], vec![Some(3)], vec![0])?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![None, Some(1)], + vec![1, 2], + )?, + make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(4), Some(7)], + vec![3, 4], + )?, + make_batch( + &left_schema, + vec![Some("B"), Some("C")], + vec![Some(2), Some(3)], + vec![5, 6], + )?, + ]; + let left = TestMemoryExec::try_new_exec( + &[left_batches], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let right_batches = vec![ + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![None, Some("A")], + vec![Some(2), None], + vec![999, 777], + )?, + make_batch(&right_schema, vec![Some("A")], vec![Some(2)], vec![20])?, + make_batch(&right_schema, vec![Some("A")], vec![Some(4)], vec![40])?, + RecordBatch::new_empty(Arc::clone(&right_schema)), + make_batch( + &right_schema, + vec![Some("A"), Some("B")], + vec![Some(6), Some(1)], + vec![60, 101], + )?, + ]; + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let on: JoinOn = vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )]; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?)) + } + + #[test] + fn eligibility_matches_public_semantics() -> Result<()> { + let left = ScalarValue::Int64(Some(10)); + let lower = ScalarValue::Int64(Some(9)); + let equal = ScalarValue::Int64(Some(10)); + let higher = ScalarValue::Int64(Some(11)); + assert!(is_eligible(Operator::Gt, &left, &lower)?); + assert!(!is_eligible(Operator::Gt, &left, &equal)?); + assert!(is_eligible(Operator::GtEq, &left, &equal)?); + assert!(is_eligible(Operator::Lt, &left, &higher)?); + assert!(!is_eligible(Operator::Lt, &left, &equal)?); + assert!(is_eligible(Operator::LtEq, &left, &equal)?); + Ok(()) + } + + #[tokio::test] + async fn state_survives_empty_input_batches_and_output_flushes() -> Result<()> { + let exec = test_exec()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let batches = collect(Arc::clone(&exec) as _, context).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 2, 2, 1] + ); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!( + ids, + vec![ + Some(0), + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + ] + ); + assert_eq!( + prices, + vec![None, None, None, Some(40), Some(60), Some(101), None] + ); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!(metrics.output_rows(), Some(7)); + assert_eq!( + metrics + .sum_by_name("matched_rows") + .map(|value| value.as_usize()), + Some(3) + ); + assert_eq!( + metrics + .sum_by_name("unmatched_left_rows") + .map(|value| value.as_usize()), + Some(4) + ); + assert!(metrics.elapsed_compute().is_some()); + assert!( + metrics.iter().any(|metric| { + matches!(metric.value(), MetricValue::ElapsedCompute(_)) + }) + ); + Ok(()) + } + + #[tokio::test] + async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { + let dictionary_type = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut left_payload = StringDictionaryBuilder::::new(); + for _ in 0..129 { + left_payload.append_value("left"); + } + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(StringArray::from(vec!["A"; 129])), + Arc::new(Int64Array::from_iter_values(-1..128)), + Arc::new(left_payload.finish()), + ], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", dictionary_type.clone(), false), + ])); + let mut right_payload = StringDictionaryBuilder::::new(); + right_payload.append_value("right"); + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(Int64Array::from(vec![0])), + Arc::new(right_payload.finish()), + ], + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch]], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(256)), + ); + let batches = collect(exec, context).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 129); + assert_eq!(batches[0].column(2).data_type(), &dictionary_type); + assert_eq!(batches[0].column(3).data_type(), &dictionary_type); + + let right_output = batches[0] + .column(3) + .as_any() + .downcast_ref::>() + .expect("right output must remain Dictionary(Int8, Utf8)"); + assert!(right_output.is_null(0)); + assert_eq!(right_output.null_count(), 1); + let values = right_output + .values() + .as_any() + .downcast_ref::() + .expect("dictionary values must be Utf8"); + for row in 1..129 { + assert_eq!( + values.value(right_output.keys().value(row) as usize), + "right" + ); + } + Ok(()) + } + + #[test] + fn rejects_volatile_physical_expressions() -> Result<()> { + let exec = test_exec()?; + let volatile = Arc::new(VolatileExpr) as PhysicalExprRef; + let match_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + exec.on().clone(), + AsOfMatchExpr::new( + Arc::clone(&volatile), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + ) + .expect_err("volatile match expression must be rejected"); + assert!(match_error.to_string().contains("must be deterministic")); + + let equality_error = AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![(volatile, Arc::new(PhysicalColumn::new("key", 0)))], + exec.match_condition().clone(), + vec![2], + ) + .expect_err("volatile equality expression must be rejected"); + assert!(equality_error.to_string().contains("must be deterministic")); + Ok(()) + } + + #[test] + fn properties_and_statistics_follow_left_preserving_contract() -> Result<()> { + let exec = test_exec()?; + let exec_plan: Arc = Arc::clone(&exec) as _; + assert_eq!(exec.maintains_input_order(), vec![true, false]); + assert_eq!(exec_plan.pipeline_behavior(), EmissionType::Incremental); + assert_eq!(exec_plan.boundedness(), Boundedness::Bounded); + assert!(matches!( + &exec.input_distribution_requirements().into_per_child()[..], + [ + Distribution::KeyPartitioned(_), + Distribution::KeyPartitioned(_) + ] + )); + for ordering in exec.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 2); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + assert_eq!( + requirement[1].options, + Some(SortOptions { + descending: false, + nulls_first: true, + }) + ); + } + + let no_keys: Arc = Arc::new(AsOfJoinExec::try_new( + Arc::clone(exec.left()), + Arc::clone(exec.right()), + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::Lt, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + assert_eq!(no_keys.output_partitioning().partition_count(), 1); + assert!(matches!( + &no_keys.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + for ordering in no_keys.required_input_ordering() { + let requirement = ordering.expect("ASOF ordering is required").into_single(); + assert_eq!(requirement.len(), 1); + assert_eq!( + requirement[0].options, + Some(SortOptions { + descending: true, + nulls_first: true, + }) + ); + } + + let mut key_stats = ColumnStatistics::new_unknown(); + key_stats.null_count = Precision::Exact(1); + key_stats.distinct_count = Precision::Exact(4); + let mut ts_stats = ColumnStatistics::new_unknown(); + ts_stats.min_value = Precision::Exact(ScalarValue::Int64(Some(1))); + ts_stats.max_value = Precision::Exact(ScalarValue::Int64(Some(7))); + let mut id_stats = ColumnStatistics::new_unknown(); + id_stats.null_count = Precision::Exact(0); + id_stats.distinct_count = Precision::Exact(7); + let left_column_statistics = vec![key_stats, ts_stats, id_stats]; + let left_stats = Arc::new(Statistics { + num_rows: Precision::Exact(7), + total_byte_size: Precision::Exact(128), + column_statistics: left_column_statistics.clone(), + }); + let right_stats = Arc::new(Statistics::new_unknown(&exec.right().schema())); + let stats = exec + .statistics_from_inputs(&[left_stats, right_stats], &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(7)); + assert_eq!(stats.total_byte_size, Precision::Absent); + assert_eq!(stats.column_statistics.len(), 4); + assert_eq!( + &stats.column_statistics[..3], + left_column_statistics.as_slice() + ); + assert_eq!(stats.column_statistics[3], ColumnStatistics::new_unknown()); + assert_eq!( + exec.child_stats_requests(None), + vec![ChildStats::At(None), ChildStats::Skip] + ); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index bbb25dda65165..820f60b09b3ba 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -18,6 +18,7 @@ //! DataFusion Join implementations use arrow::array::BooleanBufferBuilder; +pub use asof_join::{AsOfJoinExec, AsOfMatchExpr}; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ @@ -29,6 +30,7 @@ use parking_lot::Mutex; pub use piecewise_merge_join::PiecewiseMergeJoinExec; pub use sort_merge_join::SortMergeJoinExec; pub use symmetric_hash_join::SymmetricHashJoinExec; +mod asof_join; pub mod chain; mod cross_join; mod hash_join; From bd58e1debfe4622e93cb9fb52d08c35abad06799 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:28:02 +0800 Subject: [PATCH 04/18] feat: add ASOF join logical semantics --- datafusion/core/src/physical_planner.rs | 69 ++++- .../physical_optimizer/filter_pushdown.rs | 42 +++ datafusion/expr/src/logical_plan/builder.rs | 107 ++++++- datafusion/expr/src/logical_plan/display.rs | 23 +- datafusion/expr/src/logical_plan/mod.rs | 16 +- datafusion/expr/src/logical_plan/plan.rs | 279 +++++++++++++++++- datafusion/expr/src/logical_plan/tree_node.rs | 57 +++- .../optimizer/src/analyzer/type_coercion.rs | 37 ++- .../optimizer/src/common_subexpr_eliminate.rs | 1 + .../optimizer/src/optimize_projections/mod.rs | 43 ++- datafusion/optimizer/src/optimizer.rs | 5 + datafusion/optimizer/src/push_down_filter.rs | 29 ++ datafusion/proto/src/logical_plan/mod.rs | 3 + datafusion/sql/src/unparser/plan.rs | 1 + .../src/logical_plan/producer/rel/mod.rs | 3 + datafusion/substrait/tests/cases/serialize.rs | 27 +- 16 files changed, 701 insertions(+), 41 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aef8036c749a8..58b4da0f91a94 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -43,7 +43,8 @@ use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::filter::FilterExecBuilder; use crate::physical_plan::joins::utils as join_utils; use crate::physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -93,8 +94,8 @@ use datafusion_expr::physical_planning_context::{ use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, - FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, - WindowFrame, WindowFrameBound, WriteOp, + FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType, + StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, @@ -1860,6 +1861,67 @@ 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, + )?, + ); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect(); + Arc::new(AsOfJoinExec::try_new( + physical_left, + physical_right, + join_on, + match_condition, + right_output_indices, + )?) + } LogicalPlan::RecursiveQuery(RecursiveQuery { name, is_distinct, @@ -2359,6 +2421,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/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 909b80cadaae3..2f613264f45d9 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -55,6 +55,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{AsOfJoinExec, AsOfMatchExpr}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -123,6 +124,47 @@ fn test_pushdown_volatile_functions_not_allowed() { ); } +#[test] +fn test_asof_join_pushes_only_left_filters() { + let left = TestScanBuilder::new(schema()).with_support(true).build(); + let right = TestScanBuilder::new(schema()).with_support(true).build(); + let join = Arc::new( + AsOfJoinExec::try_new( + left, + right, + vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())], + AsOfMatchExpr::new( + col("b", &schema()).unwrap(), + Operator::GtEq, + col("b", &schema()).unwrap(), + ), + vec![0, 1, 2], + ) + .unwrap(), + ); + let left_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let right_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 5)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))), + )) as Arc; + let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter)); + let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap()); + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); + let formatted = format_plan_for_test(&optimized); + + assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}"); + assert!(formatted.contains("AsOfJoinExec:"), "{formatted}"); + assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}"); + assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}"); +} + /// Show that we can use config options to determine how to do pushdown. #[test] fn test_pushdown_into_scan_with_config_options() { diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 2ecb12c30afad..d5dda08bf6a45 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,43 @@ pub fn build_join_schema( dfschema.with_functional_dependencies(func_dependencies) } +/// Creates the schema for a left-preserving ASOF join. +/// +/// `ON` emits all left fields followed by nullable right fields. `USING` emits +/// each equality key once by omitting the corresponding right field. +pub fn build_asof_join_schema( + left: &DFSchema, + right: &DFSchema, + on: &[(Expr, Expr)], + join_constraint: JoinConstraint, +) -> Result { + let omitted_right_indices = if join_constraint == JoinConstraint::Using { + on.iter() + .map(|(_, right_expr)| { + let column = right_expr.get_as_join_column().ok_or_else(|| { + plan_datafusion_err!("ASOF USING keys must be columns") + })?; + right.index_of_column(column) + }) + .collect::>>()? + } else { + HashSet::new() + }; + + let full_schema = build_join_schema(left, right, &JoinType::Left)?; + let left_len = left.fields().len(); + let fields = full_schema + .iter() + .enumerate() + .filter(|(index, _)| { + *index < left_len || !omitted_right_indices.contains(&(*index - left_len)) + }) + .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field))) + .collect(); + DFSchema::new_with_metadata(fields, full_schema.metadata().clone())? + .with_functional_dependencies(left.functional_dependencies().clone()) +} + /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise /// conflict with the columns from the other. /// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 09f41c94f64fa..c5cd003d1f34e 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -21,10 +21,10 @@ use std::collections::HashMap; use std::fmt; use crate::{ - Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, - Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, - Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values, - Window, expr_vec_fmt, + Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, + Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, + Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, + Values, Window, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -493,6 +493,21 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Filter": format!("{}", filter_expr) }) } + LogicalPlan::AsOfJoin(AsOfJoin { + on, + match_condition, + join_constraint, + .. + }) => { + let join_expr: Vec = + on.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + json!({ + "Node Type": "AsOf Join", + "Join Constraint": format!("{join_constraint:?}"), + "Join Keys": join_expr.join(", "), + "Match Condition": match_condition.to_string(), + }) + } LogicalPlan::Repartition(Repartition { partitioning_scheme, .. diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 4766c3f33379f..98113d12c1b4a 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -28,8 +28,8 @@ pub mod tree_node; pub use builder::{ LogicalPlanBuilder, LogicalPlanBuilderOptions, LogicalTableSource, UNNAMED_TABLE, - build_join_schema, requalify_sides_if_needed, table_scan, union, - wrap_projection_for_join_if_necessary, + build_asof_join_schema, build_join_schema, requalify_sides_if_needed, table_scan, + union, wrap_projection_for_join_if_necessary, }; pub use ddl::{ CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction, @@ -41,12 +41,12 @@ pub use dml::{ WriteOp, }; pub use plan::{ - Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, - EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, - JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, - Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, - Unnest, Values, Window, projection_schema, + Aggregate, Analyze, AsOfJoin, AsOfMatch, ColumnUnnestList, DescribeTable, Distinct, + DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, + Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, + Projection, RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, + StringifiedPlan, Subquery, SubqueryAlias, TableScan, TableScanBuilder, + ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9cfab21a0395e..cac596b8bd58a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -41,13 +41,15 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement}; use crate::utils::{ - enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, - grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, + 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; @@ -238,6 +240,9 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), + /// Match each left row with at most one ordered row from the right input. + /// This is used to implement SQL `ASOF JOIN`. + AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -341,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(), @@ -369,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()) @@ -459,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], @@ -494,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)| { @@ -567,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() } @@ -690,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, @@ -985,6 +1024,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, @@ -1410,6 +1488,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| { @@ -1460,6 +1539,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), @@ -1498,6 +1578,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), @@ -2126,6 +2207,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, .. @@ -4237,6 +4337,165 @@ 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. + 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(), &on, join_constraint)?; + 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. /// diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c10ac92eef4f5..daf31489b08c7 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -41,11 +41,11 @@ 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, 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, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -150,6 +150,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 +464,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 @@ -610,6 +634,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 afd4e980b5424..5fb234ace4ac7 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,9 +57,9 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + AsOfJoin, AsOfMatch, Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, }; /// Performs type coercion by determining the schema @@ -175,6 +175,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), _ => Ok(plan), @@ -218,6 +219,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 2775d62144c56..1f1bcbe60c53f 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -566,6 +566,7 @@ impl OptimizerRule for CommonSubexprEliminate { LogicalPlan::Window(window) => self.try_optimize_window(window, config)?, LogicalPlan::Aggregate(agg) => self.try_optimize_aggregate(agg, config)?, LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 80aceb8cad44c..da5c715d280cb 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,8 +24,9 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, - get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, + Column, DFSchema, HashMap, JoinConstraint, JoinType, Result, + assert_eq_or_internal_err, get_required_group_by_exprs_indices, + internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; use datafusion_expr::{ @@ -407,6 +408,44 @@ fn optimize_projections( right_indices.with_projection_beneficial(), ] } + LogicalPlan::AsOfJoin(join) => { + let left_len = join.left.schema().fields().len(); + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .map(|(_, right)| { + let column = right.get_as_join_column().ok_or_else(|| { + internal_datafusion_err!("ASOF USING key is not a column") + })?; + join.right.schema().index_of_column(column) + }) + .collect::>>()? + } else { + std::collections::HashSet::new() + }; + let right_output_indices = (0..join.right.schema().fields().len()) + .filter(|index| !omitted_right.contains(index)) + .collect::>(); + 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 if let Some(right_index) = + right_output_indices.get(*index - left_len) + { + right_required.push(*right_index); + } + } + let left_indices = RequiredIndices::new_from_indices(left_required) + .with_plan_exprs(&plan, join.left.schema())?; + let right_indices = RequiredIndices::new_from_indices(right_required) + .with_plan_exprs(&plan, join.right.schema())?; + vec![ + left_indices.with_projection_beneficial(), + right_indices.with_projection_beneficial(), + ] + } // these nodes are explicitly rewritten in the match statement above LogicalPlan::Projection(_) | LogicalPlan::Aggregate(_) diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index db7ad8475273a..0abb09bea8768 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -411,6 +411,11 @@ fn map_children_mut Result>( let r = f(Arc::make_mut(right))?; l || r } + LogicalPlan::AsOfJoin(join) => { + let l = f(Arc::make_mut(&mut join.left))?; + let r = f(Arc::make_mut(&mut join.right))?; + l || r + } LogicalPlan::Union(Union { inputs, .. }) => { let mut changed = false; for input in inputs { diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index f30b1187b7bca..44170d6702eb2 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1104,6 +1104,35 @@ impl OptimizerRule for PushDownFilter { result.map_data(|plan| Ok(with_filters(keep_predicates, plan))) } LogicalPlan::Join(join) => push_down_join(join, Some(filter.predicate)), + LogicalPlan::AsOfJoin(mut join) => { + let (push, keep): (Vec<_>, Vec<_>) = + split_conjunction_owned(filter.predicate) + .into_iter() + .partition(|predicate| { + !predicate.is_volatile() + && predicate.column_refs().iter().all(|column| { + join.left.schema().is_column_from_schema(column) + }) + }); + if push.is_empty() { + let Some(predicate) = conjunction(keep) else { + return internal_err!("ASOF join filter predicates are empty"); + }; + filter.predicate = predicate; + filter.input = Arc::new(LogicalPlan::AsOfJoin(join)); + Ok(Transformed::no(LogicalPlan::Filter(filter))) + } else { + let Some(predicate) = conjunction(push) else { + return internal_err!("ASOF join push-down predicates are empty"); + }; + join.left = + Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.left))); + Ok(Transformed::yes(with_filters( + keep, + LogicalPlan::AsOfJoin(join), + ))) + } + } LogicalPlan::TableScan(mut scan) => { let filter_predicates = split_conjunction(&filter.predicate); // Filters containing scalar subqueries cannot be pushed to diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 732676a3c0a0f..353d44dc3536e 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -2190,6 +2190,9 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlan::DescribeTable(_) => Err(proto_error( "LogicalPlan serde is not yet implemented for DescribeTable", )), + LogicalPlan::AsOfJoin(_) => Err(proto_error( + "LogicalPlan serde is not yet implemented for AsOfJoin", + )), LogicalPlan::RecursiveQuery(recursive) => { let static_term = LogicalPlanNode::try_from_logical_plan( recursive.static_term.as_ref(), diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 5eef9b82d975e..036e4b34f3ac2 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -135,6 +135,7 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs index c3599a2635ffa..15f59919a2a95 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/mod.rs @@ -51,6 +51,9 @@ pub fn to_substrait_rel( LogicalPlan::Aggregate(plan) => producer.handle_aggregate(plan), LogicalPlan::Sort(plan) => producer.handle_sort(plan), LogicalPlan::Join(plan) => producer.handle_join(plan), + LogicalPlan::AsOfJoin(plan) => { + not_impl_err!("Substrait ASOF join is not supported: {plan:?}")? + } LogicalPlan::Repartition(plan) => producer.handle_repartition(plan), LogicalPlan::Union(plan) => producer.handle_union(plan), LogicalPlan::TableScan(plan) => producer.handle_table_scan(plan), diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 1981ef66db377..75e1fe251ac1b 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -18,7 +18,7 @@ #[cfg(test)] mod tests { use datafusion::datasource::provider_as_source; - use datafusion::logical_expr::LogicalPlanBuilder; + use datafusion::logical_expr::{AsOfMatch, LogicalPlanBuilder, Operator}; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::logical_plan::producer::to_substrait_plan; use datafusion_substrait::serializer; @@ -27,7 +27,7 @@ mod tests { use datafusion::prelude::*; use insta::assert_snapshot; - use std::fs; + use std::{fs, sync::Arc}; use substrait::proto::expression::field_reference::{ReferenceType, RootType}; use substrait::proto::expression::reference_segment; use substrait::proto::expression::{ReferenceSegment, RexType}; @@ -103,6 +103,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn asof_join_fails_closed_until_substrait_has_an_extension() -> Result<()> { + let ctx = create_context().await?; + let table = provider_as_source(ctx.table_provider("data").await?); + let left = LogicalPlanBuilder::scan("l", Arc::clone(&table), None)?.build()?; + let right = LogicalPlanBuilder::scan("r", table, None)?.build()?; + let plan = LogicalPlanBuilder::from(left) + .asof_join( + right, + vec![(col("l.b"), col("r.b"))], + AsOfMatch::new(col("l.a"), Operator::GtEq, col("r.a")), + )? + .build()?; + let error = to_substrait_plan(&plan, &ctx.state()) + .expect_err("ASOF must not be lowered to a generic Substrait join"); + assert!( + error + .to_string() + .contains("Substrait ASOF join is not supported") + ); + Ok(()) + } + #[tokio::test] async fn include_remaps_for_projects() -> Result<()> { let ctx = create_context().await?; From fe3b1827c9154d2a136299f7f4955c06bf500aba Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:45:16 +0800 Subject: [PATCH 05/18] feat: support ASOF JOIN SQL --- datafusion/core/tests/sql/joins.rs | 344 ++++++++++++++++++ datafusion/sql/src/relation/join.rs | 136 ++++++- datafusion/sql/src/unparser/plan.rs | 245 ++++++++++++- .../sqllogictest/test_files/asof_join.slt | 184 ++++++++++ docs/source/user-guide/sql/select.md | 34 ++ 5 files changed, 923 insertions(+), 20 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/asof_join.slt diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 7c0e89ee96418..15db970ad3283 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -20,6 +20,8 @@ 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_sql::unparser::plan_to_sql; @@ -299,3 +301,345 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } + +async 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, vec![trades])?), + )?; + + 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, vec![prices])?), + )?; + 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).await?; + + 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_without_equality_keys_is_single_partition() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + register_asof_test_tables(&ctx).await?; + let df = ctx + .sql( + "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices 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"); + assert_eq!(asof.output_partitioning().partition_count(), 1); + assert!(asof.output_ordering().is_some()); + assert!(matches!( + &asof.input_distribution_requirements().into_per_child()[..], + [Distribution::SinglePartition, Distribution::SinglePartition] + )); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); + Ok(()) +} + +#[tokio::test] +async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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_merges_key_and_unparser_round_trips() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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!["symbol", "ts", "trade_id", "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?; + Ok(()) +} + +#[tokio::test] +async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { + let ctx = SessionContext::new(); + register_asof_test_tables(&ctx).await?; + 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).await?; + 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/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 475d9a5b38099..70c9572598d1b 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 036e4b34f3ac2..d81cc735011d2 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, }; @@ -115,6 +115,7 @@ impl Unparser<'_> { | LogicalPlan::Aggregate(_) | LogicalPlan::Sort(_) | LogicalPlan::Join(_) + | LogicalPlan::AsOfJoin(_) | LogicalPlan::Repartition(_) | LogicalPlan::Union(_) | LogicalPlan::TableScan(_) @@ -135,7 +136,6 @@ impl Unparser<'_> { | LogicalPlan::Copy(_) | LogicalPlan::DescribeTable(_) | LogicalPlan::RecursiveQuery(_) - | LogicalPlan::AsOfJoin(_) | LogicalPlan::Unnest(_) => not_impl_err!("Unsupported plan: {plan:?}"), } } @@ -1299,11 +1299,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 { @@ -1437,6 +1434,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); @@ -1734,6 +1734,145 @@ 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") + })?; + let omitted_right = if join.join_constraint == JoinConstraint::Using { + join.on + .iter() + .filter_map(|(_, right)| right.get_as_join_column()) + .collect::>() + } else { + vec![] + }; + let right_projection = right_projection.into_iter().filter(|item| { + let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(ids)) = + item + else { + return true; + }; + let Some(name) = ids.last() else { + return true; + }; + !omitted_right.iter().any(|column| column.name == name.value) + }); + 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. /// @@ -2032,6 +2171,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 @@ -2052,26 +2259,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, @@ -2082,11 +2291,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 0000000000000..ec300fca74f24 --- /dev/null +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -0,0 +1,184 @@ +# 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 INT) AS VALUES + (1, 'A', 1), + (2, 'A', 4), + (3, 'A', 7), + (4, 'B', 2), + (5, 'B', 8), + (6, NULL, 3), + (7, 'A', NULL); + +statement ok +CREATE TABLE asof_right(grp TEXT, ts INT, val TEXT) AS VALUES + ('A', 2, 'a2'), + ('A', 4, 'a4'), + ('A', 6, 'a6'), + ('B', 1, 'b1'), + ('B', 6, 'b6'), + (NULL, 2, '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 IIIT +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 1 NULL NULL +2 4 4 a4 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Strict predecessor per equality group. +query IIIT +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 1 NULL NULL +2 4 2 a2 +3 7 6 a6 +4 2 1 b1 +5 8 6 b6 +6 3 NULL NULL +7 NULL NULL NULL + +# Inclusive successor per equality group. +query IIIT +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 1 2 a2 +2 4 4 a4 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# Strict successor per equality group. +query IIIT +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 1 2 a2 +2 4 6 a6 +3 7 NULL NULL +4 2 6 b6 +5 8 NULL NULL +6 3 NULL NULL +7 NULL NULL NULL + +# USING merges the equality key into one output column. +query TIIT +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 4 a4 +A 3 6 a6 +B 4 1 b1 +B 5 6 b6 +NULL 6 NULL NULL +A 7 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 + +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/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index ea96f6ae4528d..50235c169fb20 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -296,6 +296,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] @@ -377,6 +378,39 @@ 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. +Without equality keys, all rows belong to one group and DataFusion executes the +join in a single partition. + +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 From ca0e943ffc9e4fe511fa319bac5399433c65bd5b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 13:59:06 +0800 Subject: [PATCH 06/18] feat: add ASOF joins to DataFrame API --- datafusion/core/src/dataframe/mod.rs | 43 +++++++++++++++++++- datafusion/core/src/prelude.rs | 2 +- datafusion/core/tests/dataframe/mod.rs | 54 +++++++++++++++++++++++--- 3 files changed, 91 insertions(+), 8 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 325ae91d27bbf..faf7cd94fce14 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -58,8 +58,8 @@ use datafusion_common::{ }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ - ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case, - dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, + AsOfMatch, ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, + UNNAMED_TABLE, case, dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, }; use datafusion_functions::core::coalesce; use datafusion_functions::math::nanvl; @@ -1379,6 +1379,45 @@ impl DataFrame { }) } + /// Join this `DataFrame` to the closest eligible row in `right`. + /// + /// Every left row is emitted exactly once. `on` contains optional equality + /// expressions and `match_condition` selects the ordered predecessor or + /// successor from the matching right group. + pub fn join_asof( + self, + right: DataFrame, + on: Vec<(Expr, Expr)>, + match_condition: AsOfMatch, + ) -> Result { + let plan = LogicalPlanBuilder::from(self.plan) + .asof_join(right.plan, on, match_condition)? + .build()?; + Ok(DataFrame { + session_state: self.session_state, + plan, + projection_requires_validation: true, + }) + } + + /// Join this `DataFrame` to the closest eligible row in `right` using + /// same-named equality keys. + pub fn join_asof_using( + self, + right: DataFrame, + using_keys: Vec, + match_condition: AsOfMatch, + ) -> Result { + let plan = LogicalPlanBuilder::from(self.plan) + .asof_join_using(right.plan, using_keys, match_condition)? + .build()?; + Ok(DataFrame { + session_state: self.session_state, + plan, + projection_requires_validation: true, + }) + } + /// Repartition a DataFrame based on a logical partitioning scheme. /// /// # Example diff --git a/datafusion/core/src/prelude.rs b/datafusion/core/src/prelude.rs index 31d9d7eb471f0..12a4f393ca4ac 100644 --- a/datafusion/core/src/prelude.rs +++ b/datafusion/core/src/prelude.rs @@ -34,7 +34,7 @@ pub use crate::execution::options::{ pub use datafusion_common::Column; pub use datafusion_expr::{ - Expr, + AsOfMatch, Expr, Operator, expr_fn::*, lit, lit_timestamp_nano, logical_plan::{JoinType, Partitioning}, diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index b9fecb5fdd732..374964c842ad3 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -58,7 +58,7 @@ use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::logical_expr::{ColumnarValue, Volatility}; -use datafusion::prelude::{CsvReadOptions, JoinType, ParquetReadOptions}; +use datafusion::prelude::{AsOfMatch, CsvReadOptions, JoinType, ParquetReadOptions}; use datafusion::test_util::{ parquet_test_data, populate_csv_partitions, register_aggregate_csv, test_table, test_table_with_cache_factory, test_table_with_name, @@ -77,10 +77,10 @@ use datafusion_expr::expr::{GroupingSet, NullTreatment, Sort, WindowFunction}; use datafusion_expr::var_provider::{VarProvider, VarType}; use datafusion_expr::{ CreateMemoryTable, CreateView, DdlStatement, Expr, ExprFunctionExt, ExprSchemable, - LogicalPlan, LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, TableType, - WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, cast, col, - create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col, placeholder, - scalar_subquery, when, wildcard, + LogicalPlan, LogicalPlanBuilder, Operator, ScalarFunctionImplementation, SortExpr, + TableType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, + cast, col, create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col, + placeholder, scalar_subquery, when, wildcard, }; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -1503,6 +1503,50 @@ async fn join() -> Result<()> { Ok(()) } +#[tokio::test] +async fn join_asof() -> Result<()> { + let ctx = SessionContext::new(); + let left = ctx + .read_batch(record_batch!( + ("symbol", Utf8, ["A", "A", "B"]), + ("ts", Int64, [1, 4, 2]), + ("trade_id", Int32, [1, 2, 3]) + )?)? + .alias("trades")?; + let right = ctx + .read_batch(record_batch!( + ("symbol", Utf8, ["A", "A", "B"]), + ("ts", Int64, [2, 4, 1]), + ("price", Int32, [20, 40, 101]) + )?)? + .alias("prices")?; + + let results = left + .join_asof( + right, + vec![(col("symbol"), col("symbol"))], + AsOfMatch::new(col("ts"), Operator::GtEq, col("ts")), + )? + .select(vec![col("trade_id"), col("price")])? + .sort(vec![col("trade_id").sort(true, true)])? + .collect() + .await?; + + assert_batches_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 101 |", + "+----------+-------+", + ], + &results + ); + Ok(()) +} + #[tokio::test] async fn join_coercion_unnamed() -> Result<()> { let ctx = SessionContext::new(); From ed2ce92b0a7f822da9dd80c1b886ad7dc124a636 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 14:11:56 +0800 Subject: [PATCH 07/18] feat: serialize ASOF join plans --- .../proto-models/proto/datafusion.proto | 31 + .../proto-models/src/generated/pbjson.rs | 530 ++++++++++++++++++ .../proto-models/src/generated/prost.rs | 83 ++- .../proto/src/logical_plan/from_proto.rs | 16 + datafusion/proto/src/logical_plan/mod.rs | 121 +++- datafusion/proto/src/logical_plan/to_proto.rs | 20 +- datafusion/proto/src/physical_plan/mod.rs | 140 ++++- .../tests/cases/roundtrip_logical_plan.rs | 40 +- .../tests/cases/roundtrip_physical_plan.rs | 37 +- 9 files changed, 998 insertions(+), 20 deletions(-) diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 205cf89abed1b..f61a59f058cbd 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -63,6 +63,7 @@ message LogicalPlanNode { CteWorkTableScanNode cte_work_table_scan = 32; DmlNode dml = 33; EmptyTableScanNode empty_table_scan = 34; + AsOfJoinNode as_of_join = 35; } } @@ -276,6 +277,25 @@ message JoinNode { bool null_aware = 9; } +enum AsOfMatchOperator { + AS_OF_MATCH_OPERATOR_UNSPECIFIED = 0; + AS_OF_MATCH_OPERATOR_LT = 1; + AS_OF_MATCH_OPERATOR_LT_EQ = 2; + AS_OF_MATCH_OPERATOR_GT = 3; + AS_OF_MATCH_OPERATOR_GT_EQ = 4; +} + +message AsOfJoinNode { + LogicalPlanNode left = 1; + LogicalPlanNode right = 2; + repeated LogicalExprNode left_join_key = 3; + repeated LogicalExprNode right_join_key = 4; + LogicalExprNode left_match_expr = 5; + LogicalExprNode right_match_expr = 6; + AsOfMatchOperator match_operator = 7; + datafusion_common.JoinConstraint join_constraint = 8; +} + message DistinctNode { LogicalPlanNode input = 1; } @@ -881,6 +901,7 @@ message PhysicalPlanNode { BufferExecNode buffer = 37; ArrowScanExecNode arrow_scan = 38; ScalarSubqueryExecNode scalar_subquery = 39; + AsOfJoinExecNode as_of_join = 40; } } @@ -1637,6 +1658,16 @@ message SortMergeJoinExecNode { datafusion_common.NullEquality null_equality = 7; } +message AsOfJoinExecNode { + PhysicalPlanNode left = 1; + PhysicalPlanNode right = 2; + repeated JoinOn on = 3; + PhysicalExprNode left_match_expr = 4; + PhysicalExprNode right_match_expr = 5; + AsOfMatchOperator match_operator = 6; + repeated uint32 right_output_indices = 7; +} + message AsyncFuncExecNode { PhysicalPlanNode input = 1; repeated PhysicalExprNode async_exprs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d23f8eee5fd2c..d1f5c8d6985fa 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1520,6 +1520,508 @@ impl<'de> serde::Deserialize<'de> for ArrowScanExecNode { deserializer.deserialize_struct("datafusion.ArrowScanExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for AsOfJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.on.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if !self.right_output_indices.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.on.is_empty() { + struct_ser.serialize_field("on", &self.on)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if !self.right_output_indices.is_empty() { + struct_ser.serialize_field("rightOutputIndices", &self.right_output_indices)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "on", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "right_output_indices", + "rightOutputIndices", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + On, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + RightOutputIndices, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "on" => Ok(GeneratedField::On), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "rightOutputIndices" | "right_output_indices" => Ok(GeneratedField::RightOutputIndices), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut on__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut right_output_indices__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::RightOutputIndices => { + if right_output_indices__.is_some() { + return Err(serde::de::Error::duplicate_field("rightOutputIndices")); + } + right_output_indices__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(AsOfJoinExecNode { + left: left__, + right: right__, + on: on__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + right_output_indices: right_output_indices__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfJoinNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if !self.left_join_key.is_empty() { + len += 1; + } + if !self.right_join_key.is_empty() { + len += 1; + } + if self.left_match_expr.is_some() { + len += 1; + } + if self.right_match_expr.is_some() { + len += 1; + } + if self.match_operator != 0 { + len += 1; + } + if self.join_constraint != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.AsOfJoinNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if !self.left_join_key.is_empty() { + struct_ser.serialize_field("leftJoinKey", &self.left_join_key)?; + } + if !self.right_join_key.is_empty() { + struct_ser.serialize_field("rightJoinKey", &self.right_join_key)?; + } + if let Some(v) = self.left_match_expr.as_ref() { + struct_ser.serialize_field("leftMatchExpr", v)?; + } + if let Some(v) = self.right_match_expr.as_ref() { + struct_ser.serialize_field("rightMatchExpr", v)?; + } + if self.match_operator != 0 { + let v = AsOfMatchOperator::try_from(self.match_operator) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.match_operator)))?; + struct_ser.serialize_field("matchOperator", &v)?; + } + if self.join_constraint != 0 { + let v = super::datafusion_common::JoinConstraint::try_from(self.join_constraint) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_constraint)))?; + struct_ser.serialize_field("joinConstraint", &v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for AsOfJoinNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "left_join_key", + "leftJoinKey", + "right_join_key", + "rightJoinKey", + "left_match_expr", + "leftMatchExpr", + "right_match_expr", + "rightMatchExpr", + "match_operator", + "matchOperator", + "join_constraint", + "joinConstraint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + LeftJoinKey, + RightJoinKey, + LeftMatchExpr, + RightMatchExpr, + MatchOperator, + JoinConstraint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "leftJoinKey" | "left_join_key" => Ok(GeneratedField::LeftJoinKey), + "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), + "leftMatchExpr" | "left_match_expr" => Ok(GeneratedField::LeftMatchExpr), + "rightMatchExpr" | "right_match_expr" => Ok(GeneratedField::RightMatchExpr), + "matchOperator" | "match_operator" => Ok(GeneratedField::MatchOperator), + "joinConstraint" | "join_constraint" => Ok(GeneratedField::JoinConstraint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = AsOfJoinNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.AsOfJoinNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut left_join_key__ = None; + let mut right_join_key__ = None; + let mut left_match_expr__ = None; + let mut right_match_expr__ = None; + let mut match_operator__ = None; + let mut join_constraint__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::LeftJoinKey => { + if left_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("leftJoinKey")); + } + left_join_key__ = Some(map_.next_value()?); + } + GeneratedField::RightJoinKey => { + if right_join_key__.is_some() { + return Err(serde::de::Error::duplicate_field("rightJoinKey")); + } + right_join_key__ = Some(map_.next_value()?); + } + GeneratedField::LeftMatchExpr => { + if left_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("leftMatchExpr")); + } + left_match_expr__ = map_.next_value()?; + } + GeneratedField::RightMatchExpr => { + if right_match_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("rightMatchExpr")); + } + right_match_expr__ = map_.next_value()?; + } + GeneratedField::MatchOperator => { + if match_operator__.is_some() { + return Err(serde::de::Error::duplicate_field("matchOperator")); + } + match_operator__ = Some(map_.next_value::()? as i32); + } + GeneratedField::JoinConstraint => { + if join_constraint__.is_some() { + return Err(serde::de::Error::duplicate_field("joinConstraint")); + } + join_constraint__ = Some(map_.next_value::()? as i32); + } + } + } + Ok(AsOfJoinNode { + left: left__, + right: right__, + left_join_key: left_join_key__.unwrap_or_default(), + right_join_key: right_join_key__.unwrap_or_default(), + left_match_expr: left_match_expr__, + right_match_expr: right_match_expr__, + match_operator: match_operator__.unwrap_or_default(), + join_constraint: join_constraint__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.AsOfJoinNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for AsOfMatchOperator { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for AsOfMatchOperator { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + "AS_OF_MATCH_OPERATOR_LT", + "AS_OF_MATCH_OPERATOR_LT_EQ", + "AS_OF_MATCH_OPERATOR_GT", + "AS_OF_MATCH_OPERATOR_GT_EQ", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = AsOfMatchOperator; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Ok(AsOfMatchOperator::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Ok(AsOfMatchOperator::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Ok(AsOfMatchOperator::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Ok(AsOfMatchOperator::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Ok(AsOfMatchOperator::GtEq), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for AsyncFuncExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -13674,6 +14176,9 @@ impl serde::Serialize for LogicalPlanNode { logical_plan_node::LogicalPlanType::EmptyTableScan(v) => { struct_ser.serialize_field("emptyTableScan", v)?; } + logical_plan_node::LogicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -13735,6 +14240,8 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "dml", "empty_table_scan", "emptyTableScan", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -13772,6 +14279,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { CteWorkTableScan, Dml, EmptyTableScan, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -13826,6 +14334,7 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { "cteWorkTableScan" | "cte_work_table_scan" => Ok(GeneratedField::CteWorkTableScan), "dml" => Ok(GeneratedField::Dml), "emptyTableScan" | "empty_table_scan" => Ok(GeneratedField::EmptyTableScan), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14077,6 +14586,13 @@ impl<'de> serde::Deserialize<'de> for LogicalPlanNode { return Err(serde::de::Error::duplicate_field("emptyTableScan")); } logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::EmptyTableScan) +; + } + GeneratedField::AsOfJoin => { + if logical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + logical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_plan_node::LogicalPlanType::AsOfJoin) ; } } @@ -20331,6 +20847,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::ScalarSubquery(v) => { struct_ser.serialize_field("scalarSubquery", v)?; } + physical_plan_node::PhysicalPlanType::AsOfJoin(v) => { + struct_ser.serialize_field("asOfJoin", v)?; + } } } struct_ser.end() @@ -20403,6 +20922,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "arrowScan", "scalar_subquery", "scalarSubquery", + "as_of_join", + "asOfJoin", ]; #[allow(clippy::enum_variant_names)] @@ -20445,6 +20966,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { Buffer, ArrowScan, ScalarSubquery, + AsOfJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20504,6 +21026,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "buffer" => Ok(GeneratedField::Buffer), "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), + "asOfJoin" | "as_of_join" => Ok(GeneratedField::AsOfJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -20790,6 +21313,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("scalarSubquery")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::ScalarSubquery) +; + } + GeneratedField::AsOfJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("asOfJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::AsOfJoin) ; } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 6baabbf37a41c..0ece21da6dcb8 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -5,7 +5,7 @@ pub struct LogicalPlanNode { #[prost( oneof = "logical_plan_node::LogicalPlanType", - tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34" + tags = "1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35" )] pub logical_plan_type: ::core::option::Option, } @@ -79,6 +79,8 @@ pub mod logical_plan_node { Dml(::prost::alloc::boxed::Box), #[prost(message, tag = "34")] EmptyTableScan(super::EmptyTableScanNode), + #[prost(message, tag = "35")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -421,6 +423,29 @@ pub struct JoinNode { pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub left_join_key: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "4")] + pub right_join_key: ::prost::alloc::vec::Vec, + #[prost(message, optional, boxed, tag = "5")] + pub left_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(message, optional, boxed, tag = "6")] + pub right_match_expr: ::core::option::Option< + ::prost::alloc::boxed::Box, + >, + #[prost(enumeration = "AsOfMatchOperator", tag = "7")] + pub match_operator: i32, + #[prost(enumeration = "super::datafusion_common::JoinConstraint", tag = "8")] + pub join_constraint: i32, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1292,7 +1317,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40" )] pub physical_plan_type: ::core::option::Option, } @@ -1378,6 +1403,8 @@ pub mod physical_plan_node { ArrowScan(super::ArrowScanExecNode), #[prost(message, tag = "39")] ScalarSubquery(::prost::alloc::boxed::Box), + #[prost(message, tag = "40")] + AsOfJoin(::prost::alloc::boxed::Box), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2473,6 +2500,23 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AsOfJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub left: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub right: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "3")] + pub on: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "4")] + pub left_match_expr: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub right_match_expr: ::core::option::Option, + #[prost(enumeration = "AsOfMatchOperator", tag = "6")] + pub match_operator: i32, + #[prost(uint32, repeated, tag = "7")] + pub right_output_indices: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -2504,6 +2548,41 @@ pub struct PhysicalScalarSubqueryExprNode { #[prost(uint32, tag = "3")] pub index: u32, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum AsOfMatchOperator { + Unspecified = 0, + Lt = 1, + LtEq = 2, + Gt = 3, + GtEq = 4, +} +impl AsOfMatchOperator { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Unspecified => "AS_OF_MATCH_OPERATOR_UNSPECIFIED", + Self::Lt => "AS_OF_MATCH_OPERATOR_LT", + Self::LtEq => "AS_OF_MATCH_OPERATOR_LT_EQ", + Self::Gt => "AS_OF_MATCH_OPERATOR_GT", + Self::GtEq => "AS_OF_MATCH_OPERATOR_GT_EQ", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "AS_OF_MATCH_OPERATOR_UNSPECIFIED" => Some(Self::Unspecified), + "AS_OF_MATCH_OPERATOR_LT" => Some(Self::Lt), + "AS_OF_MATCH_OPERATOR_LT_EQ" => Some(Self::LtEq), + "AS_OF_MATCH_OPERATOR_GT" => Some(Self::Gt), + "AS_OF_MATCH_OPERATOR_GT_EQ" => Some(Self::GtEq), + _ => None, + } + } +} /// Identifies a built-in file format supported by DataFusion. /// Used by DefaultLogicalExtensionCodec to serialize/deserialize /// FileFormatFactory instances (e.g. in CopyTo plans). diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 6d9a73e06ff45..8e6fe2f3d9c29 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -234,6 +234,22 @@ impl FromProto for JoinConstraint { } } +impl TryFromProto for Operator { + type Error = Error; + + fn try_from_proto(value: protobuf::AsOfMatchOperator) -> Result { + match value { + protobuf::AsOfMatchOperator::Lt => Ok(Self::Lt), + protobuf::AsOfMatchOperator::LtEq => Ok(Self::LtEq), + protobuf::AsOfMatchOperator::Gt => Ok(Self::Gt), + protobuf::AsOfMatchOperator::GtEq => Ok(Self::GtEq), + protobuf::AsOfMatchOperator::Unspecified => Err(Error::General( + "ASOF match operator must be specified".to_string(), + )), + } + } +} + impl FromProto for NullEquality { fn from_proto(t: protobuf::NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 353d44dc3536e..0fa9d501cd804 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -60,17 +60,17 @@ use datafusion_datasource_json::file_format::{ use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, Operator, RangePartitioning, RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, }; use datafusion_expr::{ DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ - Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection, - Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window, - builder::project, + Aggregate, AsOfJoin, AsOfMatch, CreateCatalog, CreateCatalogSchema, + CreateExternalTable, CreateView, DdlStatement, Distinct, EmptyRelation, + Extension, Join, Prepare, Projection, Repartition, Sort, SubqueryAlias, + TableScan, TableScanBuilder, Values, Window, builder::project, }, }; use datafusion_proto_common::protobuf_common; @@ -1033,6 +1033,61 @@ impl AsLogicalPlan for LogicalPlanNode { join.null_aware, )?)) } + LogicalPlanType::AsOfJoin(join) => { + let left_keys = + from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; + let right_keys = + from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + if left_keys.len() != right_keys.len() { + return Err(proto_error(format!( + "Received an AsOfJoinNode with left_join_key and right_join_key of different lengths: {} and {}", + left_keys.len(), + right_keys.len() + ))); + } + let left_match = from_proto::parse_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode left_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let right_match = from_proto::parse_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinNode right_match_expr is missing") + })?, + ctx, + extension_codec, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from( + join.match_operator, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + let join_constraint = protobuf::JoinConstraint::try_from( + join.join_constraint, + ) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF JoinConstraint {}", + join.join_constraint + )) + })?; + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + Ok(LogicalPlan::AsOfJoin(AsOfJoin::try_new( + Arc::new(left), + Arc::new(right), + left_keys.into_iter().zip(right_keys).collect(), + AsOfMatch::new(left_match, op, right_match), + JoinConstraint::from_proto(join_constraint), + )?)) + } LogicalPlanType::Union(union) => { assert_or_internal_err!( union.inputs.len() >= 2, @@ -1685,6 +1740,59 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } + LogicalPlan::AsOfJoin(AsOfJoin { + left, + right, + on, + match_condition, + join_constraint, + .. + }) => { + let left = LogicalPlanNode::try_from_logical_plan( + left.as_ref(), + extension_codec, + )?; + let right = LogicalPlanNode::try_from_logical_plan( + right.as_ref(), + extension_codec, + )?; + let (left_join_key, right_join_key) = on + .iter() + .map(|(left, right)| { + Ok(( + serialize_expr(left, extension_codec)?, + serialize_expr(right, extension_codec)?, + )) + }) + .collect::, ToProtoError>>()? + .into_iter() + .unzip(); + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(match_condition.op)?; + Ok(LogicalPlanNode { + logical_plan_type: Some(LogicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + left_join_key, + right_join_key, + left_match_expr: Some(Box::new(serialize_expr( + &match_condition.left, + extension_codec, + )?)), + right_match_expr: Some(Box::new(serialize_expr( + &match_condition.right, + extension_codec, + )?)), + match_operator: match_operator.into(), + join_constraint: protobuf::JoinConstraint::from_proto( + *join_constraint, + ) + .into(), + }, + ))), + }) + } LogicalPlan::Subquery(subquery) => { // Serialize the inner subquery plan directly — the // LogicalPlan::Subquery wrapper is reconstructed during @@ -2190,9 +2298,6 @@ 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/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 23ce254e99a40..1fc0a03944c71 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -32,8 +32,8 @@ use datafusion_expr::expr::{ }; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ - Expr, JoinConstraint, JoinType, SortExpr, TryCast, WindowFrame, WindowFrameBound, - WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, + Expr, JoinConstraint, JoinType, Operator, SortExpr, TryCast, WindowFrame, + WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, logical_plan::StringifiedPlan, }; @@ -782,6 +782,22 @@ impl FromProto for protobuf::JoinConstraint { } } +impl TryFromProto for protobuf::AsOfMatchOperator { + type Error = Error; + + fn try_from_proto(value: Operator) -> Result { + match value { + Operator::Lt => Ok(Self::Lt), + Operator::LtEq => Ok(Self::LtEq), + Operator::Gt => Ok(Self::Gt), + Operator::GtEq => Ok(Self::GtEq), + op => Err(Error::General(format!( + "Unsupported ASOF match operator {op}" + ))), + } + } +} + impl FromProto for protobuf::NullEquality { fn from_proto(t: NullEquality) -> Self { match t { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index cea334e42aace..2c27fd23f1019 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -54,7 +54,7 @@ use datafusion_datasource_parquet::source::ParquetSource; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, Operator, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; @@ -83,8 +83,8 @@ use datafusion_physical_plan::expressions::PhysicalSortExpr; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, CrossJoinExec, HashJoinExec, NestedLoopJoinExec, + PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -865,6 +865,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::SortMergeJoin(_) => { SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } + PhysicalPlanType::AsOfJoin(asof_join) => { + self.try_into_asof_join(asof_join, ctx, proto_converter) + } PhysicalPlanType::AsyncFunc(async_func) => { self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) } @@ -949,6 +952,14 @@ pub trait PhysicalPlanNodeExt: Sized { ); } + if let Some(exec) = plan.downcast_ref::() { + return protobuf::PhysicalPlanNode::try_from_asof_join_exec( + exec, + codec, + proto_converter, + ); + } + if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_cross_join_exec( exec, @@ -2561,6 +2572,73 @@ pub trait PhysicalPlanNodeExt: Sized { SortMergeJoinExec::try_from_proto(&node, &decode_ctx) } + fn try_into_asof_join( + &self, + join: &protobuf::AsOfJoinExecNode, + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let left = into_physical_plan(&join.left, ctx, proto_converter)?; + let right = into_physical_plan(&join.right, ctx, proto_converter)?; + let on = join + .on + .iter() + .map(|pair| { + let left_expr = proto_converter.proto_to_physical_expr( + pair.left.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode left equality expression is missing", + ) + })?, + left.schema().as_ref(), + ctx, + )?; + let right_expr = proto_converter.proto_to_physical_expr( + pair.right.as_ref().ok_or_else(|| { + proto_error( + "AsOfJoinExecNode right equality expression is missing", + ) + })?, + right.schema().as_ref(), + ctx, + )?; + Ok((left_expr, right_expr)) + }) + .collect::>()?; + let left_match = proto_converter.proto_to_physical_expr( + join.left_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode left_match_expr is missing") + })?, + left.schema().as_ref(), + ctx, + )?; + let right_match = proto_converter.proto_to_physical_expr( + join.right_match_expr.as_ref().ok_or_else(|| { + proto_error("AsOfJoinExecNode right_match_expr is missing") + })?, + right.schema().as_ref(), + ctx, + )?; + let match_operator = protobuf::AsOfMatchOperator::try_from(join.match_operator) + .map_err(|_| { + proto_error(format!( + "Unknown ASOF match operator {}", + join.match_operator + )) + })?; + let op = Operator::try_from_proto(match_operator)?; + Ok(Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, op, right_match), + join.right_output_indices + .iter() + .map(|index| *index as usize) + .collect(), + )?)) + } + fn try_into_generate_series_physical_plan( &self, generate_series: &protobuf::GenerateSeriesNode, @@ -3138,6 +3216,62 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + fn try_from_asof_join_exec( + exec: &AsOfJoinExec, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.left()), + codec, + proto_converter, + )?; + let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(exec.right()), + codec, + proto_converter, + )?; + let on = exec + .on() + .iter() + .map(|(left, right)| { + Ok::<_, DataFusionError>(protobuf::JoinOn { + left: Some(proto_converter.physical_expr_to_proto(left, codec)?), + right: Some(proto_converter.physical_expr_to_proto(right, codec)?), + }) + }) + .collect::>()?; + let match_operator = + protobuf::AsOfMatchOperator::try_from_proto(exec.match_condition().op)?; + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AsOfJoin(Box::new( + protobuf::AsOfJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + left_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().left, + codec, + )?, + ), + right_match_expr: Some( + proto_converter.physical_expr_to_proto( + &exec.match_condition().right, + codec, + )?, + ), + match_operator: match_operator.into(), + right_output_indices: exec + .right_output_indices() + .iter() + .map(|index| *index as u32) + .collect(), + }, + ))), + }) + } + fn try_from_cross_join_exec( exec: &CrossJoinExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 74f7253386764..f6113a807e8a7 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -74,8 +74,9 @@ use datafusion_common::format::{ }; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, - TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, + Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, + SplitPoint, TableReference, internal_datafusion_err, internal_err, not_impl_err, + plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -90,7 +91,7 @@ use datafusion_expr::logical_plan::{ ExplainOption, Extension, UserDefinedLogicalNodeCore, }; use datafusion_expr::{ - Accumulator, AggregateUDF, ColumnarValue, DmlStatement, ExprFunctionExt, + Accumulator, AggregateUDF, AsOfMatch, ColumnarValue, DmlStatement, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, @@ -3773,6 +3774,39 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_asof_join() -> Result<()> { + let ctx = SessionContext::new(); + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + ctx.register_table("trades", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("prices", Arc::new(EmptyTable::new(right_schema)))?; + + let left = ctx.table("trades").await?.into_optimized_plan()?; + let right = ctx.table("prices").await?.into_optimized_plan()?; + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + let plan = LogicalPlanBuilder::from(left.clone()) + .asof_join_using( + right.clone(), + vec![Column::from_name("symbol")], + AsOfMatch::new(col("trades.ts"), op, col("prices.ts")), + )? + .build()?; + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{round_trip:?}")); + } + Ok(()) +} + // Single column, single split point range partitioning #[tokio::test] async fn roundtrip_range_partitioning_single_col() -> Result<()> { diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 2671f4d0152f7..682b8471320ab 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -72,8 +72,8 @@ use datafusion::physical_plan::expressions::{ }; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + AsOfJoinExec, AsOfMatchExpr, HashJoinExec, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; @@ -478,6 +478,39 @@ fn roundtrip_hash_join() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_asof_join() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("id", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("symbol", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("price", DataType::Int32, false), + ])); + let on = vec![( + Arc::new(Column::new("symbol", 0)) as _, + Arc::new(Column::new("symbol", 0)) as _, + )]; + + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + roundtrip_test(Arc::new(AsOfJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&left_schema))), + Arc::new(EmptyExec::new(Arc::clone(&right_schema))), + on.clone(), + AsOfMatchExpr::new( + Arc::new(Column::new("ts", 1)), + op, + Arc::new(Column::new("ts", 1)), + ), + vec![2], + )?))?; + } + Ok(()) +} + #[test] fn roundtrip_nested_loop_join() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); From b3ca2b6dcfb59e2c86d981c801c8b50087f7683e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 14:26:43 +0800 Subject: [PATCH 08/18] bench: add ASOF join benchmarks --- benchmarks/src/asof.rs | 194 +++++++++++++++++ benchmarks/src/bin/dfbench.rs | 4 +- benchmarks/src/lib.rs | 1 + datafusion/physical-plan/Cargo.toml | 5 + datafusion/physical-plan/benches/asof_join.rs | 196 ++++++++++++++++++ 5 files changed, 399 insertions(+), 1 deletion(-) create mode 100644 benchmarks/src/asof.rs create mode 100644 datafusion/physical-plan/benches/asof_join.rs diff --git a/benchmarks/src/asof.rs b/benchmarks/src/asof.rs new file mode 100644 index 0000000000000..d74607550f4ed --- /dev/null +++ b/benchmarks/src/asof.rs @@ -0,0 +1,194 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; +use clap::Args; +use datafusion::physical_plan::execute_stream; +use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; +use futures::StreamExt; + +/// Run end-to-end ASOF join benchmarks. +/// +/// The cases cover ordered-input reuse, optimizer-inserted sort/repartition, +/// wide payload materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 4). If not specified, runs all queries + #[arg(short, long)] + query: Option, + + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option, +} + +const ASOF_QUERIES: &[&str] = &[ + // Q1: predecessor without equality keys, ordered range input can be reused + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + "#, + // Q2: grouped predecessor, optimizer supplies repartitioning and ordering + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + value AS payload + FROM range(1000000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + value AS payload + FROM range(1000000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q3: grouped predecessor with a wide payload + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + repeat('x', 256) AS payload + FROM range(250000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + repeat('y', 256) AS payload + FROM range(250000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q4: successor matching requires descending input order + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(500000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(500000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts <= r.ts) + "#, +]; + +impl RunOpt { + pub async fn run(self) -> Result<()> { + println!("Running ASOF benchmarks with the following options: {self:#?}\n"); + + let query_range = match self.query { + Some(query_id) if (1..=ASOF_QUERIES.len()).contains(&query_id) => { + query_id..=query_id + } + Some(query_id) => { + return exec_err!( + "Query {query_id} not found. Available queries: 1 to {}", + ASOF_QUERIES.len() + ); + } + None => 1..=ASOF_QUERIES.len(), + }; + + let config = self.common.config()?; + let runtime = self.common.build_runtime()?; + let ctx = SessionContext::new_with_config_rt(config, runtime); + let mut benchmark_run = BenchmarkRun::new(); + + for query_id in query_range { + let sql = ASOF_QUERIES[query_id - 1]; + benchmark_run.start_new_case(&format!("Query {query_id}")); + match self.benchmark_query(sql, &query_id.to_string(), &ctx).await { + Ok(results) => { + for result in results { + benchmark_run.write_iter(result.elapsed, result.row_count); + } + } + Err(error) => { + return Err(DataFusionError::Context( + format!("ASOF benchmark Q{query_id} failed with error:"), + Box::new(error), + )); + } + } + } + + benchmark_run.maybe_write_json(self.output_path.as_ref())?; + Ok(()) + } + + async fn benchmark_query( + &self, + sql: &str, + query_name: &str, + ctx: &SessionContext, + ) -> Result> { + let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?; + let plan_string = format!("{physical_plan:#?}"); + if !plan_string.contains("AsOfJoinExec") { + return Err(exec_datafusion_err!( + "Query {query_name} does not use AsOfJoinExec. Physical plan: {plan_string}" + )); + } + + let mut query_results = Vec::with_capacity(self.common.iterations); + for iteration in 0..self.common.iterations { + let start = Instant::now(); + let row_count = Self::execute_sql_without_result_buffering(sql, ctx).await?; + let elapsed = start.elapsed(); + println!( + "Query {query_name} iteration {iteration} returned {row_count} rows in {elapsed:?}" + ); + query_results.push(QueryResult { elapsed, row_count }); + } + Ok(query_results) + } + + async fn execute_sql_without_result_buffering( + sql: &str, + ctx: &SessionContext, + ) -> Result { + let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?; + let mut stream = execute_stream(physical_plan, ctx.task_ctx())?; + let mut row_count = 0; + while let Some(batch) = stream.next().await { + row_count += batch?.num_rows(); + } + Ok(row_count) + } +} diff --git a/benchmarks/src/bin/dfbench.rs b/benchmarks/src/bin/dfbench.rs index 50dd99368b7f0..1e478747ceea1 100644 --- a/benchmarks/src/bin/dfbench.rs +++ b/benchmarks/src/bin/dfbench.rs @@ -32,7 +32,7 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; use datafusion_benchmarks::{ - cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch, + asof, cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch, }; #[derive(Debug, Parser)] @@ -44,6 +44,7 @@ struct Cli { #[derive(Debug, Subcommand)] enum Options { + Asof(asof::RunOpt), Cancellation(cancellation::RunOpt), Clickbench(clickbench::RunOpt), Dict(dict::RunOpt), @@ -65,6 +66,7 @@ pub async fn main() -> Result<()> { let cli = Cli::parse(); match cli.command { + Options::Asof(opt) => opt.run().await, Options::Cancellation(opt) => opt.run().await, Options::Clickbench(opt) => opt.run().await, Options::Dict(opt) => opt.run().await, diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 8d24d44a174e3..74b10ae86d78e 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -16,6 +16,7 @@ // under the License. //! DataFusion benchmark runner +pub mod asof; pub mod cancellation; pub mod clickbench; pub mod dict; diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 58c2f0d7da537..d243915c205cf 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -121,6 +121,11 @@ harness = false name = "sort_merge_join" required-features = ["test_utils"] +[[bench]] +harness = false +name = "asof_join" +required-features = ["test_utils"] + [[bench]] harness = false name = "aggregate_vectorized" diff --git a/datafusion/physical-plan/benches/asof_join.rs b/datafusion/physical-plan/benches/asof_join.rs new file mode 100644 index 0000000000000..da39080c94f90 --- /dev/null +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -0,0 +1,196 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmarks for the pre-sorted ASOF join kernel. + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, Int64Array, RecordBatch, StringArray, StringDictionaryBuilder, +}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::{TaskContext, config::SessionConfig}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::joins::{AsOfJoinExec, AsOfMatchExpr, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; +use tokio::runtime::Runtime; + +#[derive(Clone, Copy)] +enum Payload { + Int64, + WideUtf8, + Dictionary, +} + +impl Payload { + fn name(self) -> &'static str { + match self { + Self::Int64 => "int64", + Self::WideUtf8 => "wide_utf8", + Self::Dictionary => "dictionary", + } + } + + fn data_type(self) -> DataType { + match self { + Self::Int64 => DataType::Int64, + Self::WideUtf8 => DataType::Utf8, + Self::Dictionary => { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + } + } + + fn array(self, rows: &[(i64, i64, usize)]) -> ArrayRef { + match self { + Self::Int64 => Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, _, row)| *row as i64), + )), + Self::WideUtf8 => Arc::new(StringArray::from_iter_values( + rows.iter() + .map(|(_, _, row)| format!("row_{row:08}_{}", "x".repeat(244))), + )), + Self::Dictionary => { + let mut builder = StringDictionaryBuilder::::new(); + for (_, _, row) in rows { + builder.append_value(format!("category_{}", row % 64)); + } + Arc::new(builder.finish()) + } + } + } +} + +fn schema(payload: Payload) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", payload.data_type(), false), + ])) +} + +fn build_sorted_batches( + num_rows: usize, + num_groups: usize, + time_offset: i64, + payload: Payload, + schema: &SchemaRef, +) -> Vec { + let mut rows = (0..num_rows) + .map(|row| { + ( + (row % num_groups) as i64, + (row / num_groups) as i64 + time_offset, + row, + ) + }) + .collect::>(); + rows.sort_unstable_by_key(|(key, ts, _)| (*key, *ts)); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(key, _, _)| *key), + )), + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, ts, _)| *ts), + )), + payload.array(&rows), + ], + ) + .unwrap(); + + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(8192); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec(batches: &[RecordBatch], schema: &SchemaRef) -> Arc { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn do_join( + left: Arc, + right: Arc, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let left_match = col("ts", &left.schema()).unwrap(); + let right_match = col("ts", &right.schema()).unwrap(); + let join = AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, Operator::GtEq, right_match), + vec![2], + ) + .unwrap(); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(8192)), + ); + rt.block_on(async { + collect(Arc::new(join), task_ctx) + .await + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum() + }) +} + +fn bench_asof_join(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let num_rows = 100_000; + let num_groups = 10_000; + let mut group = c.benchmark_group("asof_join"); + + for payload in [Payload::Int64, Payload::WideUtf8, Payload::Dictionary] { + let schema = schema(payload); + let left_batches = + build_sorted_batches(num_rows, num_groups, 1, payload, &schema); + let right_batches = + build_sorted_batches(num_rows, num_groups, 0, payload, &schema); + group.bench_function( + BenchmarkId::new(payload.name(), format!("{num_rows}_rows_10_per_key")), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &schema); + let right = make_exec(&right_batches, &schema); + do_join(left, right, &rt) + }) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_asof_join); +criterion_main!(benches); From d36fd277c1deb42d44b9d445c09795ed3a84ab2c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 23 Jul 2026 14:39:34 +0800 Subject: [PATCH 09/18] fix: keep generated protobuf code in sync --- datafusion/proto-models/src/generated/pbjson.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index d1f5c8d6985fa..eb2f554095b2d 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1701,7 +1701,7 @@ impl<'de> serde::Deserialize<'de> for AsOfJoinExecNode { if right_output_indices__.is_some() { return Err(serde::de::Error::duplicate_field("rightOutputIndices")); } - right_output_indices__ = + right_output_indices__ = Some(map_.next_value::>>()? .into_iter().map(|x| x.0).collect()) ; From 43110a3260548ad87b73fb5670c05c24839444bb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:38:22 +0800 Subject: [PATCH 10/18] feat: broadcast ASOF right input --- .../physical-plan/src/joins/asof_join.rs | 339 ++++++++++++++---- 1 file changed, 273 insertions(+), 66 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 5945c8d750716..b55a21260de8f 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -15,7 +15,42 @@ // specific language governing permissions and limitations // under the License. -//! Ordered, left-preserving ASOF join execution. +//! Broadcast, left-preserving ASOF join execution. +//! +//! An ASOF join emits exactly one output row for every left row. Within an +//! optional equality-key group, it selects the closest right row that satisfies +//! one ordered comparison: +//! +//! ```text +//! left.ts >= right.ts => greatest eligible right.ts +//! left.ts <= right.ts => smallest eligible right.ts +//! ``` +//! +//! The right input is collected and shared by all output partitions. The left +//! input remains partitioned, and each partition performs an independent +//! monotonic scan over the ordered right input: +//! +//! ```text +//! AsOfJoinExec +//! SortExec(left equality keys, left match key) +//! RepartitionExec(RoundRobinBatch) +//! left +//! SortExec(right equality keys, right match key) +//! CoalescePartitionsExec +//! right +//! ``` +//! +//! Both inputs must be ordered by their equality keys followed by the match +//! key. For `<` and `<=`, the match ordering is reversed so all directions use +//! the same forward-only state machine. Each left partition owns its cursors, +//! equality-group state, and current candidate, while the collected right +//! batches are immutable and shared. +//! +//! This mode preserves probe-side parallelism when there are no equality keys +//! or when equality keys have low cardinality or skew. It retains the complete +//! right input in the memory pool and may scan it once per left partition, so a +//! repartitioned streaming mode remains a useful future alternative for large +//! right inputs. use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -35,27 +70,30 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, plan_err, }; use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column as PhysicalColumn; use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{Partitioning, PhysicalSortExpr}; use datafusion_physical_expr_common::physical_expr::{ PhysicalExprRef, fmt_sql, is_volatile, }; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; -use futures::{StreamExt, stream}; +use futures::{StreamExt, TryStreamExt, future::poll_fn, stream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; -use crate::joins::utils::{JoinOn, build_join_schema}; +use crate::joins::utils::{JoinOn, OnceAsync, build_join_schema}; +use crate::memory::MemoryStream; use crate::metrics::{ - BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - MetricsSet, RecordOutput, Time, + BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, + MetricCategory, MetricsSet, RecordOutput, Time, }; +use crate::spill::get_record_batch_memory_size; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ @@ -82,8 +120,8 @@ impl AsOfMatchExpr { } } -/// A sort-merge ASOF join that emits exactly one row for every left row. -#[derive(Debug, Clone)] +/// A broadcast sort-merge ASOF join that emits one row for every left row. +#[derive(Debug)] pub struct AsOfJoinExec { left: Arc, right: Arc, @@ -94,6 +132,7 @@ pub struct AsOfJoinExec { metrics: ExecutionPlanMetricsSet, left_ordering: LexOrdering, right_ordering: LexOrdering, + right_fut: OnceAsync, cache: Arc, } @@ -216,7 +255,7 @@ impl AsOfJoinExec { "ASOF right ordering must not be empty" ) })?; - let cache = Arc::new(Self::compute_properties(&left, &schema, on.is_empty())?); + let cache = Arc::new(Self::compute_properties(&left, &schema)?); Ok(Self { left, @@ -228,6 +267,7 @@ impl AsOfJoinExec { metrics: ExecutionPlanMetricsSet::new(), left_ordering, right_ordering, + right_fut: Default::default(), cache, }) } @@ -235,7 +275,6 @@ impl AsOfJoinExec { fn compute_properties( left: &Arc, schema: &SchemaRef, - single_partition: bool, ) -> Result { let left_schema = left.schema(); let mapping = ProjectionMapping::try_new( @@ -254,12 +293,9 @@ impl AsOfJoinExec { )?; let input_eq_properties = left.equivalence_properties(); let eq_properties = input_eq_properties.project(&mapping, Arc::clone(schema)); - let output_partitioning = if single_partition { - Partitioning::UnknownPartitioning(1) - } else { - left.output_partitioning() - .project(&mapping, input_eq_properties) - }; + let output_partitioning = left + .output_partitioning() + .project(&mapping, input_eq_properties); Ok(PlanProperties::new( eq_properties, output_partitioning, @@ -364,22 +400,10 @@ impl ExecutionPlan for AsOfJoinExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - if self.on.is_empty() { - InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - Distribution::SinglePartition, - ]) - } else { - let (left, right) = self - .on - .iter() - .map(|(left, right)| (Arc::clone(left), Arc::clone(right))) - .unzip(); - InputDistributionRequirements::co_partitioned(vec![ - Distribution::KeyPartitioned(left), - Distribution::KeyPartitioned(right), - ]) - } + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition, + ]) } fn required_input_ordering(&self) -> Vec> { @@ -428,8 +452,15 @@ impl ExecutionPlan for AsOfJoinExec { Ok(Arc::new(Self { left, right, + on: self.on.clone(), + match_condition: self.match_condition.clone(), + right_output_indices: self.right_output_indices.clone(), + schema: Arc::clone(&self.schema), metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&self) + left_ordering: self.left_ordering.clone(), + right_ordering: self.right_ordering.clone(), + right_fut: Default::default(), + cache: Arc::clone(&self.cache), })) } @@ -438,41 +469,62 @@ impl ExecutionPlan for AsOfJoinExec { partition: usize, context: Arc, ) -> Result { - let left_partitions = self.left.output_partitioning().partition_count(); let right_partitions = self.right.output_partitioning().partition_count(); assert_eq_or_internal_err!( - left_partitions, right_partitions, - "AsOfJoinExec partition count mismatch: {left_partitions} != {right_partitions}" + 1, + "AsOfJoinExec requires one right partition, found {right_partitions}" ); let left_stream = self.left.execute(partition, Arc::clone(&context))?; - let right_stream = self.right.execute(partition, Arc::clone(&context))?; - let (left_keys, right_keys) = self.on.iter().cloned().unzip(); - let state = AsOfJoinStreamState::new( - Arc::clone(&self.schema), - InputCursor::new( - left_stream, - left_keys, - Arc::clone(&self.match_condition.left), - ), - InputCursor::new( + let metrics = AsOfJoinMetrics::new(partition, &self.metrics); + let build_metrics = metrics.clone(); + let right_fut = self.right_fut.try_once(|| { + let right_stream = self.right.execute(0, Arc::clone(&context))?; + let reservation = + MemoryConsumer::new("AsOfJoinInput").register(context.memory_pool()); + Ok(collect_right_input( right_stream, - right_keys, - Arc::clone(&self.match_condition.right), - ), - self.match_condition.op, - self.right_output_indices.clone(), - context.session_config().batch_size(), - AsOfJoinMetrics::new(partition, &self.metrics), - ); - let stream = stream::try_unfold(state, |mut state| async move { - match state.next_batch().await? { - Some(batch) => Ok(Some((batch, state))), - None => Ok(None), - } - }); + reservation, + build_metrics, + )) + })?; + let (left_keys, right_keys) = self.on.iter().cloned().unzip(); + let output_schema = Arc::clone(&self.schema); + let stream_schema = Arc::clone(&output_schema); + let left_match = Arc::clone(&self.match_condition.left); + let right_match = Arc::clone(&self.match_condition.right); + let match_op = self.match_condition.op; + let right_output_indices = self.right_output_indices.clone(); + let batch_size = context.session_config().batch_size(); + let stream = stream::once(async move { + let mut right_fut = right_fut; + let right_input = poll_fn(|cx| right_fut.get_shared(cx)).await?; + let right_stream = right_input.stream()?; + let state = AsOfJoinStreamState::new( + Arc::clone(&stream_schema), + InputCursor::new(left_stream, left_keys, left_match), + InputCursor::new(right_stream, right_keys, right_match), + match_op, + right_output_indices, + batch_size, + metrics, + ); + let stream = stream::try_unfold( + (state, right_input), + |(mut state, right_input)| async { + match state.next_batch().await? { + Some(batch) => Ok(Some((batch, (state, right_input)))), + None => Ok(None), + } + }, + ); + Ok::(Box::pin( + RecordBatchStreamAdapter::new(stream_schema, stream), + )) + }) + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&self.schema), + output_schema, stream, ))) } @@ -535,6 +587,47 @@ impl ExecutionPlan for AsOfJoinExec { } } +struct BroadcastRightInput { + schema: SchemaRef, + batches: Vec, + _reservation: MemoryReservation, +} + +impl BroadcastRightInput { + fn stream(&self) -> Result { + Ok(Box::pin(MemoryStream::try_new( + self.batches.clone(), + Arc::clone(&self.schema), + None, + )?)) + } +} + +async fn collect_right_input( + input: SendableRecordBatchStream, + reservation: MemoryReservation, + metrics: AsOfJoinMetrics, +) -> Result { + let schema = input.schema(); + let batches = input + .try_fold(Vec::new(), |mut batches, batch| { + let batch_size = get_record_batch_memory_size(&batch); + futures::future::ready(reservation.try_grow(batch_size).map(|_| { + metrics.build_mem_used.add(batch_size); + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + batches.push(batch); + batches + })) + }) + .await?; + Ok(BroadcastRightInput { + schema, + batches, + _reservation: reservation, + }) +} + #[derive(Clone)] struct Candidate { batch: Arc, @@ -632,10 +725,14 @@ impl InputCursor { } } +#[derive(Clone)] struct AsOfJoinMetrics { baseline: BaselineMetrics, matched_rows: Count, unmatched_left_rows: Count, + build_input_batches: Count, + build_input_rows: Count, + build_mem_used: Gauge, } impl AsOfJoinMetrics { @@ -648,6 +745,14 @@ impl AsOfJoinMetrics { unmatched_left_rows: MetricBuilder::new(metrics) .with_category(MetricCategory::Rows) .counter("unmatched_left_rows", partition), + build_input_batches: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_batches", partition), + build_input_rows: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_rows", partition), + build_mem_used: MetricBuilder::new(metrics) + .peak_memory_usage("build_mem_used", partition), } } } @@ -958,8 +1063,8 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< #[cfg(test)] mod tests { use super::*; - use crate::collect; use crate::test::TestMemoryExec; + use crate::{collect, collect_partitioned}; use arrow::array::{ DictionaryArray, Int32Array, Int64Array, StringArray, StringDictionaryBuilder, }; @@ -1200,6 +1305,102 @@ mod tests { Ok(()) } + #[tokio::test] + async fn broadcasts_right_input_to_all_left_partitions() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("id", DataType::Int32, false), + ])); + let left = TestMemoryExec::try_new_exec( + &[ + vec![make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(1), Some(4)], + vec![0, 1], + )?], + vec![make_batch( + &left_schema, + vec![Some("A"), Some("A")], + vec![Some(2), Some(5)], + vec![2, 3], + )?], + ], + Arc::clone(&left_schema), + None, + )?; + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ])); + let right = TestMemoryExec::try_new_exec( + &[vec![ + make_batch(&right_schema, vec![Some("A")], vec![Some(1)], vec![10])?, + make_batch(&right_schema, vec![Some("A")], vec![Some(3)], vec![30])?, + ]], + Arc::clone(&right_schema), + None, + )?; + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + assert_eq!(exec.properties().output_partitioning().partition_count(), 2); + assert!(matches!( + &exec.input_distribution_requirements().into_per_child()[..], + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] + )); + + let partitions = collect_partitioned( + Arc::clone(&exec) as Arc, + Arc::new(TaskContext::default()), + ) + .await?; + assert_eq!(partitions.len(), 2); + for batches in partitions { + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(10), Some(30)]); + } + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!( + metrics + .sum_by_name("build_input_batches") + .map(|value| value.as_usize()), + Some(2) + ); + assert_eq!( + metrics + .sum_by_name("build_input_rows") + .map(|value| value.as_usize()), + Some(2) + ); + assert_eq!(metrics.output_rows(), Some(4)); + Ok(()) + } + #[tokio::test] async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { let dictionary_type = @@ -1333,8 +1534,8 @@ mod tests { assert!(matches!( &exec.input_distribution_requirements().into_per_child()[..], [ - Distribution::KeyPartitioned(_), - Distribution::KeyPartitioned(_) + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition ] )); for ordering in exec.required_input_ordering() { @@ -1367,10 +1568,16 @@ mod tests { ), vec![2], )?); - assert_eq!(no_keys.output_partitioning().partition_count(), 1); + assert_eq!( + no_keys.output_partitioning().partition_count(), + exec.left().output_partitioning().partition_count() + ); assert!(matches!( &no_keys.input_distribution_requirements().into_per_child()[..], - [Distribution::SinglePartition, Distribution::SinglePartition] + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] )); for ordering in no_keys.required_input_ordering() { let requirement = ordering.expect("ASOF ordering is required").into_single(); From 9556ed3ca2b0b3aa087ccd368e092a2206baf3bb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:57:02 +0800 Subject: [PATCH 11/18] fix: remove obsolete ASOF proto import --- datafusion/proto/src/physical_plan/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 4199ae49d0ccc..eec91e33d8e0f 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -50,7 +50,7 @@ use datafusion_datasource_parquet::source::ParquetSource; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, HigherOrderUDF, Operator, ScalarUDF, WindowUDF}; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; From 23e606017377bb212aaac6681877cc1bc68d2e9a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 17:04:40 +0800 Subject: [PATCH 12/18] test: cover broadcast ASOF SQL execution --- datafusion/core/tests/sql/joins.rs | 40 ++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 15db970ad3283..9668d90162818 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -302,7 +302,7 @@ async fn unparse_cross_join() -> Result<()> { Ok(()) } -async fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { +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), @@ -328,7 +328,10 @@ async fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { ]; ctx.register_table( "trades", - Arc::new(MemTable::try_new(trades_schema, vec![trades])?), + Arc::new(MemTable::try_new( + trades_schema, + trades.into_iter().map(|batch| vec![batch]).collect(), + )?), )?; let prices_schema = Arc::new(Schema::new(vec![ @@ -374,7 +377,7 @@ async fn asof_join_all_match_directions_across_batches() -> Result<()> { .with_batch_size(2) .with_target_partitions(2); let ctx = SessionContext::new_with_config(config); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; for (op, expected) in [ ( @@ -480,10 +483,10 @@ async fn asof_join_coerces_equality_and_match_types() -> Result<()> { } #[tokio::test] -async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { +async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> { let config = SessionConfig::new().with_target_partitions(4); let ctx = SessionContext::new_with_config(config); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let df = ctx .sql( "SELECT t.trade_id, p.price FROM trades t ASOF JOIN prices p \ @@ -496,11 +499,26 @@ async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { ctx.sql(&sql).await?; let plan = df.create_physical_plan().await?; let asof = find_asof_exec(&plan).expect("physical ASOF join must be present"); - assert_eq!(asof.output_partitioning().partition_count(), 1); + 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 + ); assert!(asof.output_ordering().is_some()); assert!(matches!( &asof.input_distribution_requirements().into_per_child()[..], - [Distribution::SinglePartition, Distribution::SinglePartition] + [ + Distribution::UnspecifiedDistribution, + Distribution::SinglePartition + ] )); let batches = collect(plan, ctx.task_ctx()).await?; assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); @@ -510,7 +528,7 @@ async fn asof_join_without_equality_keys_is_single_partition() -> Result<()> { #[tokio::test] async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let batches = ctx .sql( "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ @@ -568,7 +586,7 @@ async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result #[tokio::test] async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + register_asof_test_tables(&ctx)?; let df = ctx .sql( "SELECT * FROM trades t ASOF JOIN prices p \ @@ -594,7 +612,7 @@ async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { #[tokio::test] async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + 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 \ @@ -630,7 +648,7 @@ async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { #[tokio::test] async fn asof_join_rejects_invalid_contracts() -> Result<()> { let ctx = SessionContext::new(); - register_asof_test_tables(&ctx).await?; + 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", From f7324a6284816dde1f1395026560bd6b99585252 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 17:16:57 +0800 Subject: [PATCH 13/18] bench: cover broadcast ASOF execution --- benchmarks/src/asof.rs | 9 ++-- datafusion/physical-plan/benches/asof_join.rs | 49 ++++++++++++++----- 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/benchmarks/src/asof.rs b/benchmarks/src/asof.rs index d74607550f4ed..4d824c9139476 100644 --- a/benchmarks/src/asof.rs +++ b/benchmarks/src/asof.rs @@ -25,8 +25,9 @@ use futures::StreamExt; /// Run end-to-end ASOF join benchmarks. /// -/// The cases cover ordered-input reuse, optimizer-inserted sort/repartition, -/// wide payload materialization, and descending successor matching. +/// The cases cover broadcast build reuse, left-side parallelism, +/// optimizer-inserted ordering, wide payload materialization, and descending +/// successor matching. #[derive(Debug, Args, Clone)] #[command(verbatim_doc_comment)] pub struct RunOpt { @@ -44,7 +45,7 @@ pub struct RunOpt { } const ASOF_QUERIES: &[&str] = &[ - // Q1: predecessor without equality keys, ordered range input can be reused + // Q1: predecessor without equality keys broadcasts right across left partitions r#" WITH left_input AS ( SELECT value AS ts, value AS payload FROM range(1000000) @@ -56,7 +57,7 @@ const ASOF_QUERIES: &[&str] = &[ FROM left_input l ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) "#, - // Q2: grouped predecessor, optimizer supplies repartitioning and ordering + // Q2: grouped predecessor, optimizer partitions left and coalesces right r#" WITH left_input AS ( SELECT value % 10000 AS key, diff --git a/datafusion/physical-plan/benches/asof_join.rs b/datafusion/physical-plan/benches/asof_join.rs index da39080c94f90..b33c4fbcf9d8b 100644 --- a/datafusion/physical-plan/benches/asof_join.rs +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Criterion benchmarks for the pre-sorted ASOF join kernel. +//! Criterion benchmarks for the pre-sorted broadcast ASOF join kernel. use std::sync::Arc; @@ -128,8 +128,22 @@ fn build_sorted_batches( batches } -fn make_exec(batches: &[RecordBatch], schema: &SchemaRef) -> Arc { - TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +fn partition_batches( + batches: &[RecordBatch], + partition_count: usize, +) -> Vec> { + let mut partitions = vec![Vec::new(); partition_count]; + for (index, batch) in batches.iter().enumerate() { + partitions[index % partition_count].push(batch.clone()); + } + partitions +} + +fn make_exec( + partitions: &[Vec], + schema: &SchemaRef, +) -> Arc { + TestMemoryExec::try_new_exec(partitions, Arc::clone(schema), None).unwrap() } fn do_join( @@ -177,16 +191,25 @@ fn bench_asof_join(c: &mut Criterion) { build_sorted_batches(num_rows, num_groups, 1, payload, &schema); let right_batches = build_sorted_batches(num_rows, num_groups, 0, payload, &schema); - group.bench_function( - BenchmarkId::new(payload.name(), format!("{num_rows}_rows_10_per_key")), - |b| { - b.iter(|| { - let left = make_exec(&left_batches, &schema); - let right = make_exec(&right_batches, &schema); - do_join(left, right, &rt) - }) - }, - ); + let right_partitions = vec![right_batches]; + for left_partition_count in [1, 4] { + let left_partitions = partition_batches(&left_batches, left_partition_count); + group.bench_function( + BenchmarkId::new( + payload.name(), + format!( + "{num_rows}_rows_10_per_key_{left_partition_count}_left_partitions" + ), + ), + |b| { + b.iter(|| { + let left = make_exec(&left_partitions, &schema); + let right = make_exec(&right_partitions, &schema); + do_join(left, right, &rt) + }) + }, + ); + } } group.finish(); From 339d7857107c2e87bbffe87d89d2c647dbc45c63 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 01:40:50 +0800 Subject: [PATCH 14/18] fix: account shared ASOF build buffers once --- .../physical-plan/src/joins/asof_join.rs | 90 ++++++++++++++++++- 1 file changed, 88 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index b55a21260de8f..2daa5e5c9fe67 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -62,6 +62,7 @@ use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; +use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::utils::{ compare_rows, get_row_at_idx, normalize_float_zero_scalar, }; @@ -93,7 +94,6 @@ use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricCategory, MetricsSet, RecordOutput, Time, }; -use crate::spill::get_record_batch_memory_size; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ @@ -609,9 +609,10 @@ async fn collect_right_input( metrics: AsOfJoinMetrics, ) -> Result { let schema = input.schema(); + let mut memory_counter = RecordBatchMemoryCounter::new(); let batches = input .try_fold(Vec::new(), |mut batches, batch| { - let batch_size = get_record_batch_memory_size(&batch); + let batch_size = memory_counter.count_batch(&batch); futures::future::ready(reservation.try_grow(batch_size).map(|_| { metrics.build_mem_used.add(batch_size); metrics.build_input_batches.add(1); @@ -1070,6 +1071,7 @@ mod tests { }; use arrow::datatypes::{DataType, Field, Int8Type}; use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::metrics::MetricValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -1401,6 +1403,90 @@ mod tests { Ok(()) } + #[tokio::test] + async fn shared_right_buffers_are_reserved_once() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("id", DataType::Int32, false), + ])); + let left = TestMemoryExec::try_new_exec( + &[vec![make_batch( + &left_schema, + vec![Some("A")], + vec![Some(4095)], + vec![0], + )?]], + Arc::clone(&left_schema), + None, + )?; + + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("ts", DataType::Int64, false), + Field::new("price", DataType::Int32, false), + ])); + let row_count = 4096; + let parent = make_batch( + &right_schema, + vec![Some("A"); row_count], + (0..row_count).map(|value| Some(value as i64)).collect(), + (0..row_count as i32).collect(), + )?; + let mut memory_counter = RecordBatchMemoryCounter::new(); + let retained_size = memory_counter.count_batch(&parent); + let right_batches = (0..16) + .map(|index| parent.slice(index * 256, 256)) + .collect(); + let right = TestMemoryExec::try_new_exec( + &[right_batches], + Arc::clone(&right_schema), + None, + )?; + + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + vec![( + Arc::new(PhysicalColumn::new("key", 0)), + Arc::new(PhysicalColumn::new("key", 0)), + )], + AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), + ), + vec![2], + )?); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(retained_size, 1.0) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + + let batches = collect(Arc::clone(&exec) as _, context).await?; + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(4095)]); + + let metrics = exec.metrics().expect("ASOF metrics must be present"); + assert_eq!( + metrics + .sum_by_name("build_mem_used") + .map(|value| value.as_usize()), + Some(retained_size) + ); + Ok(()) + } + #[tokio::test] async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> { let dictionary_type = From 3a4f995a75a7a47f502dcb2b65c45e124c080914 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 01:52:40 +0800 Subject: [PATCH 15/18] fix: align ASOF logical contracts --- datafusion/core/src/physical_planner.rs | 22 +++--------- datafusion/expr/src/logical_plan/builder.rs | 36 +++---------------- datafusion/expr/src/logical_plan/plan.rs | 35 +++++++++++++++--- .../optimizer/src/optimize_projections/mod.rs | 27 +++----------- 4 files changed, 42 insertions(+), 78 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d852d56b1e52d..43ce2e6d71b33 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -94,8 +94,8 @@ use datafusion_expr::physical_planning_context::{ use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, - FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType, - StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, + FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, + WindowFrame, WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, @@ -1903,22 +1903,8 @@ impl DefaultPhysicalPlanner { planning_ctx, )?, ); - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .map(|(_, right)| { - let column = right.get_as_join_column().ok_or_else(|| { - internal_datafusion_err!("ASOF USING key is not a column") - })?; - join.right.schema().index_of_column(column) - }) - .collect::>>()? - } else { - HashSet::new() - }; - let right_output_indices = (0..join.right.schema().fields().len()) - .filter(|index| !omitted_right.contains(index)) - .collect(); + let right_output_indices = + (0..join.right.schema().fields().len()).collect(); Arc::new(AsOfJoinExec::try_new( physical_left, physical_right, diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 14603c3f62e24..67436332a123f 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1840,38 +1840,10 @@ pub fn build_join_schema( /// Creates the schema for a left-preserving ASOF join. /// -/// `ON` emits all left fields followed by nullable right fields. `USING` emits -/// each equality key once by omitting the corresponding right field. -pub fn build_asof_join_schema( - left: &DFSchema, - right: &DFSchema, - on: &[(Expr, Expr)], - join_constraint: JoinConstraint, -) -> Result { - let omitted_right_indices = if join_constraint == JoinConstraint::Using { - on.iter() - .map(|(_, right_expr)| { - let column = right_expr.get_as_join_column().ok_or_else(|| { - plan_datafusion_err!("ASOF USING keys must be columns") - })?; - right.index_of_column(column) - }) - .collect::>>()? - } else { - HashSet::new() - }; - - let full_schema = build_join_schema(left, right, &JoinType::Left)?; - let left_len = left.fields().len(); - let fields = full_schema - .iter() - .enumerate() - .filter(|(index, _)| { - *index < left_len || !omitted_right_indices.contains(&(*index - left_len)) - }) - .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field))) - .collect(); - DFSchema::new_with_metadata(fields, full_schema.metadata().clone())? +/// 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)? .with_functional_dependencies(left.functional_dependencies().clone()) } diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 5cee0ad97bee2..3635aeb83b6e0 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -240,9 +240,6 @@ pub enum LogicalPlan { /// Join two logical plans on one or more join columns. /// This is used to implement SQL `JOIN` Join(Join), - /// Match each left row with at most one ordered row from the right input. - /// This is used to implement SQL `ASOF JOIN`. - AsOfJoin(AsOfJoin), /// Repartitions the input based on a partitioning scheme. This is /// used to add parallelism and is sometimes referred to as an /// "exchange" operator in other systems @@ -299,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 { @@ -4455,8 +4455,7 @@ impl AsOfJoin { return plan_err!("ASOF USING keys must be columns"); } - let schema = - build_asof_join_schema(left.schema(), right.schema(), &on, join_constraint)?; + let schema = build_asof_join_schema(left.schema(), right.schema())?; Ok(Self { left, right, @@ -6867,6 +6866,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/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index da5c715d280cb..74ce2ad374af8 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,9 +24,8 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinConstraint, JoinType, Result, - assert_eq_or_internal_err, get_required_group_by_exprs_indices, - internal_datafusion_err, internal_err, + Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, + get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; use datafusion_expr::{ @@ -410,31 +409,13 @@ fn optimize_projections( } LogicalPlan::AsOfJoin(join) => { let left_len = join.left.schema().fields().len(); - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .map(|(_, right)| { - let column = right.get_as_join_column().ok_or_else(|| { - internal_datafusion_err!("ASOF USING key is not a column") - })?; - join.right.schema().index_of_column(column) - }) - .collect::>>()? - } else { - std::collections::HashSet::new() - }; - let right_output_indices = (0..join.right.schema().fields().len()) - .filter(|index| !omitted_right.contains(index)) - .collect::>(); 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 if let Some(right_index) = - right_output_indices.get(*index - left_len) - { - right_required.push(*right_index); + } else { + right_required.push(*index - left_len); } } let left_indices = RequiredIndices::new_from_indices(left_required) From fd8fddb87ca3e31485b0b778477f0b60984f390a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 02:09:23 +0800 Subject: [PATCH 16/18] fix: align ASOF USING and broadcast docs --- datafusion/core/tests/sql/joins.rs | 30 +++++++++++++++++-- datafusion/sql/src/unparser/plan.rs | 19 ------------ .../sqllogictest/test_files/asof_join.slt | 19 +++++++++++- docs/source/user-guide/sql/select.md | 10 +++++-- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 9668d90162818..2504c0077e6ff 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -584,7 +584,7 @@ async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result } #[tokio::test] -async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { +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 @@ -599,13 +599,39 @@ async fn asof_join_using_merges_key_and_unparser_round_trips() -> Result<()> { .iter() .map(|field| field.name()) .collect::>(), - vec!["symbol", "ts", "trade_id", "ts", "price"] + 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(()) } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 12d9b65be1a7d..9c0e7786fd0e1 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -1885,25 +1885,6 @@ impl Unparser<'_> { let right_projection = right_projection.ok_or_else(|| { internal_datafusion_err!("ASOF right projection is missing") })?; - let omitted_right = if join.join_constraint == JoinConstraint::Using { - join.on - .iter() - .filter_map(|(_, right)| right.get_as_join_column()) - .collect::>() - } else { - vec![] - }; - let right_projection = right_projection.into_iter().filter(|item| { - let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(ids)) = - item - else { - return true; - }; - let Some(name) = ids.last() else { - return true; - }; - !omitted_right.iter().any(|column| column.name == name.value) - }); select.projection( left_projection .into_iter() diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index ec300fca74f24..62a9e3cf35ceb 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -104,7 +104,7 @@ ORDER BY l.id; 6 3 NULL NULL 7 NULL NULL NULL -# USING merges the equality key into one output column. +# USING exposes one unqualified equality key. query TIIT SELECT grp, l.id, r.ts, r.val FROM asof_left l @@ -121,6 +121,23 @@ B 5 6 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 diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index 50235c169fb20..817f98a8b236f 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -404,8 +404,14 @@ expression from the right input using one of the following operators: An optional `ON` clause containing equality conditions combined with `AND`, or a `USING` clause, divides rows into equality groups before the ordered match. -Without equality keys, all rows belong to one group and DataFusion executes the -join in a single partition. +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 From 34552d158eef2737c9022de570dc6af5d54e5aa2 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:50:01 +0800 Subject: [PATCH 17/18] perf: compare ASOF keys without scalar materialization --- .../physical-plan/src/joins/asof_join.rs | 178 ++++++++++++------ 1 file changed, 124 insertions(+), 54 deletions(-) diff --git a/datafusion/physical-plan/src/joins/asof_join.rs b/datafusion/physical-plan/src/joins/asof_join.rs index 2daa5e5c9fe67..6da25bf50c26a 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -58,16 +58,15 @@ use std::fmt::Formatter; use std::sync::Arc; use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::buffer::NullBuffer; use arrow::compute::{SortOptions, interleave}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::utils::memory::RecordBatchMemoryCounter; -use datafusion_common::utils::{ - compare_rows, get_row_at_idx, normalize_float_zero_scalar, -}; +use datafusion_common::utils::normalize_float_zero_scalar; use datafusion_common::{ - ColumnStatistics, JoinType, Result, ScalarValue, Statistics, + ColumnStatistics, JoinType, NullEquality, Result, ScalarValue, Statistics, assert_eq_or_internal_err, internal_err, plan_err, }; use datafusion_execution::TaskContext; @@ -88,7 +87,9 @@ use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; -use crate::joins::utils::{JoinOn, OnceAsync, build_join_schema}; +use crate::joins::utils::{ + JoinKeyComparator, JoinOn, OnceAsync, build_join_schema, matchable_join_keys, +}; use crate::memory::MemoryStream; use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, @@ -633,7 +634,8 @@ async fn collect_right_input( struct Candidate { batch: Arc, row: usize, - group: Vec, + key_arrays: Arc<[ArrayRef]>, + key_batch_id: usize, } struct InputCursor { @@ -641,8 +643,10 @@ struct InputCursor { key_exprs: Vec, match_expr: PhysicalExprRef, batch: Option>, - key_arrays: Vec, + key_arrays: Arc<[ArrayRef]>, + key_validity: Option, match_array: Option, + key_batch_id: usize, row: usize, eof: bool, } @@ -658,8 +662,10 @@ impl InputCursor { key_exprs, match_expr, batch: None, - key_arrays: vec![], + key_arrays: Arc::from([]), + key_validity: None, match_array: None, + key_batch_id: 0, row: 0, eof: false, } @@ -673,7 +679,8 @@ impl InputCursor { return Ok(true); } self.batch = None; - self.key_arrays.clear(); + self.key_arrays = Arc::from([]); + self.key_validity = None; self.match_array = None; self.row = 0; if self.eof { @@ -688,23 +695,28 @@ impl InputCursor { } let batch = Arc::new(batch); let _timer = elapsed_compute.timer(); - self.key_arrays = self + let key_arrays = self .key_exprs .iter() .map(|expr| expr.evaluate(&batch)?.into_array(batch.num_rows())) - .collect::>()?; + .collect::>>()?; + self.key_validity = + matchable_join_keys(&key_arrays, NullEquality::NullEqualsNothing); + self.key_arrays = key_arrays.into(); self.match_array = Some( self.match_expr .evaluate(&batch)? .into_array(batch.num_rows())?, ); + self.key_batch_id += 1; self.batch = Some(batch); } } - fn group(&self) -> Result> { - get_row_at_idx(&self.key_arrays, self.row) - .map(|row| row.into_iter().map(normalize_float_zero_scalar).collect()) + fn group_has_null(&self) -> bool { + self.key_validity + .as_ref() + .is_some_and(|validity| validity.is_null(self.row)) } fn match_value(&self) -> Result { @@ -850,6 +862,8 @@ struct AsOfJoinStreamState { right_output_indices: Vec, candidate: Option, group_sort_options: Vec, + input_group_comparator: Option<(usize, usize, JoinKeyComparator)>, + candidate_group_comparator: Option<(usize, usize, JoinKeyComparator)>, pending_left: PendingRows, pending_right: PendingRows, batch_size: usize, @@ -883,11 +897,76 @@ impl AsOfJoinStreamState { right_output_indices, candidate: None, group_sort_options, + input_group_comparator: None, + candidate_group_comparator: None, batch_size: batch_size.max(1), metrics, } } + fn compare_input_groups(&mut self) -> Result { + if self.group_sort_options.is_empty() { + return Ok(Ordering::Equal); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_batch_id = self.right.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .input_group_comparator + .as_ref() + .is_none_or(|(right, left, _)| { + *right != right_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + self.right.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.input_group_comparator = + Some((right_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .input_group_comparator + .as_ref() + .expect("ASOF input group comparator must be initialized"); + Ok(comparator.compare(self.right.row, self.left.row)) + } + + fn candidate_is_other_group(&mut self) -> Result { + let Some(candidate) = &self.candidate else { + return Ok(false); + }; + if self.group_sort_options.is_empty() { + return Ok(false); + } + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let candidate_batch_id = candidate.key_batch_id; + let left_batch_id = self.left.key_batch_id; + if self + .candidate_group_comparator + .as_ref() + .is_none_or(|(candidate, left, _)| { + *candidate != candidate_batch_id || *left != left_batch_id + }) + { + let comparator = JoinKeyComparator::new( + candidate.key_arrays.as_ref(), + self.left.key_arrays.as_ref(), + &self.group_sort_options, + NullEquality::NullEqualsNothing, + )?; + self.candidate_group_comparator = + Some((candidate_batch_id, left_batch_id, comparator)); + } + let (_, _, comparator) = self + .candidate_group_comparator + .as_ref() + .expect("ASOF candidate group comparator must be initialized"); + Ok(comparator.compare(candidate.row, self.left.row) != Ordering::Equal) + } + async fn next_batch(&mut self) -> Result> { loop { if self.pending_left.len() >= self.batch_size { @@ -905,25 +984,20 @@ impl AsOfJoinStreamState { return Ok(None); } - let (left_group, left_match) = { + let left_match = { let _timer = self.metrics.baseline.elapsed_compute().timer(); - (self.left.group()?, self.left.match_value()?) + self.left.match_value()? }; - if left_match.is_null() || left_group.iter().any(ScalarValue::is_null) { + if left_match.is_null() || self.left.group_has_null() { self.candidate = None; + self.candidate_group_comparator = None; self.push_current_left(None)?; self.left.advance(); continue; } - let candidate_is_other_group = if let Some(candidate) = &self.candidate { - let _timer = self.metrics.baseline.elapsed_compute().timer(); - compare_rows(&candidate.group, &left_group, &self.group_sort_options)? - != Ordering::Equal - } else { - false - }; - if candidate_is_other_group { + if self.candidate_is_other_group()? { self.candidate = None; + self.candidate_group_comparator = None; } loop { @@ -934,34 +1008,27 @@ impl AsOfJoinStreamState { { break; } - let action = { - let _timer = self.metrics.baseline.elapsed_compute().timer(); - let right_group = self.right.group()?; - if right_group.iter().any(ScalarValue::is_null) { - RightAction::Advance - } else { - match compare_rows( - &right_group, - &left_group, - &self.group_sort_options, - )? { - Ordering::Less => RightAction::Advance, - Ordering::Greater => RightAction::Stop, - Ordering::Equal => { - let right_match = self.right.match_value()?; - if right_match.is_null() { - RightAction::Advance - } else if is_eligible(self.op, &left_match, &right_match)? - { - let (batch, row) = self.right.batch_row()?; - RightAction::Candidate(Candidate { - batch, - row, - group: right_group, - }) - } else { - RightAction::Stop - } + let action = if self.right.group_has_null() { + RightAction::Advance + } else { + match self.compare_input_groups()? { + Ordering::Less => RightAction::Advance, + Ordering::Greater => RightAction::Stop, + Ordering::Equal => { + let _timer = self.metrics.baseline.elapsed_compute().timer(); + let right_match = self.right.match_value()?; + if right_match.is_null() { + RightAction::Advance + } else if is_eligible(self.op, &left_match, &right_match)? { + let (batch, row) = self.right.batch_row()?; + RightAction::Candidate(Candidate { + batch, + row, + key_arrays: Arc::clone(&self.right.key_arrays), + key_batch_id: self.right.key_batch_id, + }) + } else { + RightAction::Stop } } } @@ -969,6 +1036,9 @@ impl AsOfJoinStreamState { match action { RightAction::Advance => self.right.advance(), RightAction::Candidate(candidate) => { + // Replacing the candidate selects the nearest eligible row. + // Equal match values have no secondary ordering, so which + // tied row wins is intentionally nondeterministic. self.candidate = Some(candidate); self.right.advance(); } @@ -1057,7 +1127,7 @@ fn is_eligible(op: Operator, left: &ScalarValue, right: &ScalarValue) -> Result< Operator::GtEq => ordering != Ordering::Greater, Operator::Lt => ordering == Ordering::Greater, Operator::LtEq => ordering != Ordering::Less, - _ => false, + _ => unreachable!("ASOF match operator is validated by try_new"), }) } From d618879332b4c0fea6fc981650541f3c2874cc55 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:50:13 +0800 Subject: [PATCH 18/18] test: cover temporal and multi-partition ASOF joins --- datafusion/core/tests/sql/joins.rs | 33 ++++++- .../sqllogictest/test_files/asof_join.slt | 94 +++++++++---------- 2 files changed, 75 insertions(+), 52 deletions(-) diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 2504c0077e6ff..842f5690e3828 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -17,12 +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::*; @@ -359,7 +359,10 @@ fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> { ]; ctx.register_table( "prices", - Arc::new(MemTable::try_new(prices_schema, vec![prices])?), + Arc::new(MemTable::try_new( + prices_schema, + prices.into_iter().map(|batch| vec![batch]).collect(), + )?), )?; Ok(()) } @@ -483,13 +486,14 @@ async fn asof_join_coerces_equality_and_match_types() -> Result<()> { } #[tokio::test] -async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> { +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 prices p \ + "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?; @@ -512,6 +516,11 @@ async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> 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()[..], @@ -521,7 +530,21 @@ async fn asof_join_without_equality_keys_broadcasts_right_input() -> Result<()> ] )); let batches = collect(plan, ctx.task_ctx()).await?; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 6); + assert_batches_sorted_eq!( + [ + "+----------+-------+", + "| trade_id | price |", + "+----------+-------+", + "| 1 | |", + "| 2 | 40 |", + "| 3 | 60 |", + "| 4 | 20 |", + "| 5 | 60 |", + "| 6 | 20 |", + "+----------+-------+", + ], + &batches + ); Ok(()) } diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index 62a9e3cf35ceb..46cf6cdf10bf7 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -16,28 +16,28 @@ # under the License. statement ok -CREATE TABLE asof_left(id INT, grp TEXT, ts INT) AS VALUES - (1, 'A', 1), - (2, 'A', 4), - (3, 'A', 7), - (4, 'B', 2), - (5, 'B', 8), - (6, NULL, 3), +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 INT, val TEXT) AS VALUES - ('A', 2, 'a2'), - ('A', 4, 'a4'), - ('A', 6, 'a6'), - ('B', 1, 'b1'), - ('B', 6, 'b6'), - (NULL, 2, 'null-group'), +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 IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -45,16 +45,16 @@ MATCH_CONDITION (l.ts >= r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 NULL NULL -2 4 4 a4 -3 7 6 a6 -4 2 1 b1 -5 8 6 b6 -6 3 NULL NULL +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 IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -62,16 +62,16 @@ MATCH_CONDITION (l.ts > r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 NULL NULL -2 4 2 a2 -3 7 6 a6 -4 2 1 b1 -5 8 6 b6 -6 3 NULL NULL +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 IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -79,16 +79,16 @@ MATCH_CONDITION (l.ts <= r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 2 a2 -2 4 4 a4 -3 7 NULL NULL -4 2 6 b6 -5 8 NULL NULL -6 3 NULL NULL +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 IIIT +query IPPT SELECT l.id, l.ts, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -96,16 +96,16 @@ MATCH_CONDITION (l.ts < r.ts) ON l.grp = r.grp ORDER BY l.id; ---- -1 1 2 a2 -2 4 6 a6 -3 7 NULL NULL -4 2 6 b6 -5 8 NULL NULL -6 3 NULL NULL +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 TIIT +query TIPT SELECT grp, l.id, r.ts, r.val FROM asof_left l ASOF JOIN asof_right r @@ -114,10 +114,10 @@ USING (grp) ORDER BY l.id; ---- A 1 NULL NULL -A 2 4 a4 -A 3 6 a6 -B 4 1 b1 -B 5 6 b6 +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