Repository: narwhals-dev/narwhals
Pull request #1034 (2024-09-21 13:20:54)

perf: low-hanging overhead
## What type of PR is this? (check all applicable)

- [x] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Checklist

- [x] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.

Inspired by #993, I took a very high level overview and run [perflint](https://github.com/tonybaloney/perflint) (careful with this one!!!). Tests are all passing locally and I thought to give it a go :)


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index fe0c4025..f593293f 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -37,6 +37,7 @@
         - tail
         - to_arrow
         - to_dict
+        - to_native
         - to_numpy
         - to_pandas
         - unique
diff --git a/narwhals/_arrow/selectors.py b/narwhals/_arrow/selectors.py
index f757ffed..569724c4 100644
--- a/narwhals/_arrow/selectors.py
+++ b/narwhals/_arrow/selectors.py
@@ -99,7 +99,7 @@ class ArrowSelector(ArrowExpr):
             def call(df: ArrowDataFrame) -> list[ArrowSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]]
+                return [x for x in lhs if x.name not in {x.name for x in rhs}]
 
             return ArrowSelector(
                 call,
@@ -118,7 +118,7 @@ class ArrowSelector(ArrowExpr):
             def call(df: ArrowDataFrame) -> list[ArrowSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]] + rhs
+                return [x for x in lhs if x.name not in {x.name for x in rhs}] + rhs
 
             return ArrowSelector(
                 call,
@@ -137,7 +137,7 @@ class ArrowSelector(ArrowExpr):
             def call(df: ArrowDataFrame) -> list[ArrowSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name in [x.name for x in rhs]]
+                return [x for x in lhs if x.name in {x.name for x in rhs}]
 
             return ArrowSelector(
                 call,
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index 29bc3564..ddf7a863 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -265,10 +265,11 @@ def broadcast_series(series: list[ArrowSeries]) -> list[Any]:
 
     import pyarrow as pa  # ignore-banned-import
 
+    is_max_length_gt_1 = max_length > 1
     reshaped = []
     for s, length in zip(series, lengths):
         s_native = s._native_series
-        if max_length > 1 and length == 1:
+        if is_max_length_gt_1 and length == 1:
             value = s_native[0]
             if s._backend_version < (13,) and hasattr(value, "as_py"):  # pragma: no cover
                 value = value.as_py()
diff --git a/narwhals/_dask/selectors.py b/narwhals/_dask/selectors.py
index d3525f71..54131a8a 100644
--- a/narwhals/_dask/selectors.py
+++ b/narwhals/_dask/selectors.py
@@ -102,7 +102,7 @@ class DaskSelector(DaskExpr):
             def call(df: DaskLazyFrame) -> list[Any]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]]
+                return [x for x in lhs if x.name not in {x.name for x in rhs}]
 
             return DaskSelector(
                 call,
@@ -122,7 +122,7 @@ class DaskSelector(DaskExpr):
             def call(df: DaskLazyFrame) -> list[dask_expr.Series]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]] + rhs
+                return [x for x in lhs if x.name not in {x.name for x in rhs}] + rhs
 
             return DaskSelector(
                 call,
@@ -142,7 +142,7 @@ class DaskSelector(DaskExpr):
             def call(df: DaskLazyFrame) -> list[Any]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name in [x.name for x in rhs]]
+                return [x for x in lhs if x.name in {x.name for x in rhs}]
 
             return DaskSelector(
                 call,
diff --git a/narwhals/_dask/utils.py b/narwhals/_dask/utils.py
index 1f5cda4b..27404497 100644
--- a/narwhals/_dask/utils.py
+++ b/narwhals/_dask/utils.py
@@ -45,21 +45,18 @@ def parse_exprs_and_named_exprs(
         else:  # pragma: no cover
             msg = f"Expected expression or column name, got: {expr}"
             raise TypeError(msg)
+        return_scalar = getattr(expr, "_returns_scalar", False)
         for _result in _results:
-            if getattr(expr, "_returns_scalar", False):
-                results[_result.name] = _result[0]
-            else:
-                results[_result.name] = _result
+            results[_result.name] = _result[0] if return_scalar else _result
+
     for name, value in named_exprs.items():
         _results = value._call(df)
         if len(_results) != 1:  # pragma: no cover
             msg = "Named expressions must return a single column"
             raise AssertionError(msg)
+        return_scalar = getattr(value, "_returns_scalar", False)
         for _result in _results:
-            if getattr(value, "_returns_scalar", False):
-                results[name] = _result[0]
-            else:
-                results[name] = _result
+            results[name] = _result[0] if return_scalar else _result
     return results
 
 
diff --git a/narwhals/_expression_parsing.py b/narwhals/_expression_parsing.py
index a74ca3c6..10d686ab 100644
--- a/narwhals/_expression_parsing.py
+++ b/narwhals/_expression_parsing.py
@@ -95,7 +95,7 @@ def evaluate_into_exprs(
     """Evaluate each expr into Series."""
     series: ListOfCompliantSeries = [  # type: ignore[assignment]
         item
-        for sublist in [evaluate_into_expr(df, into_expr) for into_expr in flatten(exprs)]
+        for sublist in (evaluate_into_expr(df, into_expr) for into_expr in flatten(exprs))
         for item in sublist
     ]
     for name, expr in named_exprs.items():
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index 892291d5..a0ee03f7 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -161,8 +161,9 @@ def agg_pandas(  # noqa: PLR0915
             function_name = POLARS_TO_PANDAS_AGGREGATIONS.get(
                 function_name, function_name
             )
+            is_n_unique = function_name == "nunique"
             for root_name, output_name in zip(expr._root_names, expr._output_names):
-                if function_name == "nunique":
+                if is_n_unique:
                     nunique_aggs[output_name] = root_name
                 else:
                     simple_aggregations[output_name] = (root_name, function_name)
diff --git a/narwhals/_pandas_like/selectors.py b/narwhals/_pandas_like/selectors.py
index 8dd9ceae..1214e12f 100644
--- a/narwhals/_pandas_like/selectors.py
+++ b/narwhals/_pandas_like/selectors.py
@@ -102,7 +102,7 @@ class PandasSelector(PandasLikeExpr):
             def call(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]]
+                return [x for x in lhs if x.name not in {x.name for x in rhs}]
 
             return PandasSelector(
                 call,
@@ -122,7 +122,7 @@ class PandasSelector(PandasLikeExpr):
             def call(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name not in [x.name for x in rhs]] + rhs
+                return [x for x in lhs if x.name not in {x.name for x in rhs}] + rhs
 
             return PandasSelector(
                 call,
@@ -142,7 +142,7 @@ class PandasSelector(PandasLikeExpr):
             def call(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                 lhs = self._call(df)
                 rhs = other._call(df)
-                return [x for x in lhs if x.name in [x.name for x in rhs]]
+                return [x for x in lhs if x.name in {x.name for x in rhs}]
 
             return PandasSelector(
                 call,
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index bdcb71ef..a65e361a 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -221,63 +221,54 @@ def set_axis(
 def translate_dtype(column: Any) -> DType:
     from narwhals import dtypes
 
-    dtype = column.dtype
-    if str(dtype) in ("int64", "Int64", "Int64[pyarrow]", "int64[pyarrow]"):
+    dtype = str(column.dtype)
+    if dtype in {"int64", "Int64", "Int64[pyarrow]", "int64[pyarrow]"}:
         return dtypes.Int64()
-    if str(dtype) in ("int32", "Int32", "Int32[pyarrow]", "int32[pyarrow]"):
+    if dtype in {"int32", "Int32", "Int32[pyarrow]", "int32[pyarrow]"}:
         return dtypes.Int32()
-    if str(dtype) in ("int16", "Int16", "Int16[pyarrow]", "int16[pyarrow]"):
+    if dtype in {"int16", "Int16", "Int16[pyarrow]", "int16[pyarrow]"}:
         return dtypes.Int16()
-    if str(dtype) in ("int8", "Int8", "Int8[pyarrow]", "int8[pyarrow]"):
+    if dtype in {"int8", "Int8", "Int8[pyarrow]", "int8[pyarrow]"}:
         return dtypes.Int8()
-    if str(dtype) in ("uint64", "UInt64", "UInt64[pyarrow]", "uint64[pyarrow]"):
+    if dtype in {"uint64", "UInt64", "UInt64[pyarrow]", "uint64[pyarrow]"}:
         return dtypes.UInt64()
-    if str(dtype) in ("uint32", "UInt32", "UInt32[pyarrow]", "uint32[pyarrow]"):
+    if dtype in {"uint32", "UInt32", "UInt32[pyarrow]", "uint32[pyarrow]"}:
         return dtypes.UInt32()
-    if str(dtype) in ("uint16", "UInt16", "UInt16[pyarrow]", "uint16[pyarrow]"):
+    if dtype in {"uint16", "UInt16", "UInt16[pyarrow]", "uint16[pyarrow]"}:
         return dtypes.UInt16()
-    if str(dtype) in ("uint8", "UInt8", "UInt8[pyarrow]", "uint8[pyarrow]"):
+    if dtype in {"uint8", "UInt8", "UInt8[pyarrow]", "uint8[pyarrow]"}:
         return dtypes.UInt8()
-    if str(dtype) in (
+    if dtype in {
         "float64",
         "Float64",
         "Float64[pyarrow]",
         "float64[pyarrow]",
         "double[pyarrow]",
-    ):
+    }:
         return dtypes.Float64()
-    if str(dtype) in (
+    if dtype in {
         "float32",
         "Float32",
         "Float32[pyarrow]",
         "float32[pyarrow]",
         "float[pyarrow]",
-    ):
+    }:
         return dtypes.Float32()
-    if str(dtype) in (
-        "string",
-        "string[python]",
-        "string[pyarrow]",
-        "large_string[pyarrow]",
-    ):
+    if dtype in {"string", "string[python]", "string[pyarrow]", "large_string[pyarrow]"}:
         return dtypes.String()
-    if str(dtype) in ("bool", "boolean", "boolean[pyarrow]", "bool[pyarrow]"):
+    if dtype in {"bool", "boolean", "boolean[pyarrow]", "bool[pyarrow]"}:
         return dtypes.Boolean()
-    if str(dtype) in ("category",) or str(dtype).startswith("dictionary<"):
+    if dtype == "category" or dtype.startswith("dictionary<"):
         return dtypes.Categorical()
-    if str(dtype).startswith("datetime64"):
+    if dtype.startswith(("datetime64", "timestamp[")):
         # TODO(Unassigned): different time units and time zones
         return dtypes.Datetime()
-    if str(dtype).startswith("timedelta64") or str(dtype).startswith("duration"):
+    if dtype.startswith(("timedelta64", "duration")):
         # TODO(Unassigned): different time units
         return dtypes.Duration()
-    if str(dtype).startswith("timestamp["):
-        # pyarrow-backed datetime
-        # TODO(Unassigned): different time units and time zones
-        return dtypes.Datetime()
-    if str(dtype) == "date32[day][pyarrow]":
+    if dtype == "date32[day][pyarrow]":
         return dtypes.Date()
-    if str(dtype) == "object":
+    if dtype == "object":
         if (  # pragma: no cover  TODO(unassigned): why does this show as uncovered?
             idx := getattr(column, "first_valid_index", lambda: None)()
         ) is not None and isinstance(column.loc[idx], str):
@@ -455,10 +446,10 @@ def broadcast_series(series: list[PandasLikeSeries]) -> list[Any]:
 
     idx = series[lengths.index(max_length)]._native_series.index
     reindexed = []
-
+    max_length_gt_1 = max_length > 1
     for s, length in zip(series, lengths):
         s_native = s._native_series
-        if max_length > 1 and length == 1:
+        if max_length_gt_1 and length == 1:
             reindexed.append(
                 native_namespace.Series(
                     [s_native.iloc[0]] * max_length,
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index d100e8ef..2a7a94cf 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -30,10 +30,10 @@ if TYPE_CHECKING:
     from narwhals.group_by import GroupBy
     from narwhals.group_by import LazyGroupBy
     from narwhals.series import Series
-    from narwhals.typing import IntoDataFrame
     from narwhals.typing import IntoExpr
+    from narwhals.typing import IntoFrame
 
-FrameT = TypeVar("FrameT", bound="IntoDataFrame")
+FrameT = TypeVar("FrameT", bound="IntoFrame")
 
 
 class BaseFrame(Generic[FrameT]):
@@ -390,6 +390,54 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().lazy()
 
+    def to_native(self) -> FrameT:
+        """
+        Convert Narwhals DataFrame to native one.
+
+        Returns:
+            Object of class that user started with.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> data = {"foo": [1, 2, 3], "bar": [6.0, 7.0, 8.0], "ham": ["a", "b", "c"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+            >>> df_pa = pa.table(data)
+
+            Calling `to_native` on a Narwhals DataFrame returns the native object:
+
+            >>> nw.from_native(df_pd).to_native()
+               foo  bar ham
+            0    1  6.0   a
+            1    2  7.0   b
+            2    3  8.0   c
+            >>> nw.from_native(df_pl).to_native()
+            shape: (3, 3)
+            ┌─────┬─────┬─────┐
+            │ foo ┆ bar ┆ ham │
+            │ --- ┆ --- ┆ --- │
+            │ i64 ┆ f64 ┆ str │
+            ╞═════╪═════╪═════╡
+            │ 1   ┆ 6.0 ┆ a   │
+            │ 2   ┆ 7.0 ┆ b   │
+            │ 3   ┆ 8.0 ┆ c   │
+            └─────┴─────┴─────┘
+            >>> nw.from_native(df_pa).to_native()
+            pyarrow.Table
+            foo: int64
+            bar: double
+            ham: string
+            ----
+            foo: [[1,2,3]]
+            bar: [[6,7,8]]
+            ham: [["a","b","c"]]
+        """
+
+        return self._compliant_frame._native_frame  # type: ignore[no-any-return]
+
     def to_pandas(self) -> pd.DataFrame:
         """
         Convert this DataFrame to a pandas DataFrame.
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 7b7d09de..0e7706fb 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -703,7 +703,7 @@ def narwhalify(
 
             backends = {
                 b()
-                for v in [*args, *kwargs.values()]
+                for v in (*args, *kwargs.values())
                 if (b := getattr(v, "__native_namespace__", None))
             }
 
diff --git a/narwhals/utils.py b/narwhals/utils.py
index f9d34099..15919a92 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import re
-import secrets
 from enum import Enum
 from enum import auto
+from secrets import token_hex
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
@@ -433,7 +433,7 @@ def generate_unique_token(n_bytes: int, columns: list[str]) -> str:  # pragma: n
     """
     counter = 0
     while True:
-        token = secrets.token_hex(n_bytes)
+        token = token_hex(n_bytes)
         if token not in columns:
             return token
 


Test patch:
diff --git a/tests/frame/to_native_test.py b/tests/frame/to_native_test.py
new file mode 100644
index 00000000..236403e2
--- /dev/null
+++ b/tests/frame/to_native_test.py
@@ -0,0 +1,11 @@
+from typing import Any
+
+import narwhals.stable.v1 as nw
+
+
+def test_to_native(constructor_eager: Any) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.1, 8, 9]}
+    df_raw = constructor_eager(data)
+    df = nw.from_native(df_raw, eager_only=True)
+
+    assert isinstance(df.to_native(), df_raw.__class__)