Repository: narwhals-dev/narwhals
Pull request #1389 (2024-11-16 14:39:37)

feat: consistently raise `ColumnNotFoundError` for missing columns in `select` and `drop`
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [x] 🐳 Other

## Related issues

- Related issue #1373

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below



Patch:
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 39e39c13..6d5adc6a 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -6,6 +6,7 @@ from typing import Callable
 from typing import Literal
 from typing import Sequence
 
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals._expression_parsing import reuse_series_namespace_implementation
 from narwhals.dependencies import get_numpy
@@ -64,15 +65,22 @@ class ArrowExpr:
         from narwhals._arrow.series import ArrowSeries
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
-            return [
-                ArrowSeries(
-                    df._native_frame[column_name],
-                    name=column_name,
-                    backend_version=df._backend_version,
-                    dtypes=df._dtypes,
-                )
-                for column_name in column_names
-            ]
+            try:
+                return [
+                    ArrowSeries(
+                        df._native_frame[column_name],
+                        name=column_name,
+                        backend_version=df._backend_version,
+                        dtypes=df._dtypes,
+                    )
+                    for column_name in column_names
+                ]
+            except KeyError as e:
+                missing_columns = [x for x in column_names if x not in df.columns]
+                raise ColumnNotFoundError.from_missing_and_available_column_names(
+                    missing_columns=missing_columns,
+                    available_columns=df.columns,
+                ) from e
 
         return cls(
             func,
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index c2ff8592..2b50a962 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -11,6 +11,7 @@ from typing import Sequence
 from narwhals._dask.utils import add_row_index
 from narwhals._dask.utils import maybe_evaluate
 from narwhals._dask.utils import narwhals_to_native_dtype
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals._pandas_like.utils import calculate_timestamp_date
 from narwhals._pandas_like.utils import calculate_timestamp_datetime
 from narwhals._pandas_like.utils import native_to_narwhals_dtype
@@ -67,7 +68,14 @@ class DaskExpr:
         dtypes: DTypes,
     ) -> Self:
         def func(df: DaskLazyFrame) -> list[dask_expr.Series]:
-            return [df._native_frame[column_name] for column_name in column_names]
+            try:
+                return [df._native_frame[column_name] for column_name in column_names]
+            except KeyError as e:
+                missing_columns = [x for x in column_names if x not in df.columns]
+                raise ColumnNotFoundError.from_missing_and_available_column_names(
+                    missing_columns=missing_columns,
+                    available_columns=df.columns,
+                ) from e
 
         return cls(
             func,
diff --git a/narwhals/_exceptions.py b/narwhals/_exceptions.py
index a3d6aa72..ee5ef655 100644
--- a/narwhals/_exceptions.py
+++ b/narwhals/_exceptions.py
@@ -1,7 +1,36 @@
 from __future__ import annotations
 
 
-class ColumnNotFoundError(Exception): ...
+class FormattedKeyError(KeyError):
+    """KeyError with formatted error message.
+
+    Python's `KeyError` has special casing around formatting
+    (see https://bugs.python.org/issue2651). Use this class when the error
+    message has newlines and other special format characters.
+    Needed by https://github.com/tensorflow/tensorflow/issues/36857.
+    """
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class ColumnNotFoundError(FormattedKeyError):
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(self.message)
+
+    @classmethod
+    def from_missing_and_available_column_names(
+        cls, missing_columns: list[str], available_columns: list[str]
+    ) -> ColumnNotFoundError:
+        message = (
+            f"The following columns were not found: {missing_columns}"
+            f"\n\nHint: Did you mean one of these columns: {available_columns}?"
+        )
+        return ColumnNotFoundError(message)
 
 
 class InvalidOperationError(Exception): ...
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 98f5352d..323517db 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -6,6 +6,7 @@ from typing import Callable
 from typing import Literal
 from typing import Sequence
 
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals._expression_parsing import reuse_series_namespace_implementation
 from narwhals._pandas_like.series import PandasLikeSeries
@@ -72,15 +73,22 @@ class PandasLikeExpr:
         dtypes: DTypes,
     ) -> Self:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
-            return [
-                PandasLikeSeries(
-                    df._native_frame[column_name],
-                    implementation=df._implementation,
-                    backend_version=df._backend_version,
-                    dtypes=df._dtypes,
-                )
-                for column_name in column_names
-            ]
+            try:
+                return [
+                    PandasLikeSeries(
+                        df._native_frame[column_name],
+                        implementation=df._implementation,
+                        backend_version=df._backend_version,
+                        dtypes=df._dtypes,
+                    )
+                    for column_name in column_names
+                ]
+            except KeyError as e:
+                missing_columns = [x for x in column_names if x not in df.columns]
+                raise ColumnNotFoundError.from_missing_and_available_column_names(
+                    missing_columns=missing_columns,
+                    available_columns=df.columns,
+                ) from e
 
         return cls(
             func,
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 2b31ce9c..c718850c 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -11,6 +11,7 @@ from typing import TypeVar
 from narwhals._arrow.utils import (
     native_to_narwhals_dtype as arrow_native_to_narwhals_dtype,
 )
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals.dependencies import get_polars
 from narwhals.utils import Implementation
 from narwhals.utils import isinstance_or_issubclass
@@ -659,5 +660,18 @@ 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
+            )
         return df.loc[:, column_names]  # type: ignore[no-any-return, attr-defined]
-    return df[column_names]  # type: ignore[no-any-return, index]
+    try:
+        return df[column_names]  # type: ignore[no-any-return, 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
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
index fffe5092..0f55dea6 100644
--- a/narwhals/_polars/dataframe.py
+++ b/narwhals/_polars/dataframe.py
@@ -5,6 +5,7 @@ from typing import Any
 from typing import Literal
 from typing import Sequence
 
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals._polars.namespace import PolarsNamespace
 from narwhals._polars.utils import convert_str_slice_to_int_slice
 from narwhals._polars.utils import extract_args_kwargs
@@ -79,10 +80,17 @@ class PolarsDataFrame:
             }
 
         def func(*args: Any, **kwargs: Any) -> Any:
+            import polars as pl  # ignore-banned-import()
+
             args, kwargs = extract_args_kwargs(args, kwargs)  # type: ignore[assignment]
-            return self._from_native_object(
-                getattr(self._native_frame, attr)(*args, **kwargs)
-            )
+            try:
+                return self._from_native_object(
+                    getattr(self._native_frame, attr)(*args, **kwargs)
+                )
+            except pl.exceptions.ColumnNotFoundError as e:
+                msg = str(e)
+                msg += f"\n\nHint: Did you mean one of these columns: {self.columns}?"
+                raise ColumnNotFoundError(str(e)) from e
 
         return func
 
@@ -219,12 +227,10 @@ class PolarsDataFrame:
         return self._from_native_frame(self._native_frame.with_row_index(name))
 
     def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
-        if self._backend_version < (1, 0, 0):
-            to_drop = parse_columns_to_drop(
-                compliant_frame=self, columns=columns, strict=strict
-            )
-            return self._from_native_frame(self._native_frame.drop(to_drop))
-        return self._from_native_frame(self._native_frame.drop(columns, strict=strict))
+        to_drop = parse_columns_to_drop(
+            compliant_frame=self, columns=columns, strict=strict
+        )
+        return self._from_native_frame(self._native_frame.drop(to_drop))
 
     def unpivot(
         self: Self,
@@ -341,8 +347,15 @@ class PolarsLazyFrame:
             }
 
     def collect(self) -> PolarsDataFrame:
+        import polars as pl  # ignore-banned-import
+
+        try:
+            result = self._native_frame.collect()
+        except pl.exceptions.ColumnNotFoundError as e:
+            raise ColumnNotFoundError(str(e)) from e
+
         return PolarsDataFrame(
-            self._native_frame.collect(),
+            result,
             backend_version=self._backend_version,
             dtypes=self._dtypes,
         )
diff --git a/narwhals/utils.py b/narwhals/utils.py
index 43ad11a1..b3127b7d 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -638,16 +638,16 @@ def parse_columns_to_drop(
     columns: Iterable[str],
     strict: bool,  # noqa: FBT001
 ) -> list[str]:
-    cols = set(compliant_frame.columns)
+    cols = compliant_frame.columns
     to_drop = list(columns)
-
     if strict:
-        for d in to_drop:
-            if d not in cols:
-                msg = f'"{d}" not found'
-                raise ColumnNotFoundError(msg)
+        missing_columns = [x for x in to_drop if x not in cols]
+        if missing_columns:
+            raise ColumnNotFoundError.from_missing_and_available_column_names(
+                missing_columns=missing_columns, available_columns=cols
+            )
     else:
-        to_drop = list(cols.intersection(set(to_drop)))
+        to_drop = list(set(cols).intersection(set(to_drop)))
     return to_drop
 
 


Test patch:
diff --git a/tests/frame/select_test.py b/tests/frame/select_test.py
index 93f28ace..222db8a9 100644
--- a/tests/frame/select_test.py
+++ b/tests/frame/select_test.py
@@ -5,7 +5,9 @@ import pyarrow as pa
 import pytest
 
 import narwhals.stable.v1 as nw
+from narwhals._exceptions import ColumnNotFoundError
 from tests.utils import PANDAS_VERSION
+from tests.utils import POLARS_VERSION
 from tests.utils import Constructor
 from tests.utils import assert_equal_data
 
@@ -53,3 +55,43 @@ def test_comparison_with_list_error_message() -> None:
         nw.from_native(pa.chunked_array([[1, 2, 3]]), series_only=True) == [1, 2, 3]  # noqa: B015
     with pytest.raises(ValueError, match=msg):
         nw.from_native(pd.Series([[1, 2, 3]]), series_only=True) == [1, 2, 3]  # noqa: B015
+
+
+def test_missing_columns(constructor: Constructor) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+    df = nw.from_native(constructor(data))
+    selected_columns = ["a", "e", "f"]
+    msg = (
+        r"The following columns were not found: \[.*\]"
+        r"\n\nHint: Did you mean one of these columns: \['a', 'b', 'z'\]?"
+    )
+    if "polars" in str(constructor):
+        # In the lazy case, Polars only errors when we call `collect`,
+        # and we have no way to recover exactly which columns the user
+        # tried selecting. So, we just emit their message (which varies
+        # across versions...)
+        msg = "e|f"
+        if isinstance(df, nw.LazyFrame):
+            with pytest.raises(ColumnNotFoundError, match=msg):
+                df.select(selected_columns).collect()
+        else:
+            with pytest.raises(ColumnNotFoundError, match=msg):
+                df.select(selected_columns)
+        if POLARS_VERSION >= (1,):
+            # Old Polars versions wouldn't raise an error
+            # at all here
+            if isinstance(df, nw.LazyFrame):
+                with pytest.raises(ColumnNotFoundError, match=msg):
+                    df.drop(selected_columns, strict=True).collect()
+            else:
+                with pytest.raises(ColumnNotFoundError, match=msg):
+                    df.drop(selected_columns, strict=True)
+        else:  # pragma: no cover
+            pass
+    else:
+        with pytest.raises(ColumnNotFoundError, match=msg):
+            df.select(selected_columns)
+        with pytest.raises(ColumnNotFoundError, match=msg):
+            df.drop(selected_columns, strict=True)
+        with pytest.raises(ColumnNotFoundError, match=msg):
+            df.select(nw.col("fdfa"))