perf: preserve left functional dependencies for ASOF joins

datafusion/core/src/physical_planner.rs

@@ -41,7 +41,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};
@@ -1790,6 +1791,51 @@ 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::<Result<join_utils::JoinOn>>()?;
+                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,
+                    )?,
+                );
+                Arc::new(AsOfJoinExec::try_new(
+                    physical_left,
+                    physical_right,
+                    join_on,
+                    match_condition,
+                    None,
+                )?)
+            }
             LogicalPlan::RecursiveQuery(RecursiveQuery {
                 name,
                 is_distinct,
@@ -2291,6 +2337,7 @@ fn extract_dml_filters(
             | LogicalPlan::Sort(_)
             | LogicalPlan::Union(_)
             | LogicalPlan::Join(_)
+            | LogicalPlan::AsOfJoin(_)
             | LogicalPlan::Repartition(_)
             | LogicalPlan::Aggregate(_)
             | LogicalPlan::Window(_)

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> {
+        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<Column>,
+        match_condition: AsOfMatch,
+    ) -> Result<Self> {
+        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::<Result<_>>()?;
+        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<Self> {
+        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::<Result<_>>()?;
+        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<Column> {
         if column.relation.is_some() {
             // column is already normalized
@@ -1776,6 +1838,19 @@ pub fn build_join_schema(
     dfschema.with_functional_dependencies(func_dependencies)
 }
 
+/// Creates the schema for a left-preserving ASOF join.
+///
+/// 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> {
+    // 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
 /// conflict with the columns from the other.
 /// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying
@@ -2311,7 +2386,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;
 
@@ -2969,6 +3045,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<String, String> =

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<String> =
+                    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,
                 ..

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,