Support multi-partition Select operations with 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/ir.py

@@ -57,6 +57,7 @@
     "Filter",
     "GroupBy",
     "HConcat",
+    "HConcatBcast",
     "HStack",
     "Join",
     "MapFunction",
@@ -2023,3 +2024,22 @@ def do_evaluate(cls, *dfs: DataFrame) -> DataFrame:
                 )
             )
         )
+
+
+class HConcatBcast(IR):
+    """Concatenate dataframes horizontally with broadcasting."""
+
+    __slots__ = ()
+    _non_child = ("schema",)
+
+    def __init__(self, schema: Schema, *children: IR):
+        self.schema = schema
+        self._non_child_args = ()
+        self.children = children
+
+    @classmethod
+    def do_evaluate(cls, *dfs: DataFrame) -> DataFrame:
+        """Evaluate and return a dataframe."""
+        return DataFrame(
+            broadcast(*itertools.chain.from_iterable(df.columns for df in dfs))
+        )

python/cudf_polars/cudf_polars/dsl/utils/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""DSL utilities."""
+
+from __future__ import annotations
+
+__all__: list[str] = []

python/cudf_polars/cudf_polars/dsl/utils/naming.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Name generation utilities."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterable
+
+
+__all__ = ["unique_names"]
+
+
+def unique_names(names: Iterable[str]) -> Generator[str, None, None]:
+    """
+    Generate unique names relative to some known names.
+
+    Parameters
+    ----------
+    names
+        Names we should be unique with respect to.
+
+    Yields
+    ------
+    Unique names (just using sequence numbers)
+    """
+    prefix = "_" * max(map(len, names))
+    i = 0
+    while True:
+        yield f"{prefix}{i}"
+        i += 1