Add partition-wise Select support to cuDF-Polars

python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py

@@ -40,6 +40,7 @@ def __init__(
         self.dtype = dtype
         self.name = name
         self.options = options
+        self.is_pointwise = False
         self.children = children
         if name not in Agg._SUPPORTED:
             raise NotImplementedError(

python/cudf_polars/cudf_polars/dsl/expressions/base.py

@@ -36,9 +36,11 @@ class ExecutionContext(IntEnum):
 class Expr(Node["Expr"]):
     """An abstract expression object."""
 
-    __slots__ = ("dtype",)
+    __slots__ = ("dtype", "is_pointwise")
     dtype: plc.DataType
     """Data type of the expression."""
+    is_pointwise: bool
+    """Whether row output is independent."""
     # This annotation is needed because of https://github.com/python/mypy/issues/17981
     _non_child: ClassVar[tuple[str, ...]] = ("dtype",)
     """Names of non-child data (not Exprs) for reconstruction."""
@@ -164,6 +166,7 @@ def __init__(self, dtype: plc.DataType, error: str) -> None:
         self.dtype = dtype
         self.error = error
         self.children = ()
+        self.is_pointwise = True
 
 
 class NamedExpr:
@@ -178,6 +181,11 @@ def __init__(self, name: str, value: Expr) -> None:
         self.name = name
         self.value = value
 
+    @property
+    def is_pointwise(self) -> bool:
+        """Whether row output is independent."""
+        return self.value.is_pointwise
+
     def __hash__(self) -> int:
         """Hash of the expression."""
         return hash((type(self), self.name, self.value))
@@ -243,6 +251,7 @@ class Col(Expr):
     def __init__(self, dtype: plc.DataType, name: str) -> None:
         self.dtype = dtype
         self.name = name
+        self.is_pointwise = True
         self.children = ()
 
     def do_evaluate(
@@ -280,6 +289,7 @@ def __init__(
         self.dtype = dtype
         self.index = index
         self.table_ref = table_ref
+        self.is_pointwise = True
         self.children = (column,)
 
     def do_evaluate(

python/cudf_polars/cudf_polars/dsl/expressions/binaryop.py

@@ -42,6 +42,7 @@ def __init__(
             op = BinOp._BOOL_KLEENE_MAPPING.get(op, op)
         self.op = op
         self.children = (left, right)
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
         if not plc.binaryop.is_supported_operation(
             self.dtype, left.dtype, right.dtype, op
         ):

python/cudf_polars/cudf_polars/dsl/expressions/boolean.py

@@ -81,6 +81,7 @@ def __init__(
         self.options = options
         self.name = name
         self.children = children
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
         if self.name is BooleanFunction.Name.IsIn and not all(
             c.dtype == self.children[0].dtype for c in self.children
         ):

python/cudf_polars/cudf_polars/dsl/expressions/datetime.py

@@ -114,6 +114,7 @@ def __init__(
         self.options = options
         self.name = name
         self.children = children
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
         if self.name not in self._COMPONENT_MAP:
             raise NotImplementedError(f"Temporal function {self.name}")
 

python/cudf_polars/cudf_polars/dsl/expressions/literal.py

@@ -38,6 +38,7 @@ def __init__(self, dtype: plc.DataType, value: pa.Scalar[Any]) -> None:
         assert value.type == plc.interop.to_arrow(dtype)
         self.value = value
         self.children = ()
+        self.is_pointwise = True
 
     def do_evaluate(
         self,
@@ -65,6 +66,7 @@ def __init__(self, dtype: plc.DataType, value: pl.Series) -> None:
         data = value.to_arrow()
         self.value = data.cast(dtypes.downcast_arrow_lists(data.type))
         self.children = ()
+        self.is_pointwise = True
 
     def get_hashable(self) -> Hashable:
         """Compute a hash of the column."""

python/cudf_polars/cudf_polars/dsl/expressions/rolling.py

@@ -24,6 +24,7 @@ def __init__(self, dtype: plc.DataType, options: Any, agg: Expr) -> None:
         self.dtype = dtype
         self.options = options
         self.children = (agg,)
+        self.is_pointwise = False
         raise NotImplementedError("Rolling window not implemented")
 
 
@@ -35,4 +36,5 @@ def __init__(self, dtype: plc.DataType, options: Any, agg: Expr, *by: Expr) -> N
         self.dtype = dtype
         self.options = options
         self.children = (agg, *by)
+        self.is_pointwise = False
         raise NotImplementedError("Grouped rolling window not implemented")

python/cudf_polars/cudf_polars/dsl/expressions/selection.py

@@ -30,6 +30,7 @@ class Gather(Expr):
     def __init__(self, dtype: plc.DataType, values: Expr, indices: Expr) -> None:
         self.dtype = dtype
         self.children = (values, indices)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,
@@ -71,6 +72,7 @@ class Filter(Expr):
     def __init__(self, dtype: plc.DataType, values: Expr, indices: Expr):
         self.dtype = dtype
         self.children = (values, indices)
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
 
     def do_evaluate(
         self,

python/cudf_polars/cudf_polars/dsl/expressions/sorting.py

@@ -32,6 +32,7 @@ def __init__(
         self.dtype = dtype
         self.options = options
         self.children = (column,)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,
@@ -71,6 +72,7 @@ def __init__(
         self.dtype = dtype
         self.options = options
         self.children = (column, *by)
+        self.is_pointwise = False
 
     def do_evaluate(
         self,

python/cudf_polars/cudf_polars/dsl/expressions/string.py

@@ -106,6 +106,7 @@ def __init__(
         self.options = options
         self.name = name
         self.children = children
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
         self._validate_input()
 
     def _validate_input(self):

python/cudf_polars/cudf_polars/dsl/expressions/ternary.py

@@ -34,6 +34,7 @@ def __init__(
     ) -> None:
         self.dtype = dtype
         self.children = (when, then, otherwise)
+        self.is_pointwise = all(c.is_pointwise for c in self.children)
 
     def do_evaluate(
         self,

python/cudf_polars/cudf_polars/experimental/select.py

@@ -7,7 +7,6 @@
 from typing import TYPE_CHECKING
 
 from cudf_polars.dsl.ir import Select
-from cudf_polars.dsl.traversal import traversal
 from cudf_polars.experimental.base import PartitionInfo
 from cudf_polars.experimental.dispatch import lower_ir_node
 
@@ -18,22 +17,6 @@
     from cudf_polars.experimental.parallel import LowerIRTransformer
 
 
-_PARTWISE = (
-    "Literal",
-    "LiteralColumn",
-    "Col",
-    "ColRef",
-    "BooleanFunction",
-    "StringFunction",
-    "TemporalFunction",
-    "Filter",
-    "Cast",
-    "Ternary",
-    "BinOp",
-    "UnaryFunction",
-)
-
-
 @lower_ir_node.register(Select)
 def _(
     ir: Select, rec: LowerIRTransformer
@@ -43,17 +26,15 @@ def _(
     new_node = ir.reconstruct([child])
 
     # Search the expression graph for "complex" operations
-    for ne in ir.exprs:
-        for expr in traversal(ne.value):
-            if type(expr).__name__ not in _PARTWISE:
-                # TODO: Handle non-partition-wise expressions.
-                if partition_info[child].count > 1:
-                    raise NotImplementedError(
-                        f"Expr {type(expr)} does not support multiple partitions."
-                    )
-                else:  # pragma: no cover
-                    partition_info[new_node] = PartitionInfo(count=1)
-                    return new_node, partition_info
+    if not all(e.is_pointwise for e in ir.exprs):
+        if partition_info[child].count > 1:
+            # TODO: Handle non-pointwise expressions.
+            raise NotImplementedError(
+                f"Selection {ir} does not support multiple partitions."
+            )
+        else:  # pragma: no cover
+            partition_info[new_node] = PartitionInfo(count=1)
+            return new_node, partition_info
 
     # Remaining Select ops are partition-wise
     partition_info[new_node] = PartitionInfo(count=partition_info[child].count)