perf: preserve left functional dependencies for ASOF joins

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<DFSchema> {
-    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<String, String> =

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