feat(expr-ir): Add date_range

Mostly just re-purposing `int_range` *so far*, as mentioned in (#2895 (comment))

Need to think about `Interval` some more

Author: dangotbanned

narwhals/_plan/__init__.py

@@ -9,6 +9,7 @@
     any_horizontal,
     col,
     concat_str,
+    date_range,
     exclude,
     int_range,
     len,
@@ -38,6 +39,7 @@
     "any_horizontal",
     "col",
     "concat_str",
+    "date_range",
     "exclude",
     "int_range",
     "len",

narwhals/_plan/arrow/functions.py

@@ -20,6 +20,7 @@
 from narwhals._utils import Implementation
 
 if TYPE_CHECKING:
+    import datetime as dt
     from collections.abc import Iterable, Mapping
 
     from typing_extensions import TypeAlias, TypeIs
@@ -43,6 +44,7 @@
         DataType,
         DataTypeRemap,
         DataTypeT,
+        DateScalar,
         IntegerScalar,
         IntegerType,
         LargeStringType,
@@ -328,12 +330,38 @@ def int_range(
     *,
     dtype: IntegerType = pa.int64(),  # noqa: B008
 ) -> ChunkedArray[IntegerScalar]:
-    import numpy as np  # ignore-banned-import
-
     if end is None:
         end = start
         start = 0
-    return pa.chunked_array([pa.array(np.arange(start, end, step), dtype)])
+    if BACKEND_VERSION < (21, 0, 0):  # pragma: no cover
+        import numpy as np  # ignore-banned-import
+
+        arr = pa.array(np.arange(start=start, stop=end, step=step), type=dtype)
+    else:
+        pa_arange: Incomplete = t.cast("Incomplete", pa.arange)  # type: ignore[attr-defined]
+        arr = t.cast("ArrayAny", pa_arange(start=start, stop=end, step=step)).cast(dtype)
+    return pa.chunked_array([arr])
+
+
+def date_range(
+    start: dt.date,
+    end: dt.date,
+    interval: int,  # (* assuming the `Interval` part is solved)
+    *,
+    closed: ClosedInterval = "both",
+) -> ChunkedArray[DateScalar]:
+    start_i = pa.scalar(start).cast(pa.int32()).as_py()
+    end_i = pa.scalar(end).cast(pa.int32()).as_py()
+    ca = int_range(start_i, end_i + 1, interval, dtype=pa.int32())
+    if closed == "both":
+        return ca.cast(pa.date32())
+    if closed == "left":
+        ca = ca.slice(length=ca.length() - 1)
+    elif closed == "none":
+        ca = ca.slice(1, length=ca.length() - 1)
+    else:
+        ca = ca.slice(1)
+    return ca.cast(pa.date32())
 
 
 def nulls_like(n: int, native: ArrowAny) -> ArrayAny:

narwhals/_plan/arrow/namespace.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime as dt
 from functools import reduce
 from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
@@ -11,6 +12,7 @@
 from narwhals._plan.arrow import functions as fn
 from narwhals._plan.compliant.namespace import EagerNamespace
 from narwhals._plan.expressions.literal import is_literal_scalar
+from narwhals._typing_compat import TypeVar
 from narwhals._utils import Version
 from narwhals.exceptions import InvalidOperationError
 
@@ -24,12 +26,15 @@
     from narwhals._plan.expressions import expr, functions as F
     from narwhals._plan.expressions.boolean import AllHorizontal, AnyHorizontal
     from narwhals._plan.expressions.expr import FunctionExpr, RangeExpr
-    from narwhals._plan.expressions.ranges import IntRange
+    from narwhals._plan.expressions.ranges import DateRange, IntRange
     from narwhals._plan.expressions.strings import ConcatStr
     from narwhals._plan.series import Series as NwSeries
     from narwhals.typing import ConcatMethod, NonNestedLiteral, PythonLiteral
 
 
+PythonLiteralT = TypeVar("PythonLiteralT", bound="PythonLiteral")
+
+
 class ArrowNamespace(EagerNamespace["Frame", "Series", "Expr", "Scalar"]):
     def __init__(self, version: Version = Version.MAIN) -> None:
         self._version = version
@@ -155,12 +160,12 @@ def concat_str(
             return self._scalar.from_native(result, name, self.version)
         return self._expr.from_native(result, name, self.version)
 
-    def int_range(self, node: RangeExpr[IntRange], frame: Frame, name: str) -> Expr:
+    def _range_function_inputs(
+        self, node: RangeExpr, frame: Frame, valid_type: type[PythonLiteralT]
+    ) -> tuple[PythonLiteralT, PythonLiteralT]:
         start_: PythonLiteral
         end_: PythonLiteral
         start, end = node.function.unwrap_input(node)
-        step = node.function.step
-        dtype = node.function.dtype
         if is_literal_scalar(start) and is_literal_scalar(end):
             start_, end_ = start.unwrap(), end.unwrap()
         else:
@@ -172,22 +177,29 @@ def int_range(self, node: RangeExpr[IntRange], frame: Frame, name: str) -> Expr:
                 start_, end_ = scalar_start.to_python(), scalar_end.to_python()
             else:
                 msg = (
-                    f"All inputs for `int_range()` must be scalar or aggregations, but got \n"
+                    f"All inputs for `{node.function}()` must be scalar or aggregations, but got \n"
                     f"{scalar_start.native!r}\n{scalar_end.native!r}"
                 )
                 raise InvalidOperationError(msg)
-        if isinstance(start_, int) and isinstance(end_, int):
-            pa_dtype = narwhals_to_native_dtype(dtype, self.version)
-            if not pa.types.is_integer(pa_dtype):
-                raise TypeError(pa_dtype)
-            native = fn.int_range(start_, end_, step, dtype=pa_dtype)
-            return self._expr.from_native(native, name, self.version)
-
-        msg = (
-            f"All inputs for `int_range()` resolve to int, but got \n{start_!r}\n{end_!r}"
-        )
+        if isinstance(start_, valid_type) and isinstance(end_, valid_type):
+            return start_, end_
+        msg = f"All inputs for `{node.function}()` resolve to {valid_type.__name__}, but got \n{start_!r}\n{end_!r}"
         raise InvalidOperationError(msg)
 
+    def int_range(self, node: RangeExpr[IntRange], frame: Frame, name: str) -> Expr:
+        start, end = self._range_function_inputs(node, frame, int)
+        dtype = narwhals_to_native_dtype(node.function.dtype, self.version)
+        if not pa.types.is_integer(dtype):
+            raise TypeError(dtype)
+        native = fn.int_range(start, end, node.function.step, dtype=dtype)
+        return self._expr.from_native(native, name, self.version)
+
+    def date_range(self, node: RangeExpr[DateRange], frame: Frame, name: str) -> Expr:
+        start, end = self._range_function_inputs(node, frame, dt.date)
+        func = node.function
+        native = fn.date_range(start, end, func.interval, closed=func.closed)
+        return self._expr.from_native(native, name, self.version)
+
     @overload
     def concat(self, items: Iterable[Frame], *, how: ConcatMethod) -> Frame: ...
     @overload

narwhals/_plan/arrow/typing.py

@@ -10,6 +10,7 @@
     import pyarrow as pa
     import pyarrow.compute as pc
     from pyarrow.lib import (
+        Date32Type,
         Int8Type,
         Int16Type,
         Int32Type,
@@ -28,6 +29,7 @@
     StringScalar: TypeAlias = "Scalar[StringType | LargeStringType]"
     IntegerType: TypeAlias = "Int8Type | Int16Type | Int32Type | Int64Type | Uint8Type | Uint16Type | Uint32Type | Uint64Type"
     IntegerScalar: TypeAlias = "Scalar[IntegerType]"
+    DateScalar: TypeAlias = "Scalar[Date32Type]"
 
     class NativeArrowSeries(NativeSeries, Protocol):
         @property

narwhals/_plan/compliant/namespace.py

@@ -24,7 +24,7 @@
 
     from narwhals._plan import expressions as ir
     from narwhals._plan.expressions import FunctionExpr, boolean, functions as F
-    from narwhals._plan.expressions.ranges import IntRange
+    from narwhals._plan.expressions.ranges import DateRange, IntRange
     from narwhals._plan.expressions.strings import ConcatStr
     from narwhals._plan.series import Series
     from narwhals.typing import ConcatMethod, NonNestedLiteral
@@ -47,6 +47,9 @@ def col(self, node: ir.Column, frame: FrameT, name: str) -> ExprT_co: ...
     def concat_str(
         self, node: FunctionExpr[ConcatStr], frame: FrameT, name: str
     ) -> ExprT_co | ScalarT_co: ...
+    def date_range(
+        self, node: ir.RangeExpr[DateRange], frame: FrameT, name: str
+    ) -> ExprT_co: ...
     def int_range(
         self, node: ir.RangeExpr[IntRange], frame: FrameT, name: str
     ) -> ExprT_co: ...

narwhals/_plan/expressions/ranges.py

@@ -10,6 +10,7 @@
 
     from narwhals._plan.expressions import ExprIR, RangeExpr
     from narwhals.dtypes import IntegerType
+    from narwhals.typing import ClosedInterval
 
 
 class RangeFunction(Function, config=FEOptions.namespaced()):
@@ -18,6 +19,10 @@ def to_function_expr(self, *inputs: ExprIR) -> RangeExpr[Self]:
 
         return RangeExpr(input=inputs, function=self, options=self.function_options)
 
+    def unwrap_input(self, node: RangeExpr[Self], /) -> tuple[ExprIR, ExprIR]:
+        start, end = node.input
+        return start, end
+
 
 class IntRange(RangeFunction, options=FunctionOptions.row_separable):
     """N-ary (start, end)."""
@@ -26,6 +31,10 @@ class IntRange(RangeFunction, options=FunctionOptions.row_separable):
     step: int
     dtype: IntegerType
 
-    def unwrap_input(self, node: RangeExpr[Self], /) -> tuple[ExprIR, ExprIR]:
-        start, end = node.input
-        return start, end
+
+class DateRange(RangeFunction, options=FunctionOptions.row_separable):
+    """N-ary (start, end)."""
+
+    __slots__ = ("interval", "closed")  # noqa: RUF023
+    interval: int
+    closed: ClosedInterval

narwhals/_plan/functions.py

@@ -1,23 +1,26 @@
 from __future__ import annotations
 
 import builtins
+import datetime as dt
 import typing as t
 from typing import TYPE_CHECKING
 
+from narwhals._duration import Interval
 from narwhals._plan import _guards, _parse, common, expressions as ir, selectors as cs
 from narwhals._plan.expressions import functions as F
 from narwhals._plan.expressions.literal import ScalarLiteral, SeriesLiteral
-from narwhals._plan.expressions.ranges import IntRange
+from narwhals._plan.expressions.ranges import DateRange, IntRange
 from narwhals._plan.expressions.strings import ConcatStr
 from narwhals._plan.when_then import When
 from narwhals._utils import Version, flatten
+from narwhals.exceptions import ComputeError
 
 if TYPE_CHECKING:
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
     from narwhals._plan.typing import IntoExpr, IntoExprColumn, NativeSeriesT
     from narwhals.dtypes import IntegerType
-    from narwhals.typing import IntoDType, NonNestedLiteral
+    from narwhals.typing import ClosedInterval, IntoDType, NonNestedLiteral
 
 
 def col(*names: str | t.Iterable[str]) -> Expr:
@@ -161,3 +164,36 @@ def int_range(
         .to_function_expr(*_parse.parse_into_seq_of_expr_ir(start, end))
         .to_narwhals()
     )
+
+
+def date_range(
+    start: dt.date | IntoExprColumn,
+    end: dt.date | IntoExprColumn,
+    interval: str | dt.timedelta = "1d",
+    *,
+    closed: ClosedInterval = "both",
+    eager: bool = False,
+) -> Expr:
+    if eager:
+        msg = f"{eager=}"
+        raise NotImplementedError(msg)
+    return (
+        DateRange(interval=_interval_days(interval), closed=closed)
+        .to_function_expr(*_parse.parse_into_seq_of_expr_ir(start, end))
+        .to_narwhals()
+    )
+
+
+def _interval_days(interval: str | dt.timedelta, /) -> int:
+    if not isinstance(interval, dt.timedelta):
+        if interval == "1d":
+            return 1
+        parsed = Interval.parse_no_constraints(interval)
+        if parsed.unit not in {"d", "mo", "q", "y"}:
+            msg = f"`interval` input for `date_range` must consist of full days, got: {parsed.multiple}{parsed.unit}"
+            raise ComputeError(msg)
+        if parsed.unit in {"mo", "q", "y"}:
+            msg = f"`interval` input for `date_range` does not support {parsed.unit!r} yet, got: {parsed.multiple}{parsed.unit}"
+            raise NotImplementedError(msg)
+        interval = parsed.to_timedelta()
+    return interval.days