Repository: narwhals-dev/narwhals
Pull request #773 (2024-08-10 21:03:00)

feat: extend dataframe `drop` method
<!--
# 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)

- [x] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #742 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.



FBruzzesi: Aside: I noticed that polars allows selectors other than column names. I think it should not be too complex to integrate. 
Thoughts?
MarcoGorelli: > Aside: I noticed that polars allows selectors other than column names. I think it should not be too complex to integrate.
Thoughts?

if it's not too complex - sure, I do like these 👍 
FBruzzesi: > > Aside: I noticed that polars allows selectors other than column names. I think it should not be too complex to integrate.
> > Thoughts?
> 
> if it's not too complex - sure, I do like these 👍

Happy to keep it as a follow up honestly. Other methods would benefit from that as well and it could be worth to factor out some sort of "parser" for selectors and column names

Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 7497c21f..880e6d95 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -19,6 +19,7 @@ from narwhals.dependencies import get_pyarrow_parquet
 from narwhals.utils import Implementation
 from narwhals.utils import flatten
 from narwhals.utils import generate_unique_token
+from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -264,14 +265,16 @@ class ArrowDataFrame:
             )
 
             return self._from_native_frame(
-                self.with_columns(**{key_token: plx.lit(0, None)})._native_frame.join(
+                self.with_columns(**{key_token: plx.lit(0, None)})
+                ._native_frame.join(
                     other.with_columns(**{key_token: plx.lit(0, None)})._native_frame,
                     keys=key_token,
                     right_keys=key_token,
                     join_type="inner",
                     right_suffix="_right",
-                ),
-            ).drop(key_token)
+                )
+                .drop([key_token]),
+            )
 
         return self._from_native_frame(
             self._native_frame.join(
@@ -283,8 +286,11 @@ class ArrowDataFrame:
             ),
         )
 
-    def drop(self, *columns: str) -> Self:
-        return self._from_native_frame(self._native_frame.drop(list(columns)))
+    def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
+        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 drop_nulls(self: Self, subset: str | list[str] | None) -> Self:
         if subset is None:
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index 33f3c8b8..cf3c6cc1 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -14,6 +14,7 @@ from narwhals.dependencies import get_pandas
 from narwhals.utils import Implementation
 from narwhals.utils import flatten
 from narwhals.utils import generate_unique_token
+from narwhals.utils import parse_columns_to_drop
 from narwhals.utils import parse_version
 
 if TYPE_CHECKING:
@@ -129,8 +130,12 @@ class DaskLazyFrame:
     def collect_schema(self) -> dict[str, DType]:
         return self.schema
 
-    def drop(self: Self, *columns: str) -> Self:
-        return self._from_native_frame(self._native_frame.drop(columns=list(columns)))
+    def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
+        to_drop = parse_columns_to_drop(
+            compliant_frame=self, columns=columns, strict=strict
+        )
+
+        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:
@@ -204,14 +209,16 @@ class DaskLazyFrame:
             )
 
             return self._from_native_frame(
-                self._native_frame.assign(**{key_token: 0}).merge(
+                self._native_frame.assign(**{key_token: 0})
+                .merge(
                     other._native_frame.assign(**{key_token: 0}),
                     how="inner",
                     left_on=key_token,
                     right_on=key_token,
                     suffixes=("", "_right"),
-                ),
-            ).drop(key_token)
+                )
+                .drop(columns=key_token),
+            )
 
         if how == "anti":
             indicator_token = generate_unique_token(
@@ -297,5 +304,5 @@ class DaskLazyFrame:
                 pln.col(row_index_token) >= offset,  # type: ignore[operator]
                 (pln.col(row_index_token) - offset) % n == 0,  # type: ignore[arg-type]
             )
-            .drop(row_index_token)
+            .drop([row_index_token], strict=False)
         )
diff --git a/narwhals/_exceptions.py b/narwhals/_exceptions.py
new file mode 100644
index 00000000..18995451
--- /dev/null
+++ b/narwhals/_exceptions.py
@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+
+class ColumnNotFoundError(Exception): ...
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index f6fb2489..a8729eee 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -24,6 +24,7 @@ from narwhals.dependencies import get_pyarrow
 from narwhals.utils import Implementation
 from narwhals.utils import flatten
 from narwhals.utils import generate_unique_token
+from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -321,8 +322,11 @@ class PandasLikeDataFrame:
     def rename(self, mapping: dict[str, str]) -> Self:
         return self._from_native_frame(self._native_frame.rename(columns=mapping))
 
-    def drop(self, *columns: str) -> Self:
-        return self._from_native_frame(self._native_frame.drop(columns=list(columns)))
+    def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
+        to_drop = parse_columns_to_drop(
+            compliant_frame=self, columns=columns, strict=strict
+        )
+        return self._from_native_frame(self._native_frame.drop(columns=to_drop))
 
     # --- transform ---
     def sort(
@@ -382,14 +386,16 @@ class PandasLikeDataFrame:
                 )
 
                 return self._from_native_frame(
-                    self._native_frame.assign(**{key_token: 0}).merge(
+                    self._native_frame.assign(**{key_token: 0})
+                    .merge(
                         other._native_frame.assign(**{key_token: 0}),
                         how="inner",
                         left_on=key_token,
                         right_on=key_token,
                         suffixes=("", "_right"),
-                    ),
-                ).drop(key_token)
+                    )
+                    .drop(columns=key_token),
+                )
             else:
                 return self._from_native_frame(
                     self._native_frame.merge(
@@ -420,7 +426,7 @@ class PandasLikeDataFrame:
                     right_on=left_on,
                 )
                 .loc[lambda t: t[indicator_token] == "left_only"]
-                .drop(columns=[indicator_token])
+                .drop(columns=indicator_token)
             )
 
         if how == "semi":
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
index afef8187..2137f891 100644
--- a/narwhals/_polars/dataframe.py
+++ b/narwhals/_polars/dataframe.py
@@ -8,6 +8,7 @@ from narwhals._polars.utils import extract_args_kwargs
 from narwhals._polars.utils import translate_dtype
 from narwhals.dependencies import get_polars
 from narwhals.utils import Implementation
+from narwhals.utils import parse_columns_to_drop
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -124,6 +125,14 @@ class PolarsDataFrame:
             return self._from_native_frame(self._native_frame.with_row_count(name))
         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):  # pragma: no cover
+            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))
+
 
 class PolarsLazyFrame:
     def __init__(self, df: Any, *, backend_version: tuple[int, ...]) -> None:
@@ -184,3 +193,8 @@ class PolarsLazyFrame:
         if self._backend_version < (0, 20, 4):  # pragma: no cover
             return self._from_native_frame(self._native_frame.with_row_count(name))
         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):  # pragma: no cover
+            return self._from_native_frame(self._native_frame.drop(columns))
+        return self._from_native_frame(self._native_frame.drop(columns, strict=strict))
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 487767c3..bd6ef115 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -137,9 +137,9 @@ class BaseFrame(Generic[FrameT]):
     def tail(self, n: int) -> Self:
         return self._from_compliant_dataframe(self._compliant_frame.tail(n))
 
-    def drop(self, *columns: str | Iterable[str]) -> Self:
+    def drop(self, *columns: Iterable[str], strict: bool) -> Self:
         return self._from_compliant_dataframe(
-            self._compliant_frame.drop(*flatten(columns))
+            self._compliant_frame.drop(columns, strict=strict)
         )
 
     def unique(
@@ -1286,12 +1286,14 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().tail(n)
 
-    def drop(self, *columns: str | Iterable[str]) -> Self:
+    def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self:
         """
         Remove columns from the dataframe.
 
         Arguments:
             *columns: Names of the columns that should be removed from the dataframe.
+            strict: Validate that all column names exist in the schema and throw an
+                exception if a column name does not exist in the schema.
 
         Examples:
             >>> import pandas as pd
@@ -1349,7 +1351,7 @@ class DataFrame(BaseFrame[FrameT]):
             │ 8.0 │
             └─────┘
         """
-        return super().drop(*columns)
+        return super().drop(*flatten(columns), strict=strict)
 
     def unique(
         self,
@@ -2743,13 +2745,19 @@ class LazyFrame(BaseFrame[FrameT]):
         """
         return super().tail(n)
 
-    def drop(self, *columns: str | Iterable[str]) -> Self:
+    def drop(self, *columns: str | Iterable[str], strict: bool = True) -> Self:
         r"""
         Remove columns from the LazyFrame.
 
         Arguments:
-            *columns: Names of the columns that should be removed from the
-                      dataframe. Accepts column selector input.
+            *columns: Names of the columns that should be removed from the dataframe.
+            strict: Validate that all column names exist in the schema and throw an
+                exception if a column name does not exist in the schema.
+
+        Warning:
+            `strict` argument is ignored for `polars<1.0.0`.
+
+            Please consider upgrading to a newer version or pass to eager mode.
 
         Examples:
             >>> import pandas as pd
@@ -2807,7 +2815,7 @@ class LazyFrame(BaseFrame[FrameT]):
             │ 8.0 │
             └─────┘
         """
-        return super().drop(*flatten(columns))
+        return super().drop(*flatten(columns), strict=strict)
 
     def unique(
         self,
diff --git a/narwhals/utils.py b/narwhals/utils.py
index 78299ba4..512099cc 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -12,6 +12,7 @@ from typing import TypeVar
 from typing import cast
 
 from narwhals import dtypes
+from narwhals._exceptions import ColumnNotFoundError
 from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
 from narwhals.dependencies import get_pandas
@@ -395,3 +396,21 @@ def generate_unique_token(n_bytes: int, columns: list[str]) -> str:  # pragma: n
                 "join operation"
             )
             raise AssertionError(msg)
+
+
+def parse_columns_to_drop(
+    compliant_frame: Any,
+    columns: Iterable[str],
+    strict: bool,  # noqa: FBT001
+) -> list[str]:
+    cols = set(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)
+    else:
+        to_drop = list(cols.intersection(set(to_drop)))
+    return to_drop


Test patch:
diff --git a/tests/frame/drop_test.py b/tests/frame/drop_test.py
index f22d148b..547ddc74 100644
--- a/tests/frame/drop_test.py
+++ b/tests/frame/drop_test.py
@@ -1,21 +1,58 @@
 from __future__ import annotations
 
+from contextlib import nullcontext as does_not_raise
 from typing import Any
 
+import polars as pl
 import pytest
+from polars.exceptions import ColumnNotFoundError as PlColumnNotFoundError
 
 import narwhals.stable.v1 as nw
+from narwhals._exceptions import ColumnNotFoundError
+from narwhals.utils import parse_version
 
 
 @pytest.mark.parametrize(
-    ("drop", "left"),
+    ("to_drop", "expected"),
     [
-        (["a"], ["b", "z"]),
-        (["a", "b"], ["z"]),
+        ("abc", ["b", "z"]),
+        (["abc"], ["b", "z"]),
+        (["abc", "b"], ["z"]),
     ],
 )
-def test_drop(constructor: Any, drop: list[str], left: list[str]) -> None:
-    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+def test_drop(constructor: Any, to_drop: list[str], expected: list[str]) -> None:
+    data = {"abc": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
     df = nw.from_native(constructor(data))
-    assert df.drop(drop).collect_schema().names() == left
-    assert df.drop(*drop).collect_schema().names() == left
+    assert df.drop(to_drop).collect_schema().names() == expected
+    if not isinstance(to_drop, str):
+        assert df.drop(*to_drop).collect_schema().names() == expected
+
+
+@pytest.mark.parametrize(
+    ("strict", "context"),
+    [
+        (
+            True,
+            pytest.raises(
+                (ColumnNotFoundError, PlColumnNotFoundError), match='"z" not found'
+            ),
+        ),
+        (False, does_not_raise()),
+    ],
+)
+def test_drop_strict(request: Any, constructor: Any, strict: bool, context: Any) -> None:  # noqa: FBT001
+    if (
+        "polars_lazy" in str(request)
+        and parse_version(pl.__version__) < (1, 0, 0)
+        and strict
+    ):
+        request.applymarker(pytest.mark.xfail)
+
+    data = {"a": [1, 3, 2], "b": [4, 4, 6]}
+    to_drop = ["a", "z"]
+
+    df = nw.from_native(constructor(data))
+
+    with context:
+        names_out = df.drop(to_drop, strict=strict).collect_schema().names()
+        assert names_out == ["b"]