Support equi joins that Polars plans with a fused predicate (NVIDIA#24858)

### What Polars gives us

Polars plans an equi join with an extra, non-equality condition as a single join with a "fused predicate". For example `SQL ... WHERE a.key = b.key AND x <> y`, `join_where` with an equality and an inequality, and correlated `EXISTS` / `NOT EXISTS` subqueries such as the ones in TPC-H queries 7, 17, 19, 20 and 21. With Polars 2.0 the visitor did not expose that condition, so cudf-polars fell back to the CPU. Polars 2.1 reports the join as `("Inner", "fused_predicate")`, with the equality keys in `left_on`/`right_on` and the extra condition in `node.predicate`.

```python
import polars as pl

orders = pl.LazyFrame({"key": [1, 2, 3], "supp_nation": ["FRANCE", "GERMANY", "FRANCE"]})
customers = pl.LazyFrame({"key": [1, 2, 3], "cust_nation": ["GERMANY", "FRANCE", "FRANCE"]})
q = pl.SQLContext(orders=orders, customers=customers).execute(
    "SELECT * FROM orders, customers WHERE orders.key = customers.key AND supp_nation <> cust_nation"
)
q.collect(engine=pl.GPUEngine(raise_on_fail=True))  # before: NotImplementedError, now: runs on the GPU
```

### The obvious approach does not scale

The first version of this PR translated the join as an equi join followed by a filter. That is simple, but the join writes out every pair of rows that match on the keys before the filter sees any of them. If 30,000 rows on each side share one key, that is 900 million pairs, even when the predicate rejects all of them.

```python
n = 30_000
left = pl.LazyFrame({"k": [1] * n, "a": range(n)})
right = pl.LazyFrame({"k": [1] * n, "b": range(n)})
q = pl.SQLContext(left=left, right=right).execute(
    "SELECT * FROM left, right WHERE left.k = right.k AND left.a + right.b < 0"
)
q.collect(engine=pl.GPUEngine(raise_on_fail=True))
# join then filter: out of memory (a 7.2 GB allocation) on a 32 GB GPU
# this PR: 0 rows
```

### What this PR does

libcudf already has a mixed join. It builds a hash table on the keys, finds the rows with the same key, and tests the extra condition on the spot. It only writes out the pairs that pass, so memory scales with the answer, not with the number of key matches. `Join` now takes an optional predicate and evaluates it with `plc.join.mixed_inner_join`.

The catch is that the condition must be written in libcudf's AST. The AST is strict: no implicit casts, both operands of an operation must have the same type, decimal128 only supports comparisons, and there are no string functions. Polars conditions are arbitrary expressions, such as `price > 1.2 * avg_price` (TPC-DS query 6) or `lower(a) <> lower(b)`.

So the condition is split. Any part that uses only one input is computed once per row on that input, and added as an extra column to the table the AST looks at. The condition is rewritten so the AST only compares columns of the two inputs. That is 60,000 evaluations in the example above, not 900 million, and it should also be cheaper than evaluating the same expression for every pair (I have not benchmarked that). Conditions that cannot be split this way fall back to the CPU: a non-pointwise expression on one input, or operands that still differ in type.

## Changes

- `dsl/ir.py`: `Join` takes an optional predicate (inner joins only) and calls `plc.join.mixed_inner_join` when there is one. Gather maps, row ordering and the streaming paths are unchanged.
- `dsl/to_ast.py`: `split_join_predicate` does the split described above.
- `dsl/translate.py`: `_translate_fused_predicate` builds that `Join`. Other join kinds fall back to the CPU (I only saw a predicate fused into an inner join, including when the original query used an outer join and Polars rewrote it).
- `streaming/join.py`, `streaming/filter_hint.py`, `streaming/actor_graph/join.py`: pass the predicate along when a join is rebuilt. Prefilters are not used for a join with a predicate.
- `FUSED_PREDICATE_XFAIL` in `tests/test_join.py` is now conditional on Polars 2.0, where the visitor still hides the predicate (uses `POLARS_VERSION_LT_21` from NVIDIA#24825).
- Tests, all skipped on Polars 2.0:
  - `tests/test_join_fused_predicate.py`: the SQL query above, an inner join followed by a filter (`nulls_equal` both ways, with NULL keys that pass the predicate so the two settings differ), `join_where`, suffixed columns, correlated `EXISTS` and `NOT EXISTS`, an empty result, a decimal compared with a computed float (the TPC-DS query 6 shape), expressions of each input, a repeated expression, a string function, mismatched integer widths, a check that no `Filter` is left in the plan, and the 30,000 row case above.
  - `tests/dsl/test_join_predicate.py`: unit tests for the split and for the cases that fall back.

## Notes for reviewers

- This is partly a workaround for limits of libcudf's AST. The hoisting of one-sided work is worth doing anyway, but doing it for type and string reasons is not ideal. From reading the libcudf source (not tested), it has AST casts to int64, uint64 and float64 and the float64 cast accepts decimal sources, but pylibcudf does not expose them and they do not cover narrower types, strings or decimal128 arithmetic. Exposing them, and better type promotion in the AST, would be worth a follow-up. It would shrink the split but not replace it. `ConditionalJoin` already works around the same limits by pre-casting decimal and float columns.
- `Join` gains a positional `predicate` argument, so every `Join(...)` call site changed.
- The 30,000 row test is skipped on spmd because it takes 36 s with many small partitions.

## Verification

| Check | Result |
| --- | --- |
| `tests/test_join_fused_predicate.py` on Polars 2.1.0, in-memory, dask and spmd | 59 passed, 2 skipped |
| `tests/test_join_fused_predicate.py`, `tests/dsl/test_join_predicate.py`, `tests/test_join.py` | 633 passed, 173 skipped, 5 xfailed |
| New tests against the join-then-filter version | `test_predicate_is_evaluated_inside_the_join` and the 30,000 row test fail (out of memory, 7.2 GB) |
| 30,000 rows per side, one shared key, predicate rejects all (V100, 32 GB) | before: out of memory. now: 0 rows on the GPU |
| `test_join.py`, `test_join_fused_predicate.py`, `test_sql.py`, `tests/dsl`, `tests/streaming` on this branch and on the join-then-filter version | the same 318 failures in both, with identical test ids. They are environment failures on my host (TPC-DS and TPC-H data, hive scans, metadata, a libcudf JIT cache directory). 2,683 passed here, 2,616 before |
| TPC-DS query 6 (it broke with a first version of the predicate join, fixed by the split) | passes |
| Coverage of the added lines | all covered except `JoinWithPrefilter.predicate = None`, which only runs with dynamic planning (rapidsmpf) |
| ruff and mypy | `check-style` passed in CI after the mypy fix; later commits are small refactors, CI re-running |
| Polars 2.0, upstream Polars tests, TPC-H 7, 17, 19, 20, 21 and 22, dynamic planning, speed against join then filter | not run after this change |

Authors:
  - Matthew Murray (https://github.com/Matt711)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)
  - Matthew Roeschke (https://github.com/mroeschke)

URL: NVIDIA#24858

Author: Matt711

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -50,6 +50,7 @@
 from cudf_polars.dsl.nodebase import Node
 from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter
 from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars
+from cudf_polars.dsl.traversal import traversal
 from cudf_polars.dsl.utils.naming import unique_names
 from cudf_polars.dsl.utils.per_path import PerPathValues
 from cudf_polars.dsl.utils.reshape import broadcast
@@ -80,6 +81,7 @@
     from rmm.pylibrmm.stream import Stream
 
     from cudf_polars.containers.dataframe import NamedColumn
+    from cudf_polars.dsl.to_ast import JoinPredicate
     from cudf_polars.dsl.utils.io import CachedParquetInfo
     from cudf_polars.quent._context import (
         ActorBoundQuentIRExecutionState,
@@ -2979,6 +2981,37 @@ def find(key: _ColumnKey) -> _ColumnKey:
     return left_casts, right_casts
 
 
+def _check_ast_operand_types(predicate: expr.Expr) -> None:
+    """
+    Check that libcudf's AST can evaluate a predicate.
+
+    The AST requires both operands of a binary operation to share a type id,
+    and it only fails once the join runs, so reject a predicate whose operands
+    differ in type (after the decimal and float casts) when translating.
+    """
+    left_casts, right_casts = _collect_decimal_binop_casts(predicate)
+
+    def type_id(node: expr.Expr) -> plc.TypeId:
+        if isinstance(node, expr.ColRef):
+            (col,) = node.children
+            assert isinstance(col, expr.Col)
+            casts = (
+                left_casts
+                if node.table_ref == plc_expr.TableReference.LEFT
+                else right_casts
+            )
+            if col.name in casts:
+                return casts[col.name].plc_type.id()
+        return node.dtype.plc_type.id()
+
+    for node in traversal([predicate]):
+        if (
+            isinstance(node, expr.BinOp)
+            and len({type_id(c) for c in node.children}) > 1
+        ):
+            raise NotImplementedError("Join predicate with operands of different types")
+
+
 def _apply_casts(df: DataFrame, casts: dict[str, DataType]) -> DataFrame:
     if not casts:
         return df
@@ -3127,18 +3160,24 @@ def do_evaluate(
 class Join(IR):
     """A join of two dataframes."""
 
-    __slots__ = ("left_on", "options", "right_on")
+    __slots__ = ("left_on", "options", "predicate", "right_on")
     _non_child: ClassVar[tuple[str, ...]] = (
         "schema",
         "left_on",
         "right_on",
         "options",
+        "predicate",
     )
-    _n_non_child_args = 3
+    _n_non_child_args = 4
     left_on: tuple[expr.NamedExpr, ...]
     """List of expressions used as keys in the left frame."""
     right_on: tuple[expr.NamedExpr, ...]
     """List of expressions used as keys in the right frame."""
+    predicate: JoinPredicate | None
+    """
+    Optional extra condition (inner joins only) that a pair of rows must
+    satisfy to match, evaluated while joining.
+    """
     options: tuple[
         Literal["Inner", "Left", "Right", "Full", "Semi", "Anti", "Cross"],
         bool,
@@ -3176,6 +3215,7 @@ def __init__(
         left_on: Sequence[expr.NamedExpr],
         right_on: Sequence[expr.NamedExpr],
         options: Any,
+        predicate: JoinPredicate | None,
         left: IR,
         right: IR,
     ):
@@ -3184,7 +3224,24 @@ def __init__(
         self.right_on = tuple(right_on)
         self.options = options
         self.children = (left, right)
-        self._non_child_args = (self.left_on, self.right_on, self.options)
+        predicate_args = None
+        if predicate is not None:
+            if options[0] != "Inner":
+                raise NotImplementedError(f"Join predicate with join type {options[0]}")
+            ast_predicate = _strip_predicate_casts(predicate.predicate)
+            _check_ast_operand_types(ast_predicate)
+            predicate_args = (
+                ConditionalJoin.Predicate(ast_predicate),
+                predicate.left_exprs,
+                predicate.right_exprs,
+            )
+        self.predicate = predicate
+        self._non_child_args = (
+            self.left_on,
+            self.right_on,
+            self.options,
+            predicate_args,
+        )
 
     @staticmethod
     @cache
@@ -3370,6 +3427,12 @@ def do_evaluate(
             bool,
             Literal["none", "left", "right", "left_right", "right_left"],
         ],
+        predicate_args: tuple[
+            ConditionalJoin.Predicate,
+            tuple[expr.NamedExpr, ...],
+            tuple[expr.NamedExpr, ...],
+        ]
+        | None,
         left: DataFrame,
         right: DataFrame,
         *,
@@ -3442,6 +3505,19 @@ def do_evaluate(
                 ),
                 stream=right.stream,
             )
+            if predicate_args is not None:
+                predicate_wrapper, left_exprs, right_exprs = predicate_args
+                left_casts, right_casts = _collect_decimal_binop_casts(
+                    predicate_wrapper.predicate
+                )
+                left_conditional = _apply_casts(
+                    HStack.do_evaluate(left_exprs, True, left, context=context),  # noqa: FBT003
+                    left_casts,
+                )
+                right_conditional = _apply_casts(
+                    HStack.do_evaluate(right_exprs, True, right, context=context),  # noqa: FBT003
+                    right_casts,
+                )
             with context.stream_ordered_after(left, right) as stream:
                 null_equality = (
                     plc.types.NullEquality.EQUAL
@@ -3464,9 +3540,20 @@ def do_evaluate(
                         left, right = right, left
                         left_on, right_on = right_on, left_on
                         maintain_order = Join.SWAPPED_ORDER[maintain_order]
-                    lg, rg = join_fn(
-                        left_on.table, right_on.table, null_equality, stream=stream
-                    )
+                    if predicate_args is None:
+                        lg, rg = join_fn(
+                            left_on.table, right_on.table, null_equality, stream=stream
+                        )
+                    else:
+                        lg, rg = plc.join.mixed_inner_join(
+                            left_on.table,
+                            right_on.table,
+                            left_conditional.table,
+                            right_conditional.table,
+                            predicate_wrapper.ast,
+                            null_equality,
+                            stream=stream,
+                        )
                     if (
                         how in ("Inner", "Left", "Right", "Full")
                         and maintain_order != "none"

python/cudf_polars/cudf_polars/dsl/to_ast.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 from functools import partial, reduce, singledispatch
-from typing import TYPE_CHECKING, TypedDict, cast
+from typing import TYPE_CHECKING, NamedTuple, TypedDict, cast
 
 import polars as pl
 
@@ -16,10 +16,12 @@
 from cudf_polars.containers import DataType
 from cudf_polars.dsl import expr
 from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged, traversal
+from cudf_polars.dsl.utils.naming import unique_names
+from cudf_polars.dsl.utils.replace import replace
 from cudf_polars.typing import GenericTransformer
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Generator, Mapping, Sequence
 
     from rmm.pylibrmm.stream import Stream
 
@@ -112,9 +114,35 @@ class ExprTransformerState(TypedDict):
     table_ref: plc.expressions.TableReference
 
 
+class JoinPredicateState(TypedDict):
+    """
+    State used for splitting the extra condition of a join.
+
+    Parameters
+    ----------
+    columns
+        For each output column of the join, the table it comes from and its
+        name there.
+    indices
+        For each table, the index of each of its columns.
+    names
+        For each table, a generator of names for columns to append to it.
+    extras
+        For each table, the expressions to append to it and their names.
+    """
+
+    columns: Mapping[str, tuple[plc_expr.TableReference, str]]
+    indices: Mapping[plc_expr.TableReference, Mapping[str, int]]
+    names: Mapping[plc_expr.TableReference, Generator[str, None, None]]
+    extras: Mapping[plc_expr.TableReference, dict[expr.Expr, str]]
+
+
 type Transformer = GenericTransformer[expr.Expr, plc_expr.Expression, ASTState]
 type ExprTransformer = GenericTransformer[expr.Expr, expr.Expr, ExprTransformerState]
 """Protocol for transformation of Expr nodes."""
+type JoinPredicateTransformer = GenericTransformer[
+    expr.Expr, expr.Expr, JoinPredicateState
+]
 
 
 @singledispatch
@@ -441,3 +469,109 @@ def insert_colrefs(
         state={"name_to_index": name_to_index, "table_ref": table_ref},
     )
     return mapper(node)
+
+
+class JoinPredicate(NamedTuple):
+    """
+    An extra condition of an equi join, split for libcudf's mixed join.
+
+    Parts of the condition that use columns of one input only are evaluated
+    on that input before the join. Only the remaining comparison between the
+    inputs is left to the AST.
+    """
+
+    predicate: expr.Expr
+    """Expression with column references into the left and right tables."""
+    left_exprs: tuple[expr.NamedExpr, ...]
+    """Columns appended to the left table, evaluated on the left input."""
+    right_exprs: tuple[expr.NamedExpr, ...]
+    """Columns appended to the right table, evaluated on the right input."""
+
+
+def _split_join_predicate(node: expr.Expr, rec: JoinPredicateTransformer) -> expr.Expr:
+    state = rec.state
+    if isinstance(node, expr.Col):
+        table_ref, name = state["columns"][node.name]
+        return expr.ColRef(
+            node.dtype,
+            state["indices"][table_ref][name],
+            table_ref,
+            expr.Col(node.dtype, name),
+        )
+    sides = {
+        state["columns"][col.name][0]
+        for col in traversal([node])
+        if isinstance(col, expr.Col)
+    }
+    if len(sides) == 1:
+        (table_ref,) = sides
+        if not node.all_pointwise():
+            raise NotImplementedError("Join predicate with a non-pointwise expression")
+        renames: dict[expr.Expr, expr.Expr] = {
+            col: expr.Col(col.dtype, state["columns"][col.name][1])
+            for col in traversal([node])
+            if isinstance(col, expr.Col)
+        }
+        (local,) = replace([node], renames)
+        extras = state["extras"][table_ref]
+        if local not in extras:
+            extras[local] = next(state["names"][table_ref])
+        index = len(state["indices"][table_ref]) + list(extras).index(local)
+        return expr.ColRef(
+            node.dtype, index, table_ref, expr.Col(node.dtype, extras[local])
+        )
+    return reuse_if_unchanged(node, rec)
+
+
+def split_join_predicate(
+    node: expr.Expr,
+    *,
+    columns: Mapping[str, tuple[plc_expr.TableReference, str]],
+    left_names: Sequence[str],
+    right_names: Sequence[str],
+) -> JoinPredicate:
+    """
+    Split an equi join's extra condition for libcudf's mixed join.
+
+    Parameters
+    ----------
+    node
+        Condition over the columns of the join's output.
+    columns
+        For each output column, the table it comes from and its name there.
+    left_names, right_names
+        Column names of the left and right inputs.
+
+    Returns
+    -------
+    The condition with column references inserted, and the columns to add
+    to each input so that every comparison in the condition is between
+    columns of the two inputs.
+
+    Raises
+    ------
+    NotImplementedError
+        If a part of the condition that uses one input cannot be evaluated
+        on that input.
+    """
+    left = plc_expr.TableReference.LEFT
+    right = plc_expr.TableReference.RIGHT
+    extras: dict[plc_expr.TableReference, dict[expr.Expr, str]] = {left: {}, right: {}}
+    mapper: JoinPredicateTransformer = CachingVisitor(
+        _split_join_predicate,
+        state={
+            "columns": columns,
+            "indices": {
+                left: {name: i for i, name in enumerate(left_names)},
+                right: {name: i for i, name in enumerate(right_names)},
+            },
+            "names": {left: unique_names(left_names), right: unique_names(right_names)},
+            "extras": extras,
+        },
+    )
+    predicate = mapper(node)
+    return JoinPredicate(
+        predicate,
+        tuple(expr.NamedExpr(name, e) for e, name in extras[left].items()),
+        tuple(expr.NamedExpr(name, e) for e, name in extras[right].items()),
+    )