feat(expr-ir): Add meta.as_selector, parse_into_selector_ir

Still have some translations missing
`by_index` will mean updating `matches_column` to *also* pass in the schema index

Author: dangotbanned

narwhals/_plan/_expr_ir.py

@@ -8,6 +8,7 @@
 from narwhals._plan.common import replace
 from narwhals._plan.options import ExprIROptions
 from narwhals._plan.typing import ExprIRT
+from narwhals.exceptions import InvalidOperationError
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
@@ -59,6 +60,10 @@ def to_narwhals(self, version: Version = Version.MAIN) -> Expr:
         tp = expr.Expr if version is Version.MAIN else expr.ExprV1
         return tp._from_ir(self)
 
+    def to_selector_ir(self) -> SelectorIR:
+        msg = f"cannot turn `{self!r}` into a selector"
+        raise InvalidOperationError(msg)
+
     @property
     def is_scalar(self) -> bool:
         return False
@@ -201,6 +206,9 @@ def matches_column(self, name: str, dtype: DType) -> bool:
         """
         raise NotImplementedError(type(self))
 
+    def to_selector_ir(self) -> Self:
+        return self
+
 
 class NamedIR(Immutable, Generic[ExprIRT]):
     """Post-projection expansion wrapper for `ExprIR`.

narwhals/_plan/_guards.py

@@ -11,7 +11,7 @@
 if TYPE_CHECKING:
     from typing_extensions import TypeIs
 
-    from narwhals._plan import expressions as ir
+    from narwhals._plan import expr, expressions as ir
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
@@ -58,6 +58,10 @@ def is_expr(obj: Any) -> TypeIs[Expr]:
     return isinstance(obj, _expr().Expr)
 
 
+def is_selector(obj: Any) -> TypeIs[expr.Selector]:
+    return isinstance(obj, _expr().Selector)
+
+
 def is_column(obj: Any) -> TypeIs[Expr]:
     """Indicate if the given object is a basic/unaliased column."""
     return is_expr(obj) and obj.meta.is_column()

narwhals/_plan/_parse.py

@@ -6,12 +6,18 @@
 from itertools import chain
 from typing import TYPE_CHECKING
 
-from narwhals._plan._guards import is_expr, is_into_expr_column, is_iterable_reject
+from narwhals._plan._guards import (
+    is_expr,
+    is_into_expr_column,
+    is_iterable_reject,
+    is_selector,
+)
 from narwhals._plan.exceptions import (
     invalid_into_expr_error,
     is_iterable_pandas_error,
     is_iterable_polars_error,
 )
+from narwhals._utils import qualified_type_name
 from narwhals.dependencies import get_polars, is_pandas_dataframe, is_pandas_series
 from narwhals.exceptions import InvalidOperationError
 
@@ -22,8 +28,10 @@
     import polars as pl
     from typing_extensions import TypeAlias, TypeIs
 
-    from narwhals._plan.expressions import ExprIR
+    from narwhals._plan.expr import Expr
+    from narwhals._plan.expressions import ExprIR, SelectorIR
     from narwhals._plan.typing import (
+        ColumnNameOrSelector,
         IntoExpr,
         IntoExprColumn,
         OneOrIterable,
@@ -124,6 +132,21 @@ def parse_into_expr_ir(
     return expr._ir
 
 
+def parse_into_selector_ir(input: ColumnNameOrSelector | Expr, /) -> SelectorIR:
+    if is_selector(input):
+        selector = input
+    elif isinstance(input, str):
+        from narwhals._plan import selectors as cs
+
+        selector = cs.by_name(input)
+    elif is_expr(input):
+        selector = input.meta.as_selector()
+    else:
+        msg = f"cannot turn {qualified_type_name(input)!r} into selector"
+        raise TypeError(msg)
+    return selector._ir
+
+
 def parse_into_seq_of_expr_ir(
     first_input: OneOrIterable[IntoExpr] = (),
     *more_inputs: IntoExpr | _RaisesInvalidIntoExprError,

narwhals/_plan/expressions/expr.py

@@ -9,6 +9,7 @@
 from narwhals._plan._expr_ir import ExprIR, SelectorIR
 from narwhals._plan.common import flatten_hash_safe
 from narwhals._plan.exceptions import function_expr_invalid_operation_error
+from narwhals._plan.expressions import selectors as cs
 from narwhals._plan.options import ExprIROptions
 from narwhals._plan.typing import (
     FunctionT_co,
@@ -32,7 +33,6 @@
     from narwhals._plan.compliant.typing import Ctx, FrameT_contra, R_co
     from narwhals._plan.expressions.functions import MapBatches  # noqa: F401
     from narwhals._plan.expressions.literal import LiteralValue
-    from narwhals._plan.expressions.selectors import Selector
     from narwhals._plan.expressions.window import Window
     from narwhals._plan.options import FunctionOptions, SortMultipleOptions, SortOptions
     from narwhals.dtypes import DType
@@ -100,6 +100,9 @@ class Column(ExprIR, config=ExprIROptions.namespaced("col")):
     def __repr__(self) -> str:
         return f"col({self.name!r})"
 
+    def to_selector_ir(self) -> RootSelector:
+        return cs.ByName.from_name(self.name).to_selector_ir()
+
 
 class _ColumnSelection(ExprIR, config=ExprIROptions.no_dispatch()):
     """Nodes which can resolve to `Column`(s) with a `Schema`."""
@@ -112,7 +115,11 @@ class Columns(_ColumnSelection):
     def __repr__(self) -> str:
         return f"cols({list(self.names)!r})"
 
+    def to_selector_ir(self) -> RootSelector:
+        return cs.ByName.from_names(*self.names).to_selector_ir()
+
 
+# TODO @dangotbanned: Add `selectors.by_index`
 class Nth(_ColumnSelection):
     __slots__ = ("index",)
     index: int
@@ -121,6 +128,7 @@ def __repr__(self) -> str:
         return f"nth({self.index})"
 
 
+# TODO @dangotbanned: Add `selectors.by_index`
 class IndexColumns(_ColumnSelection):
     __slots__ = ("indices",)
     indices: Seq[int]
@@ -133,7 +141,11 @@ class All(_ColumnSelection):
     def __repr__(self) -> str:
         return "all()"
 
+    def to_selector_ir(self) -> RootSelector:
+        return cs.All().to_selector_ir()
+
 
+# TODO @dangotbanned: Add `selectors.exclude`
 class Exclude(_ColumnSelection, child=("expr",)):
     __slots__ = ("expr", "names")
     expr: ExprIR
@@ -450,7 +462,7 @@ class RootSelector(SelectorIR):
     """A single selector expression."""
 
     __slots__ = ("selector",)
-    selector: Selector
+    selector: cs.Selector
 
     def __repr__(self) -> str:
         return f"{self.selector!r}"

narwhals/_plan/expressions/selectors.py

@@ -38,7 +38,7 @@
 
 
 class Selector(Immutable):
-    def to_selector(self) -> RootSelector:
+    def to_selector_ir(self) -> RootSelector:
         from narwhals._plan.expressions.expr import RootSelector
 
         return RootSelector(selector=self)
@@ -248,14 +248,14 @@ def matches_column(self, name: str, dtype: DType) -> bool:
 
 
 def all() -> expr.Selector:
-    return All().to_selector().to_narwhals()
+    return All().to_selector_ir().to_narwhals()
 
 
 def array(
     inner: expr.Selector | None = None, *, size: int | None = None
 ) -> expr.Selector:
     s_ir = inner._ir if inner is not None else None
-    return Array(inner=s_ir, size=size).to_selector().to_narwhals()
+    return Array(inner=s_ir, size=size).to_selector_ir().to_narwhals()
 
 
 def by_dtype(*dtypes: OneOrIterable[DType | type[DType]]) -> expr.Selector:
@@ -267,15 +267,15 @@ def by_name(*names: OneOrIterable[str]) -> expr.Selector:
         sel = ByName.from_name(names[0])
     else:
         sel = ByName.from_names(*names)
-    return sel.to_selector().to_narwhals()
+    return sel.to_selector_ir().to_narwhals()
 
 
 def boolean() -> expr.Selector:
-    return Boolean().to_selector().to_narwhals()
+    return Boolean().to_selector_ir().to_narwhals()
 
 
 def categorical() -> expr.Selector:
-    return Categorical().to_selector().to_narwhals()
+    return Categorical().to_selector_ir().to_narwhals()
 
 
 def datetime(
@@ -284,38 +284,38 @@ def datetime(
 ) -> expr.Selector:
     return (
         Datetime.from_time_unit_and_time_zone(time_unit, time_zone)
-        .to_selector()
+        .to_selector_ir()
         .to_narwhals()
     )
 
 
 def list(inner: expr.Selector | None = None) -> expr.Selector:
     s_ir = inner._ir if inner is not None else None
-    return List(inner=s_ir).to_selector().to_narwhals()
+    return List(inner=s_ir).to_selector_ir().to_narwhals()
 
 
 def duration(time_unit: OneOrIterable[TimeUnit] | None = None) -> expr.Selector:
-    return Duration.from_time_unit(time_unit).to_selector().to_narwhals()
+    return Duration.from_time_unit(time_unit).to_selector_ir().to_narwhals()
 
 
 def enum() -> expr.Selector:
-    return Enum().to_selector().to_narwhals()
+    return Enum().to_selector_ir().to_narwhals()
 
 
 def matches(pattern: str) -> expr.Selector:
-    return Matches.from_string(pattern).to_selector().to_narwhals()
+    return Matches.from_string(pattern).to_selector_ir().to_narwhals()
 
 
 def numeric() -> expr.Selector:
-    return Numeric().to_selector().to_narwhals()
+    return Numeric().to_selector_ir().to_narwhals()
 
 
 def string() -> expr.Selector:
-    return String().to_selector().to_narwhals()
+    return String().to_selector_ir().to_narwhals()
 
 
 def struct() -> expr.Selector:
-    return Struct().to_selector().to_narwhals()
+    return Struct().to_selector_ir().to_narwhals()
 
 
 _HASH_SENSITIVE_TO_SELECTOR: Mapping[type[DType], Callable[[], expr.Selector]] = {
@@ -343,7 +343,7 @@ def _from_dtypes(*by_dtypes: OneOrIterable[DType | type[DType]]) -> expr.Selecto
         else:
             dtypes.append(dtype)  # type: ignore[arg-type]
     if dtypes:
-        dtype_selector = ByDType(dtypes=frozenset(dtypes)).to_selector().to_narwhals()
+        dtype_selector = ByDType(dtypes=frozenset(dtypes)).to_selector_ir().to_narwhals()
         selectors.appendleft(dtype_selector)
     it = iter(selectors)
     return reduce(operator.or_, it, next(it))

narwhals/_plan/meta.py

@@ -14,12 +14,14 @@
 from narwhals._plan._guards import is_literal
 from narwhals._plan.expressions.literal import is_literal_scalar
 from narwhals._plan.expressions.namespace import IRNamespace
-from narwhals.exceptions import ComputeError
+from narwhals.exceptions import ComputeError, InvalidOperationError
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator
 
+    from narwhals._plan import expr
+
 
 class MetaNamespace(IRNamespace):
     """Methods to modify and traverse existing expressions."""
@@ -75,6 +77,16 @@ def root_names(self) -> list[str]:
         """Get the root column names."""
         return list(_expr_to_leaf_column_names_iter(self._ir))
 
+    def as_selector(self) -> expr.Selector:
+        """Try to turn this expression into a selector.
+
+        Raises if the underlying expressions is not a column or selector.
+        """
+        if not self.is_column_selection():
+            msg = f"cannot turn `{self._ir!r}` into a selector"
+            raise InvalidOperationError(msg)
+        return self._ir.to_selector_ir().to_narwhals()
+
 
 def _expr_to_leaf_column_names_iter(expr: ir.ExprIR, /) -> Iterator[str]:
     for e in _expr_to_leaf_column_exprs_iter(expr):