Enable multi-partition Select operations containing basic aggregations

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

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 # TODO: remove need for this
 # ruff: noqa: D101
@@ -18,6 +18,8 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+    from typing_extensions import Self
+
     from cudf_polars.containers import Column, DataFrame
 
 __all__ = ["AggInfo", "Col", "ColRef", "ExecutionContext", "Expr", "NamedExpr"]
@@ -237,6 +239,24 @@ def collect_agg(self, *, depth: int) -> AggInfo:
         """Collect information about aggregations in groupbys."""
         return self.value.collect_agg(depth=depth)
 
+    def reconstruct(self, expr: Expr) -> Self:
+        """
+        Rebuild with a new `Expr` value.
+
+        Parameters
+        ----------
+        expr
+            New `Expr` value
+
+        Returns
+        -------
+        New `NamedExpr` with `expr` as the underlying expression.
+        The name of the original `NamedExpr` is preserved.
+        """
+        if expr is self.value:
+            return self
+        return type(self)(self.name, expr)
+
 
 class Col(Expr):
     __slots__ = ("name",)

python/cudf_polars/cudf_polars/dsl/traversal.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """Traversal and visitor utilities for nodes."""
@@ -23,14 +23,19 @@
 ]
 
 
-def traversal(nodes: Sequence[NodeT]) -> Generator[NodeT, None, None]:
+def traversal(
+    nodes: Sequence[NodeT], *, cutoff_types: tuple[type[NodeT], ...] = ()
+) -> Generator[NodeT, None, None]:
     """
     Pre-order traversal of nodes in an expression.
 
     Parameters
     ----------
     nodes
         Roots of expressions to traverse.
+    cutoff_types
+        Types to terminate traversal at. If a type is in this tuple
+        then we do not yield any of its children.
 
     Yields
     ------
@@ -43,6 +48,8 @@ def traversal(nodes: Sequence[NodeT]) -> Generator[NodeT, None, None]:
     while lifo:
         node = lifo.pop()
         yield node
+        if cutoff_types and isinstance(node, cutoff_types):
+            continue
         for child in reversed(node.children):
             if child not in seen:
                 seen.add(child)

python/cudf_polars/cudf_polars/experimental/base.py

@@ -36,6 +36,6 @@ def keys(self, node: Node) -> Iterator[tuple[str, int]]:
         yield from ((name, i) for i in range(self.count))
 
 
-def get_key_name(node: Node) -> str:
+def get_key_name(node: Node, *other: Node) -> str:
     """Generate the key name for a Node."""
-    return f"{type(node).__name__.lower()}-{hash(node)}"
+    return f"{type(node).__name__.lower()}-{hash((node, *other))}"