chore: always use check_columns_exist where possible (#2495)

EdAbati · dangotbanned · web-flow · commit c952d5868a5f · 2025-05-26T12:49:10.000+01:00
Co-authored-by: dangotbanned &lt;125183946+dangotbanned@users.noreply.github.com&gt;
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
@@ -25,7 +25,6 @@
 from narwhals.utils import (
     Implementation,
     Version,
-    check_column_exists,
     check_column_names_are_unique,
     convert_str_slice_to_int_slice,
     generate_temporary_column_name,
@@ -440,9 +439,7 @@ def join(
     join_asof = not_implemented()
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        to_drop = parse_columns_to_drop(
-            compliant_frame=self, columns=columns, strict=strict
-        )
+        to_drop = parse_columns_to_drop(self, columns, strict=strict)
         return self._with_native(self.native.drop(to_drop), validate_column_names=False)
 
     def drop_nulls(self: ArrowDataFrame, subset: Sequence[str] | None) -> ArrowDataFrame:
@@ -693,7 +690,8 @@ def unique(
         # and has no effect on the output.
         import numpy as np  # ignore-banned-import
 
-        check_column_exists(self.columns, subset)
+        if subset and (error := self._check_columns_exist(subset)):
+            raise error
         subset = list(subset or self.columns)
 
         if keep in {"any", "first", "last"}:
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
@@ -7,7 +7,6 @@
 from narwhals._arrow.series import ArrowSeries
 from narwhals._compliant import EagerExpr
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
-from narwhals.exceptions import ColumnNotFoundError
 from narwhals.utils import Implementation, generate_temporary_column_name, not_implemented
 
 if TYPE_CHECKING:
@@ -69,12 +68,9 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
                     for column_name in evaluate_column_names(df)
                 ]
             except KeyError as e:
-                missing_columns = [
-                    x for x in evaluate_column_names(df) if x not in df.columns
-                ]
-                raise ColumnNotFoundError.from_missing_and_available_column_names(
-                    missing_columns=missing_columns, available_columns=df.columns
-                ) from e
+                if error := df._check_columns_exist(evaluate_column_names(df)):
+                    raise error from e
+                raise
 
         return cls(
             func,
diff --git a/narwhals/_compliant/dataframe.py b/narwhals/_compliant/dataframe.py
@@ -36,6 +36,7 @@
 from narwhals.utils import (
     Version,
     _StoresNative,
+    check_columns_exist,
     is_compliant_series,
     is_index_selector,
     is_range,
@@ -59,6 +60,7 @@
     from narwhals._translate import IntoArrowTable
     from narwhals.dataframe import DataFrame
     from narwhals.dtypes import DType
+    from narwhals.exceptions import ColumnNotFoundError
     from narwhals.schema import Schema
     from narwhals.typing import (
         AsofJoinStrategy,
@@ -270,6 +272,9 @@ def _evaluate_aliases(self, *exprs: CompliantExprT_contra) -> list[str]:
         it = (expr._evaluate_aliases(self) for expr in exprs)
         return list(chain.from_iterable(it))
 
+    def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:
+        return check_columns_exist(subset, available=self.columns)
+
 
 class CompliantLazyFrame(
     _StoresNative[NativeFrameT],
@@ -377,6 +382,9 @@ def _evaluate_aliases(self, *exprs: CompliantExprT_contra) -> list[str]:
         it = (expr._evaluate_aliases(self) for expr in exprs)
         return list(chain.from_iterable(it))
 
+    def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:
+        return check_columns_exist(subset, available=self.columns)
+
 
 class EagerDataFrame(
     CompliantDataFrame[EagerSeriesT, EagerExprT, NativeFrameT, "DataFrame[NativeFrameT]"],
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
@@ -11,7 +11,6 @@
 from narwhals.utils import (
     Implementation,
     _remap_full_join_keys,
-    check_column_exists,
     check_column_names_are_unique,
     generate_temporary_column_name,
     not_implemented,
@@ -200,9 +199,7 @@ def collect_schema(self) -> dict[str, DType]:
         return self.schema
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        to_drop = parse_columns_to_drop(
-            compliant_frame=self, columns=columns, strict=strict
-        )
+        to_drop = parse_columns_to_drop(self, columns, strict=strict)
 
         return self._with_native(self.native.drop(columns=to_drop))
 
@@ -222,7 +219,8 @@ def head(self, n: int) -> Self:
     def unique(
         self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
     ) -> Self:
-        check_column_exists(self.columns, subset)
+        if subset and (error := self._check_columns_exist(subset)):
+            raise error
         if keep == "none":
             subset = subset or self.columns
             token = generate_temporary_column_name(n_bytes=8, columns=subset)
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
@@ -14,7 +14,7 @@
 )
 from narwhals._expression_parsing import ExprKind, evaluate_output_names_and_aliases
 from narwhals._pandas_like.utils import native_to_narwhals_dtype
-from narwhals.exceptions import ColumnNotFoundError, InvalidOperationError
+from narwhals.exceptions import InvalidOperationError
 from narwhals.utils import Implementation, generate_temporary_column_name, not_implemented
 
 if TYPE_CHECKING:
@@ -106,12 +106,9 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
                     for column_name in evaluate_column_names(df)
                 ]
             except KeyError as e:
-                missing_columns = [
-                    x for x in evaluate_column_names(df) if x not in df.columns
-                ]
-                raise ColumnNotFoundError.from_missing_and_available_column_names(
-                    missing_columns=missing_columns, available_columns=df.columns
-                ) from e
+                if error := df._check_columns_exist(evaluate_column_names(df)):
+                    raise error from e
+                raise
 
         return cls(
             func,
diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py
@@ -17,7 +17,7 @@
     native_to_narwhals_dtype,
 )
 from narwhals.dependencies import get_duckdb
-from narwhals.exceptions import ColumnNotFoundError, InvalidOperationError
+from narwhals.exceptions import InvalidOperationError
 from narwhals.typing import CompliantLazyFrame
 from narwhals.utils import (
     Implementation,
@@ -182,7 +182,7 @@ def select(self, *exprs: DuckDBExpr) -> Self:
         return self._with_native(self.native.select(*selection))
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        columns_to_drop = parse_columns_to_drop(self, columns=columns, strict=strict)
+        columns_to_drop = parse_columns_to_drop(self, columns, strict=strict)
         selection = (name for name in self.columns if name not in columns_to_drop)
         return self._with_native(self.native.select(*selection))
 
@@ -387,9 +387,8 @@ def unique(
                 )
                 raise NotImplementedError(msg)
             # Sanitise input
-            if any(x not in self.columns for x in subset_):
-                msg = f"Columns {set(subset_).difference(self.columns)} not found in {self.columns}."
-                raise ColumnNotFoundError(msg)
+            if error := self._check_columns_exist(subset_):
+                raise error
             idx_name = generate_temporary_column_name(8, self.columns)
             count_name = generate_temporary_column_name(8, [*self.columns, idx_name])
             partition_by_sql = generate_partition_by_sql(*(subset_))
diff --git a/narwhals/_ibis/dataframe.py b/narwhals/_ibis/dataframe.py
@@ -176,7 +176,7 @@ def select(self, *exprs: IbisExpr) -> Self:
         return self._with_native(t)
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        columns_to_drop = parse_columns_to_drop(self, columns=columns, strict=strict)
+        columns_to_drop = parse_columns_to_drop(self, columns, strict=strict)
         selection = (col for col in self.columns if col not in columns_to_drop)
         return self._with_native(self.native.select(*selection))
 
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
@@ -35,7 +35,6 @@
     Implementation,
     _into_arrow_table,
     _remap_full_join_keys,
-    check_column_exists,
     exclude_column_names,
     generate_temporary_column_name,
     parse_columns_to_drop,
@@ -485,9 +484,7 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
         )
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        to_drop = parse_columns_to_drop(
-            compliant_frame=self, columns=columns, strict=strict
-        )
+        to_drop = parse_columns_to_drop(self, columns, strict=strict)
         return self._with_native(
             self.native.drop(columns=to_drop), validate_column_names=False
         )
@@ -753,7 +750,8 @@ def unique(
         # The param `maintain_order` is only here for compatibility with the Polars API
         # and has no effect on the output.
         mapped_keep = {"none": False, "any": "first"}.get(keep, keep)
-        check_column_exists(self.columns, subset)
+        if subset and (error := self._check_columns_exist(subset)):
+            raise error
         return self._with_native(
             self.native.drop_duplicates(subset=subset, keep=mapped_keep),
             validate_column_names=False,
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
@@ -6,7 +6,6 @@
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
 from narwhals._pandas_like.group_by import PandasLikeGroupBy
 from narwhals._pandas_like.series import PandasLikeSeries
-from narwhals.exceptions import ColumnNotFoundError
 from narwhals.utils import generate_temporary_column_name
 
 if TYPE_CHECKING:
@@ -138,12 +137,9 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                     for column_name in evaluate_column_names(df)
                 ]
             except KeyError as e:
-                missing_columns = [
-                    x for x in evaluate_column_names(df) if x not in df.columns
-                ]
-                raise ColumnNotFoundError.from_missing_and_available_column_names(
-                    missing_columns=missing_columns, available_columns=df.columns
-                ) from e
+                if error := df._check_columns_exist(evaluate_column_names(df)):
+                    raise error from e
+                raise
 
         return cls(
             func,
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
@@ -8,11 +8,12 @@
 import pandas as pd
 
 from narwhals._compliant.series import EagerSeriesNamespace
-from narwhals.exceptions import ColumnNotFoundError, DuplicateError, ShapeError
+from narwhals.exceptions import DuplicateError, ShapeError
 from narwhals.utils import (
     Implementation,
     Version,
     _DeferredIterable,
+    check_columns_exist,
     isinstance_or_issubclass,
 )
 
@@ -622,21 +623,21 @@ def select_columns_by_name(
     ):
         # See https://github.com/narwhals-dev/narwhals/issues/1349#issuecomment-2470118122
         # for why we need this
-        available_columns = df.columns.tolist()  # type: ignore[attr-defined]
-        missing_columns = [x for x in column_names if x not in available_columns]
-        if missing_columns:  # pragma: no cover
-            raise ColumnNotFoundError.from_missing_and_available_column_names(
-                missing_columns, available_columns
-            )
+        if error := check_columns_exist(
+            column_names,  # type: ignore[arg-type]
+            available=df.columns.tolist(),  # type: ignore[attr-defined]
+        ):
+            raise error
         return df.loc[:, column_names]  # type: ignore[attr-defined]
     try:
         return df[column_names]  # type: ignore[index]
     except KeyError as e:
-        available_columns = df.columns.tolist()  # type: ignore[attr-defined]
-        missing_columns = [x for x in column_names if x not in available_columns]
-        raise ColumnNotFoundError.from_missing_and_available_column_names(
-            missing_columns, available_columns
-        ) from e
+        if error := check_columns_exist(
+            column_names,  # type: ignore[arg-type]
+            available=df.columns.tolist(),  # type: ignore[attr-defined]
+        ):
+            raise error from e
+        raise
 
 
 def check_column_names_are_unique(columns: pd.Index[str]) -> None:
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
@@ -26,6 +26,7 @@
 from narwhals.utils import (
     Implementation,
     _into_arrow_table,
+    check_columns_exist,
     convert_str_slice_to_int_slice,
     is_compliant_series,
     is_index_selector,
@@ -454,9 +455,7 @@ def with_row_index(self, name: str) -> Self:
         return self._with_native(self.native.with_row_index(name))
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        to_drop = parse_columns_to_drop(
-            compliant_frame=self, columns=columns, strict=strict
-        )
+        to_drop = parse_columns_to_drop(self, columns, strict=strict)
         return self._with_native(self.native.drop(to_drop))
 
     def unpivot(
@@ -533,6 +532,9 @@ def join(
         except Exception as e:  # noqa: BLE001
             raise catch_polars_exception(e, self._backend_version) from None
 
+    def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:
+        return check_columns_exist(subset, available=self.columns)
+
 
 class PolarsLazyFrame:
     drop_nulls: Method[Self]
@@ -760,3 +762,8 @@ def join(
                 suffix=suffix,
             )
         )
+
+    def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:
+        return check_columns_exist(  # pragma: no cover
+            subset, available=self.columns
+        )
diff --git a/narwhals/_spark_like/dataframe.py b/narwhals/_spark_like/dataframe.py
@@ -17,7 +17,6 @@
 from narwhals.typing import CompliantLazyFrame
 from narwhals.utils import (
     Implementation,
-    check_column_exists,
     find_stacklevel,
     generate_temporary_column_name,
     not_implemented,
@@ -294,9 +293,7 @@ def collect_schema(self) -> dict[str, DType]:
         return self.schema
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
-        columns_to_drop = parse_columns_to_drop(
-            compliant_frame=self, columns=columns, strict=strict
-        )
+        columns_to_drop = parse_columns_to_drop(self, columns, strict=strict)
         return self._with_native(self.native.drop(*columns_to_drop))
 
     def head(self, n: int) -> Self:
@@ -344,7 +341,8 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
     def unique(
         self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
     ) -> Self:
-        check_column_exists(self.columns, subset)
+        if subset and (error := self._check_columns_exist(subset)):
+            raise error
         subset = list(subset) if subset else None
         if keep == "none":
             tmp = generate_temporary_column_name(8, self.columns)
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
@@ -25,7 +25,6 @@
 )
 from narwhals.dependencies import get_polars, is_numpy_array
 from narwhals.exceptions import (
-    ColumnNotFoundError,
     InvalidIntoExprError,
     LengthChangingExprError,
     OrderDependentExprError,
@@ -174,11 +173,9 @@ def select(
                 )
             except Exception as e:
                 # Column not found is the only thing that can realistically be raised here.
-                available_columns = self.columns
-                missing_columns = [x for x in flat_exprs if x not in available_columns]
-                raise ColumnNotFoundError.from_missing_and_available_column_names(
-                    missing_columns, available_columns
-                ) from e
+                if error := self._compliant_frame._check_columns_exist(flat_exprs):
+                    raise error from e
+                raise
         compliant_exprs, kinds = self._flatten_and_extract(*flat_exprs, **named_exprs)
         if compliant_exprs and all_exprs_are_scalar_like(*flat_exprs, **named_exprs):
             return self._with_compliant(self._compliant_frame.aggregate(*compliant_exprs))
diff --git a/narwhals/exceptions.py b/narwhals/exceptions.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Iterable, Sequence
+
 
 class NarwhalsError(ValueError):
     """Base class for all Narwhals exceptions."""
@@ -30,11 +32,11 @@ def __init__(self, message: str) -> None:
 
     @classmethod
     def from_missing_and_available_column_names(
-        cls: type, missing_columns: list[str], available_columns: list[str]
+        cls, missing_columns: Iterable[str], available_columns: Sequence[str], /
     ) -> ColumnNotFoundError:
         message = (
-            f"The following columns were not found: {missing_columns}"
-            f"\n\nHint: Did you mean one of these columns: {available_columns}?"
+            f"The following columns were not found: {sorted(missing_columns)}"
+            f"\n\nHint: Did you mean one of these columns: {list(available_columns)}?"
         )
         return ColumnNotFoundError(message)
 
diff --git a/narwhals/utils.py b/narwhals/utils.py
diff --git a/tests/utils_test.py b/tests/utils_test.py
