diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index a44505b15ab45..4ae86f9318342 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -56,8 +56,8 @@ use datafusion_common::display::ToStringifiedPlan; use datafusion_common::file_options::file_type::FileType; use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ - Column, Constraints, DFSchema, DFSchemaRef, NullEquality, Result, ScalarValue, - TableReference, ToDFSchema, UnnestOptions, exec_err, + Column, Constraints, DFSchema, DFSchemaRef, FunctionalDependencies, NullEquality, + Result, ScalarValue, TableReference, ToDFSchema, UnnestOptions, exec_err, get_target_functional_dependencies, internal_datafusion_err, plan_datafusion_err, plan_err, }; @@ -1894,7 +1894,20 @@ 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 remain valid but not unique because one + // right row may match multiple left rows. + let schema = build_join_schema(left, right, &JoinType::Left)?; + // `build_join_schema` places the downgraded left dependencies before the + // nullable, non-unique right dependencies. Replace only the left prefix. + let left_dependencies_len = left.functional_dependencies().len(); + let right_dependencies = + schema.functional_dependencies()[left_dependencies_len..].to_vec(); + let mut dependencies = left.functional_dependencies().clone(); + dependencies.extend_target_indices(schema.fields().len()); + dependencies.extend(FunctionalDependencies::new(right_dependencies)); + schema.with_functional_dependencies(dependencies) } /// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise @@ -2432,7 +2445,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; @@ -3198,6 +3212,45 @@ mod tests { Ok(()) } + #[test] + fn asof_join_schema_preserves_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(), + &FunctionalDependencies::new(vec![ + FunctionalDependence::new(vec![0], vec![0, 1, 2, 3], false) + .with_mode(Dependency::Single), + FunctionalDependence::new(vec![2], vec![2, 3], true) + .with_mode(Dependency::Multi), + ]) + ); + Ok(()) + } + #[test] fn test_values_metadata() -> Result<()> { let metadata: HashMap = diff --git a/datafusion/sqllogictest/test_files/asof_join.slt b/datafusion/sqllogictest/test_files/asof_join.slt index bcfa41b9023b3..9071424a019dd 100644 --- a/datafusion/sqllogictest/test_files/asof_join.slt +++ b/datafusion/sqllogictest/test_files/asof_join.slt @@ -16,7 +16,7 @@ # under the License. statement ok -CREATE TABLE asof_left(id INT, grp TEXT, ts TIMESTAMP) AS VALUES +CREATE TABLE asof_left(id INT PRIMARY KEY, 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'), @@ -256,6 +256,61 @@ physical_plan 04)--SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false] 05)----DataSourceExec: partitions=1, partition_sizes=[1] +# An ASOF join emits each left row exactly once, so the primary key from +# asof_left remains unique. Because the outer LEFT JOIN does not use columns +# from the ASOF side, it can be removed entirely. +statement ok +CREATE TABLE asof_probe(id INT, payload TEXT) AS VALUES + (1, 'p1'), + (8, 'p8'); + +query TT +EXPLAIN SELECT p.payload +FROM asof_probe p +LEFT JOIN ( + SELECT l.id + FROM asof_left l + ASOF JOIN asof_right r + MATCH_CONDITION (l.ts >= r.ts) + ON l.grp = r.grp +) matched +ON p.id = matched.id; +---- +logical_plan +01)SubqueryAlias: p +02)--TableScan: asof_probe projection=[payload] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +# A right-side primary key still determines the other right columns after an +# ASOF join, even though one right row may match multiple left rows. The +# dependency is therefore retained as nullable and non-unique. +statement ok +CREATE TABLE asof_fd_left(id INT PRIMARY KEY, grp TEXT, ts INT) AS VALUES + (1, 'A', 1), + (2, 'A', 4), + (3, 'A', 7), + (4, 'B', 2), + (5, 'C', 1); + +statement ok +CREATE TABLE asof_fd_right(id INT PRIMARY KEY, grp TEXT, ts INT, v INT) AS VALUES + (10, 'A', 0, 100), + (11, 'A', 5, 110), + (12, 'B', 1, 120); + +query II rowsort +SELECT r.id, r.v +FROM asof_fd_left l +ASOF JOIN asof_fd_right r +MATCH_CONDITION (l.ts >= r.ts) +ON l.grp = r.grp +GROUP BY r.id; +---- +10 100 +11 110 +12 120 +NULL NULL + query error ASOF MATCH_CONDITION requires <, <=, >, or >= SELECT * FROM asof_left l