Implement streaming window functions in cudf-polars

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

@@ -615,13 +615,17 @@ def _grouped_window_scan_setup(
         local = self._sorted_grouper(by_cols_for_scan)
         return order_index, by_cols_for_scan, local
 
+    # TODO: this is an ordered left-join that drops the join keys.
+    # Rename it and replace the manual scatter+gather with the Join IR's
+    # _reorder_maps helper (lifted to a shared utility) so the streaming
+    # and in-memory over paths share the same primitive.
     def _broadcast_agg_results(
         self,
         by_tbl: plc.Table,
         group_keys_tbl: plc.Table,
-        value_tbls: list[plc.Table],
-        names: list[str],
-        dtypes: list[DataType],
+        value_tbl: plc.Table,
+        names: Sequence[str],
+        dtypes: Sequence[DataType],
         stream: Stream,
     ) -> list[Column]:
         # We do a left-join between the input keys to group-keys
@@ -647,7 +651,7 @@ def _broadcast_agg_results(
 
         # Broadcast each scalar aggregated result back to row-shape using
         # the aligned mapping between row indices and group indices.
-        out_cols = (t.columns()[0] for t in value_tbls)
+        out_cols = value_tbl.columns()
         return [
             Column(
                 plc.copying.gather(
@@ -773,7 +777,7 @@ def do_evaluate(  # noqa: D102
         broadcasted_cols = self._broadcast_agg_results(
             by_tbl,
             group_keys_tbl,
-            value_tables,
+            plc.Table([col for t in value_tables for col in t.columns()]),
             out_names,
             out_dtypes,
             df.stream,
@@ -807,7 +811,7 @@ def do_evaluate(  # noqa: D102
                 self._broadcast_agg_results(
                     by_tbl,
                     group_keys_tbl_local,
-                    value_tables_local,
+                    plc.Table([col for t in value_tables_local for col in t.columns()]),
                     out_names,
                     out_dtypes,
                     df.stream,

python/cudf_polars/cudf_polars/dsl/translate.py

@@ -38,6 +38,8 @@
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+
     from polars import GPUEngine
 
     from cudf_polars.typing import NodeTraverser
@@ -90,6 +92,7 @@ def __init__(self, visitor: NodeTraverser, engine: GPUEngine):
         self.errors: list[Exception] = []
         self._cache_nodes: dict[int, ir.Cache] = {}
         self._expr_context: ExecutionContext = ExecutionContext.FRAME
+        self._internal_name_gen: Generator[str, None, None] | None = None
 
     def translate_ir(self, *, n: int | None = None) -> ir.IR:
         """
@@ -247,6 +250,24 @@ def __exit__(self, *args: Any) -> None:
         self.translator._expr_context = self._prev
 
 
+class set_internal_name_gen(AbstractContextManager[None]):
+    """Share one internal-name generator across sibling expression translations."""
+
+    __slots__ = ("_prev", "schema", "translator")
+
+    def __init__(self, translator: Translator, schema: Schema) -> None:
+        self.translator = translator
+        self.schema = schema
+        self._prev: Generator[str, None, None] | None = None
+
+    def __enter__(self) -> None:
+        self._prev = self.translator._internal_name_gen
+        self.translator._internal_name_gen = unique_names(self.schema)
+
+    def __exit__(self, *args: Any) -> None:
+        self.translator._internal_name_gen = self._prev
+
+
 @singledispatch
 def _translate_ir(node: Any, translator: Translator, schema: Schema) -> ir.IR:
     raise NotImplementedError(
@@ -362,9 +383,11 @@ def _(
 def _(node: plrs._ir_nodes.Select, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        exprs = [
-            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.expr
-        ]
+        with set_internal_name_gen(translator, inp.schema):
+            exprs = [
+                translate_named_expr(translator, n=e, schema=inp.schema)
+                for e in node.expr
+            ]
     return ir.Select(schema, exprs, node.should_broadcast, inp)
 
 
@@ -478,9 +501,11 @@ def _(node: plrs._ir_nodes.Join, translator: Translator, schema: Schema) -> ir.I
 def _(node: plrs._ir_nodes.HStack, translator: Translator, schema: Schema) -> ir.IR:
     with set_node(translator.visitor, node.input):
         inp = translator.translate_ir(n=None)
-        exprs = [
-            translate_named_expr(translator, n=e, schema=inp.schema) for e in node.exprs
-        ]
+        with set_internal_name_gen(translator, inp.schema):
+            exprs = [
+                translate_named_expr(translator, n=e, schema=inp.schema)
+                for e in node.exprs
+            ]
     return ir.HStack(schema, exprs, node.should_broadcast, inp)
 
 
@@ -830,7 +855,7 @@ def _(
         # pl.col("a").rolling(...)
         with set_expr_context(translator, ExecutionContext.ROLLING):
             agg = translator.translate_expr(n=node.function, schema=schema)
-        name_generator = unique_names(schema)
+        name_generator = translator._internal_name_gen or unique_names(schema)
         aggs, named_post_agg = decompose_single_agg(
             expr.NamedExpr(next(name_generator), agg),
             name_generator,
@@ -875,7 +900,7 @@ def _(
         # not exposed until polars 1.39.
         with set_expr_context(translator, ExecutionContext.WINDOW):
             agg = translator.translate_expr(n=node.function, schema=schema)
-        name_gen = unique_names(schema)
+        name_gen = translator._internal_name_gen or unique_names(schema)
         aggs, post = decompose_single_agg(
             expr.NamedExpr(next(name_gen), agg),
             name_gen,

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

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
 """Name generation utilities."""
@@ -7,11 +7,15 @@
 
 from typing import TYPE_CHECKING
 
+from cudf_polars.dsl.expr import NamedExpr
+
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterable
 
+    from cudf_polars.typing import Schema
+
 
-__all__ = ["unique_names"]
+__all__ = ["names_to_indices", "unique_names"]
 
 
 def unique_names(names: Iterable[str]) -> Generator[str, None, None]:
@@ -32,3 +36,28 @@ def unique_names(names: Iterable[str]) -> Generator[str, None, None]:
     while True:
         yield f"{prefix}{i}"
         i += 1
+
+
+def names_to_indices(
+    names: tuple[str | NamedExpr, ...], schema: Schema
+) -> tuple[int, ...]:
+    """
+    Return column indices for the given names in schema order.
+
+    Accepts either column names (str) or NamedExpr, so it can be used with
+    e.g. ir.left_on, ir.right_on as well as plain name tuples.
+
+    Parameters
+    ----------
+    names
+        The names to get indices for.
+    schema
+        The schema to get indices from.
+
+    Returns
+    -------
+    The column indices for each name in schema order.
+    """
+    keys = list(schema.keys())
+    str_names = [n.name if isinstance(n, NamedExpr) else n for n in names]
+    return tuple(keys.index(n) for n in str_names)

python/cudf_polars/cudf_polars/experimental/expressions.py

@@ -44,13 +44,15 @@
 from cudf_polars.dsl.expressions.base import Col, ExecutionContext, NamedExpr
 from cudf_polars.dsl.expressions.binaryop import BinOp
 from cudf_polars.dsl.expressions.literal import Literal
+from cudf_polars.dsl.expressions.rolling import GroupedWindow
 from cudf_polars.dsl.expressions.ternary import Ternary
 from cudf_polars.dsl.expressions.unary import Cast, Len, UnaryFunction
 from cudf_polars.dsl.ir import Distinct, Empty, HConcat, Select
 from cudf_polars.dsl.traversal import (
     CachingVisitor,
 )
 from cudf_polars.experimental.base import PartitionInfo
+from cudf_polars.experimental.over import _decompose_grouped_window_node
 from cudf_polars.experimental.repartition import Repartition
 from cudf_polars.experimental.utils import _dynamic_planning_on
 
@@ -462,6 +464,10 @@ def _decompose_expr_node(
         )
         (expr,) = columns
         return expr, input_ir, partition_info
+    elif isinstance(expr, GroupedWindow) and _dynamic_planning_on(config_options):
+        return _decompose_grouped_window_node(
+            expr, input_ir, partition_info, config_options, names=names
+        )
     else:
         # This is an un-supported expression - raise.
         raise NotImplementedError(