Add higher level join_between method

Author: adamreeve

py-polars/polars/dataframe/frame.py

@@ -7124,6 +7124,58 @@ def inequality_join(
             .collect(_eager=True)
         )
 
+    def join_between(
+        self,
+        other: DataFrame,
+        *,
+        left_on: str | Expr,
+        right_on_lower: str | Expr,
+        right_on_upper: str | Expr,
+        exclusive_lower: bool = False,
+        exclusive_upper: bool = True,
+        suffix: str = "_right",
+    ) -> DataFrame:
+        """
+        Join by matching values from this table with an interval in another table.
+
+        A row from this table may be included in zero or multiple rows in the result,
+        and the relative order of rows may differ between the input and output tables.
+
+        Parameters
+        ----------
+        other
+            DataFrame to join with.
+        left_on
+            Join column of the left table.
+        right_on_lower
+            Lower bound of the interval in the other table
+        right_on_upper
+            Upper bound of the interval in the other table
+        exclusive_lower
+            Whether the lower bound of the interval is an exclusive bound
+        exclusive_upper
+            Whether the upper bound of the interval is an exclusive bound
+        suffix
+            Suffix to append to columns with a duplicate name.
+        """
+        if not isinstance(other, DataFrame):
+            msg = f"expected `other` join table to be a DataFrame, got {type(other).__name__!r}"
+            raise TypeError(msg)
+
+        return (
+            self.lazy()
+            .join_between(
+                other.lazy(),
+                left_on=left_on,
+                right_on_lower=right_on_lower,
+                right_on_upper=right_on_upper,
+                exclusive_lower=exclusive_lower,
+                exclusive_upper=exclusive_upper,
+                suffix=suffix,
+            )
+            .collect(_eager=True)
+        )
+
     def map_rows(
         self,
         function: Callable[[tuple[Any, ...]], Any],

py-polars/polars/lazyframe/frame.py

@@ -4607,6 +4607,63 @@ def join(
             )
         )
 
+    def join_between(
+        self,
+        other: LazyFrame,
+        *,
+        left_on: str | Expr,
+        right_on_lower: str | Expr,
+        right_on_upper: str | Expr,
+        exclusive_lower: bool = False,
+        exclusive_upper: bool = True,
+        suffix: str = "_right",
+    ) -> LazyFrame:
+        """
+        Join by matching values from this table with an interval in another table.
+
+        A row from this table may be included in zero or multiple rows in the result,
+        and the relative order of rows may differ between the input and output tables.
+
+        Parameters
+        ----------
+        other
+            LazyFrame to join with.
+        left_on
+            Join column of the left table.
+        right_on_lower
+            Lower bound of the interval in the other table
+        right_on_upper
+            Upper bound of the interval in the other table
+        exclusive_lower
+            Whether the lower bound of the interval is an exclusive bound
+        exclusive_upper
+            Whether the upper bound of the interval is an exclusive bound
+        suffix
+            Suffix to append to columns with a duplicate name.
+        """
+        if not isinstance(other, LazyFrame):
+            msg = f"expected `other` join table to be a LazyFrame, not a {type(other).__name__!r}"
+            raise TypeError(msg)
+
+        left_on = wrap_expr(parse_into_expression(left_on))
+        right_on_lower = wrap_expr(parse_into_expression(right_on_lower))
+        right_on_upper = wrap_expr(parse_into_expression(right_on_upper))
+
+        lower_expr = (
+            left_on > right_on_lower if exclusive_lower else left_on >= right_on_lower
+        )
+        upper_expr = (
+            left_on < right_on_upper if exclusive_upper else left_on <= right_on_upper
+        )
+
+        return self._from_pyldf(
+            self._ldf.inequality_join(
+                other._ldf,
+                [lower_expr._pyexpr, upper_expr._pyexpr],
+                suffix,
+            )
+        )
+
     def with_columns(
         self,
         *exprs: IntoExpr | Iterable[IntoExpr],

py-polars/tests/unit/operations/test_inequality_join.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
 import hypothesis.strategies as st
@@ -167,6 +168,52 @@ def test_ie_join_with_expressions() -> None:
     assert_frame_equal(actual, expected, check_row_order=False, check_exact=True)
 
 
+def test_join_between() -> None:
+    left = pl.DataFrame(
+        {
+            "id": [0, 1, 2, 3, 4, 5],
+            "time": [
+                datetime(2024, 8, 26, 15, 34, 30),
+                datetime(2024, 8, 26, 15, 35, 30),
+                datetime(2024, 8, 26, 15, 36, 30),
+                datetime(2024, 8, 26, 15, 37, 30),
+                datetime(2024, 8, 26, 15, 38, 0),
+                datetime(2024, 8, 26, 15, 39, 0),
+            ],
+        }
+    )
+    right = pl.DataFrame(
+        {
+            "id": [0, 1, 2],
+            "start_time": [
+                datetime(2024, 8, 26, 15, 34, 0),
+                datetime(2024, 8, 26, 15, 35, 0),
+                datetime(2024, 8, 26, 15, 38, 0),
+            ],
+            "end_time": [
+                datetime(2024, 8, 26, 15, 36, 0),
+                datetime(2024, 8, 26, 15, 37, 0),
+                datetime(2024, 8, 26, 15, 39, 0),
+            ],
+        }
+    )
+
+    actual = left.join_between(
+        right,
+        left_on="time",
+        right_on_lower="start_time",
+        right_on_upper="end_time",
+    ).select("id", "id_right")
+
+    expected = pl.DataFrame(
+        {
+            "id": [0, 1, 1, 2, 4],
+            "id_right": [0, 0, 1, 1, 2],
+        }
+    )
+    assert_frame_equal(actual, expected, check_row_order=False, check_exact=True)
+
+
 def _inequality_expression(col1: str, op: str, col2: str) -> pl.Expr:
     if op == "<":
         return pl.col(col1) < pl.col(col2)