feat: Support drop(*columns: OneOrIterable[ColumnNameOrSelector])

This also simplifies the compliant-level, since resolving names is part of expansion
Heavily based on what `polars` does

Author: dangotbanned

narwhals/_plan/_parse.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+import operator
+from collections import deque
+from collections.abc import Collection, Iterable, Sequence
 
 # ruff: noqa: A002
+from functools import reduce
 from itertools import chain
 from typing import TYPE_CHECKING
 
@@ -14,6 +17,7 @@
     is_iterable_reject,
     is_selector,
 )
+from narwhals._plan.common import flatten_hash_safe
 from narwhals._plan.exceptions import invalid_into_expr_error, is_iterable_error
 from narwhals._utils import qualified_type_name
 from narwhals.dependencies import get_polars
@@ -27,6 +31,7 @@
 
     from narwhals._plan.expr import Expr
     from narwhals._plan.expressions import ExprIR, SelectorIR
+    from narwhals._plan.selectors import Selector
     from narwhals._plan.typing import (
         ColumnNameOrSelector,
         IntoExpr,
@@ -129,19 +134,55 @@ def parse_into_expr_ir(
     return expr._ir
 
 
-def parse_into_selector_ir(input: ColumnNameOrSelector | Expr, /) -> SelectorIR:
+def parse_into_selector_ir(
+    input: ColumnNameOrSelector | Expr, /, *, require_all: bool = True
+) -> SelectorIR:
+    return _parse_into_selector(input, require_all=require_all)._ir
+
+
+def _parse_into_selector(
+    input: ColumnNameOrSelector | Expr, /, *, require_all: bool = True
+) -> Selector:
     if is_selector(input):
         selector = input
     elif isinstance(input, str):
-        from narwhals._plan import selectors as cs
+        import narwhals._plan.selectors as cs
 
-        selector = cs.by_name(input)
+        selector = cs.by_name(input, require_all=require_all)
     elif is_expr(input):
         selector = input.meta.as_selector()
     else:
         msg = f"cannot turn {qualified_type_name(input)!r} into a selector"
         raise TypeError(msg)
-    return selector._ir
+    return selector
+
+
+def parse_into_combined_selector_ir(
+    *inputs: OneOrIterable[ColumnNameOrSelector], require_all: bool = True
+) -> SelectorIR:
+    import narwhals._plan.selectors as cs
+
+    flat = tuple(flatten_hash_safe(inputs))
+    selectors = deque["Selector"]()
+    if names := tuple(el for el in flat if isinstance(el, str)):
+        selector = cs.by_name(names, require_all=require_all)
+        if len(names) == len(flat):
+            return selector._ir
+        selectors.append(selector)
+    selectors.extend(_parse_into_selector(el) for el in flat if not isinstance(el, str))
+    return _any_of(selectors)._ir
+
+
+def _any_of(selectors: Iterable[Selector], /) -> Selector:
+    import narwhals._plan.selectors as cs
+
+    if isinstance(selectors, Collection):
+        if not selectors:
+            return cs.empty()
+        if len(selectors) == 1:
+            return next(iter(selectors))  # type: ignore[no-any-return]
+    s: Selector = reduce(operator.or_, selectors)
+    return s
 
 
 def parse_into_seq_of_expr_ir(

narwhals/_plan/arrow/dataframe.py

@@ -17,7 +17,7 @@
 from narwhals._plan.compliant.typing import namespace
 from narwhals._plan.expressions import NamedIR
 from narwhals._plan.typing import Seq
-from narwhals._utils import Implementation, Version, parse_columns_to_drop
+from narwhals._utils import Implementation, Version
 from narwhals.schema import Schema
 
 if TYPE_CHECKING:
@@ -106,9 +106,8 @@ def get_column(self, name: str) -> Series:
         chunked = self.native.column(name)
         return Series.from_native(chunked, name, version=self.version)
 
-    def drop(self, columns: Sequence[str], *, strict: bool = True) -> Self:
-        to_drop = parse_columns_to_drop(self, columns, strict=strict)
-        return self._with_native(self.native.drop(to_drop))
+    def drop(self, columns: Sequence[str]) -> Self:
+        return self._with_native(self.native.drop(list(columns)))
 
     def drop_nulls(self, subset: Sequence[str] | None) -> Self:
         if subset is None:

narwhals/_plan/common.py

@@ -22,7 +22,13 @@
 
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.series import Series
-    from narwhals._plan.typing import DTypeT, NonNestedDTypeT, OneOrIterable, Seq
+    from narwhals._plan.typing import (
+        ColumnNameOrSelector,
+        DTypeT,
+        NonNestedDTypeT,
+        OneOrIterable,
+        Seq,
+    )
     from narwhals._utils import _StoresColumns
     from narwhals.typing import NonNestedDType, NonNestedLiteral
 
@@ -85,8 +91,12 @@ def flatten_hash_safe(
     iterable: Iterable[OneOrIterable[CompliantSeries]], /
 ) -> Iterator[CompliantSeries]: ...
 @overload
+def flatten_hash_safe(
+    iterable: Iterable[OneOrIterable[ColumnNameOrSelector]], /
+) -> Iterator[ColumnNameOrSelector]: ...
+@overload
 def flatten_hash_safe(iterable: Iterable[OneOrIterable[T]], /) -> Iterator[T]: ...
-def flatten_hash_safe(iterable: Iterable[OneOrIterable[T]], /) -> Iterator[T]:
+def flatten_hash_safe(iterable: Iterable[OneOrIterable[Any]], /) -> Iterator[Any]:
     """Fully unwrap all levels of nesting.
 
     Aiming to reduce the chances of passing an unhashable argument.
@@ -95,7 +105,7 @@ def flatten_hash_safe(iterable: Iterable[OneOrIterable[T]], /) -> Iterator[T]:
         if isinstance(element, Iterable) and not is_iterable_reject(element):
             yield from flatten_hash_safe(element)
         else:
-            yield element  # type: ignore[misc]
+            yield element
 
 
 def _not_one_or_iterable_str_error(obj: Any, /) -> TypeError:  # pragma: no cover

narwhals/_plan/compliant/dataframe.py

@@ -55,7 +55,7 @@ def native(self) -> NativeFrameT_co: ...
     def to_narwhals(self) -> BaseFrame[NativeFrameT_co]: ...
     @property
     def columns(self) -> list[str]: ...
-    def drop(self, columns: Sequence[str], *, strict: bool = True) -> Self: ...
+    def drop(self, columns: Sequence[str]) -> Self: ...
     def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
     # Shouldn't *need* to be `NamedIR`, but current impl depends on a name being passed around
     def filter(self, predicate: NamedIR, /) -> Self: ...

narwhals/_plan/dataframe.py

@@ -110,8 +110,12 @@ def sort(
         named_irs, _ = prepare_projection(sort, schema=self)
         return self._with_compliant(self._compliant.sort(named_irs, opts))
 
-    def drop(self, *columns: str, strict: bool = True) -> Self:
-        return self._with_compliant(self._compliant.drop(columns, strict=strict))
+    def drop(
+        self, *columns: OneOrIterable[ColumnNameOrSelector], strict: bool = True
+    ) -> Self:
+        s_ir = _parse.parse_into_combined_selector_ir(*columns, require_all=strict)
+        names = expand_selector_irs_names((s_ir,), schema=self)
+        return self._with_compliant(self._compliant.drop(names))
 
     def drop_nulls(
         self, subset: OneOrIterable[ColumnNameOrSelector] | None = None

narwhals/_plan/expressions/selectors.py

@@ -257,6 +257,8 @@ class ByDType(DTypeSelector, dtype=DType):
     dtypes: frozenset[DType | type[DType]]
 
     def __repr__(self) -> str:
+        if not self.dtypes:
+            return "ncs.empty()"
         return f"ncs.by_dtype([{', '.join(sorted(map(repr, self.dtypes)))}])"
 
     def _matches(self, dtype: DType | type[DType]) -> bool:

narwhals/_plan/selectors.py

@@ -34,6 +34,7 @@
     "categorical",
     "datetime",
     "duration",
+    "empty",
     "enum",
     "first",
     "float",
@@ -180,7 +181,7 @@ def by_dtype(*dtypes: OneOrIterable[DType | type[DType]]) -> Selector:
     it = iter(selectors)
     if first := next(it, None):
         return reduce(operator.or_, it, first)
-    return s_ir.ByDType.empty().to_selector_ir().to_narwhals()
+    return empty()
 
 
 def by_index(*indices: OneOrIterable[int], require_all: bool = True) -> Selector:
@@ -222,6 +223,10 @@ def duration(time_unit: OneOrIterable[TimeUnit] | None = None) -> Selector:
     return s_ir.Duration.from_time_unit(time_unit).to_selector_ir().to_narwhals()
 
 
+def empty() -> Selector:
+    return s_ir.ByDType.empty().to_selector_ir().to_narwhals()
+
+
 def enum() -> Selector:
     return s_ir.Enum().to_selector_ir().to_narwhals()