perf: Rewrite with_row_index so it's faster for pandas/pyarrow and doesn't use rank (#3239)

* wip

* wip

* shiny ci

Author: MarcoGorelli

narwhals/_arrow/dataframe.py

@@ -502,17 +502,18 @@ def to_dict(
         return {ser.name: ser.to_list() for ser in it}
 
     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
+        import numpy as np  # ignore-banned-import
 
-            data = pa.array(np.arange(len(self), dtype=np.int64))
+        plx = self.__narwhals_namespace__()
+        data = pa.array(np.arange(len(self)))
+        row_index_s = plx._series.from_iterable(data, context=self, name=name)
+        row_index = plx._expr._from_series(row_index_s)
+        if order_by:
             row_index = plx._expr._from_series(
-                plx._series.from_iterable(data, context=self, name=name)
+                self.select(row_index, *(plx.col(x) for x in order_by))
+                .sort(*order_by, descending=False, nulls_last=False)
+                .get_column(name)
             )
-        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())
 
     def filter(self, predicate: ArrowExpr) -> Self:

narwhals/_pandas_like/dataframe.py

@@ -444,18 +444,17 @@ 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_s = plx._series.from_iterable(
+            data, context=self, index=self.native.index, name=name
+        )
+        row_index = plx._expr._from_series(row_index_s)
+        if order_by:
             row_index = plx._expr._from_series(
-                plx._series.from_iterable(
-                    data, context=self, index=self.native.index, name=name
-                )
+                self.select(row_index, *(plx.col(x) for x in order_by))
+                .sort(*order_by, descending=False, nulls_last=False)
+                .get_column(name)
             )
-        else:
-            rank = plx.col(order_by[0]).rank(method="ordinal", descending=False)
-            row_index = (rank.over(partition_by=[], order_by=order_by) - 1).alias(name)
         return self.select(row_index, plx.all())
 
     def row(self, index: int) -> tuple[Any, ...]: