perf: only validate duplicate column names when collecting for duckdb/pyspark/dask (#2092)

Author: MarcoGorelli

narwhals/_arrow/dataframe.py

@@ -356,7 +356,7 @@ def select(self: Self, *exprs: ArrowExpr) -> Self:
         names = [s.name for s in new_series]
         new_series = align_series_full_broadcast(*new_series)
         df = pa.Table.from_arrays([s._native_series for s in new_series], names=names)
-        return self._from_native_frame(df, validate_column_names=False)
+        return self._from_native_frame(df, validate_column_names=True)
 
     def with_columns(self: Self, *exprs: ArrowExpr) -> Self:
         native_frame = self._native_frame

narwhals/_dask/dataframe.py

@@ -10,7 +10,6 @@
 
 from narwhals._dask.utils import add_row_index
 from narwhals._dask.utils import evaluate_exprs
-from narwhals._pandas_like.utils import check_column_names_are_unique
 from narwhals._pandas_like.utils import native_to_narwhals_dtype
 from narwhals._pandas_like.utils import select_columns_by_name
 from narwhals.typing import CompliantDataFrame
@@ -41,15 +40,14 @@ def __init__(
         *,
         backend_version: tuple[int, ...],
         version: Version,
-        validate_column_names: bool,
+        # Unused, just for compatibility. We only validate when collecting.
+        validate_column_names: bool = False,
     ) -> None:
         self._native_frame: dd.DataFrame = native_dataframe
         self._backend_version = backend_version
         self._implementation = Implementation.DASK
         self._version = version
         validate_backend_version(self._implementation, self._backend_version)
-        if validate_column_names:
-            check_column_names_are_unique(native_dataframe.columns)
 
     def __native_namespace__(self: Self) -> ModuleType:
         if self._implementation is Implementation.DASK:
@@ -71,23 +69,19 @@ def _change_version(self: Self, version: Version) -> Self:
             self._native_frame,
             backend_version=self._backend_version,
             version=version,
-            validate_column_names=False,
         )
 
-    def _from_native_frame(
-        self: Self, df: Any, *, validate_column_names: bool = True
-    ) -> Self:
+    def _from_native_frame(self: Self, df: Any) -> Self:
         return self.__class__(
             df,
             backend_version=self._backend_version,
             version=self._version,
-            validate_column_names=validate_column_names,
         )
 
     def with_columns(self: Self, *exprs: DaskExpr) -> Self:
         df = self._native_frame
         new_series = evaluate_exprs(self, *exprs)
-        df = df.assign(**new_series)
+        df = df.assign(**dict(new_series))
         return self._from_native_frame(df)
 
     def collect(
@@ -107,7 +101,7 @@ def collect(
                 implementation=Implementation.PANDAS,
                 backend_version=parse_version(pd),
                 version=self._version,
-                validate_column_names=False,
+                validate_column_names=True,
             )
 
         if backend is Implementation.POLARS:
@@ -130,7 +124,7 @@ def collect(
                 pa.Table.from_pandas(result),
                 backend_version=parse_version(pa),
                 version=self._version,
-                validate_column_names=False,
+                validate_column_names=True,
             )
 
         msg = f"Unsupported `backend` value: {backend}"  # pragma: no cover
@@ -144,9 +138,7 @@ def filter(self: Self, predicate: DaskExpr) -> Self:
         # `[0]` is safe as the predicate's expression only returns a single column
         mask = predicate._call(self)[0]
 
-        return self._from_native_frame(
-            self._native_frame.loc[mask], validate_column_names=False
-        )
+        return self._from_native_frame(self._native_frame.loc[mask])
 
     def simple_select(self: Self, *column_names: str) -> Self:
         return self._from_native_frame(
@@ -156,13 +148,12 @@ def simple_select(self: Self, *column_names: str) -> Self:
                 self._backend_version,
                 self._implementation,
             ),
-            validate_column_names=False,
         )
 
     def aggregate(self: Self, *exprs: DaskExpr) -> Self:
         new_series = evaluate_exprs(self, *exprs)
-        df = dd.concat([val.rename(name) for name, val in new_series.items()], axis=1)
-        return self._from_native_frame(df, validate_column_names=False)
+        df = dd.concat([val.rename(name) for name, val in new_series], axis=1)
+        return self._from_native_frame(df)
 
     def select(self: Self, *exprs: DaskExpr) -> Self:
         new_series = evaluate_exprs(self, *exprs)
@@ -173,22 +164,19 @@ def select(self: Self, *exprs: DaskExpr) -> Self:
                 dd.from_pandas(
                     pd.DataFrame(), npartitions=self._native_frame.npartitions
                 ),
-                validate_column_names=False,
             )
 
         df = select_columns_by_name(
-            self._native_frame.assign(**new_series),
-            list(new_series.keys()),
+            self._native_frame.assign(**dict(new_series)),
+            [s[0] for s in new_series],
             self._backend_version,
             self._implementation,
         )
-        return self._from_native_frame(df, validate_column_names=False)
+        return self._from_native_frame(df)
 
     def drop_nulls(self: Self, subset: list[str] | None) -> Self:
         if subset is None:
-            return self._from_native_frame(
-                self._native_frame.dropna(), validate_column_names=False
-            )
+            return self._from_native_frame(self._native_frame.dropna())
         plx = self.__narwhals_namespace__()
         return self.filter(~plx.any_horizontal(plx.col(*subset).is_null()))
 
@@ -210,9 +198,7 @@ def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
             compliant_frame=self, columns=columns, strict=strict
         )
 
-        return self._from_native_frame(
-            self._native_frame.drop(columns=to_drop), validate_column_names=False
-        )
+        return self._from_native_frame(self._native_frame.drop(columns=to_drop))
 
     def with_row_index(self: Self, name: str) -> Self:
         # Implementation is based on the following StackOverflow reply:
@@ -228,8 +214,7 @@ def rename(self: Self, mapping: dict[str, str]) -> Self:
 
     def head(self: Self, n: int) -> Self:
         return self._from_native_frame(
-            self._native_frame.head(n=n, compute=False, npartitions=-1),
-            validate_column_names=False,
+            self._native_frame.head(n=n, compute=False, npartitions=-1)
         )
 
     def unique(
@@ -250,7 +235,7 @@ def unique(
         else:
             mapped_keep = {"any": "first"}.get(keep, keep)
             result = native_frame.drop_duplicates(subset=subset, keep=mapped_keep)
-        return self._from_native_frame(result, validate_column_names=False)
+        return self._from_native_frame(result)
 
     def sort(
         self: Self,
@@ -265,8 +250,7 @@ def sort(
             ascending = [not d for d in descending]
         na_position = "last" if nulls_last else "first"
         return self._from_native_frame(
-            df.sort_values(list(by), ascending=ascending, na_position=na_position),
-            validate_column_names=False,
+            df.sort_values(list(by), ascending=ascending, na_position=na_position)
         )
 
     def join(
@@ -413,9 +397,7 @@ def tail(self: Self, n: int) -> Self:  # pragma: no cover
         n_partitions = native_frame.npartitions
 
         if n_partitions == 1:
-            return self._from_native_frame(
-                self._native_frame.tail(n=n, compute=False), validate_column_names=False
-            )
+            return self._from_native_frame(self._native_frame.tail(n=n, compute=False))
         else:
             msg = "`LazyFrame.tail` is not supported for Dask backend with multiple partitions."
             raise NotImplementedError(msg)

narwhals/_dask/group_by.py

@@ -99,7 +99,6 @@ def _from_native_frame(self: Self, df: dd.DataFrame) -> DaskLazyFrame:
             df,
             backend_version=self._df._backend_version,
             version=self._df._version,
-            validate_column_names=True,
         )
 
 

narwhals/_dask/namespace.py

@@ -192,7 +192,6 @@ def concat(
                 dd.concat(dfs, axis=0, join="inner"),
                 backend_version=self._backend_version,
                 version=self._version,
-                validate_column_names=True,
             )
         if how == "horizontal":
             all_column_names: list[str] = [
@@ -211,14 +210,12 @@ def concat(
                 dd.concat(dfs, axis=1, join="outer"),
                 backend_version=self._backend_version,
                 version=self._version,
-                validate_column_names=True,
             )
         if how == "diagonal":
             return DaskLazyFrame(
                 dd.concat(dfs, axis=0, join="outer"),
                 backend_version=self._backend_version,
                 version=self._version,
-                validate_column_names=True,
             )
 
         raise NotImplementedError

narwhals/_dask/utils.py

@@ -39,20 +39,15 @@ def maybe_evaluate_expr(df: DaskLazyFrame, obj: DaskExpr | object) -> dx.Series
     return obj
 
 
-def evaluate_exprs(df: DaskLazyFrame, /, *exprs: DaskExpr) -> dict[str, dx.Series]:
-    native_results: dict[str, dx.Series] = {}
+def evaluate_exprs(df: DaskLazyFrame, /, *exprs: DaskExpr) -> list[tuple[str, dx.Series]]:
+    native_results: list[tuple[str, dx.Series]] = []
     for expr in exprs:
         native_series_list = expr._call(df)
         _, aliases = evaluate_output_names_and_aliases(expr, df, [])
         if len(aliases) != len(native_series_list):  # pragma: no cover
             msg = f"Internal error: got aliases {aliases}, but only got {len(native_series_list)} results"
             raise AssertionError(msg)
-        native_results.update(
-            {
-                alias: native_series
-                for native_series, alias in zip(native_series_list, aliases)
-            }
-        )
+        native_results.extend(zip(aliases, native_series_list))
     return native_results