diff --git a/benchmarks/src/asof.rs b/benchmarks/src/asof.rs new file mode 100644 index 0000000000000..3b7acc1bfa4a3 --- /dev/null +++ b/benchmarks/src/asof.rs @@ -0,0 +1,226 @@ +// 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 broadcast-side size asymmetry, equality-key cardinality and +/// skew, left-side parallelism, optimizer-inserted ordering, wide payload +/// materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 6). 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: small broadcast input and a large, equality-free probe input + 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(10000) + ) + 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 partitions left and coalesces right + 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) + "#, + // Q5: a large broadcast input exposes the opposite size asymmetry + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(100000) + ), + 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) + "#, + // Q6: low-cardinality, heavily skewed equality keys + r#" + WITH left_input AS ( + SELECT CASE WHEN value % 100 < 95 THEN 0 ELSE value % 16 + 1 END AS key, + value / 100 + 1 AS ts, + value AS payload + FROM range(1000000) + ), + right_input AS ( + SELECT CASE WHEN value % 100 < 95 THEN 0 ELSE value % 16 + 1 END AS key, + value / 100 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 + "#, +]; + +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 29cc8d63d2d8d..1607a52323e90 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, statistics, + asof, cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, statistics, tpcds, tpch, }; @@ -45,6 +45,7 @@ struct Cli { #[derive(Debug, Subcommand)] enum Options { + Asof(asof::RunOpt), Cancellation(cancellation::RunOpt), Clickbench(clickbench::RunOpt), Dict(dict::RunOpt), @@ -67,6 +68,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 0b3783421f840..a9bc8c31b359c 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/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/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index dd32830dd5eb8..7e57b8e330020 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -456,33 +456,6 @@ async fn asof_join_all_match_directions_across_batches() -> Result<()> { Ok(()) } -#[tokio::test] -async fn asof_join_coerces_equality_and_match_types() -> Result<()> { - let ctx = SessionContext::new(); - let batches = ctx - .sql( - "SELECT t.id, p.price \ - FROM (VALUES (CAST(1 AS INT), CAST(4 AS INT), 7)) t(k, ts, id) \ - ASOF JOIN \ - (VALUES (CAST(1 AS BIGINT), CAST(2 AS BIGINT), 20)) p(k, ts, price) \ - MATCH_CONDITION (t.ts >= p.ts) ON t.k = p.k", - ) - .await? - .collect() - .await?; - assert_batches_eq!( - [ - "+----+-------+", - "| id | price |", - "+----+-------+", - "| 7 | 20 |", - "+----+-------+", - ], - &batches - ); - Ok(()) -} - #[tokio::test] async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { let config = SessionConfig::new().with_target_partitions(4); @@ -519,7 +492,10 @@ async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { .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!(asof.children()[0].output_ordering().is_some()); + // The embedded output projection drops the match column, so the join no + // longer exposes its input ordering. + assert!(asof.output_ordering().is_none()); assert!(matches!( &asof.input_distribution_requirements().into_per_child()[..], [ @@ -546,28 +522,6 @@ async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> { Ok(()) } -#[tokio::test] -async fn asof_join_explain_names_equality_and_match_conditions() -> Result<()> { - let ctx = SessionContext::new(); - register_asof_test_tables(&ctx)?; - let batches = ctx - .sql( - "EXPLAIN SELECT t.trade_id, p.price FROM trades t \ - ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \ - ON t.symbol = p.symbol", - ) - .await? - .collect() - .await?; - let explain = arrow::util::pretty::pretty_format_batches(&batches)?.to_string(); - assert_contains!(explain.as_str(), "AsOf Join: match=[t.ts >= p.ts]"); - assert_contains!(explain.as_str(), "on=[t.symbol = p.symbol]"); - assert_contains!(explain.as_str(), "AsOfJoinExec:"); - assert_contains!(explain.as_str(), "on=[(symbol = symbol)]"); - assert_contains!(explain.as_str(), "match=[ts >= ts]"); - Ok(()) -} - #[tokio::test] async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result<()> { let ctx = SessionContext::new(); @@ -605,7 +559,7 @@ async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result } #[tokio::test] -async fn asof_join_using_preserves_key_access_and_unparser_round_trips() -> Result<()> { +async fn asof_join_using_unparser_round_trips() -> Result<()> { let ctx = SessionContext::new(); register_asof_test_tables(&ctx)?; let df = ctx @@ -614,45 +568,11 @@ async fn asof_join_using_preserves_key_access_and_unparser_round_trips() -> Resu MATCH_CONDITION (t.ts >= p.ts) USING (symbol)", ) .await?; - assert_eq!( - df.schema() - .fields() - .iter() - .map(|field| field.name()) - .collect::>(), - vec!["ts", "trade_id", "symbol", "ts", "price"] - ); let sql = plan_to_sql(df.logical_plan())?.to_string(); assert!(sql.contains("ASOF JOIN")); assert!(sql.contains("MATCH_CONDITION")); assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}"); ctx.sql(&sql).await?; - - let batches = ctx - .sql( - "SELECT t.trade_id, t.symbol AS left_symbol, p.symbol AS right_symbol \ - FROM trades t ASOF JOIN prices p \ - MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \ - ORDER BY t.trade_id", - ) - .await? - .collect() - .await?; - assert_batches_eq!( - [ - "+----------+-------------+--------------+", - "| trade_id | left_symbol | right_symbol |", - "+----------+-------------+--------------+", - "| 1 | A | |", - "| 2 | A | A |", - "| 3 | A | A |", - "| 4 | B | B |", - "| 5 | B | B |", - "| 6 | | |", - "+----------+-------------+--------------+", - ], - &batches - ); Ok(()) } @@ -691,20 +611,3 @@ async fn asof_join_unparser_preserves_right_preselection() -> Result<()> { } Ok(()) } - -#[tokio::test] -async fn asof_join_rejects_invalid_contracts() -> Result<()> { - let ctx = SessionContext::new(); - register_asof_test_tables(&ctx)?; - for sql in [ - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts = p.ts) ON t.symbol = p.symbol", - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (p.ts >= t.ts) ON t.symbol = p.symbol", - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON t.symbol > p.symbol", - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (1 >= p.ts) ON t.symbol = p.symbol", - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) ON 1 = 1", - "SELECT * FROM trades t ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts AND t.ts > p.ts) ON t.symbol = p.symbol", - ] { - assert!(ctx.sql(sql).await.is_err(), "query should fail: {sql}"); - } - Ok(()) -} diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 36aa67bbe7e3d..7d21eac8ee95f 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1898,7 +1898,12 @@ pub fn build_join_schema( /// 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) + // ASOF emits exactly one output row for each left row. Unlike a general + // left join, it cannot duplicate left rows, so left dependencies retain + // their modes. Right dependencies do not hold because one right row may + // match multiple left rows. + build_join_schema(left, right, &JoinType::Left)? + .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 @@ -2436,7 +2441,8 @@ mod tests { use crate::test::function_stub::sum; use datafusion_common::{ - Constraint, DataFusionError, RecursionUnnestOption, SchemaError, + Constraint, DataFusionError, Dependency, FunctionalDependence, + FunctionalDependencies, RecursionUnnestOption, SchemaError, }; use insta::assert_snapshot; @@ -3202,6 +3208,40 @@ mod tests { Ok(()) } + #[test] + fn asof_join_schema_preserves_only_left_dependencies() -> Result<()> { + let left = DFSchema::try_from_qualified_schema( + "left", + &Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ]), + )? + .with_functional_dependencies(FunctionalDependencies::new(vec![ + FunctionalDependence::new(vec![0], vec![0, 1], false) + .with_mode(Dependency::Single), + ]))?; + let right = DFSchema::try_from_qualified_schema( + "right", + &Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ]), + )? + .with_functional_dependencies(FunctionalDependencies::new(vec![ + FunctionalDependence::new(vec![0], vec![0, 1], false) + .with_mode(Dependency::Single), + ]))?; + + let schema = build_asof_join_schema(&left, &right)?; + + assert_eq!( + schema.functional_dependencies(), + left.functional_dependencies() + ); + Ok(()) + } + #[test] fn test_values_metadata() -> Result<()> { let metadata: HashMap = diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 8a1dcc12ef874..754724250fd27 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1115,6 +1115,39 @@ 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) => { + // ASOF emits exactly one output row per left row without + // changing left values, so deterministic left-only predicates + // can run before matching. Right or mixed predicates must stay + // above the join because unmatched right fields are NULL-padded. + 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 @@ -1444,7 +1477,7 @@ mod tests { use datafusion_expr::expr::ScalarFunction; use datafusion_expr::logical_plan::table_scan; use datafusion_expr::{ - ColumnarValue, ExprFunctionExt, Extension, LogicalPlanBuilder, + ColumnarValue, ExprFunctionExt, Extension, LogicalPlanBuilder, Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableScan, TableSource, TableType, UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, in_list, in_subquery, lit, @@ -2490,6 +2523,67 @@ mod tests { ) } + #[test] + fn asof_join_pushes_only_deterministic_left_filters() -> Result<()> { + let left = test_table_scan_with_name("test1")?; + let right = test_table_scan_with_name("test2")?; + let plan = LogicalPlanBuilder::from(left) + .asof_join_on( + right, + Some(col("test1.a").eq(col("test2.a"))), + col("test1.b").gt_eq(col("test2.b")), + )? + .filter( + col("test1.a") + .gt(lit(1u32)) + .and(col("test2.c").gt(lit(2u32))) + .and(col("test1.c").lt(col("test2.c"))), + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Filter: test2.c > UInt32(2) AND test1.c < test2.c + AsOf Join: match=[test1.b >= test2.b], constraint=On, on=[test1.a = test2.a] + TableScan: test1, full_filters=[test1.a > UInt32(1)] + TableScan: test2 + " + ) + } + + #[test] + fn asof_join_keeps_volatile_left_filter_above_join() -> Result<()> { + let left = test_table_scan_with_name("test1")?; + let right = test_table_scan_with_name("test2")?; + let fun = ScalarUDF::new_from_impl(TestScalarUDF { + signature: Signature::exact(vec![DataType::UInt32], Volatility::Volatile), + }); + let predicate = Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::new(fun), + vec![col("test1.c")], + )) + .gt(lit(0)); + let plan = LogicalPlanBuilder::from(left) + .asof_join_on( + right, + Some(col("test1.a").eq(col("test2.a"))), + col("test1.b").gt_eq(col("test2.b")), + )? + .filter(predicate)? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Filter: TestScalarUDF(test1.c) > Int32(0) + AsOf Join: match=[test1.b >= test2.b], constraint=On, on=[test1.a = test2.a] + TableScan: test1 + TableScan: test2 + " + ) + } + /// post-on-join predicates on a column common to both sides is pushed to both sides #[test] fn filter_on_join_on_common_independent() -> Result<()> { diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index f2f9285de560a..676a7165a6a93 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -32,6 +32,7 @@ mod like; mod literal; mod negative; mod no_op; +mod normalize_float_zero; mod not; mod similar_to_pattern; mod try_cast; @@ -58,6 +59,7 @@ pub use like::{LikeExpr, like}; pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; +pub use normalize_float_zero::NormalizeFloatZeroExpr; pub use not::{NotExpr, not}; pub(crate) use similar_to_pattern::translate_scalar; pub use similar_to_pattern::{SqlSimilarToPattern, sql_similar_to_regex}; diff --git a/datafusion/physical-expr/src/expressions/normalize_float_zero.rs b/datafusion/physical-expr/src/expressions/normalize_float_zero.rs new file mode 100644 index 0000000000000..7e4dc3db150d1 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/normalize_float_zero.rs @@ -0,0 +1,255 @@ +// 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. + +//! Floating-point signed-zero normalization expression. + +use std::hash::Hash; +use std::sync::Arc; + +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; +use datafusion_expr::ColumnarValue; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::ExprProperties; + +use crate::PhysicalExpr; + +/// Replaces floating-point `-0.0` values with `+0.0`. +/// +/// Other values and data types are returned unchanged. This expression is +/// order-preserving but not strictly order-preserving because it collapses the +/// two signed-zero representations. +#[derive(Debug, Eq)] +pub struct NormalizeFloatZeroExpr { + arg: Arc, +} + +impl PartialEq for NormalizeFloatZeroExpr { + fn eq(&self, other: &Self) -> bool { + self.arg.eq(&other.arg) + } +} + +impl Hash for NormalizeFloatZeroExpr { + fn hash(&self, state: &mut H) { + self.arg.hash(state); + } +} + +impl NormalizeFloatZeroExpr { + /// Creates a signed-zero normalization expression. + pub fn new(arg: Arc) -> Self { + Self { arg } + } + + /// Returns the input expression. + pub fn arg(&self) -> &Arc { + &self.arg + } +} + +impl std::fmt::Display for NormalizeFloatZeroExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "normalize_float_zero({})", self.arg) + } +} + +impl PhysicalExpr for NormalizeFloatZeroExpr { + fn data_type(&self, input_schema: &Schema) -> Result { + self.arg.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.arg.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + Ok(match self.arg.evaluate(batch)? { + ColumnarValue::Array(array) => { + ColumnarValue::Array(normalize_float_zero(&array)) + } + ColumnarValue::Scalar(scalar) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(scalar)) + } + }) + } + + fn return_field(&self, input_schema: &Schema) -> Result { + self.arg.return_field(input_schema) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.arg] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } + + fn evaluate_bounds(&self, children: &[&Interval]) -> Result { + Interval::try_new( + normalize_float_zero_scalar(children[0].lower().clone()), + normalize_float_zero_scalar(children[0].upper().clone()), + ) + } + + fn get_properties(&self, children: &[ExprProperties]) -> Result { + let range = self.evaluate_bounds(&[&children[0].range])?; + Ok(children[0] + .clone() + .with_range(range) + .with_strictly_order_preserving(false)) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "normalize_float_zero(")?; + self.arg.fmt_sql(f)?; + write!(f, ")") + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::NormalizeFloatZero( + Box::new(protobuf::PhysicalNormalizeFloatZeroNode { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }), + )), + })) + } +} + +#[cfg(feature = "proto")] +impl NormalizeFloatZeroExpr { + /// Reconstructs a [`NormalizeFloatZeroExpr`] from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::NormalizeFloatZero, + "NormalizeFloatZero", + ); + let arg = ctx.decode_required_expression( + node.expr.as_deref(), + "NormalizeFloatZeroExpr", + "expr", + )?; + Ok(Arc::new(Self::new(arg))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ArrayRef, AsArray, Float64Array}; + use datafusion_common::ScalarValue; + use datafusion_expr::sort_properties::SortProperties; + use half::f16; + + use crate::expressions::{Column, Literal}; + + #[test] + fn normalizes_array_and_scalar_signed_zero() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Float64Array::from(vec![-0.0, 0.0, 1.0])) as ArrayRef, + )])?; + let expr = NormalizeFloatZeroExpr::new(Arc::new(Column::new("a", 0))); + let ColumnarValue::Array(array) = expr.evaluate(&batch)? else { + panic!("column evaluation must return an array"); + }; + let array = array.as_primitive::(); + assert_eq!(array.value(0).to_bits(), 0.0_f64.to_bits()); + assert_eq!(array.value(1).to_bits(), 0.0_f64.to_bits()); + assert_eq!(array.value(2), 1.0); + + let expr = NormalizeFloatZeroExpr::new(Arc::new(Literal::new( + ScalarValue::Float64(Some(-0.0)), + ))); + let ColumnarValue::Scalar(ScalarValue::Float64(Some(value))) = + expr.evaluate(&RecordBatch::new_empty(Arc::new(Schema::empty())))? + else { + panic!("literal evaluation must return a Float64 scalar"); + }; + assert_eq!(value.to_bits(), 0.0_f64.to_bits()); + Ok(()) + } + + #[test] + fn normalizes_signed_zero_bounds_and_properties() -> Result<()> { + let batch = RecordBatch::new_empty(Arc::new(Schema::empty())); + let cases = [ + ( + ScalarValue::Float16(Some(f16::NEG_ZERO)), + ScalarValue::Float16(Some(f16::ZERO)), + ), + ( + ScalarValue::Float32(Some(-0.0)), + ScalarValue::Float32(Some(0.0)), + ), + ( + ScalarValue::Float64(Some(-0.0)), + ScalarValue::Float64(Some(0.0)), + ), + ]; + + for (negative_zero, positive_zero) in cases { + let child_range = + Interval::try_new(negative_zero.clone(), negative_zero.clone())?; + let expected_range = + Interval::try_new(positive_zero.clone(), positive_zero.clone())?; + let child_properties = ExprProperties::new_unknown() + .with_order(SortProperties::Singleton) + .with_range(child_range.clone()) + .with_preserves_lex_ordering(true) + .with_strictly_order_preserving(true); + let expr = NormalizeFloatZeroExpr::new(Arc::new(Literal::new(negative_zero))); + + let ColumnarValue::Scalar(value) = expr.evaluate(&batch)? else { + panic!("literal evaluation must return a scalar"); + }; + let bounds = expr.evaluate_bounds(&[&child_range])?; + assert_eq!(value, positive_zero); + assert_eq!(bounds, expected_range); + assert!(bounds.contains_value(&value)?); + + let properties = expr.get_properties(&[child_properties])?; + assert_eq!(properties.sort_properties, SortProperties::Singleton); + assert_eq!(properties.range, expected_range); + assert!(properties.preserves_lex_ordering); + assert!(!properties.strictly_order_preserving); + } + Ok(()) + } +} diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index c16a8de97e15a..d8fdc8248f988 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -149,6 +149,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..65f57dc43753d --- /dev/null +++ b/datafusion/physical-plan/benches/asof_join.rs @@ -0,0 +1,219 @@ +// 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 broadcast 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 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( + 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), + Some(vec![0, 1, 2, 5]), + ) + .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); + 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(); +} + +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 index abf1b37f1f8e4..8fdc576bebfc3 100644 --- a/datafusion/physical-plan/src/joins/asof_join.rs +++ b/datafusion/physical-plan/src/joins/asof_join.rs @@ -91,7 +91,9 @@ 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::expressions::{ + Column as PhysicalColumn, NormalizeFloatZeroExpr, +}; use datafusion_physical_expr::projection::{ProjectionMapping, ProjectionRef}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr_common::physical_expr::{ @@ -110,6 +112,7 @@ use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, RecordOutput, Time, }; +use crate::projection::{EmbeddedProjection, ProjectionExec, try_embed_projection}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ @@ -165,10 +168,8 @@ impl AsOfJoinExec { /// /// The match operator must be `<`, `<=`, `>`, or `>=`. Equality and match /// expressions must be deterministic, reference only their corresponding - /// input, and have matching input types. Equality types must support hashing; - /// floating-point equality keys are not supported because Arrow sorting - /// distinguishes signed zero while SQL equality does not. Projection indices - /// refer to the full left-then-right join schema. + /// input, and have matching input types. Equality types must support hashing. + /// Projection indices refer to the full left-then-right join schema. /// /// The logical ASOF constructor validates the corresponding pre-coercion /// contract. Keep the shared operator, side-ownership, and determinism checks @@ -196,24 +197,33 @@ impl AsOfJoinExec { descending, nulls_first: true, }; - let mut left_sort_exprs = on - .iter() - .map(|(left, _)| PhysicalSortExpr { - expr: Arc::clone(left), + let mut left_sort_exprs = Vec::with_capacity(on.len() + 1); + let mut right_sort_exprs = Vec::with_capacity(on.len() + 1); + for (left, right) in &on { + let left_expr = if left.data_type(&left_schema)?.is_floating() { + Arc::new(NormalizeFloatZeroExpr::new(Arc::clone(left))) as PhysicalExprRef + } else { + Arc::clone(left) + }; + let right_expr = if right.data_type(&right_schema)?.is_floating() { + Arc::new(NormalizeFloatZeroExpr::new(Arc::clone(right))) + as PhysicalExprRef + } else { + Arc::clone(right) + }; + left_sort_exprs.push(PhysicalSortExpr { + expr: left_expr, options: equality_options, - }) - .collect::>(); + }); + right_sort_exprs.push(PhysicalSortExpr { + expr: right_expr, + options: equality_options, + }); + } 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, @@ -250,6 +260,18 @@ impl AsOfJoinExec { }) } + /// Returns this join emitting only the columns in `projection`, in that order. + /// The indices address the join's own schema, before any projection. + pub fn with_projection(&self, projection: Option>) -> Result { + Self::try_new( + Arc::clone(&self.left), + Arc::clone(&self.right), + self.on.clone(), + self.match_condition.clone(), + projection, + ) + } + fn compute_properties( left: &Arc, join_schema: &SchemaRef, @@ -293,6 +315,12 @@ impl AsOfJoinExec { } } +impl EmbeddedProjection for AsOfJoinExec { + fn with_projection(&self, projection: Option>) -> Result { + self.with_projection(projection) + } +} + impl DisplayAs for AsOfJoinExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { let on = self @@ -386,6 +414,16 @@ impl ExecutionPlan for AsOfJoinExec { vec![&self.left, &self.right] } + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + if self.projection.is_some() { + return Ok(None); + } + try_embed_projection(projection, self) + } + fn apply_expressions( &self, f: &mut dyn FnMut(&Arc) -> Result, @@ -1433,11 +1471,6 @@ fn validate_asof_join( "AsOfJoinExec equality expressions have unsupported hash type {left_type}" ); } - if left_type.is_floating() { - return plan_err!( - "AsOfJoinExec equality expressions do not support floating-point 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)?; @@ -1484,6 +1517,7 @@ mod tests { use super::*; use crate::collect; + use crate::sorts::sort::SortExec; use crate::test::TestMemoryExec; use arrow::array::{Float64Array, Int32Array, Int64Array, StringArray}; use arrow::datatypes::{DataType, Field}; @@ -1493,6 +1527,7 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr}; + use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use insta::assert_snapshot; @@ -1670,6 +1705,95 @@ mod tests { Ok(()) } + #[tokio::test] + async fn embeds_output_projection() -> Result<()> { + let exec = Arc::new(test_exec()?.with_projection(None)?); + let input = Arc::clone(&exec) as Arc; + let projection = ProjectionExec::try_new( + [ + ProjectionExpr { + expr: Arc::new(PhysicalColumn::new("id", 2)), + alias: "id".to_string(), + }, + ProjectionExpr { + expr: Arc::new(PhysicalColumn::new("price", 5)), + alias: "price".to_string(), + }, + ], + input, + )?; + + let embedded = exec + .try_swapping_with_projection(&projection)? + .expect("projection should be embedded"); + let embedded_exec = embedded + .downcast_ref::() + .expect("identity projection should be removed"); + assert_eq!(embedded_exec.projection.as_deref(), Some(&[2, 5][..])); + assert_eq!( + embedded_exec + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec!["id", "price"] + ); + + let batches = collect(embedded, Arc::new(TaskContext::default())).await?; + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+ + | id | price | + +----+-------+ + | 0 | | + | 1 | | + | 2 | | + | 3 | 40 | + | 4 | 60 | + | 5 | 101 | + | 6 | | + +----+-------+ + "); + Ok(()) + } + + #[tokio::test] + async fn empty_projection_preserves_row_count() -> Result<()> { + let exec = Arc::new(test_exec()?.with_projection(None)?); + let input = Arc::clone(&exec) as Arc; + let projection = ProjectionExec::try_new(Vec::::new(), input)?; + + let embedded = exec + .try_swapping_with_projection(&projection)? + .expect("empty projection should be embedded"); + let embedded_exec = embedded + .downcast_ref::() + .expect("empty projection should remove ProjectionExec"); + assert_eq!(embedded_exec.projection.as_deref(), Some(&[][..])); + assert!(embedded_exec.schema().fields().is_empty()); + + let batches = collect(embedded, Arc::new(TaskContext::default())).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 7); + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + Ok(()) + } + + #[test] + fn declines_projection_when_already_embedded() -> Result<()> { + let exec = test_exec()?; + let input = Arc::clone(&exec) as Arc; + let projection = ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(PhysicalColumn::new("id", 2)), + alias: "id".to_string(), + }], + input, + )?; + + assert!(exec.try_swapping_with_projection(&projection)?.is_none()); + Ok(()) + } + fn exec_without_equality_keys( left_times: Vec, right_times: Vec, @@ -1994,34 +2118,88 @@ mod tests { Ok(()) } - #[test] - fn rejects_floating_equality_expressions() -> Result<()> { - let exec = test_exec()?; + #[tokio::test] + async fn floating_equality_keys_treat_signed_zero_as_equal() -> Result<()> { + let left_batch = RecordBatch::try_from_iter(vec![ + ("key", Arc::new(Float64Array::from(vec![0.0])) as ArrayRef), + ("ts", Arc::new(Int64Array::from(vec![5])) as ArrayRef), + ("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef), + ])?; + let right_batch = RecordBatch::try_from_iter(vec![ + ( + "key", + Arc::new(Float64Array::from(vec![-0.0, 0.0])) as ArrayRef, + ), + ("ts", Arc::new(Int64Array::from(vec![10, 1])) as ArrayRef), + ( + "price", + Arc::new(Int32Array::from(vec![100, 10])) as ArrayRef, + ), + ])?; for data_type in [DataType::Float16, DataType::Float32, DataType::Float64] { - let left = Arc::new(CastExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), - data_type.clone(), + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch.clone()]], + left_batch.schema(), None, - )); - let right = Arc::new(CastExpr::new( - Arc::new(PhysicalColumn::new("ts", 1)), - data_type.clone(), + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()]], + right_batch.schema(), None, - )); - let error = AsOfJoinExec::try_new( - Arc::clone(&exec.left), - Arc::clone(&exec.right), - vec![(left, right)], - exec.match_condition.clone(), - Some(vec![0, 1, 2, 5]), - ) - .expect_err("floating equality expressions must be rejected"); - assert!( - error.to_string().contains(&format!( - "equality expressions do not support floating-point type {data_type}" + )?; + let on: JoinOn = vec![( + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("key", 0)), + data_type.clone(), + None, + )), + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("key", 0)), + data_type.clone(), + None, )), - "unexpected error: {error}" + )]; + let match_condition = AsOfMatchExpr::new( + Arc::new(PhysicalColumn::new("ts", 1)), + Operator::GtEq, + Arc::new(PhysicalColumn::new("ts", 1)), ); + let unsorted = AsOfJoinExec::try_new( + left, + right, + on.clone(), + match_condition.clone(), + Some(vec![2, 5]), + )?; + let left = Arc::new(SortExec::new( + unsorted.left_ordering.clone(), + Arc::clone(&unsorted.left), + )); + let right = Arc::new(SortExec::new( + unsorted.right_ordering.clone(), + Arc::clone(&unsorted.right), + )); + let exec = Arc::new(AsOfJoinExec::try_new( + left, + right, + on, + match_condition, + Some(vec![2, 5]), + )?); + + let batches = collect(exec, Arc::new(TaskContext::default())).await?; + let prices = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(prices, vec![Some(10)], "data type: {data_type}"); } Ok(()) } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index fac5ff27191cd..e0337b4149205 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1074,6 +1074,7 @@ message PhysicalExprNode { PhysicalLambdaVariableExprNode lambda_variable = 26; PhysicalRangeExprNode range_expr = 27; PhysicalSqlSimilarToPatternNode sql_similar_to_pattern = 28; + PhysicalNormalizeFloatZeroNode normalize_float_zero = 29; } } @@ -1219,6 +1220,10 @@ message PhysicalNegativeNode { PhysicalExprNode expr = 1; } +message PhysicalNormalizeFloatZeroNode { + PhysicalExprNode expr = 1; +} + message PhysicalExtensionExprNode { bytes expr = 1; repeated PhysicalExprNode inputs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index e06f5b7011504..02267fb5aa3c4 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -19253,6 +19253,9 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::SqlSimilarToPattern(v) => { struct_ser.serialize_field("sqlSimilarToPattern", v)?; } + physical_expr_node::ExprType::NormalizeFloatZero(v) => { + struct_ser.serialize_field("normalizeFloatZero", v)?; + } } } struct_ser.end() @@ -19312,6 +19315,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "rangeExpr", "sql_similar_to_pattern", "sqlSimilarToPattern", + "normalize_float_zero", + "normalizeFloatZero", ]; #[allow(clippy::enum_variant_names)] @@ -19343,6 +19348,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { LambdaVariable, RangeExpr, SqlSimilarToPattern, + NormalizeFloatZero, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -19391,6 +19397,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), "rangeExpr" | "range_expr" => Ok(GeneratedField::RangeExpr), "sqlSimilarToPattern" | "sql_similar_to_pattern" => Ok(GeneratedField::SqlSimilarToPattern), + "normalizeFloatZero" | "normalize_float_zero" => Ok(GeneratedField::NormalizeFloatZero), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -19602,6 +19609,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("sqlSimilarToPattern")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::SqlSimilarToPattern) +; + } + GeneratedField::NormalizeFloatZero => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("normalizeFloatZero")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::NormalizeFloatZero) ; } } @@ -20972,6 +20986,97 @@ impl<'de> serde::Deserialize<'de> for PhysicalNegativeNode { deserializer.deserialize_struct("datafusion.PhysicalNegativeNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalNormalizeFloatZeroNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.expr.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalNormalizeFloatZeroNode", len)?; + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalNormalizeFloatZeroNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "expr", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Expr, + } + 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 { + "expr" => Ok(GeneratedField::Expr), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalNormalizeFloatZeroNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalNormalizeFloatZeroNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut expr__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + } + } + Ok(PhysicalNormalizeFloatZeroNode { + expr: expr__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalNormalizeFloatZeroNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalNot { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 149839dacc967..aa906abdf5be4 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1588,7 +1588,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option, #[prost( oneof = "physical_expr_node::ExprType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29" )] pub expr_type: ::core::option::Option, } @@ -1657,6 +1657,10 @@ pub mod physical_expr_node { SqlSimilarToPattern( ::prost::alloc::boxed::Box, ), + #[prost(message, tag = "29")] + NormalizeFloatZero( + ::prost::alloc::boxed::Box, + ), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1887,6 +1891,11 @@ pub struct PhysicalNegativeNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalNormalizeFloatZeroNode { + #[prost(message, optional, boxed, tag = "1")] + pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalExtensionExprNode { #[prost(bytes = "vec", tag = "1")] pub expr: ::prost::alloc::vec::Vec, diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index bc443149df413..a9a41a15d435f 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -36,8 +36,8 @@ use datafusion_physical_expr::{ }; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, - LikeExpr, Literal, NegativeExpr, NotExpr, SqlSimilarToPattern, TryCastExpr, - UnKnownColumn, + LikeExpr, Literal, NegativeExpr, NormalizeFloatZeroExpr, NotExpr, + SqlSimilarToPattern, TryCastExpr, UnKnownColumn, }; use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; @@ -288,6 +288,9 @@ pub fn parse_physical_expr_with_converter( ExprType::IsNotNullExpr(_) => IsNotNullExpr::try_from_proto(proto, &decode_ctx)?, ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::NormalizeFloatZero(_) => { + NormalizeFloatZeroExpr::try_from_proto(proto, &decode_ctx)? + } ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?, ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, diff --git a/datafusion/proto/tests/cases/plans/sorts.rs b/datafusion/proto/tests/cases/plans/sorts.rs index 1172775b1bad7..92a4d5173b944 100644 --- a/datafusion/proto/tests/cases/plans/sorts.rs +++ b/datafusion/proto/tests/cases/plans/sorts.rs @@ -22,7 +22,9 @@ use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::LexOrdering; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::physical_plan::expressions::{ + NormalizeFloatZeroExpr, PhysicalSortExpr, col, +}; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::prelude::SessionContext; @@ -61,6 +63,20 @@ fn roundtrip_sort() -> Result<()> { ))) } +#[test] +fn roundtrip_sort_with_normalized_float_zero() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); + let sort_exprs = [PhysicalSortExpr { + expr: Arc::new(NormalizeFloatZeroExpr::new(col("a", &schema)?)), + options: SortOptions::default(), + }] + .into(); + roundtrip_test(Arc::new(SortExec::new( + sort_exprs, + Arc::new(EmptyExec::new(schema)), + ))) +} + #[test] fn roundtrip_sort_preserve_partitioning() -> Result<()> { let field_a = Field::new("a", DataType::Boolean, false); diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 18af08fc18361..9681afd5ca7e5 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -170,6 +170,17 @@ struct DerivedInputScope<'a> { schema: &'a DFSchema, } +/// Left-side joins can be chained in the surrounding `FROM`, while right-side +/// joins need nesting. ASOF inputs also preserve derived boundaries because +/// their input shaping is part of candidate matching. +#[derive(Clone, Copy)] +enum JoinInputMode { + RegularLeft, + RegularRight, + AsOfLeft, + AsOfRight, +} + impl Unparser<'_> { pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { let mut plan = normalize_union_schema(plan)?; @@ -1452,45 +1463,29 @@ impl Unparser<'_> { left_plan, &mut table_scan_filters, )?; - let left_plan = if already_projected { - Self::unwrap_qualified_passthrough_join_projection(left_plan) - } else { - left_plan - }; - - self.select_to_sql_recursively( - left_plan.as_ref(), + let left_projection = self.unparse_join_input( + &left_plan, query, select, relation, + already_projected, + JoinInputMode::RegularLeft, )?; - let left_projection: Option> = if !already_projected - { - Some(select.pop_projections()) - } else { - None - }; - let right_plan = Self::extract_join_input_table_scan_filters( right_plan, &mut table_scan_filters, )?; let mut right_relation = RelationBuilder::default(); - if already_projected - && let Some(nested_relation) = - self.join_input_to_nested_relation(right_plan.as_ref(), query)? - { - right_relation = nested_relation; - } else { - self.select_to_sql_recursively( - right_plan.as_ref(), - query, - select, - &mut right_relation, - )?; - } + let right_projection = self.unparse_join_input( + &right_plan, + query, + select, + &mut right_relation, + already_projected, + JoinInputMode::RegularRight, + )?; let (join_filters, where_filters) = Self::split_join_on_and_where_filters( join.join_type, @@ -1508,13 +1503,6 @@ impl Unparser<'_> { join_filters.as_ref(), )?; - let right_projection: Option> = if !already_projected - { - Some(select.pop_projections()) - } else { - None - }; - match join.join_type { JoinType::LeftSemi | JoinType::LeftAnti @@ -1919,78 +1907,24 @@ impl Unparser<'_> { relation: &mut RelationBuilder, ) -> Result<()> { let already_projected = select.already_projected(); - let left_plan = - Self::unwrap_qualified_passthrough_join_projection(Arc::clone(&join.left)); - let inline_left_join = matches!( - left_plan.as_ref(), - LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) - ); - let left_projection = if already_projected { - None - } else if inline_left_join { - self.select_to_sql_recursively(left_plan.as_ref(), query, select, relation)?; - select.pop_projections(); - Some(self.derived_input_projection(join.left.as_ref(), None)?) - } else if Self::asof_input_requires_derived(join.left.as_ref()) { - let qualifier = self.derive_asof_input(join.left.as_ref(), relation)?; - Some(self.derived_input_projection(join.left.as_ref(), qualifier.as_ref())?) - } else { - self.select_to_sql_recursively(join.left.as_ref(), query, select, relation)?; - Some(select.pop_projections()) - }; - if already_projected { - if inline_left_join { - self.select_to_sql_recursively( - left_plan.as_ref(), - query, - select, - relation, - )?; - } else if Self::asof_input_requires_derived(join.left.as_ref()) { - self.derive_asof_input(join.left.as_ref(), relation)?; - } else { - self.select_to_sql_recursively( - join.left.as_ref(), - query, - select, - relation, - )?; - } - } + let left_projection = self.unparse_join_input( + &join.left, + query, + select, + relation, + already_projected, + JoinInputMode::AsOfLeft, + )?; 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 right_projection = self.unparse_join_input( + &join.right, + query, + select, + &mut right_relation, + already_projected, + JoinInputMode::AsOfRight, + )?; let Ok(Some(relation)) = right_relation.build() else { return internal_err!("Failed to build ASOF right relation"); }; @@ -2330,7 +2264,76 @@ impl Unparser<'_> { ) } - fn asof_input_requires_derived(plan: &LogicalPlan) -> bool { + fn unparse_join_input( + &self, + plan: &Arc, + query: &mut Option, + select: &mut SelectBuilder, + relation: &mut RelationBuilder, + already_projected: bool, + mode: JoinInputMode, + ) -> Result>> { + let preserve_boundary = + matches!(mode, JoinInputMode::AsOfLeft | JoinInputMode::AsOfRight); + let is_left = + matches!(mode, JoinInputMode::RegularLeft | JoinInputMode::AsOfLeft); + let unwrapped_plan = if preserve_boundary || already_projected { + Self::unwrap_qualified_passthrough_join_projection(Arc::clone(plan)) + } else { + Arc::clone(plan) + }; + let inline_left_join = is_left + && matches!( + unwrapped_plan.as_ref(), + LogicalPlan::Join(_) | LogicalPlan::AsOfJoin(_) + ); + + if !is_left + && (preserve_boundary || already_projected) + && let Some(nested_relation) = + self.join_input_to_nested_relation(plan.as_ref(), query)? + { + *relation = nested_relation; + return if already_projected { + Ok(None) + } else { + Ok(Some(self.derived_input_projection(plan.as_ref(), None)?)) + }; + } + + if preserve_boundary + && !inline_left_join + && Self::join_input_requires_derived(plan.as_ref()) + { + let qualifier = self.derive_join_input(plan.as_ref(), relation)?; + return if already_projected { + Ok(None) + } else { + Ok(Some(self.derived_input_projection( + plan.as_ref(), + qualifier.as_ref(), + )?)) + }; + } + + let recursive_plan = if inline_left_join { + unwrapped_plan.as_ref() + } else { + plan.as_ref() + }; + self.select_to_sql_recursively(recursive_plan, query, select, relation)?; + + if already_projected { + Ok(None) + } else if preserve_boundary && inline_left_join { + select.pop_projections(); + Ok(Some(self.derived_input_projection(plan.as_ref(), None)?)) + } else { + Ok(Some(select.pop_projections())) + } + } + + fn join_input_requires_derived(plan: &LogicalPlan) -> bool { let simple_scan = |scan: &TableScan| scan.filters.is_empty() && scan.fetch.is_none(); match plan { @@ -2342,7 +2345,7 @@ impl Unparser<'_> { } } - fn derive_asof_input( + fn derive_join_input( &self, plan: &LogicalPlan, relation: &mut RelationBuilder, diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index 3e84cd037653e..97a99be4e7069 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -138,6 +138,23 @@ ORDER BY l.id; 6 NULL NULL 7 A NULL +# SELECT * exposes the USING key once. +query ITPPT +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp) +ORDER BY l.id; +---- +1 A 2024-01-01T09:00:01 NULL NULL +2 A 2024-01-01T09:00:04 2024-01-01T09:00:04 a4 +3 A 2024-01-01T09:00:07 2024-01-01T09:00:06 a6 +4 B 2024-01-01T09:00:02 2024-01-01T09:00:01 b1 +5 B 2024-01-01T09:00:08 2024-01-01T09:00:06 b6 +6 NULL 2024-01-01T09:00:03 NULL NULL +7 A NULL NULL NULL + # Equality keys are optional. query IT SELECT l.id, r.label @@ -169,6 +186,25 @@ ORDER BY l.id; 1 x-a2 2 y-a3 +# USING accepts multiple equality keys. +query IT +SELECT l.id, r.val +FROM (VALUES + (1, 'X', 'A', TIMESTAMP '2024-01-01 09:00:04'), + (2, 'Y', 'A', TIMESTAMP '2024-01-01 09:00:04') +) AS l(id, venue, grp, ts) +ASOF JOIN (VALUES + ('X', 'A', TIMESTAMP '2024-01-01 09:00:02', 'x-a2'), + ('Y', 'A', TIMESTAMP '2024-01-01 09:00:03', 'y-a3'), + ('X', 'B', TIMESTAMP '2024-01-01 09:00:04', 'x-b4') +) AS r(venue, grp, ts, val) +MATCH_CONDITION (l.ts >= r.ts) +USING (venue, grp) +ORDER BY l.id; +---- +1 x-a2 +2 y-a3 + # Candidate selection sees the right input after subquery filtering. query IT SELECT l.id, r.val @@ -182,6 +218,150 @@ ORDER BY l.id; 3 a4 5 b6 +# A filter on a right output column is applied after candidate selection. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +WHERE r.val <> 'a6' +ORDER BY l.id; +---- +2 a4 +4 b1 +5 b6 + +query TT +EXPLAIN SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +WHERE r.val <> 'a6'; +---- +logical_plan +01)Filter: r.val != Utf8View("a6") +02)--Projection: l.id, r.val +03)----AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +04)------SubqueryAlias: l +05)--------TableScan: asof_left projection=[id, grp, ts] +06)------SubqueryAlias: r +07)--------TableScan: asof_right projection=[grp, ts, val] +physical_plan +01)FilterExec: val@1 != a6 +02)--AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts], projection=[id@0, val@5] +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] + +# Projection pruning keeps only columns required by the ASOF contract. +query TT +EXPLAIN SELECT l.id +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- +logical_plan +01)Projection: l.id +02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp] +03)----SubqueryAlias: l +04)------TableScan: asof_left projection=[id, grp, ts] +05)----SubqueryAlias: r +06)------TableScan: asof_right projection=[grp, ts] +physical_plan +01)AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts], projection=[id@0] +02)--SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# Equality and ordered operands may be expressions. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts + INTERVAL '1 second') +ON lower(l.grp) = lower(r.grp) +WHERE l.id IN (2, 3, 4) +ORDER BY l.id; +---- +2 a2 +3 a6 +4 b1 + +# Empty inputs preserve left-join semantics. +query IT +SELECT l.id, r.val +FROM asof_left l +ASOF JOIN (SELECT * FROM asof_right WHERE false) r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id; +---- +1 NULL +2 NULL +3 NULL +4 NULL +5 NULL +6 NULL +7 NULL + +query I +SELECT l.id +FROM (SELECT * FROM asof_left WHERE false) l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp; +---- + +# Duplicate left rows are emitted independently. +query IT +SELECT l.id, r.val +FROM ( + SELECT * FROM asof_left + UNION ALL + SELECT * FROM asof_left +) l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +ORDER BY l.id, r.val; +---- +1 NULL +1 NULL +2 a4 +2 a4 +3 a6 +3 a6 +4 b1 +4 b1 +5 b6 +5 b6 +6 NULL +6 NULL +7 NULL +7 NULL + +# ASOF JOIN supports independently aliased self joins. +query II +SELECT a.id, b.id +FROM asof_left a +ASOF JOIN asof_left b +MATCH_CONDITION (a.ts > b.ts) +ON a.grp = b.grp +ORDER BY a.id; +---- +1 NULL +2 1 +3 2 +4 NULL +5 4 +6 NULL +7 NULL + # Equality and match operands use the planner's common coercion types. query II SELECT l.id, r.payload @@ -214,14 +394,13 @@ logical_plan 07)------Projection: column1 AS grp, column2 AS ts, column3 AS payload 08)--------Values: (Int64(5), Int64(9), Int64(90)) physical_plan -01)ProjectionExec: expr=[id@0 as id, payload@5 as payload] -02)--AsOfJoinExec: on=[(CAST(grp AS Int64) = grp)], match=[CAST(ts AS Int64) >= ts] -03)----SortExec: expr=[CAST(grp@1 AS Int64) ASC, CAST(ts@2 AS Int64) ASC], preserve_partitioning=[false] -04)------ProjectionExec: expr=[column1@0 as id, column2@1 as grp, column3@2 as ts] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)----ProjectionExec: expr=[column1@0 as grp, column2@1 as ts, column3@2 as payload] -07)------SortExec: expr=[column1@0 ASC, column2@1 ASC], preserve_partitioning=[false] -08)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)AsOfJoinExec: on=[(CAST(grp AS Int64) = grp)], match=[CAST(ts AS Int64) >= ts], projection=[id@0, payload@5] +02)--SortExec: expr=[CAST(grp@1 AS Int64) ASC, CAST(ts@2 AS Int64) ASC], preserve_partitioning=[false] +03)----ProjectionExec: expr=[column1@0 as id, column2@1 as grp, column3@2 as ts] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)--ProjectionExec: expr=[column1@0 as grp, column2@1 as ts, column3@2 as payload] +06)----SortExec: expr=[column1@0 ASC, column2@1 ASC], preserve_partitioning=[false] +07)------DataSourceExec: partitions=1, partition_sizes=[1] # Equality operands can name the right input first. query IT @@ -251,12 +430,11 @@ logical_plan 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] +01)AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts], projection=[id@0, val@5] +02)--SortExec: expr=[grp@1 ASC, ts@2 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] +05)----DataSourceExec: partitions=1, partition_sizes=[1] query error ASOF MATCH_CONDITION requires <, <=, >, or >= SELECT * @@ -285,3 +463,45 @@ FROM asof_left l ASOF JOIN asof_right r MATCH_CONDITION (l.ts) ON l.grp = r.grp; + +query error ASOF MATCH_CONDITION left operand must reference only the left input +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (1 >= r.ts) +ON l.grp = r.grp; + +query error Each ASOF equality condition must compare one left expression with one right expression +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON 1 = 1; + +query error ASOF MATCH_CONDITION requires <, <=, >, or >= +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts AND l.ts > r.ts) +ON l.grp = r.grp; + +query error Each ASOF equality condition must compare one left expression with one right expression +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = l.grp; + +query error ASOF ON accepts only equality conditions combined with AND +SELECT * +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp AND l.ts >= r.ts; + +query error Ambiguous reference to unqualified field ts +SELECT ts +FROM asof_left l +ASOF JOIN asof_right r +MATCH_CONDITION (l.ts >= r.ts) +USING (grp); diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 4483fe76059b2..1812d05a9442c 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -1351,6 +1351,37 @@ readers that don't recognize the new tag. See [PR #23188](https://github.com/apache/datafusion/pull/23188) for details. +### `ExprType` gained a `NormalizeFloatZero` variant + +Floating-point equality keys in ASOF joins now use a +`NormalizeFloatZeroExpr` physical expression to ensure that `-0.0` and `+0.0` +have the same ordering and equality semantics. The generated `ExprType` enum +on `PhysicalExprNode` gained a matching `NormalizeFloatZero` variant (tag 29). + +**Who is affected:** + +- Users matching exhaustively on `ExprType` (for example in a custom + physical-plan encoder/decoder). + +**Migration guide:** + +Add a `NormalizeFloatZero` arm, or fall back to a wildcard arm: + +```rust,ignore +match expr_type { + // ... + ExprType::NormalizeFloatZero(node) => { /* ... */ } + _ => { /* ... */ } +} +``` + +Plans encoded before this variant existed are unaffected: they never produced +this variant, so decoding is unchanged. Plans containing a +`NormalizeFloatZeroExpr` encoded after this change fail to decode on older +readers that don't recognize the new tag. + +See [PR #24375](https://github.com/apache/datafusion/pull/24375) for details. + ### `ParquetObjectReader` / `ParquetObjectWriter` deprecated upstream The [`parquet` crate] deprecated [`ParquetObjectReader`]