[DO NOT MERGE] PDS DS ALL

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],
+        names: Sequence[str],
+        dtypes: Sequence[DataType],
         stream: Stream,
     ) -> list[Column]:
         # We do a left-join between the input keys to group-keys

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/benchmarks/pdsds_queries/q14.py

@@ -335,65 +335,66 @@ def polars_impl(run_config: RunConfig) -> QueryResult:
     item = get_data(run_config.dataset_path, "item", run_config.suffix)
     date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix)
 
+    cross_items = build_cross_items(
+        store_sales, catalog_sales, web_sales, item, date_dim, year=year
+    )
+    average_sales = build_average_sales(
+        store_sales, catalog_sales, web_sales, date_dim, year=year
+    )
+
+    # week_dates is ≤7 rows (one calendar week), computed once as a 1-partition frame.
+    # Push the week filter into each channel before the UNION via a semi-join so that
+    # ~99% of rows (everything outside the target week) are dropped before the
+    # expensive cross_items join and groupby.
+    target_week = (
+        date_dim.filter(
+            (pl.col("d_year") == year + 1)
+            & (pl.col("d_moy") == 12)
+            & (pl.col("d_dom") == day)
+        )
+        .select("d_week_seq")
+        .unique()
+    )
+    week_dates = date_dim.join(target_week, on="d_week_seq").select("d_date_sk")
+
     all_sales = pl.concat(
         [
-            store_sales.select(
+            store_sales.join(
+                week_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi"
+            ).select(
                 [
                     pl.lit("store").alias("channel"),
                     pl.col("ss_item_sk").alias("item_sk"),
                     pl.col("ss_quantity").alias("quantity"),
                     pl.col("ss_list_price").alias("list_price"),
-                    pl.col("ss_sold_date_sk").alias("date_sk"),
                 ]
             ),
-            catalog_sales.select(
+            catalog_sales.join(
+                week_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi"
+            ).select(
                 [
                     pl.lit("catalog").alias("channel"),
                     pl.col("cs_item_sk").alias("item_sk"),
                     pl.col("cs_quantity").alias("quantity"),
                     pl.col("cs_list_price").alias("list_price"),
-                    pl.col("cs_sold_date_sk").alias("date_sk"),
                 ]
             ),
-            web_sales.select(
+            web_sales.join(
+                week_dates, left_on="ws_sold_date_sk", right_on="d_date_sk", how="semi"
+            ).select(
                 [
                     pl.lit("web").alias("channel"),
                     pl.col("ws_item_sk").alias("item_sk"),
                     pl.col("ws_quantity").alias("quantity"),
                     pl.col("ws_list_price").alias("list_price"),
-                    pl.col("ws_sold_date_sk").alias("date_sk"),
                 ]
             ),
         ]
     )
 
-    cross_items = build_cross_items(
-        store_sales, catalog_sales, web_sales, item, date_dim, year=year
-    )
-    average_sales = build_average_sales(
-        store_sales, catalog_sales, web_sales, date_dim, year=year
-    )
-
-    # d_week_seq target is the same for all 3 channels; compute it once.
-    target_week = (
-        date_dim.filter(
-            (pl.col("d_year") == year + 1)
-            & (pl.col("d_moy") == 12)
-            & (pl.col("d_dom") == day)
-        )
-        .select("d_week_seq")
-        .unique()
-    )
-    week_dates = date_dim.join(target_week, on="d_week_seq").select("d_date_sk")
-
-    # Build y: all 3 channels in a single pipeline.
-    # cross_items and average_sales each appear once — no CSE needed.
-    # After group_by the frame is tiny, so the cross join with the 1-row
-    # average_sales frame is negligible even if Polars fuses it into an IEJoin.
     y = (
         all_sales.join(cross_items, left_on="item_sk", right_on="ss_item_sk")
         .join(item, left_on="item_sk", right_on="i_item_sk")
-        .join(week_dates, left_on="date_sk", right_on="d_date_sk")
         .group_by(["channel", "i_brand_id", "i_class_id", "i_category_id"])
         .agg(
             [

python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.py

@@ -110,38 +110,103 @@ def polars_impl(run_config: RunConfig) -> QueryResult:
     sort_by = {"i_item_id": False, "i_item_desc": False, "s_state": False}
     limit = 100
 
-    store_sales_base = (
+    q1 = f"{year}Q1"
+    q1_q3 = [f"{year}Q1", f"{year}Q2", f"{year}Q3"]
+
+    # Pre-filter date_dim to only qualifying d_date_sk values.
+    d1_dates = date_dim.filter(pl.col("d_quarter_name") == q1).select("d_date_sk")
+    d_q3_dates = date_dim.filter(pl.col("d_quarter_name").is_in(q1_q3)).select(
+        "d_date_sk"
+    )
+
+    # store_returns has [6] partitions — at the broadcast limit. Filter it to Q1-Q3 dates
+    # first, then use the (customer, item) pairs it contains to pre-filter both store_sales
+    # and catalog_sales before those larger tables enter the expensive shuffle joins.
+    store_returns_filtered = store_returns.join(
+        d_q3_dates, left_on="sr_returned_date_sk", right_on="d_date_sk", how="semi"
+    ).select(["sr_customer_sk", "sr_item_sk", "sr_ticket_number", "sr_return_quantity"])
+
+    # (customer, item) pairs present in any qualifying store return; stays at [6] partitions
+    # so broadcast is free. Polars will CACHE this shared subplan.
+    sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"])
+
+    store_sales_filtered = (
         store_sales.join(
-            date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk", suffix="_d1"
+            d1_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi"
+        )
+        .join(
+            sr_customer_item,
+            left_on=["ss_customer_sk", "ss_item_sk"],
+            right_on=["sr_customer_sk", "sr_item_sk"],
+            how="semi",
+        )
+        .select(
+            [
+                "ss_customer_sk",
+                "ss_item_sk",
+                "ss_store_sk",
+                "ss_ticket_number",
+                "ss_quantity",
+            ]
+        )
+        .join(
+            item.select(["i_item_sk", "i_item_id", "i_item_desc"]),
+            left_on="ss_item_sk",
+            right_on="i_item_sk",
+        )
+        .join(
+            store.select(["s_store_sk", "s_state"]),
+            left_on="ss_store_sk",
+            right_on="s_store_sk",
+        )
+        .select(
+            [
+                "ss_customer_sk",
+                "ss_item_sk",
+                "ss_ticket_number",
+                "ss_quantity",
+                "i_item_id",
+                "i_item_desc",
+                "s_state",
+            ]
         )
-        .join(item, left_on="ss_item_sk", right_on="i_item_sk")
-        .join(store, left_on="ss_store_sk", right_on="s_store_sk")
-        .filter(pl.col("d_quarter_name") == f"{year}Q1")
     )
 
-    store_returns_base = store_returns.join(
-        date_dim, left_on="sr_returned_date_sk", right_on="d_date_sk", suffix="_d2"
-    ).filter(pl.col("d_quarter_name").is_in([f"{year}Q1", f"{year}Q2", f"{year}Q3"]))
-
-    catalog_sales_base = catalog_sales.join(
-        date_dim, left_on="cs_sold_date_sk", right_on="d_date_sk", suffix="_d3"
-    ).filter(pl.col("d_quarter_name").is_in([f"{year}Q1", f"{year}Q2", f"{year}Q3"]))
+    catalog_sales_filtered = (
+        catalog_sales.join(
+            d_q3_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi"
+        )
+        .join(
+            sr_customer_item,
+            left_on=["cs_bill_customer_sk", "cs_item_sk"],
+            right_on=["sr_customer_sk", "sr_item_sk"],
+            how="semi",
+        )
+        .select(["cs_bill_customer_sk", "cs_item_sk", "cs_quantity"])
+    )
 
     return QueryResult(
         frame=(
-            store_sales_base.join(
-                store_returns_base,
+            store_sales_filtered.join(
+                store_returns_filtered,
                 left_on=["ss_customer_sk", "ss_item_sk", "ss_ticket_number"],
                 right_on=["sr_customer_sk", "sr_item_sk", "sr_ticket_number"],
-                how="inner",
-                suffix="_sr",
+            )
+            .select(
+                [
+                    "ss_customer_sk",
+                    "ss_item_sk",
+                    "ss_quantity",
+                    "sr_return_quantity",
+                    "i_item_id",
+                    "i_item_desc",
+                    "s_state",
+                ]
             )
             .join(
-                catalog_sales_base,
+                catalog_sales_filtered,
                 left_on=["ss_customer_sk", "ss_item_sk"],
                 right_on=["cs_bill_customer_sk", "cs_item_sk"],
-                how="inner",
-                suffix="_cs",
             )
             .group_by(["i_item_id", "i_item_desc", "s_state"])
             .agg(