perf: Set pandas attributes inplace when it is safe to do so (#2559)

Author: MarcoGorelli

CONTRIBUTING.md

@@ -139,6 +139,23 @@ If you add code that should be tested, please add tests.
   - To run tests using `cudf.pandas`, run `NARWHALS_DEFAULT_CONSTRUCTORS=pandas python -m cudf.pandas -m pytest`
   - To run tests using `polars[gpu]`, run `NARWHALS_POLARS_GPU=1 pytest --constructors=polars[lazy]`
 
+### Backend-specific advice
+
+- pandas:
+
+  - Don't use `apply` or `map`. The only place we currently use `apply` is in `group_by` for operations
+    which the pandas API just doesn't support, and even then, it's accompanied by a big warning.
+  - Don't use inplace methods, unless you're creating a new object and are sure it's safe to modify
+    it. In particular, you should never ever modify the user's input data.
+
+- Polars:
+
+  - Never use `map_elements`.
+
+- DuckDB / PySpark / anything lazy-only:
+
+  - Never assume that your data is ordered in any pre-defined way.
+
 ### Test Failure Patterns
 
 We aim to use three standard patterns for handling test failures:

narwhals/_pandas_like/group_by.py

@@ -12,7 +12,6 @@
 from narwhals._compliant import EagerGroupBy
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
 from narwhals._pandas_like.utils import select_columns_by_name
-from narwhals._pandas_like.utils import set_columns
 from narwhals.utils import find_stacklevel
 
 if TYPE_CHECKING:
@@ -199,25 +198,17 @@ def agg(self, *exprs: PandasLikeExpr) -> PandasLikeDataFrame:  # noqa: C901, PLR
                 result_aggs.append(result_nunique_aggs)
 
             if std_aggs:
-                result_aggs.extend(
-                    set_columns(
-                        self._grouped[std_output_names].std(ddof=ddof),
-                        columns=std_aliases,
-                        implementation=implementation,
-                        backend_version=backend_version,
-                    )
-                    for ddof, (std_output_names, std_aliases) in std_aggs.items()
-                )
+                for ddof, (std_output_names, std_aliases) in std_aggs.items():
+                    _aggregation = self._grouped[std_output_names].std(ddof=ddof)
+                    # `_aggregation` is a new object so it's OK to operate inplace.
+                    _aggregation.columns = std_aliases
+                    result_aggs.append(_aggregation)
             if var_aggs:
-                result_aggs.extend(
-                    set_columns(
-                        self._grouped[var_output_names].var(ddof=ddof),
-                        columns=var_aliases,
-                        implementation=implementation,
-                        backend_version=backend_version,
-                    )
-                    for ddof, (var_output_names, var_aliases) in var_aggs.items()
-                )
+                for ddof, (var_output_names, var_aliases) in var_aggs.items():
+                    _aggregation = self._grouped[var_output_names].var(ddof=ddof)
+                    # `_aggregation` is a new object so it's OK to operate inplace.
+                    _aggregation.columns = var_aliases
+                    result_aggs.append(_aggregation)
 
             if result_aggs:
                 output_names_counter = collections.Counter(

narwhals/_pandas_like/series_list.py

@@ -6,7 +6,6 @@
 from narwhals._pandas_like.utils import PandasLikeSeriesNamespace
 from narwhals._pandas_like.utils import get_dtype_backend
 from narwhals._pandas_like.utils import narwhals_to_native_dtype
-from narwhals._pandas_like.utils import set_index
 
 if TYPE_CHECKING:
     from narwhals._pandas_like.series import PandasLikeSeries
@@ -20,12 +19,8 @@ def len(self) -> PandasLikeSeries:
         implementation = self.implementation
         backend_version = self.backend_version
         if implementation.is_pandas() and backend_version < (3, 0):  # pragma: no cover
-            result = set_index(
-                result,
-                self.native.index,
-                implementation=implementation,
-                backend_version=backend_version,
-            )
+            # `result` is a new object so it's safe to do this inplace.
+            result.index = self.native.index
         dtype = narwhals_to_native_dtype(
             self.version.dtypes.UInt32(),
             get_dtype_backend(result.dtype, implementation),

narwhals/_pandas_like/utils.py

@@ -166,36 +166,6 @@ def set_index(
     return obj.set_axis(index, axis=0, **kwargs)  # type: ignore[attr-defined]
 
 
-def set_columns(
-    obj: T,
-    columns: list[str],
-    *,
-    implementation: Implementation,
-    backend_version: tuple[int, ...],
-) -> T:
-    """Wrapper around pandas' set_axis to set object columns.
-
-    We can set `copy` / `inplace` based on implementation/version.
-    """
-    if implementation is Implementation.CUDF:  # pragma: no cover
-        obj = obj.copy(deep=False)  # type: ignore[attr-defined]
-        obj.columns = columns  # type: ignore[attr-defined]
-        return obj
-    if implementation is Implementation.PANDAS and (
-        backend_version < (1,)
-    ):  # pragma: no cover
-        kwargs = {"inplace": False}
-    else:
-        kwargs = {}
-    if implementation is Implementation.PANDAS and (
-        (1, 5) <= backend_version < (3,)
-    ):  # pragma: no cover
-        kwargs["copy"] = False
-    else:  # pragma: no cover
-        pass
-    return obj.set_axis(columns, axis=1, **kwargs)  # type: ignore[attr-defined]
-
-
 def rename(
     obj: T,
     *args: Any,