Repository: narwhals-dev/narwhals
Pull request #911 (2024-09-05 13:07:04)

feat: Add join_asof support for pandas and dask
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #896
- Closes #

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/.github/workflows/extremes.yml b/.github/workflows/extremes.yml
index ae9c7900..7e1a5586 100644
--- a/.github/workflows/extremes.yml
+++ b/.github/workflows/extremes.yml
@@ -104,7 +104,7 @@ jobs:
       - name: uninstall pandas
         run: uv pip uninstall pandas --system
       - name: install-pandas-nightly
-        run: uv pip install --pre --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pandas --system
+        run: uv pip install --prerelease=allow --pre --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pandas --system
       - name: uninstall numpy
         run: uv pip uninstall numpy --system
       - name: install numpy nightly
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index c144b4af..f78b4e3d 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -22,6 +22,7 @@
         - item
         - iter_rows
         - join
+        - join_asof
         - lazy
         - null_count
         - pipe
diff --git a/docs/api-reference/lazyframe.md b/docs/api-reference/lazyframe.md
index 9ca6a974..5d472bab 100644
--- a/docs/api-reference/lazyframe.md
+++ b/docs/api-reference/lazyframe.md
@@ -15,6 +15,7 @@
         - group_by
         - head
         - join
+        - join_asof
         - lazy
         - pipe
         - rename
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index f6cb4710..755a9241 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -315,6 +315,17 @@ class ArrowDataFrame:
             ),
         )
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        msg = "join_asof is not yet supported on PyArrow tables"
+        raise NotImplementedError(msg)
+
     def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
         to_drop = parse_columns_to_drop(
             compliant_frame=self, columns=columns, strict=strict
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index f11a8890..91a7e96a 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -298,6 +298,26 @@ class DaskLazyFrame:
             ),
         )
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        plx = self.__native_namespace__()
+        return self._from_native_frame(
+            plx.merge_asof(
+                self._native_frame,
+                other._native_frame,
+                left_on=left_on,
+                right_on=right_on,
+                direction=strategy,
+                suffixes=("", "_right"),
+            ),
+        )
+
     def group_by(self, *by: str) -> DaskLazyGroupBy:
         from narwhals._dask.group_by import DaskLazyGroupBy
 
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 193955cb..0425e28e 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -509,6 +509,26 @@ class PandasLikeDataFrame:
             ),
         )
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        plx = self.__native_namespace__()
+        return self._from_native_frame(
+            plx.merge_asof(
+                self._native_frame,
+                other._native_frame,
+                left_on=left_on,
+                right_on=right_on,
+                direction=strategy,
+                suffixes=("", "_right"),
+            ),
+        )
+
     # --- partial reduction ---
 
     def head(self, n: int) -> Self:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 9276fda1..da1ee1dc 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -214,6 +214,29 @@ class BaseFrame(Generic[FrameT]):
             self._compliant_frame.gather_every(n=n, offset=offset)
         )
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        _supported_strategies = ("backward", "forward", "nearest")
+
+        if strategy not in _supported_strategies:
+            msg = f"Only the following strategies are supported: {_supported_strategies}; found '{strategy}'."
+            raise NotImplementedError(msg)
+
+        return self._from_compliant_dataframe(
+            self._compliant_frame.join_asof(
+                self._extract_compliant(other),
+                left_on=left_on,
+                right_on=right_on,
+                strategy=strategy,
+            )
+        )
+
 
 class DataFrame(BaseFrame[FrameT]):
     """
@@ -1839,6 +1862,96 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().join(other, how=how, left_on=left_on, right_on=right_on)
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        """
+        Perform an asof join.
+
+        This is similar to a left-join except that we match on nearest key rather than equal keys.
+
+        Both DataFrames must be sorted by the asof_join key.
+
+        Arguments:
+            other: DataFrame to join with.
+
+            left_on: Name(s) of the left join column(s).
+
+            right_on: Name(s) of the right join column(s).
+
+            strategy: Join strategy. The default is "backward".
+
+                  * *backward*: selects the last row in the right DataFrame whose "on" key is less than or equal to the left's key.
+                  * *forward*: selects the first row in the right DataFrame whose "on" key is greater than or equal to the left's key.
+                  * *nearest*: search selects the last row in the right DataFrame whose value is nearest to the left's key.
+
+        Returns:
+            A new joined DataFrame
+
+        Examples:
+            >>> from datetime import datetime
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data_gdp = {
+            ...     "datetime": [
+            ...         datetime(2016, 1, 1),
+            ...         datetime(2017, 1, 1),
+            ...         datetime(2018, 1, 1),
+            ...         datetime(2019, 1, 1),
+            ...         datetime(2020, 1, 1),
+            ...     ],
+            ...     "gdp": [4164, 4411, 4566, 4696, 4827],
+            ... }
+            >>> data_population = {
+            ...     "datetime": [
+            ...         datetime(2016, 3, 1),
+            ...         datetime(2018, 8, 1),
+            ...         datetime(2019, 1, 1),
+            ...     ],
+            ...     "population": [82.19, 82.66, 83.12],
+            ... }
+            >>> gdp_pd = pd.DataFrame(data_gdp)
+            >>> population_pd = pd.DataFrame(data_population)
+
+            >>> gdp_pl = pl.DataFrame(data_gdp).sort("datetime")
+            >>> population_pl = pl.DataFrame(data_population).sort("datetime")
+
+            Let's define a dataframe-agnostic function in which we join over "datetime" column:
+
+            >>> @nw.narwhalify
+            ... def join_asof_date(df, other_any, strategy):
+            ...     return df.join_asof(
+            ...         other_any, left_on="datetime", right_on="datetime", strategy=strategy
+            ...     )
+            >>> # We can now pass either pandas or Polars to the function:
+            >>> join_asof_date(population_pd, gdp_pd, strategy="backward")
+                datetime  population   gdp
+            0 2016-03-01       82.19  4164
+            1 2018-08-01       82.66  4566
+            2 2019-01-01       83.12  4696
+
+            >>> join_asof_date(population_pl, gdp_pl, strategy="backward")
+            shape: (3, 3)
+            ┌─────────────────────┬────────────┬──────┐
+            │ datetime            ┆ population ┆ gdp  │
+            │ ---                 ┆ ---        ┆ ---  │
+            │ datetime[μs]        ┆ f64        ┆ i64  │
+            ╞═════════════════════╪════════════╪══════╡
+            │ 2016-03-01 00:00:00 ┆ 82.19      ┆ 4164 │
+            │ 2018-08-01 00:00:00 ┆ 82.66      ┆ 4566 │
+            │ 2019-01-01 00:00:00 ┆ 83.12      ┆ 4696 │
+            └─────────────────────┴────────────┴──────┘
+        """
+        return super().join_asof(
+            other, left_on=left_on, right_on=right_on, strategy=strategy
+        )
+
     # --- descriptive ---
     def is_duplicated(self: Self) -> Series:
         r"""
@@ -3378,6 +3491,95 @@ class LazyFrame(BaseFrame[FrameT]):
         """
         return super().join(other, how=how, left_on=left_on, right_on=right_on)
 
+    def join_asof(
+        self,
+        other: Self,
+        *,
+        left_on: str,
+        right_on: str,
+        strategy: Literal["backward", "forward", "nearest"] = "backward",
+    ) -> Self:
+        """
+        Perform an asof join.
+
+        This is similar to a left-join except that we match on nearest key rather than equal keys.
+
+        Both DataFrames must be sorted by the asof_join key.
+
+        Arguments:
+            other: DataFrame to join with.
+
+            left_on: Name(s) of the left join column(s).
+
+            right_on: Name(s) of the right join column(s).
+
+            strategy: Join strategy. The default is "backward".
+
+                  * *backward*: selects the last row in the right DataFrame whose "on" key is less than or equal to the left's key.
+                  * *forward*: selects the first row in the right DataFrame whose "on" key is greater than or equal to the left's key.
+                  * *nearest*: search selects the last row in the right DataFrame whose value is nearest to the left's key.
+
+        Returns:
+            A new joined DataFrame
+
+        Examples:
+            >>> from datetime import datetime
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data_gdp = {
+            ...     "datetime": [
+            ...         datetime(2016, 1, 1),
+            ...         datetime(2017, 1, 1),
+            ...         datetime(2018, 1, 1),
+            ...         datetime(2019, 1, 1),
+            ...         datetime(2020, 1, 1),
+            ...     ],
+            ...     "gdp": [4164, 4411, 4566, 4696, 4827],
+            ... }
+            >>> data_population = {
+            ...     "datetime": [
+            ...         datetime(2016, 3, 1),
+            ...         datetime(2018, 8, 1),
+            ...         datetime(2019, 1, 1),
+            ...     ],
+            ...     "population": [82.19, 82.66, 83.12],
+            ... }
+            >>> gdp_pd = pd.DataFrame(data_gdp)
+            >>> population_pd = pd.DataFrame(data_population)
+            >>> gdp_pl = pl.LazyFrame(data_gdp).sort("datetime")
+            >>> population_pl = pl.LazyFrame(data_population).sort("datetime")
+
+            Let's define a dataframe-agnostic function in which we join over "datetime" column:
+
+            >>> @nw.narwhalify
+            ... def join_asof_date(df, other_any, strategy):
+            ...     return df.join_asof(
+            ...         other_any, left_on="datetime", right_on="datetime", strategy=strategy
+            ...     )
+            >>> # We can now pass either pandas or Polars to the function:
+            >>> join_asof_date(population_pd, gdp_pd, strategy="backward")
+                datetime  population   gdp
+            0 2016-03-01       82.19  4164
+            1 2018-08-01       82.66  4566
+            2 2019-01-01       83.12  4696
+
+            >>> join_asof_date(population_pl, gdp_pl, strategy="backward").collect()
+            shape: (3, 3)
+            ┌─────────────────────┬────────────┬──────┐
+            │ datetime            ┆ population ┆ gdp  │
+            │ ---                 ┆ ---        ┆ ---  │
+            │ datetime[μs]        ┆ f64        ┆ i64  │
+            ╞═════════════════════╪════════════╪══════╡
+            │ 2016-03-01 00:00:00 ┆ 82.19      ┆ 4164 │
+            │ 2018-08-01 00:00:00 ┆ 82.66      ┆ 4566 │
+            │ 2019-01-01 00:00:00 ┆ 83.12      ┆ 4696 │
+            └─────────────────────┴────────────┴──────┘
+        """
+        return super().join_asof(
+            other, left_on=left_on, right_on=right_on, strategy=strategy
+        )
+
     def clone(self) -> Self:
         r"""
         Create a copy of this DataFrame.


Test patch:
diff --git a/tests/expr_and_series/dt/datetime_attributes_test.py b/tests/expr_and_series/dt/datetime_attributes_test.py
index 4d59567d..22e20590 100644
--- a/tests/expr_and_series/dt/datetime_attributes_test.py
+++ b/tests/expr_and_series/dt/datetime_attributes_test.py
@@ -42,6 +42,8 @@ def test_datetime_attributes(
         and "pyarrow" not in str(constructor)
     ):
         request.applymarker(pytest.mark.xfail)
+    if attribute == "date" and "cudf" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor(data))
     result = df.select(getattr(nw.col("a").dt, attribute)())
@@ -73,6 +75,8 @@ def test_datetime_attributes_series(
         and "pyarrow" not in str(constructor_eager)
     ):
         request.applymarker(pytest.mark.xfail)
+    if attribute == "date" and "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor_eager(data), eager_only=True)
     result = df.select(getattr(df["a"].dt, attribute)())
@@ -82,6 +86,8 @@ def test_datetime_attributes_series(
 def test_datetime_chained_attributes(request: Any, constructor_eager: Any) -> None:
     if "pandas" in str(constructor_eager) and "pyarrow" not in str(constructor_eager):
         request.applymarker(pytest.mark.xfail)
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor_eager(data), eager_only=True)
     result = df.select(df["a"].dt.date().dt.year())
diff --git a/tests/frame/join_test.py b/tests/frame/join_test.py
index e6dfad63..c9119e20 100644
--- a/tests/frame/join_test.py
+++ b/tests/frame/join_test.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from datetime import datetime
 from typing import Any
 
 import pandas as pd
@@ -8,6 +9,7 @@ import pytest
 
 import narwhals.stable.v1 as nw
 from narwhals.utils import Implementation
+from narwhals.utils import parse_version
 from tests.utils import compare_dicts
 
 
@@ -202,3 +204,113 @@ def test_left_join_overlapping_column(constructor: Any) -> None:
         "index": [0, 1, 2],
     }
     compare_dicts(result, expected)
+
+
+def test_joinasof_numeric(constructor: Any, request: Any) -> None:
+    if "pyarrow_table" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+    if parse_version(pd.__version__) < (2, 1) and (
+        ("pandas_pyarrow" in str(constructor)) or ("pandas_nullable" in str(constructor))
+    ):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor({"a": [1, 5, 10], "val": ["a", "b", "c"]})).sort("a")
+    df_right = nw.from_native(
+        constructor({"a": [1, 2, 3, 6, 7], "val": [1, 2, 3, 6, 7]})
+    ).sort("a")
+    result_backward = df.join_asof(df_right, left_on="a", right_on="a")  # type: ignore[arg-type]
+    result_forward = df.join_asof(df_right, left_on="a", right_on="a", strategy="forward")  # type: ignore[arg-type]
+    result_nearest = df.join_asof(df_right, left_on="a", right_on="a", strategy="nearest")  # type: ignore[arg-type]
+    expected_backward = {
+        "a": [1, 5, 10],
+        "val": ["a", "b", "c"],
+        "val_right": [1, 3, 7],
+    }
+    expected_forward = {
+        "a": [1, 5, 10],
+        "val": ["a", "b", "c"],
+        "val_right": [1, 6, float("nan")],
+    }
+    expected_nearest = {
+        "a": [1, 5, 10],
+        "val": ["a", "b", "c"],
+        "val_right": [1, 6, 7],
+    }
+    compare_dicts(result_backward, expected_backward)
+    compare_dicts(result_forward, expected_forward)
+    compare_dicts(result_nearest, expected_nearest)
+
+
+def test_joinasof_time(constructor: Any, request: Any) -> None:
+    if "pyarrow_table" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+    if parse_version(pd.__version__) < (2, 1) and ("pandas_pyarrow" in str(constructor)):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(
+        constructor(
+            {
+                "datetime": [
+                    datetime(2016, 3, 1),
+                    datetime(2018, 8, 1),
+                    datetime(2019, 1, 1),
+                ],
+                "population": [82.19, 82.66, 83.12],
+            }
+        )
+    ).sort("datetime")
+    df_right = nw.from_native(
+        constructor(
+            {
+                "datetime": [
+                    datetime(2016, 1, 1),
+                    datetime(2017, 1, 1),
+                    datetime(2018, 1, 1),
+                    datetime(2019, 1, 1),
+                    datetime(2020, 1, 1),
+                ],
+                "gdp": [4164, 4411, 4566, 4696, 4827],
+            }
+        )
+    ).sort("datetime")
+    result_backward = df.join_asof(df_right, left_on="datetime", right_on="datetime")  # type: ignore[arg-type]
+    result_forward = df.join_asof(
+        df_right,  # type: ignore[arg-type]
+        left_on="datetime",
+        right_on="datetime",
+        strategy="forward",
+    )
+    result_nearest = df.join_asof(
+        df_right,  # type: ignore[arg-type]
+        left_on="datetime",
+        right_on="datetime",
+        strategy="nearest",
+    )
+    expected_backward = {
+        "datetime": [datetime(2016, 3, 1), datetime(2018, 8, 1), datetime(2019, 1, 1)],
+        "population": [82.19, 82.66, 83.12],
+        "gdp": [4164, 4566, 4696],
+    }
+    expected_forward = {
+        "datetime": [datetime(2016, 3, 1), datetime(2018, 8, 1), datetime(2019, 1, 1)],
+        "population": [82.19, 82.66, 83.12],
+        "gdp": [4411, 4696, 4696],
+    }
+    expected_nearest = {
+        "datetime": [datetime(2016, 3, 1), datetime(2018, 8, 1), datetime(2019, 1, 1)],
+        "population": [82.19, 82.66, 83.12],
+        "gdp": [4164, 4696, 4696],
+    }
+    compare_dicts(result_backward, expected_backward)
+    compare_dicts(result_forward, expected_forward)
+    compare_dicts(result_nearest, expected_nearest)
+
+
+@pytest.mark.parametrize("strategy", ["back", "furthest"])
+def test_joinasof_not_implemented(constructor: Any, strategy: str) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+    df = nw.from_native(constructor(data))
+
+    with pytest.raises(
+        NotImplementedError,
+        match=rf"Only the following strategies are supported: \('backward', 'forward', 'nearest'\); found '{strategy}'.",
+    ):
+        df.join_asof(df, left_on="a", right_on="a", strategy=strategy)  # type: ignore[arg-type]