This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 4d87bafcdf feat: add ASOF joins to DataFrame API (#23831)
4d87bafcdf is described below
commit 4d87bafcdf1e0a038ffe0fa329ec5bc880e0ea07
Author: Xuanwo <[email protected]>
AuthorDate: Fri Sep 11 06:45:10 2026 +0000
feat: add ASOF joins to DataFrame API (#23831)
## Which issue does this PR close?
- Part of #318.
- Umbrella PR: #23738.
## Rationale for this change
This is the DataFrame API layer of the ASOF JOIN stack. It exposes the
logical contract already available on `main` through the fluent API
without coupling DataFrame callers to SQL.
## What changes are included in this PR?
- Add `DataFrame::join_asof` with an optional equality expression and
one ordered match expression.
- Add `DataFrame::join_asof_using` for same-named equality keys.
- Cover both entry points with a DataFrame execution test.
## Are these changes tested?
Yes:
- `cargo fmt --all`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `cargo test -p datafusion --test core_integration dataframe::join_asof
--all-features`
- The extended workspace test command from the contributor guide
## Are there any user-facing changes?
Users can build ASOF joins through `DataFrame::join_asof` and
`DataFrame::join_asof_using`. Both methods use the same expression
validation and left-preserving semantics as the logical builder and SQL
API.
---
datafusion/core/src/dataframe/mod.rs | 101 +++++++++++++++++++++++++++++++++
datafusion/core/tests/dataframe/mod.rs | 69 ++++++++++++++++++++++
2 files changed, 170 insertions(+)
diff --git a/datafusion/core/src/dataframe/mod.rs
b/datafusion/core/src/dataframe/mod.rs
index ed3dc5ea83..55d3695a68 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -1380,6 +1380,107 @@ impl DataFrame {
})
}
+ /// Join this `DataFrame` to the closest eligible row in `right`.
+ ///
+ /// Every left row is emitted exactly once, with `NULL` values for the
right
+ /// columns when no eligible row exists. `NULL` ordered values and equality
+ /// keys never match. When present, `on` must contain equality comparisons
+ /// combined with `AND`. `match_condition` must be a single `<`, `<=`, `>`,
+ /// or `>=` comparison whose left and right operands reference this
+ /// `DataFrame` and `right`, respectively.
+ ///
+ /// # Example
+ /// ```
+ /// # use datafusion::arrow::array::record_batch;
+ /// # use datafusion::error::Result;
+ /// # use datafusion::prelude::*;
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// # let ctx = SessionContext::new();
+ /// # let trades = ctx.read_batch(record_batch!(
+ /// # ("symbol", Utf8, ["A"]),
+ /// # ("ts", Int64, [4])
+ /// # )?)?.alias("trades")?;
+ /// # let prices = ctx.read_batch(record_batch!(
+ /// # ("symbol", Utf8, ["A"]),
+ /// # ("ts", Int64, [2]),
+ /// # ("price", Int32, [20])
+ /// # )?)?.alias("prices")?;
+ /// // For each trade, find the latest price at or before its timestamp.
+ /// let joined = trades.join_asof(
+ /// prices,
+ /// Some(col("trades.symbol").eq(col("prices.symbol"))),
+ /// col("trades.ts").gt_eq(col("prices.ts")),
+ /// )?;
+ /// # let _ = joined;
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn join_asof(
+ self,
+ right: DataFrame,
+ on: Option<Expr>,
+ match_condition: Expr,
+ ) -> Result<DataFrame> {
+ let plan = LogicalPlanBuilder::from(self.plan)
+ .asof_join_on(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.
+ ///
+ /// This has the same matching behavior as [`join_asof`](Self::join_asof),
+ /// but accepts columns that appear under the same name on both inputs.
+ ///
+ /// # Example
+ /// ```
+ /// # use datafusion::arrow::array::record_batch;
+ /// # use datafusion::error::Result;
+ /// # use datafusion::prelude::*;
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// # let ctx = SessionContext::new();
+ /// # let trades = ctx.read_batch(record_batch!(
+ /// # ("symbol", Utf8, ["A"]),
+ /// # ("ts", Int64, [4])
+ /// # )?)?.alias("trades")?;
+ /// # let prices = ctx.read_batch(record_batch!(
+ /// # ("symbol", Utf8, ["A"]),
+ /// # ("ts", Int64, [2]),
+ /// # ("price", Int32, [20])
+ /// # )?)?.alias("prices")?;
+ /// // Same-named equality keys can be specified once.
+ /// let joined = trades.join_asof_using(
+ /// prices,
+ /// vec![Column::from_name("symbol")],
+ /// col("trades.ts").gt_eq(col("prices.ts")),
+ /// )?;
+ /// # let _ = joined;
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn join_asof_using(
+ self,
+ right: DataFrame,
+ using_keys: Vec<Column>,
+ match_condition: Expr,
+ ) -> 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
diff --git a/datafusion/core/tests/dataframe/mod.rs
b/datafusion/core/tests/dataframe/mod.rs
index 392a6391bd..5db3d23ca1 100644
--- a/datafusion/core/tests/dataframe/mod.rs
+++ b/datafusion/core/tests/dataframe/mod.rs
@@ -1584,6 +1584,75 @@ 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
+ .clone()
+ .join_asof(
+ right.clone(),
+ Some(col("trades.symbol").eq(col("prices.symbol"))),
+ col("trades.ts").gt_eq(col("prices.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
+ );
+
+ let results = left
+ .join_asof_using(
+ right,
+ vec![datafusion_common::Column::from_name("symbol")],
+ col("trades.ts").gt_eq(col("prices.ts")),
+ )?
+ .select(vec![col("symbol"), col("trade_id"), col("price")])?
+ .sort(vec![col("trade_id").sort(true, true)])?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(
+ [
+ "+--------+----------+-------+",
+ "| symbol | trade_id | price |",
+ "+--------+----------+-------+",
+ "| A | 1 | |",
+ "| A | 2 | 40 |",
+ "| B | 3 | 101 |",
+ "+--------+----------+-------+",
+ ],
+ &results
+ );
+ Ok(())
+}
+
#[tokio::test]
async fn join_coercion_unnamed() -> Result<()> {
let ctx = SessionContext::new();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
