feat: Support int_range(eager=...)

May need tweaking pending (#2895 (comment))

Author: dangotbanned

narwhals/_plan/arrow/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from narwhals._plan.arrow.dataframe import ArrowDataFrame as DataFrame
+from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
+from narwhals._plan.arrow.namespace import ArrowNamespace as Namespace
+from narwhals._plan.arrow.series import ArrowSeries as Series
+
+__all__ = ["DataFrame", "Expr", "Namespace", "Scalar", "Series"]

narwhals/_plan/arrow/namespace.py

@@ -13,7 +13,7 @@
 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._utils import Implementation, Version
 from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
@@ -23,19 +23,24 @@
     from narwhals._plan.arrow.dataframe import ArrowDataFrame as Frame
     from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
     from narwhals._plan.arrow.series import ArrowSeries as Series
+    from narwhals._plan.arrow.typing import ChunkedArray, IntegerScalar
     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 DateRange, IntRange
     from narwhals._plan.expressions.strings import ConcatStr
     from narwhals._plan.series import Series as NwSeries
+    from narwhals.dtypes import IntegerType
     from narwhals.typing import ConcatMethod, NonNestedLiteral, PythonLiteral
 
 
 PythonLiteralT = TypeVar("PythonLiteralT", bound="PythonLiteral")
+Int64 = Version.MAIN.dtypes.Int64()
 
 
 class ArrowNamespace(EagerNamespace["Frame", "Series", "Expr", "Scalar"]):
+    implementation = Implementation.PYARROW
+
     def __init__(self, version: Version = Version.MAIN) -> None:
         self._version = version
 
@@ -186,14 +191,33 @@ def _range_function_inputs(
         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, start: int, end: int, step: int, dtype: IntegerType, /
+    ) -> ChunkedArray[IntegerScalar]:
+        if dtype is not Int64:
+            pa_dtype = narwhals_to_native_dtype(dtype, self.version)
+            if not pa.types.is_integer(pa_dtype):
+                raise TypeError(dtype)
+            return fn.int_range(start, end, step, dtype=pa_dtype)
+        return fn.int_range(start, end, step)
+
     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)
+        native = self._int_range(start, end, node.function.step, node.function.dtype)
         return self._expr.from_native(native, name, self.version)
 
+    def int_range_eager(
+        self,
+        start: int,
+        end: int,
+        step: int = 1,
+        *,
+        dtype: IntegerType = Int64,
+        name: str = "literal",
+    ) -> Series:
+        native = self._int_range(start, end, step, dtype)
+        return self._series.from_native(native, name, version=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

narwhals/_plan/compliant/namespace.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Protocol, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, overload
 
 from narwhals._plan.compliant.typing import (
     ConcatT1,
@@ -16,6 +16,7 @@
     ScalarT_co,
     SeriesT,
 )
+from narwhals._utils import Implementation, Version
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -27,10 +28,15 @@
     from narwhals._plan.expressions.ranges import DateRange, IntRange
     from narwhals._plan.expressions.strings import ConcatStr
     from narwhals._plan.series import Series
+    from narwhals.dtypes import IntegerType
     from narwhals.typing import ConcatMethod, NonNestedLiteral
 
+Int64 = Version.MAIN.dtypes.Int64()
+
 
 class CompliantNamespace(HasVersion, Protocol[FrameT, ExprT_co, ScalarT_co]):
+    implementation: ClassVar[Implementation]
+
     @property
     def _expr(self) -> type[ExprT_co]: ...
     @property
@@ -131,6 +137,15 @@ def lit(
     def lit(
         self, node: ir.Literal[Any], frame: EagerDataFrameT, name: str
     ) -> EagerExprT_co | EagerScalarT_co: ...
+    def int_range_eager(
+        self,
+        start: int,
+        end: int,
+        step: int = 1,
+        *,
+        dtype: IntegerType = Int64,
+        name: str = "literal",
+    ) -> SeriesT: ...
 
 
 class LazyNamespace(

narwhals/_plan/functions.py

@@ -12,15 +12,27 @@
 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
+from narwhals._utils import Implementation, Version, flatten, is_eager_allowed
+from narwhals.exceptions import ComputeError, InvalidOperationError
 
 if TYPE_CHECKING:
+    import pyarrow as pa
+
+    from narwhals._plan import arrow as _arrow
+    from narwhals._plan.compliant.namespace import EagerNamespace
+    from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
     from narwhals._plan.typing import IntoExpr, IntoExprColumn, NativeSeriesT
+    from narwhals._typing import Arrow
     from narwhals.dtypes import IntegerType
-    from narwhals.typing import ClosedInterval, IntoDType, NonNestedLiteral
+    from narwhals.typing import (
+        ClosedInterval,
+        EagerAllowed,
+        IntoBackend,
+        IntoDType,
+        NonNestedLiteral,
+    )
 
 
 def col(*names: str | t.Iterable[str]) -> Expr:
@@ -145,27 +157,94 @@ def when(
     return When._from_ir(condition)
 
 
+@t.overload
+def int_range(
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
+    step: int = ...,
+    *,
+    dtype: IntegerType | type[IntegerType] = ...,
+    eager: t.Literal[False] = ...,
+) -> Expr: ...
+@t.overload
+def int_range(
+    start: int = ...,
+    end: int | None = ...,
+    step: int = ...,
+    *,
+    dtype: IntegerType | type[IntegerType] = ...,
+    eager: Arrow,
+) -> Series[pa.ChunkedArray[t.Any]]: ...
+@t.overload
+def int_range(
+    start: int = ...,
+    end: int | None = ...,
+    step: int = ...,
+    *,
+    dtype: IntegerType | type[IntegerType] = ...,
+    eager: IntoBackend[EagerAllowed],
+) -> Series: ...
 def int_range(
     start: int | IntoExprColumn = 0,
     end: int | IntoExprColumn | None = None,
     step: int = 1,
     *,
     dtype: IntegerType | type[IntegerType] = Version.MAIN.dtypes.Int64,
-    eager: bool = False,
-) -> Expr:
+    eager: IntoBackend[EagerAllowed] | t.Literal[False] = False,
+) -> Expr | Series:
     if end is None:
         end = start
         start = 0
+    dtype = common.into_dtype(dtype)
     if eager:
-        msg = f"{eager=}"
-        raise NotImplementedError(msg)
+        return _int_range_eager(start, end, step, dtype=dtype, ns=_eager_namespace(eager))
     return (
-        IntRange(step=step, dtype=common.into_dtype(dtype))
+        IntRange(step=step, dtype=dtype)
         .to_function_expr(*_parse.parse_into_seq_of_expr_ir(start, end))
         .to_narwhals()
     )
 
 
+def _int_range_eager(
+    start: t.Any,
+    end: t.Any,
+    step: int,
+    *,
+    dtype: IntegerType,
+    ns: EagerNamespace[t.Any, CompliantSeries[NativeSeriesT], t.Any, t.Any],
+) -> Series[NativeSeriesT]:
+    if not (isinstance(start, int) and isinstance(end, int)):
+        msg = (
+            f"Expected `start` and `end` to be integer values since `eager={ns.implementation}`.\n"
+            f"Found: `start` of type {type(start)} and `end` of type {type(end)}\n\n"
+            "Hint: Calling `nw.int_range` with expressions requires:\n"
+            "  - `eager=False`"
+            "  - a context such as `select` or `with_columns`"
+        )
+        raise InvalidOperationError(msg)
+    return ns.int_range_eager(start, end, step, dtype=dtype).to_narwhals()
+
+
+@t.overload
+def _eager_namespace(backend: Arrow, /) -> _arrow.Namespace: ...
+@t.overload
+def _eager_namespace(
+    backend: IntoBackend[EagerAllowed], /
+) -> EagerNamespace[t.Any, t.Any, t.Any, t.Any]: ...
+def _eager_namespace(
+    backend: IntoBackend[EagerAllowed], /
+) -> EagerNamespace[t.Any, t.Any, t.Any, t.Any] | _arrow.Namespace:
+    impl = Implementation.from_backend(backend)
+    if is_eager_allowed(impl):
+        if impl is Implementation.PYARROW:
+            from narwhals._plan.arrow.namespace import ArrowNamespace
+
+            return ArrowNamespace(Version.MAIN)
+        raise NotImplementedError(impl)
+    msg = f"{impl} support in Narwhals is lazy-only"
+    raise ValueError(msg)
+
+
 def date_range(
     start: dt.date | IntoExprColumn,
     end: dt.date | IntoExprColumn,

tests/plan/compliant_test.py

@@ -658,3 +658,11 @@ def test_dataframe_from_native_overloads() -> None:
         native_bad = native_good.to_batches()[0]
         nwp.DataFrame.from_native(native_bad)  # type: ignore[call-overload]
         assert_type(native_bad, "pa.RecordBatch")
+
+    def test_int_range_overloads() -> None:
+        series = nwp.int_range(50, eager="pyarrow")
+        assert_type(series, "nwp.Series[pa.ChunkedArray[Any]]")
+        native = series.to_native()
+        assert_type(native, "pa.ChunkedArray[Any]")
+        roundtrip = nwp.Series.from_native(native)
+        assert_type(roundtrip, "nwp.Series[pa.ChunkedArray[Any]]")

tests/plan/expr_parsing_test.py

@@ -16,6 +16,7 @@
 from narwhals._plan.expressions import functions as F, operators as ops
 from narwhals._plan.expressions.literal import SeriesLiteral
 from narwhals._plan.expressions.ranges import IntRange
+from narwhals._utils import Implementation
 from narwhals.exceptions import (
     ComputeError,
     InvalidIntoExprError,
@@ -186,11 +187,21 @@ def test_int_range_invalid() -> None:
         int_range.to_function_expr(ir.col("a"))
 
 
-@pytest.mark.xfail(
-    reason="Not implemented `int_range(eager=True)`", raises=NotImplementedError
-)
-def test_int_range_series() -> None:
-    assert isinstance(nwp.int_range(50, eager=True), nwp.Series)
+def test_int_range_eager() -> None:
+    series = nwp.int_range(50, eager="pyarrow")
+    assert isinstance(series, nwp.Series)
+    assert series.to_list() == list(range(50))
+    series = nwp.int_range(50, eager=Implementation.PYARROW)
+    assert series.to_list() == list(range(50))
+
+    with pytest.raises(InvalidOperationError):
+        nwp.int_range(nwp.len(), eager="pyarrow")  # type: ignore[call-overload]
+    with pytest.raises(InvalidOperationError):
+        nwp.int_range(10, nwp.col("a").last(), eager=Implementation.PYARROW)  # type: ignore[call-overload]
+    with pytest.raises(NotImplementedError):
+        nwp.int_range(10, eager="pandas")
+    with pytest.raises(ValueError, match=r"lazy-only"):
+        nwp.int_range(10, eager="duckdb")  # type: ignore[call-overload]
 
 
 def test_over_invalid() -> None: