feat: support ASOF joins

benchmarks/src/asof.rs

@@ -0,0 +1,194 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::util::{BenchmarkRun, CommonOpt, QueryResult};
+use clap::Args;
+use datafusion::physical_plan::execute_stream;
+use datafusion::{error::Result, prelude::SessionContext};
+use datafusion_common::instant::Instant;
+use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err};
+use futures::StreamExt;
+
+/// Run end-to-end ASOF join benchmarks.
+///
+/// The cases cover ordered-input reuse, optimizer-inserted sort/repartition,
+/// wide payload materialization, and descending successor matching.
+#[derive(Debug, Args, Clone)]
+#[command(verbatim_doc_comment)]
+pub struct RunOpt {
+    /// Query number (between 1 and 4). If not specified, runs all queries
+    #[arg(short, long)]
+    query: Option<usize>,
+
+    /// Common options
+    #[command(flatten)]
+    common: CommonOpt,
+
+    /// If present, write results json here
+    #[arg(short = 'o', long = "output")]
+    output_path: Option<std::path::PathBuf>,
+}
+
+const ASOF_QUERIES: &[&str] = &[
+    // Q1: predecessor without equality keys, ordered range input can be reused
+    r#"
+        WITH left_input AS (
+            SELECT value AS ts, value AS payload FROM range(1000000)
+        ),
+        right_input AS (
+            SELECT value AS ts, value AS payload FROM range(1000000)
+        )
+        SELECT l.ts, l.payload, r.payload AS right_payload
+        FROM left_input l
+        ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts)
+    "#,
+    // Q2: grouped predecessor, optimizer supplies repartitioning and ordering
+    r#"
+        WITH left_input AS (
+            SELECT value % 10000 AS key,
+                   value / 10000 + 1 AS ts,
+                   value AS payload
+            FROM range(1000000)
+        ),
+        right_input AS (
+            SELECT value % 10000 AS key,
+                   value / 10000 AS ts,
+                   value AS payload
+            FROM range(1000000)
+        )
+        SELECT l.key, l.ts, l.payload, r.payload AS right_payload
+        FROM left_input l
+        ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts)
+        ON l.key = r.key
+    "#,
+    // Q3: grouped predecessor with a wide payload
+    r#"
+        WITH left_input AS (
+            SELECT value % 10000 AS key,
+                   value / 10000 + 1 AS ts,
+                   repeat('x', 256) AS payload
+            FROM range(250000)
+        ),
+        right_input AS (
+            SELECT value % 10000 AS key,
+                   value / 10000 AS ts,
+                   repeat('y', 256) AS payload
+            FROM range(250000)
+        )
+        SELECT l.key, l.ts, l.payload, r.payload AS right_payload
+        FROM left_input l
+        ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts)
+        ON l.key = r.key
+    "#,
+    // Q4: successor matching requires descending input order
+    r#"
+        WITH left_input AS (
+            SELECT value AS ts, value AS payload FROM range(500000)
+        ),
+        right_input AS (
+            SELECT value AS ts, value AS payload FROM range(500000)
+        )
+        SELECT l.ts, l.payload, r.payload AS right_payload
+        FROM left_input l
+        ASOF JOIN right_input r MATCH_CONDITION (l.ts <= r.ts)
+    "#,
+];
+
+impl RunOpt {
+    pub async fn run(self) -> Result<()> {
+        println!("Running ASOF benchmarks with the following options: {self:#?}\n");
+
+        let query_range = match self.query {
+            Some(query_id) if (1..=ASOF_QUERIES.len()).contains(&query_id) => {
+                query_id..=query_id
+            }
+            Some(query_id) => {
+                return exec_err!(
+                    "Query {query_id} not found. Available queries: 1 to {}",
+                    ASOF_QUERIES.len()
+                );
+            }
+            None => 1..=ASOF_QUERIES.len(),
+        };
+
+        let config = self.common.config()?;
+        let runtime = self.common.build_runtime()?;
+        let ctx = SessionContext::new_with_config_rt(config, runtime);
+        let mut benchmark_run = BenchmarkRun::new();
+
+        for query_id in query_range {
+            let sql = ASOF_QUERIES[query_id - 1];
+            benchmark_run.start_new_case(&format!("Query {query_id}"));
+            match self.benchmark_query(sql, &query_id.to_string(), &ctx).await {
+                Ok(results) => {
+                    for result in results {
+                        benchmark_run.write_iter(result.elapsed, result.row_count);
+                    }
+                }
+                Err(error) => {
+                    return Err(DataFusionError::Context(
+                        format!("ASOF benchmark Q{query_id} failed with error:"),
+                        Box::new(error),
+                    ));
+                }
+            }
+        }
+
+        benchmark_run.maybe_write_json(self.output_path.as_ref())?;
+        Ok(())
+    }
+
+    async fn benchmark_query(
+        &self,
+        sql: &str,
+        query_name: &str,
+        ctx: &SessionContext,
+    ) -> Result<Vec<QueryResult>> {
+        let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?;
+        let plan_string = format!("{physical_plan:#?}");
+        if !plan_string.contains("AsOfJoinExec") {
+            return Err(exec_datafusion_err!(
+                "Query {query_name} does not use AsOfJoinExec. Physical plan: {plan_string}"
+            ));
+        }
+
+        let mut query_results = Vec::with_capacity(self.common.iterations);
+        for iteration in 0..self.common.iterations {
+            let start = Instant::now();
+            let row_count = Self::execute_sql_without_result_buffering(sql, ctx).await?;
+            let elapsed = start.elapsed();
+            println!(
+                "Query {query_name} iteration {iteration} returned {row_count} rows in {elapsed:?}"
+            );
+            query_results.push(QueryResult { elapsed, row_count });
+        }
+        Ok(query_results)
+    }
+
+    async fn execute_sql_without_result_buffering(
+        sql: &str,
+        ctx: &SessionContext,
+    ) -> Result<usize> {
+        let physical_plan = ctx.sql(sql).await?.create_physical_plan().await?;
+        let mut stream = execute_stream(physical_plan, ctx.task_ctx())?;
+        let mut row_count = 0;
+        while let Some(batch) = stream.next().await {
+            row_count += batch?.num_rows();
+        }
+        Ok(row_count)
+    }
+}

benchmarks/src/bin/dfbench.rs

@@ -32,7 +32,7 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
 static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
 use datafusion_benchmarks::{
-    cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch,
+    asof, cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch,
 };
 
 #[derive(Debug, Parser)]
@@ -44,6 +44,7 @@ struct Cli {
 
 #[derive(Debug, Subcommand)]
 enum Options {
+    Asof(asof::RunOpt),
     Cancellation(cancellation::RunOpt),
     Clickbench(clickbench::RunOpt),
     Dict(dict::RunOpt),
@@ -65,6 +66,7 @@ pub async fn main() -> Result<()> {
 
     let cli = Cli::parse();
     match cli.command {
+        Options::Asof(opt) => opt.run().await,
         Options::Cancellation(opt) => opt.run().await,
         Options::Clickbench(opt) => opt.run().await,
         Options::Dict(opt) => opt.run().await,

benchmarks/src/lib.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 //! DataFusion benchmark runner
+pub mod asof;
 pub mod cancellation;
 pub mod clickbench;
 pub mod dict;

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};
@@ -93,8 +94,8 @@ use datafusion_expr::physical_planning_context::{
 use datafusion_expr::utils::{expr_to_columns, split_conjunction};
 use datafusion_expr::{
     Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
-    FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan,
-    WindowFrame, WindowFrameBound, WriteOp,
+    FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType,
+    StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{
     AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder,
@@ -1860,6 +1861,67 @@ 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 omitted_right = if join.join_constraint == JoinConstraint::Using {
+                    join.on
+                        .iter()
+                        .map(|(_, right)| {
+                            let column = right.get_as_join_column().ok_or_else(|| {
+                                internal_datafusion_err!("ASOF USING key is not a column")
+                            })?;
+                            join.right.schema().index_of_column(column)
+                        })
+                        .collect::<Result<HashSet<_>>>()?
+                } else {
+                    HashSet::new()
+                };
+                let right_output_indices = (0..join.right.schema().fields().len())
+                    .filter(|index| !omitted_right.contains(index))
+                    .collect();
+                Arc::new(AsOfJoinExec::try_new(
+                    physical_left,
+                    physical_right,
+                    join_on,
+                    match_condition,
+                    right_output_indices,
+                )?)
+            }
             LogicalPlan::RecursiveQuery(RecursiveQuery {
                 name,
                 is_distinct,
@@ -2359,6 +2421,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},