more partition_by prep

Author: dangotbanned

narwhals/_plan/_expansion.py

@@ -72,6 +72,7 @@
     IntoFrozenSchema,
     freeze_schema,
 )
+from narwhals._utils import check_column_names_are_unique
 from narwhals.dtypes import DType
 from narwhals.exceptions import ComputeError, InvalidOperationError
 
@@ -156,7 +157,7 @@ def with_multiple_columns(self) -> ExpansionFlags:
 def prepare_projection(
     exprs: Sequence[ExprIR], /, keys: GroupByKeys = (), *, schema: IntoFrozenSchema
 ) -> tuple[Seq[NamedIR], FrozenSchema]:
-    """Expand IRs into named column selections.
+    """Expand IRs into named column projections.
 
     **Primary entry-point**, for `select`, `with_columns`,
     and any other context that requires resolving expression names.
@@ -173,13 +174,33 @@ def prepare_projection(
     return named_irs, frozen_schema
 
 
+def expand_selector_irs_names(
+    selectors: Sequence[SelectorIR],
+    /,
+    keys: GroupByKeys = (),
+    *,
+    schema: IntoFrozenSchema,
+) -> OutputNames:
+    """Expand selector-only input into the column names that match.
+
+    Similar to `prepare_projection`, but intended for allowing a subset of `Expr` and all `Selector`s
+    to be used in more places like `DataFrame.{drop,sort,partition_by}`.
+
+    Arguments:
+        selectors: IRs that **only** contain subclasses of `SelectorIR`.
+        keys: Names of `group_by` columns.
+        schema: Scope to expand multi-column selectors in.
+    """
+    frozen_schema = freeze_schema(schema)
+    names = tuple(_iter_expand_selector_names(selectors, keys, schema=frozen_schema))
+    return _ensure_valid_output_names(names, frozen_schema)
+
+
 def into_named_irs(exprs: Seq[ExprIR], names: OutputNames) -> Seq[NamedIR]:
     if len(exprs) != len(names):
         msg = f"zip length mismatch: {len(exprs)} != {len(names)}"
         raise ValueError(msg)
-    return tuple(
-        NamedIR(expr=remove_alias(ir), name=name) for ir, name in zip(exprs, names)
-    )
+    return tuple(ir.named_ir(name, remove_alias(e)) for e, name in zip(exprs, names))
 
 
 def ensure_valid_exprs(exprs: Seq[ExprIR], schema: FrozenSchema) -> OutputNames:
@@ -191,13 +212,40 @@ def ensure_valid_exprs(exprs: Seq[ExprIR], schema: FrozenSchema) -> OutputNames:
     return output_names
 
 
+def _ensure_valid_output_names(names: Seq[str], schema: FrozenSchema) -> OutputNames:
+    """Selector-only variant of `ensure_valid_exprs`."""
+    check_column_names_are_unique(names)
+    output_names = names
+    if not (set(schema.names).issuperset(output_names)):
+        raise column_not_found_error(output_names, schema)
+    return output_names
+
+
 def _ensure_output_names_unique(exprs: Seq[ExprIR]) -> OutputNames:
     names = tuple(e.meta.output_name() for e in exprs)
     if len(names) != len(set(names)):
         raise duplicate_error(exprs)
     return names
 
 
+def _ensure_columns(expr: ExprIR, /) -> Columns:
+    if not isinstance(expr, Columns):
+        msg = f"Expected only column selections here, but got {expr!r}"
+        raise NotImplementedError(msg)
+    return expr
+
+
+def _iter_expand_selector_names(
+    selectors: Iterable[SelectorIR], /, keys: GroupByKeys = (), *, schema: FrozenSchema
+) -> Iterator[str]:
+    for selector in selectors:
+        names = _ensure_columns(replace_selector(selector, schema=schema)).names
+        if keys:
+            yield from (name for name in names if name not in keys)
+        else:
+            yield from names
+
+
 # NOTE: Recursive for all `input` expressions which themselves contain `Seq[ExprIR]`
 def rewrite_projections(
     input: Seq[ExprIR], /, keys: GroupByKeys = (), *, schema: FrozenSchema

narwhals/_plan/_guards.py

@@ -15,7 +15,12 @@
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
-    from narwhals._plan.typing import IntoExprColumn, NativeSeriesT, Seq
+    from narwhals._plan.typing import (
+        ColumnNameOrSelector,
+        IntoExprColumn,
+        NativeSeriesT,
+        Seq,
+    )
     from narwhals.typing import NonNestedLiteral
 
     T = TypeVar("T")
@@ -75,6 +80,13 @@ def is_into_expr_column(obj: Any) -> TypeIs[IntoExprColumn]:
     return isinstance(obj, (str, _expr().Expr, _series().Series))
 
 
+def is_column_name_or_selector(
+    obj: Any, *, allow_expr: bool = False
+) -> TypeIs[ColumnNameOrSelector]:
+    tps = (str, _expr().Selector) if not allow_expr else (str, _expr().Expr)
+    return isinstance(obj, tps)
+
+
 def is_compliant_series(
     obj: CompliantSeries[NativeSeriesT] | Any,
 ) -> TypeIs[CompliantSeries[NativeSeriesT]]:

narwhals/_plan/_parse.py

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING
 
 from narwhals._plan._guards import (
+    is_column_name_or_selector,
     is_expr,
     is_into_expr_column,
     is_iterable_reject,
@@ -132,6 +133,8 @@ def parse_into_expr_ir(
     return expr._ir
 
 
+# NOTE: Might need to add `require_all`, since selectors are created indirectly from `str`
+# here, but use set semantics
 def parse_into_selector_ir(input: ColumnNameOrSelector | Expr, /) -> SelectorIR:
     if is_selector(input):
         selector = input
@@ -198,6 +201,34 @@ def _parse_sort_by_into_iter_expr_ir(
         yield e
 
 
+def parse_into_seq_of_selector_ir(
+    first_input: OneOrIterable[ColumnNameOrSelector], *more_inputs: ColumnNameOrSelector
+) -> Seq[SelectorIR]:
+    return tuple(_parse_into_iter_selector_ir(first_input, more_inputs))
+
+
+def _parse_into_iter_selector_ir(
+    first_input: OneOrIterable[ColumnNameOrSelector],
+    more_inputs: tuple[ColumnNameOrSelector, ...],
+    /,
+) -> Iterator[SelectorIR]:
+    if is_column_name_or_selector(first_input) and not more_inputs:
+        yield parse_into_selector_ir(first_input)
+        return
+
+    if not _is_empty_sequence(first_input):
+        if _is_iterable(first_input) and not isinstance(first_input, str):
+            if more_inputs:
+                raise invalid_into_expr_error(first_input, more_inputs, {})
+            else:
+                for into in first_input:  # type: ignore[var-annotated]
+                    yield parse_into_selector_ir(into)
+        else:
+            yield parse_into_selector_ir(first_input)
+    for into in more_inputs:
+        yield parse_into_selector_ir(into)
+
+
 def _parse_into_iter_expr_ir(
     first_input: OneOrIterable[IntoExpr],
     *more_inputs: IntoExpr | list[Any],

narwhals/_plan/arrow/dataframe.py

@@ -171,3 +171,8 @@ def filter(self, predicate: NamedIR) -> Self:
         else:
             mask = acero.lit(resolved.native)
         return self._with_native(self.native.filter(mask))
+
+    def partition_by(self, by: Sequence[str], *, include_key: bool = True) -> list[Self]:
+        """Review https://github.com/pola-rs/polars/blob/870f0e01811b8b0cf9b846ded9d97685f143d27c/crates/polars-core/src/frame/mod.rs#L3225-L3284."""
+        msg = "TODO: `ArrowDataFrame.partition_by`"
+        raise NotImplementedError(msg)

narwhals/_plan/compliant/dataframe.py

@@ -129,6 +129,9 @@ def join(
         suffix: str = "_right",
     ) -> Self: ...
     def join_cross(self, other: Self, *, suffix: str = "_right") -> Self: ...
+    def partition_by(
+        self, by: Sequence[str], *, include_key: bool = True
+    ) -> list[Self]: ...
     def row(self, index: int) -> tuple[Any, ...]: ...
     @overload
     def to_dict(self, *, as_series: Literal[True]) -> dict[str, SeriesT]: ...

narwhals/_plan/dataframe.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, get_args, overload
 
 from narwhals._plan import _parse
-from narwhals._plan._expansion import prepare_projection
+from narwhals._plan._expansion import expand_selector_irs_names, prepare_projection
 from narwhals._plan.common import ensure_seq_str, temp
 from narwhals._plan.group_by import GroupBy, Grouped
 from narwhals._plan.options import SortMultipleOptions
@@ -258,6 +258,17 @@ def filter(
             raise ValueError(msg)
         return self._with_compliant(self._compliant.filter(named_irs[0]))
 
+    def partition_by(
+        self,
+        by: OneOrIterable[ColumnNameOrSelector],
+        *more_by: ColumnNameOrSelector,
+        include_key: bool = True,
+    ) -> list[Self]:
+        by_selectors = _parse.parse_into_seq_of_selector_ir(by, *more_by)
+        names = expand_selector_irs_names(by_selectors, schema=self)
+        partitions = self._compliant.partition_by(names, include_key=include_key)
+        return [self._with_compliant(p) for p in partitions]
+
 
 def _is_join_strategy(obj: Any) -> TypeIs[JoinStrategy]:
     return obj in {"inner", "left", "full", "cross", "anti", "semi"}