Fix silent row drops in multi-GPU joins with computed key expressions

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

@@ -7,10 +7,10 @@
 
 from typing import TYPE_CHECKING
 
-from cudf_polars.dsl.expr import NamedExpr
+from cudf_polars.dsl.expr import Col, NamedExpr
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Iterable
+    from collections.abc import Generator, Iterable, Sequence
 
     from cudf_polars.typing import Schema
 
@@ -38,8 +38,26 @@ def unique_names(names: Iterable[str]) -> Generator[str, None, None]:
         i += 1
 
 
+def _concrete_prefix(names: Sequence[str | NamedExpr]) -> tuple[str, ...]:
+    # Exclude NamedExprs that are not concrete Col references.
+    # We don't throw out the entire NamedExpr tuple if a prefix
+    # of the tuple is concrete.
+    prefix: list[str] = []
+    for name in names:
+        if isinstance(name, str):
+            prefix.append(name)
+        elif isinstance(name.value, Col):
+            prefix.append(name.value.name)
+        else:
+            break
+    return tuple(prefix)
+
+
 def names_to_indices(
-    names: tuple[str | NamedExpr, ...], schema: Schema
+    names: tuple[str | NamedExpr, ...],
+    schema: Schema,
+    *,
+    concrete_prefix: bool = False,
 ) -> tuple[int, ...]:
     """
     Return column indices for the given names in schema order.
@@ -53,11 +71,17 @@ def names_to_indices(
         The names to get indices for.
     schema
         The schema to get indices from.
+    concrete_prefix
+        If True, use only the prefix of names corresponding
+        to concrete column references. If False (default),
+        use all names.
 
     Returns
     -------
     The column indices for each name in schema order.
     """
     keys = list(schema.keys())
+    if concrete_prefix:
+        names = _concrete_prefix(names)
     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/streaming/actor_graph/groupby.py

@@ -23,7 +23,7 @@
 from cudf_polars.containers import DataType
 from cudf_polars.dsl.expr import Col, NamedExpr
 from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select
-from cudf_polars.dsl.utils.naming import unique_names
+from cudf_polars.dsl.utils.naming import names_to_indices, unique_names
 from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.actor_graph.dispatch import (
     generate_ir_sub_network,
@@ -150,7 +150,6 @@ def from_ir(cls, ir: GroupBy | Distinct) -> DecomposedGroupBy:
         else:  # pragma: no cover
             raise TypeError(f"Unsupported IR type: {type(ir)}")
 
-        # Distinguish between output and shuffle indices
         output_indices = _key_indices(ir, ir.schema)
         if isinstance(ir, Distinct):
             shuffle_indices = output_indices
@@ -485,9 +484,13 @@ def _enforce_schema(
     )
 
 
-def _key_indices(ir: GroupBy | Distinct, schema: Schema) -> tuple[int, ...]:
+def _key_indices(
+    ir: GroupBy | Distinct, schema: Schema, *, concrete_prefix: bool = False
+) -> tuple[int, ...]:
     schema_keys = {n: i for i, n in enumerate(schema.keys())}
     if isinstance(ir, GroupBy):
+        if concrete_prefix:
+            return names_to_indices(ir.keys, schema, concrete_prefix=True)
         groupby_key_names = tuple(ne.name for ne in ir.keys)
         if not all(k in schema_keys for k in groupby_key_names):
             return ()
@@ -638,7 +641,7 @@ async def groupby_actor(
         partitioning = NormalizedPartitioning.from_keys(
             metadata_in.partitioning,
             nranks,
-            keys=_key_indices(ir, ir.children[0].schema),
+            keys=_key_indices(ir, ir.children[0].schema, concrete_prefix=True),
         )
         require_tree = _require_tree(ir)
         fully_partitioned = partitioning.is_strictly_partitioned()

python/cudf_polars/cudf_polars/streaming/actor_graph/join.py

@@ -1048,12 +1048,12 @@ async def _choose_strategy(
     left_partitioning = NormalizedPartitioning.from_keys(
         left_metadata.partitioning,
         nranks,
-        keys=names_to_indices(ir.left_on, ir.children[0].schema),
+        keys=names_to_indices(ir.left_on, ir.children[0].schema, concrete_prefix=True),
     )
     right_partitioning = NormalizedPartitioning.from_keys(
         right_metadata.partitioning,
         nranks,
-        keys=names_to_indices(ir.right_on, ir.children[1].schema),
+        keys=names_to_indices(ir.right_on, ir.children[1].schema, concrete_prefix=True),
     )
 
     if left_partitioning.is_aligned_with(right_partitioning, context.br()):

python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py

@@ -386,7 +386,6 @@ def pytest_report_header(config: pytest.Config) -> str:
     "tests/unit/functions/test_concat.py::test_concat_horizontal_zero_width_height_mismatch_26876": "https://github.com/rapidsai/cudf/issues/21644",
     "tests/unit/functions/test_concat.py::test_concat_horizontally_strict": "Correct polars.exceptions.ShapeError raised but it's in a ExceptionGroup",
     "tests/unit/interop/test_interop.py::test_0_width_df_roundtrip": "https://github.com/rapidsai/cudf/issues/21644",
-    "tests/unit/lazyframe/test_projections.py::test_join_projection_pushdown_struct_field_as_key_24446": "https://github.com/rapidsai/cudf/issues/22105",
     "tests/unit/operations/test_slice.py::test_slice_pushdown_literal_projection_14349": "https://github.com/rapidsai/cudf/issues/22072",
     "tests/unit/operations/test_group_by.py::test_group_by_lit_series": "Incorrect broadcasting of literals in groupby-agg",
     "tests/unit/operations/test_group_by.py::test_group_by_series_partitioned": "https://github.com/rapidsai/cudf/issues/22072",

python/cudf_polars/tests/dsl/test_naming.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import polars as pl
+
+from cudf_polars.containers import DataType
+from cudf_polars.dsl import expr
+from cudf_polars.dsl.utils.naming import names_to_indices
+
+
+def test_names_to_indices_concrete_prefix() -> None:
+    dtype = DataType(pl.Int64())
+    schema = {"a": dtype, "b": dtype, "c": dtype}
+    names = (
+        expr.NamedExpr("a_alias", expr.Col(dtype, "a")),
+        "b",
+        expr.NamedExpr("computed", expr.Literal(dtype, 1)),
+        expr.NamedExpr("c_alias", expr.Col(dtype, "c")),
+    )
+
+    assert names_to_indices(names, schema, concrete_prefix=True) == (0, 1)

python/cudf_polars/tests/streaming/test_join.py

@@ -388,3 +388,59 @@ def test_dynamic_planning_skips_compile_time_partition_wise_join():
         right_ir: PartitionInfo(1, partitioned_on=()),
     }
     assert not _use_pwise_join(executor, partition_info, join_ir)
+
+
+def test_join_computed_expr_right_key(streaming_engine_factory) -> None:
+    """Join on a computed key expression."""
+    engine = streaming_engine_factory(
+        StreamingOptions(
+            target_partition_size=1,
+            max_rows_per_partition=4,
+            broadcast_limit=1,  # Disable broadcast joins
+        ),
+    )
+    if engine.nranks < 2:
+        pytest.skip("bug only manifests on 2+ ranks")
+
+    zip_prefixes = ["10", "20", "30", "40"]
+    full_zips = ["10001", "20001", "30001", "40001"]
+    reps = 4
+
+    # Start with joins on concrete column references
+    # to establish left and right partitioning metadata.
+    left_a = pl.LazyFrame(
+        {
+            "zip_prefix": zip_prefixes * reps,
+            "val_a": list(range(len(zip_prefixes) * reps)),
+        }
+    )
+    left_b = pl.LazyFrame(
+        {
+            "zip_prefix": zip_prefixes * reps,
+            "val_b": list(range(100, 100 + len(zip_prefixes) * reps)),
+        }
+    )
+    left = left_a.join(left_b, on="zip_prefix", how="inner")
+
+    right_a = pl.LazyFrame(
+        {
+            "full_zip": full_zips * reps,
+            "val_c": list(range(200, 200 + len(full_zips) * reps)),
+        }
+    )
+    right_b = pl.LazyFrame(
+        {
+            "full_zip": full_zips * reps,
+            "val_d": list(range(300, 300 + len(full_zips) * reps)),
+        }
+    )
+    right = right_a.join(right_b, on="full_zip", how="inner")
+
+    # Now join on a computed key expression.
+    # This should not silently drop rows across ranks
+    q = left.join(
+        right,
+        left_on="zip_prefix",
+        right_on=pl.col("full_zip").str.slice(0, 2),
+    )
+    assert_gpu_result_equal(q, engine=engine, check_row_order=False)