feat: support ASOF JOIN SQL

datafusion/expr/src/logical_plan/builder.rs

@@ -38,13 +38,15 @@ use crate::logical_plan::{
 };
 use crate::select_expr::SelectExpr;
 use crate::utils::{
-    can_hash, columnize_expr, compare_sort_expr, expand_qualified_wildcard,
-    expand_wildcard, expr_to_columns, find_valid_equijoin_key_pair,
-    group_window_expr_by_sort_keys,
+    can_hash, check_all_columns_from_schema, columnize_expr, compare_sort_expr,
+    expand_qualified_wildcard, expand_wildcard, expr_to_columns,
+    find_valid_equijoin_key_pair, group_window_expr_by_sort_keys,
+    split_conjunction_owned,
 };
 use crate::{
-    DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, RecursiveQuery,
-    Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, binary_expr, lit,
+    BinaryExpr, DmlStatement, ExplainOption, Expr, ExprSchemable, Operator,
+    RecursiveQuery, Statement, TableProviderFilterPushDown, TableSource, WriteOp, and,
+    binary_expr, lit,
 };
 
 use super::dml::InsertOp;
@@ -1007,23 +1009,67 @@ impl LogicalPlanBuilder {
         )
     }
 
-    /// Apply a left-preserving ASOF join using equality expressions and one
-    /// ordered match condition.
-    pub fn asof_join(
+    /// Apply a left-preserving ASOF join using an optional equality condition
+    /// and one ordered match condition.
+    ///
+    /// When present, `on_expr` must contain equality comparisons combined with
+    /// `AND`. Each comparison must have one operand that references only the
+    /// left input and one that references only `right`; their order does not
+    /// matter. `match_condition` must be a single `<`, `<=`, `>`, or `>=`
+    /// comparison whose left operand references only the left input and whose
+    /// right operand references only `right`.
+    pub fn asof_join_on(
         self,
         right: LogicalPlan,
-        on: Vec<(Expr, Expr)>,
-        match_condition: AsOfMatch,
+        on_expr: Option<Expr>,
+        match_condition: Expr,
     ) -> Result<Self> {
-        self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::On)
+        let on = on_expr
+            .into_iter()
+            .flat_map(split_conjunction_owned)
+            .map(|predicate| {
+                let Expr::BinaryExpr(BinaryExpr {
+                    left,
+                    op: Operator::Eq,
+                    right: right_expr,
+                }) = predicate
+                else {
+                    return plan_err!(
+                        "ASOF ON accepts only equality conditions combined with AND"
+                    );
+                };
+                find_valid_equijoin_key_pair(
+                    &left,
+                    &right_expr,
+                    self.plan.schema(),
+                    right.schema(),
+                )?
+                .ok_or_else(|| {
+                    plan_datafusion_err!(
+                        "Each ASOF equality condition must compare one left expression with one right expression"
+                    )
+                })
+            })
+            .collect::<Result<_>>()?;
+        self.asof_join_with_constraint(
+            right,
+            on,
+            AsOfMatch::try_from(match_condition)?,
+            JoinConstraint::On,
+        )
     }
 
-    /// Apply a left-preserving ASOF join using `USING` equality keys.
+    /// Apply a left-preserving ASOF join using `USING` equality keys and one
+    /// ordered match condition.
+    ///
+    /// Every key in `using_keys` must resolve in both inputs.
+    /// `match_condition` follows the same operand and operator requirements as
+    /// [`asof_join_on`](Self::asof_join_on).
     pub fn asof_join_using(
         self,
         right: LogicalPlan,
         using_keys: Vec<Column>,
-        match_condition: AsOfMatch,
+        match_condition: Expr,
     ) -> Result<Self> {
         let on = using_keys
             .into_iter()
@@ -1033,7 +1079,12 @@ impl LogicalPlanBuilder {
                 Ok((Expr::Column(left), Expr::Column(right)))
             })
             .collect::<Result<_>>()?;
-        self.asof_join_with_constraint(right, on, match_condition, JoinConstraint::Using)
+        self.asof_join_with_constraint(
+            right,
+            on,
+            AsOfMatch::try_from(match_condition)?,
+            JoinConstraint::Using,
+        )
     }
 
     fn asof_join_with_constraint(
@@ -1043,6 +1094,17 @@ impl LogicalPlanBuilder {
         match_condition: AsOfMatch,
         join_constraint: JoinConstraint,
     ) -> Result<Self> {
+        let left_columns = match_condition.left.column_refs();
+        let right_columns = match_condition.right.column_refs();
+        if left_columns.is_empty()
+            || right_columns.is_empty()
+            || !check_all_columns_from_schema(&left_columns, self.plan.schema())?
+            || !check_all_columns_from_schema(&right_columns, right.schema())?
+        {
+            return plan_err!(
+                "ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input"
+            );
+        }
         let normalize = |expr, schema: &DFSchema| {
             normalize_col_with_schemas_and_ambiguity_check(expr, &[&[schema]], &[])
         };
@@ -2914,6 +2976,70 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn asof_join_on_extracts_and_validates_conditions() -> Result<()> {
+        let values = vec![vec![lit(1), lit(2)]];
+        let left = LogicalPlanBuilder::values(values.clone())?
+            .alias("l")?
+            .build()?;
+        let right = LogicalPlanBuilder::values(values)?.alias("r")?.build()?;
+
+        let plan = LogicalPlanBuilder::from(left.clone())
+            .asof_join_on(
+                right.clone(),
+                Some(
+                    col("r.column1")
+                        .eq(col("l.column1"))
+                        .and(col("l.column2").eq(col("r.column2"))),
+                ),
+                col("l.column2").gt_eq(col("r.column2")),
+            )?
+            .build()?;
+        let LogicalPlan::AsOfJoin(join) = plan else {
+            panic!("expected ASOF join")
+        };
+        assert_eq!(
+            join.on,
+            vec![
+                (col("l.column1"), col("r.column1")),
+                (col("l.column2"), col("r.column2")),
+            ]
+        );
+        assert_eq!(
+            join.match_condition.as_ref(),
+            &AsOfMatch::new(col("l.column2"), Operator::GtEq, col("r.column2"))
+        );
+
+        let invalid_on = LogicalPlanBuilder::from(left.clone())
+            .asof_join_on(
+                right.clone(),
+                Some(col("l.column1").gt(col("r.column1"))),
+                col("l.column2").gt_eq(col("r.column2")),
+            )
+            .expect_err("non-equality ASOF ON should fail");
+        assert_snapshot!(invalid_on.strip_backtrace(), @r#"Error during planning: ASOF ON accepts only equality conditions combined with AND"#);
+
+        let invalid_match = LogicalPlanBuilder::from(left.clone())
+            .asof_join_on(
+                right.clone(),
+                Some(col("l.column1").eq(col("r.column1"))),
+                col("l.column2").eq(col("r.column2")),
+            )
+            .expect_err("equality ASOF MATCH_CONDITION should fail");
+        assert_snapshot!(invalid_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION requires <, <=, >, or >=, found ="#);
+
+        let reversed_match = LogicalPlanBuilder::from(left)
+            .asof_join_on(
+                right,
+                Some(col("l.column1").eq(col("r.column1"))),
+                col("r.column2").gt_eq(col("l.column2")),
+            )
+            .expect_err("reversed ASOF MATCH_CONDITION should fail");
+        assert_snapshot!(reversed_match.strip_backtrace(), @r#"Error during planning: ASOF MATCH_CONDITION left operand must reference only the left input and right operand only the right input"#);
+
+        Ok(())
+    }
+
     #[test]
     fn plan_builder_from_logical_plan() -> Result<()> {
         let plan =

datafusion/expr/src/logical_plan/plan.rs

@@ -4462,6 +4462,25 @@ impl AsOfMatch {
     }
 }
 
+impl TryFrom<Expr> for AsOfMatch {
+    type Error = DataFusionError;
+
+    fn try_from(condition: Expr) -> Result<Self> {
+        let Expr::BinaryExpr(BinaryExpr { left, op, right }) = condition else {
+            return plan_err!("ASOF MATCH_CONDITION must be a single comparison");
+        };
+        if !matches!(
+            op,
+            Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq
+        ) {
+            return plan_err!(
+                "ASOF MATCH_CONDITION requires <, <=, >, or >=, found {op}"
+            );
+        }
+        Ok(Self::new(*left, op, *right))
+    }
+}
+
 impl Display for AsOfMatch {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         write!(f, "{} {} {}", self.left, self.op, self.right)
@@ -6179,10 +6198,10 @@ mod tests {
         assert_eq!(cross_join.min_rows(), 2);
 
         let asof_join = LogicalPlanBuilder::from(two_rows.clone())
-            .asof_join(
+            .asof_join_on(
                 one_row.clone(),
-                vec![],
-                AsOfMatch::new(col("l.column1"), Operator::GtEq, col("r.column1")),
+                None,
+                col("l.column1").gt_eq(col("r.column1")),
             )?
             .build()?;
         assert_eq!(asof_join.min_rows(), 2);

datafusion/sql/src/relation/join.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
-use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err};
+use datafusion_common::{Column, Result, not_impl_err, plan_datafusion_err, plan_err};
 use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder};
 use sqlparser::ast::{
     Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins,
@@ -98,10 +98,75 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             JoinOperator::CrossJoin(JoinConstraint::None) => {
                 self.parse_cross_join(left, right)
             }
+            JoinOperator::AsOf {
+                match_condition,
+                constraint,
+            } => self.parse_asof_join(
+                left,
+                right,
+                match_condition,
+                constraint,
+                planner_context,
+            ),
             other => not_impl_err!("Unsupported JOIN operator {other:?}"),
         }
     }
 
+    fn parse_asof_join(
+        &self,
+        left: LogicalPlan,
+        right: LogicalPlan,
+        sql_match_condition: sqlparser::ast::Expr,
+        constraint: JoinConstraint,
+        planner_context: &mut PlannerContext,
+    ) -> Result<LogicalPlan> {
+        let join_schema = left.schema().join(right.schema())?;
+        let match_condition =
+            self.sql_to_expr(sql_match_condition, &join_schema, planner_context)?;
+
+        match constraint {
+            JoinConstraint::On(sql_on) => {
+                let on = self.sql_to_expr(sql_on, &join_schema, planner_context)?;
+                LogicalPlanBuilder::from(left)
+                    .asof_join_on(right, Some(on), match_condition)?
+                    .build()
+            }
+            JoinConstraint::Using(object_names) => {
+                let keys = object_names
+                    .into_iter()
+                    .map(|object_name| {
+                        let ObjectName(mut object_names) = object_name;
+                        if object_names.len() != 1 {
+                            return not_impl_err!(
+                                "Invalid identifier in ASOF USING clause. Expected single identifier, got {}",
+                                ObjectName(object_names)
+                            );
+                        }
+                        let id = object_names.swap_remove(0);
+                        id.as_ident()
+                            .ok_or_else(|| {
+                                plan_datafusion_err!(
+                                    "Expected identifier in ASOF USING clause"
+                                )
+                            })
+                            .map(|ident| {
+                                Column::from_name(
+                                    self.ident_normalizer.normalize(ident.clone()),
+                                )
+                            })
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+                LogicalPlanBuilder::from(left)
+                    .asof_join_using(right, keys, match_condition)?
+                    .build()
+            }
+            JoinConstraint::None => LogicalPlanBuilder::from(left)
+                .asof_join_on(right, None, match_condition)?
+                .build(),
+            JoinConstraint::Natural => plan_err!("NATURAL ASOF JOIN is not supported"),
+        }
+    }
+
     fn parse_cross_join(
         &self,
         left: LogicalPlan,