feat: add ASOF joins to DataFrame API

datafusion/core/src/dataframe/mod.rs

@@ -58,8 +58,8 @@ use datafusion_common::{
 };
 use datafusion_expr::select_expr::SelectExpr;
 use datafusion_expr::{
-    ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case,
-    dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION,
+    AsOfMatch, ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown,
+    UNNAMED_TABLE, case, dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION,
 };
 use datafusion_functions::core::coalesce;
 use datafusion_functions::math::nanvl;
@@ -1379,6 +1379,45 @@ impl DataFrame {
         })
     }
 
+    /// Join this `DataFrame` to the closest eligible row in `right`.
+    ///
+    /// Every left row is emitted exactly once. `on` contains optional equality
+    /// expressions and `match_condition` selects the ordered predecessor or
+    /// successor from the matching right group.
+    pub fn join_asof(
+        self,
+        right: DataFrame,
+        on: Vec<(Expr, Expr)>,
+        match_condition: AsOfMatch,
+    ) -> Result<DataFrame> {
+        let plan = LogicalPlanBuilder::from(self.plan)
+            .asof_join(right.plan, on, match_condition)?
+            .build()?;
+        Ok(DataFrame {
+            session_state: self.session_state,
+            plan,
+            projection_requires_validation: true,
+        })
+    }
+
+    /// Join this `DataFrame` to the closest eligible row in `right` using
+    /// same-named equality keys.
+    pub fn join_asof_using(
+        self,
+        right: DataFrame,
+        using_keys: Vec<Column>,
+        match_condition: AsOfMatch,
+    ) -> Result<DataFrame> {
+        let plan = LogicalPlanBuilder::from(self.plan)
+            .asof_join_using(right.plan, using_keys, match_condition)?
+            .build()?;
+        Ok(DataFrame {
+            session_state: self.session_state,
+            plan,
+            projection_requires_validation: true,
+        })
+    }
+
     /// Repartition a DataFrame based on a logical partitioning scheme.
     ///
     /// # Example

datafusion/core/src/physical_planner.rs

@@ -43,7 +43,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};
@@ -1865,6 +1866,53 @@ 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,
+                    )?,
+                );
+                let right_output_indices =
+                    (0..join.right.schema().fields().len()).collect();
+                Arc::new(AsOfJoinExec::try_new(
+                    physical_left,
+                    physical_right,
+                    join_on,
+                    match_condition,
+                    right_output_indices,
+                )?)
+            }
             LogicalPlan::RecursiveQuery(RecursiveQuery {
                 name,
                 is_distinct,
@@ -2364,6 +2412,7 @@ fn extract_dml_filters(
             | LogicalPlan::Sort(_)
             | LogicalPlan::Union(_)
             | LogicalPlan::Join(_)
+            | LogicalPlan::AsOfJoin(_)
             | LogicalPlan::Repartition(_)
             | LogicalPlan::Aggregate(_)
             | LogicalPlan::Window(_)

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},

datafusion/core/tests/dataframe/mod.rs

@@ -58,7 +58,7 @@ use datafusion::error::Result;
 use datafusion::execution::context::SessionContext;
 use datafusion::execution::session_state::SessionStateBuilder;
 use datafusion::logical_expr::{ColumnarValue, Volatility};
-use datafusion::prelude::{CsvReadOptions, JoinType, ParquetReadOptions};
+use datafusion::prelude::{AsOfMatch, CsvReadOptions, JoinType, ParquetReadOptions};
 use datafusion::test_util::{
     parquet_test_data, populate_csv_partitions, register_aggregate_csv, test_table,
     test_table_with_cache_factory, test_table_with_name,
@@ -77,10 +77,10 @@ use datafusion_expr::expr::{GroupingSet, NullTreatment, Sort, WindowFunction};
 use datafusion_expr::var_provider::{VarProvider, VarType};
 use datafusion_expr::{
     CreateMemoryTable, CreateView, DdlStatement, Expr, ExprFunctionExt, ExprSchemable,
-    LogicalPlan, LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, TableType,
-    WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, cast, col,
-    create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col, placeholder,
-    scalar_subquery, when, wildcard,
+    LogicalPlan, LogicalPlanBuilder, Operator, ScalarFunctionImplementation, SortExpr,
+    TableType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
+    cast, col, create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col,
+    placeholder, scalar_subquery, when, wildcard,
 };
 use datafusion_physical_expr::Partitioning;
 use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -1503,6 +1503,50 @@ async fn join() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn join_asof() -> Result<()> {
+    let ctx = SessionContext::new();
+    let left = ctx
+        .read_batch(record_batch!(
+            ("symbol", Utf8, ["A", "A", "B"]),
+            ("ts", Int64, [1, 4, 2]),
+            ("trade_id", Int32, [1, 2, 3])
+        )?)?
+        .alias("trades")?;
+    let right = ctx
+        .read_batch(record_batch!(
+            ("symbol", Utf8, ["A", "A", "B"]),
+            ("ts", Int64, [2, 4, 1]),
+            ("price", Int32, [20, 40, 101])
+        )?)?
+        .alias("prices")?;
+
+    let results = left
+        .join_asof(
+            right,
+            vec![(col("symbol"), col("symbol"))],
+            AsOfMatch::new(col("ts"), Operator::GtEq, col("ts")),
+        )?
+        .select(vec![col("trade_id"), col("price")])?
+        .sort(vec![col("trade_id").sort(true, true)])?
+        .collect()
+        .await?;
+
+    assert_batches_eq!(
+        [
+            "+----------+-------+",
+            "| trade_id | price |",
+            "+----------+-------+",
+            "| 1        |       |",
+            "| 2        | 40    |",
+            "| 3        | 101   |",
+            "+----------+-------+",
+        ],
+        &results
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn join_coercion_unnamed() -> Result<()> {
     let ctx = SessionContext::new();

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -55,6 +55,7 @@ use datafusion_physical_plan::{
     coalesce_partitions::CoalescePartitionsExec,
     collect,
     filter::{FilterExec, FilterExecBuilder},
+    joins::{AsOfJoinExec, AsOfMatchExpr},
     projection::ProjectionExec,
     repartition::RepartitionExec,
     sorts::sort::SortExec,
@@ -123,6 +124,47 @@ fn test_pushdown_volatile_functions_not_allowed() {
     );
 }
 
+#[test]
+fn test_asof_join_pushes_only_left_filters() {
+    let left = TestScanBuilder::new(schema()).with_support(true).build();
+    let right = TestScanBuilder::new(schema()).with_support(true).build();
+    let join = Arc::new(
+        AsOfJoinExec::try_new(
+            left,
+            right,
+            vec![(col("a", &schema()).unwrap(), col("a", &schema()).unwrap())],
+            AsOfMatchExpr::new(
+                col("b", &schema()).unwrap(),
+                Operator::GtEq,
+                col("b", &schema()).unwrap(),
+            ),
+            vec![0, 1, 2],
+        )
+        .unwrap(),
+    );
+    let left_filter = Arc::new(BinaryExpr::new(
+        Arc::new(Column::new("c", 2)),
+        Operator::Gt,
+        Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
+    )) as Arc<dyn PhysicalExpr>;
+    let right_filter = Arc::new(BinaryExpr::new(
+        Arc::new(Column::new("c", 5)),
+        Operator::Gt,
+        Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
+    )) as Arc<dyn PhysicalExpr>;
+    let predicate = Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter));
+    let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap());
+    let mut config = ConfigOptions::default();
+    config.execution.parquet.pushdown_filters = true;
+    let optimized = FilterPushdown::new().optimize(plan, &config).unwrap();
+    let formatted = format_plan_for_test(&optimized);
+
+    assert!(formatted.contains("FilterExec: c@5 > 0"), "{formatted}");
+    assert!(formatted.contains("AsOfJoinExec:"), "{formatted}");
+    assert!(formatted.contains("predicate=c@2 > 0"), "{formatted}");
+    assert_eq!(formatted.matches("predicate=").count(), 1, "{formatted}");
+}
+
 /// Show that we can use config options to determine how to do pushdown.
 #[test]
 fn test_pushdown_into_scan_with_config_options() {

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,15 @@ 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> {
+    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