feat: Add support for nw.int_range for eager backends

docs/api-reference/narwhals.md

@@ -21,6 +21,7 @@ Here are the top-level functions available in Narwhals.
         - generate_temporary_column_name
         - get_level
         - get_native_namespace
+        - int_range
         - is_ordered_categorical
         - len
         - lit

narwhals/__init__.py

@@ -58,6 +58,7 @@
     from_dict,
     from_numpy,
     get_level,
+    int_range,
     len_ as len,
     lit,
     max,
@@ -143,6 +144,7 @@
     "generate_temporary_column_name",
     "get_level",
     "get_native_namespace",
+    "int_range",
     "is_ordered_categorical",
     "len",
     "lit",

narwhals/_arrow/dataframe.py

@@ -472,16 +472,13 @@ def to_dict(
     def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self:
         plx = self.__narwhals_namespace__()
         if order_by is None:
-            import numpy as np  # ignore-banned-import
-
-            data = pa.array(np.arange(len(self), dtype=np.int64))
-            row_index = plx._expr._from_series(
-                plx._series.from_iterable(data, context=self, name=name)
+            row_index = plx.int_range(
+                start=0, end=len(self), step=1, dtype=self._version.dtypes.Int64()
             )
         else:
             rank = plx.col(order_by[0]).rank("ordinal", descending=False)
-            row_index = (rank.over(partition_by=[], order_by=order_by) - 1).alias(name)
-        return self.select(row_index, plx.all())
+            row_index = rank.over(partition_by=[], order_by=order_by) - 1
+        return self.select(row_index.alias(name), plx.all())
 
     def filter(self, predicate: ArrowExpr | list[bool | None]) -> Self:
         if isinstance(predicate, list):

narwhals/_arrow/namespace.py

@@ -26,6 +26,7 @@
     from narwhals._arrow.typing import ArrayOrScalar, ChunkedArrayAny, Incomplete
     from narwhals._compliant.typing import ScalarKwargs
     from narwhals._utils import Version
+    from narwhals.dtypes import IntegerType
     from narwhals.typing import IntoDType, NonNestedLiteral
 
 
@@ -284,6 +285,49 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             context=self,
         )
 
+    def int_range(
+        self,
+        start: int | ArrowExpr,
+        end: int | ArrowExpr,
+        step: int,
+        *,
+        dtype: IntegerType | type[IntegerType],
+    ) -> ArrowExpr:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            if isinstance(start, ArrowExpr):
+                start_eval = start(df)[0]
+                name = start_eval.name
+                start_value = start_eval.item()
+            else:
+                name = "literal"
+                start_value = start
+
+            end_value = end(df)[0].item() if isinstance(end, ArrowExpr) else end
+            return [
+                ArrowSeries._int_range(
+                    start=start_value,
+                    end=end_value,
+                    step=step,
+                    dtype=dtype,
+                    context=self,
+                    name=name,
+                )
+            ]
+
+        evaluate_output_names = (
+            combine_evaluate_output_names(start)
+            if isinstance(start, ArrowExpr)
+            else lambda _df: ["literal"]
+        )
+        return self._expr._from_callable(
+            func=func,
+            depth=0,
+            function_name="int_range",
+            evaluate_output_names=evaluate_output_names,
+            alias_output_names=None,
+            context=self,
+        )
+
 
 class ArrowWhen(EagerWhen[ArrowDataFrame, ArrowSeries, ArrowExpr, "ChunkedArrayAny"]):
     @property

narwhals/_arrow/series.py

@@ -59,7 +59,7 @@
     )
     from narwhals._compliant.series import HistData
     from narwhals._utils import Version, _LimitedContext
-    from narwhals.dtypes import DType
+    from narwhals.dtypes import DType, IntegerType
     from narwhals.typing import (
         ClosedInterval,
         FillNullStrategy,
@@ -161,6 +161,31 @@ def from_iterable(
             chunked_array([data], dtype_pa), name=name, context=context
         )
 
+    @classmethod
+    def _int_range(
+        cls,
+        start: int,
+        end: int,
+        step: int,
+        dtype: IntegerType | type[IntegerType],
+        context: _LimitedContext,
+        name: str,
+    ) -> Self:
+        version = context._version
+        dtype_pa = narwhals_to_native_dtype(dtype, version)
+        if cls._implementation._backend_version() < (21, 0, 0):  # pragma: no cover
+            import numpy as np  # ignore-banned-import
+
+            data = np.arange(start=start, stop=end, step=step)
+        else:
+            data = pc.cast(  # type: ignore[assignment]
+                pa.arange(start=start, stop=end, step=step),  # type: ignore[attr-defined]
+                dtype_pa,
+            )
+        return cls.from_native(
+            chunked_array([data], dtype_pa), name=name, context=context
+        )
+
     def _from_scalar(self, value: Any) -> Self:
         if hasattr(value, "as_py"):
             value = value.as_py()

narwhals/_compliant/namespace.py

@@ -19,6 +19,7 @@
 from narwhals._utils import (
     exclude_column_names,
     get_column_names,
+    not_implemented,
     passthrough_column_names,
 )
 from narwhals.dependencies import is_numpy_array_2d
@@ -31,7 +32,7 @@
     from narwhals._compliant.selectors import CompliantSelectorNamespace
     from narwhals._compliant.when_then import CompliantWhen, EagerWhen
     from narwhals._utils import Implementation, Version
-    from narwhals.dtypes import DType
+    from narwhals.dtypes import DType, IntegerType
     from narwhals.schema import Schema
     from narwhals.typing import (
         ConcatMethod,
@@ -92,6 +93,14 @@ def when(
     def concat_str(
         self, *exprs: CompliantExprT, separator: str, ignore_nulls: bool
     ) -> CompliantExprT: ...
+    def int_range(
+        self,
+        start: int | CompliantExprT,
+        end: int | CompliantExprT,
+        step: int,
+        *,
+        dtype: IntegerType | type[IntegerType],
+    ) -> CompliantExprT: ...
     @property
     def selectors(self) -> CompliantSelectorNamespace[Any, Any]: ...
     @property
@@ -139,6 +148,8 @@ def from_native(self, data: NativeFrameT_co | Any, /) -> CompliantLazyFrameT:
             msg = f"Unsupported type: {type(data).__name__!r}"
             raise TypeError(msg)
 
+    int_range: not_implemented = not_implemented()
+
 
 class EagerNamespace(
     DepthTrackingNamespace[EagerDataFrameT, EagerExprT],

narwhals/_compliant/series.py

@@ -39,7 +39,7 @@
     from narwhals._compliant.expr import CompliantExpr, EagerExpr
     from narwhals._compliant.namespace import CompliantNamespace, EagerNamespace
     from narwhals._utils import Implementation, Version, _LimitedContext
-    from narwhals.dtypes import DType
+    from narwhals.dtypes import DType, IntegerType
     from narwhals.series import Series
     from narwhals.typing import (
         ClosedInterval,
@@ -122,6 +122,16 @@ def from_iterable(
         name: str = "",
         dtype: IntoDType | None = None,
     ) -> Self: ...
+    @classmethod
+    def _int_range(
+        cls,
+        start: int,
+        end: int,
+        step: int,
+        dtype: IntegerType | type[IntegerType],
+        context: _LimitedContext,
+        name: str,
+    ) -> Self: ...
     def to_narwhals(self) -> Series[NativeSeriesT]:
         return self._version.series(self, level="full")
 

narwhals/_pandas_like/dataframe.py

@@ -413,8 +413,7 @@ def estimated_size(self, unit: SizeUnit) -> int | float:
     def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self:
         plx = self.__narwhals_namespace__()
         if order_by is None:
-            size = len(self)
-            data = self._array_funcs.arange(size)
+            data = self._array_funcs.arange(len(self))
 
             row_index = plx._expr._from_series(
                 plx._series.from_iterable(

narwhals/_pandas_like/namespace.py

@@ -25,6 +25,7 @@
 
     from narwhals._compliant.typing import ScalarKwargs
     from narwhals._utils import Implementation, Version
+    from narwhals.dtypes import IntegerType
     from narwhals.typing import IntoDType, NonNestedLiteral
 
 
@@ -370,6 +371,48 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             context=self,
         )
 
+    def int_range(
+        self,
+        start: int | PandasLikeExpr,
+        end: int | PandasLikeExpr,
+        step: int,
+        *,
+        dtype: IntegerType | type[IntegerType],
+    ) -> PandasLikeExpr:
+        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+            if isinstance(start, PandasLikeExpr):
+                start_eval = start(df)[0]
+                name = start_eval.name
+                start_value = start_eval.item()
+            else:  # pragma: no cover
+                name = "literal"
+                start_value = start
+            end_value = end(df)[0].item() if isinstance(end, PandasLikeExpr) else end
+            return [
+                PandasLikeSeries._int_range(
+                    start=start_value,
+                    end=end_value,
+                    step=step,
+                    dtype=dtype,
+                    context=self,
+                    name=name,
+                )
+            ]
+
+        evaluate_output_names = (
+            combine_evaluate_output_names(start)
+            if isinstance(start, PandasLikeExpr)
+            else lambda _df: ["literal"]
+        )
+        return self._expr._from_callable(
+            func=func,
+            depth=0,
+            function_name="int_range",
+            evaluate_output_names=evaluate_output_names,
+            alias_output_names=None,
+            context=self,
+        )
+
 
 class _NativeConcat(Protocol[NativeDataFrameT, NativeSeriesT]):
     @overload

narwhals/_pandas_like/series.py

@@ -41,7 +41,7 @@
     from narwhals._pandas_like.dataframe import PandasLikeDataFrame
     from narwhals._pandas_like.namespace import PandasLikeNamespace
     from narwhals._utils import Version, _LimitedContext
-    from narwhals.dtypes import DType
+    from narwhals.dtypes import DType, IntegerType
     from narwhals.typing import (
         ClosedInterval,
         FillNullStrategy,
@@ -182,6 +182,20 @@ def from_iterable(
                 kwds["index"] = index
         return cls.from_native(ns.Series(data, name=name, **kwds), context=context)
 
+    @classmethod
+    def _int_range(
+        cls,
+        start: int,
+        end: int,
+        step: int,
+        dtype: IntegerType | type[IntegerType],
+        context: _LimitedContext,
+        name: str,
+    ) -> Self:
+        array_funcs = import_array_module(context._implementation)
+        data = array_funcs.arange(start, end, step)
+        return cls.from_iterable(data, context=context, name=name, dtype=dtype)
+
     @staticmethod
     def _is_native(obj: Any) -> TypeIs[Any]:
         return is_pandas_like_series(obj)  # pragma: no cover
@@ -311,7 +325,7 @@ def cast(self, dtype: IntoDType) -> Self:
         )
         return self._with_native(self.native.astype(pd_dtype), preserve_broadcast=True)
 
-    def item(self, index: int | None) -> Any:
+    def item(self, index: int | None = None) -> Any:
         # cuDF doesn't have Series.item().
         if index is None:
             if len(self) != 1:

narwhals/_polars/namespace.py

@@ -16,10 +16,13 @@
     from collections.abc import Iterable, Mapping, Sequence
     from datetime import timezone
 
+    from polars.datatypes import IntegerType as PlIntegerType
+
     from narwhals._compliant import CompliantSelectorNamespace, CompliantWhen
     from narwhals._polars.dataframe import Method, PolarsDataFrame, PolarsLazyFrame
     from narwhals._polars.typing import FrameT
     from narwhals._utils import Version, _LimitedContext
+    from narwhals.dtypes import IntegerType
     from narwhals.schema import Schema
     from narwhals.typing import Into1DArray, IntoDType, TimeUnit, _2DArray
 
@@ -201,6 +204,24 @@ def concat_str(
             version=self._version,
         )
 
+    def int_range(
+        self,
+        start: int | PolarsExpr,
+        end: int | PolarsExpr,
+        step: int,
+        *,
+        dtype: IntegerType | type[IntegerType],
+    ) -> PolarsExpr:
+        start_ = start if isinstance(start, int) else start.native
+        end_ = end if isinstance(end, int) else end.native
+        pl_dtype: PlIntegerType = narwhals_to_native_dtype(
+            dtype=dtype, version=self._version
+        )  # type: ignore[assignment]
+        return self._expr(
+            pl.int_range(start=start_, end=end_, step=step, dtype=pl_dtype, eager=False),
+            version=self._version,
+        )
+
     # NOTE: Implementation is too different to annotate correctly (vs other `*SelectorNamespace`)
     # 1. Others have lots of private stuff for code reuse
     #    i. None of that is useful here

narwhals/_polars/series.py

@@ -22,13 +22,14 @@
 
     import pandas as pd
     import pyarrow as pa
+    from polars.datatypes import IntegerType as PlIntegerType
     from typing_extensions import Self, TypeAlias, TypeIs
 
     from narwhals._polars.dataframe import Method, PolarsDataFrame
     from narwhals._polars.expr import PolarsExpr
     from narwhals._polars.namespace import PolarsNamespace
     from narwhals._utils import Version, _LimitedContext
-    from narwhals.dtypes import DType
+    from narwhals.dtypes import DType, IntegerType
     from narwhals.series import Series
     from narwhals.typing import Into1DArray, IntoDType, MultiIndexSelector, _1DArray
 
@@ -164,6 +165,23 @@ def from_iterable(
         native = pl.Series(name=name, values=cast("Sequence[Any]", data), dtype=dtype_pl)
         return cls.from_native(native, context=context)
 
+    @classmethod
+    def _int_range(
+        cls,
+        start: int,
+        end: int,
+        step: int,
+        dtype: IntegerType | type[IntegerType],
+        context: _LimitedContext,
+        name: str,
+    ) -> Self:
+        version = context._version
+        dtype_pl: PlIntegerType = narwhals_to_native_dtype(dtype, version)  # type: ignore[assignment]
+        return cls.from_native(
+            pl.int_range(start=start, end=end, step=step, dtype=dtype_pl, eager=True),
+            context=context,
+        ).alias(name)
+
     @staticmethod
     def _is_native(obj: pl.Series | Any) -> TypeIs[pl.Series]:
         return isinstance(obj, pl.Series)