Add plan optimizer pass that extracts join keys and uses them to construct pre-filters for inner joins

docs/cudf/source/cudf_polars/api.md

@@ -76,6 +76,7 @@ Most users interact with them through `StreamingOptions` fields rather than dire
 .. automodule:: cudf_polars.utils.config
    :members:
       DynamicPlanningOptions,
+      JoinFilterPushdownOptions,
       MemoryResourceConfig,
       ParquetOptions,
       StreamingExecutor,

docs/cudf/source/cudf_polars/options.md

@@ -108,6 +108,7 @@ Environment variables follow these patterns:
 | `broadcast_limit`        | Maximum number of bytes for broadcast joins.                                                                                                        | auto        |
 | `target_partition_size`  | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto.                                                                   | auto        |
 | `dynamic_planning`       | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables.                                 | enabled     |
+| `join_filter_pushdown`   | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables.        | enabled     |
 | `sink_to_directory`      | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True`      |
 
 ### Category: `engine`

python/cudf_polars/cudf_polars/dsl/utils/column_domain.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for tracking column value domains between IR nodes."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import singledispatch
+from typing import TYPE_CHECKING
+
+from cudf_polars.dsl import expr
+from cudf_polars.dsl.ir import (
+    Distinct,
+    Filter,
+    GroupBy,
+    HStack,
+    Join,
+    Projection,
+    Select,
+    Slice,
+    Sort,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from cudf_polars.dsl.ir import IR
+
+__all__ = [
+    "ColumnBinding",
+    "ColumnLineage",
+    "ColumnRef",
+    "column_domain_bindings",
+]
+
+
+@dataclass(frozen=True)
+class ColumnBinding:
+    """A direct binding to a named column on a specific child edge."""
+
+    child_index: int
+    name: str
+
+
+@dataclass(frozen=True)
+class ColumnRef:
+    """A named column produced by an IR node."""
+
+    node: IR
+    name: str
+
+
+@dataclass(frozen=True)
+class ColumnLineage:
+    """Persistent value-domain lineage, sharing suffixes across DAG branches."""
+
+    column: ColumnRef
+    source: ColumnLineage | None = None
+    source_child_index: int | None = None
+    """Child edge leading to ``source``, or None if there is no source."""
+
+
+@singledispatch
+def column_domain_bindings(node: IR) -> Mapping[str, ColumnBinding]:
+    """
+    Map output columns to child columns containing their value domains.
+
+    For every ``output_name -> ColumnBinding(child_index, input_name)`` binding,
+    every value appearing in ``node[output_name]`` is guaranteed to appear in
+    ``node.children[child_index][input_name]``. Row order, multiplicity, and
+    cardinality are not preserved.
+
+    If a name in ``node.schema`` does not appear in the mapping it means
+    that it was not possible to derive a relationship between the domain of
+    the output and input values for that column.
+    """
+    return {}
+
+
+@column_domain_bindings.register(Select)
+def _(node: Select) -> Mapping[str, ColumnBinding]:
+    return {
+        item.name: ColumnBinding(0, item.value.name)
+        for item in node.exprs
+        if isinstance(item.value, expr.Col)
+    }
+
+
+@column_domain_bindings.register(HStack)
+def _(node: HStack) -> Mapping[str, ColumnBinding]:
+    child = node.children[0]
+    replaced = {item.name for item in node.columns}
+    return {
+        name: ColumnBinding(0, name) for name in child.schema if name not in replaced
+    } | {
+        item.name: ColumnBinding(0, item.value.name)
+        for item in node.columns
+        if isinstance(item.value, expr.Col)
+    }
+
+
+@column_domain_bindings.register(GroupBy)
+def _(node: GroupBy) -> Mapping[str, ColumnBinding]:
+    return {
+        key.name: ColumnBinding(0, key.value.name)
+        for key in node.keys
+        if isinstance(key.value, expr.Col)
+    }
+
+
+@column_domain_bindings.register(Join)
+def _(node: Join) -> Mapping[str, ColumnBinding]:
+    left, right = node.children
+    how = node.options[0]
+    if how in ("Semi", "Anti"):
+        return {
+            name: ColumnBinding(0, name) for name in node.schema if name in left.schema
+        }
+    if how != "Inner":
+        return {}
+
+    bindings = {name: ColumnBinding(0, name) for name in left.schema}
+    suffix = node.options[3]
+    for name in right.schema:
+        output_name = f"{name}{suffix}" if name in left.schema else name
+        if output_name in node.schema:
+            bindings[output_name] = ColumnBinding(1, name)
+    return bindings
+
+
+@column_domain_bindings.register(Distinct)
+@column_domain_bindings.register(Filter)
+@column_domain_bindings.register(Projection)
+@column_domain_bindings.register(Slice)
+@column_domain_bindings.register(Sort)
+def _(
+    node: Distinct | Filter | Projection | Slice | Sort,
+) -> Mapping[str, ColumnBinding]:
+    child = node.children[0]
+    return {
+        name: ColumnBinding(0, name) for name in node.schema if name in child.schema
+    }

python/cudf_polars/cudf_polars/engine/core.py

@@ -734,11 +734,17 @@ def evaluate_on_rank(
     """
     stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor)
 
+    lowering, node_map = lower_ir_graph_with_node_map(
+        ir, config_options, stats, rank=comm.rank, nranks=comm.nranks
+    )
+    optimized = lowering.optimized
+    ir = lowering.lowered
+    partition_info = lowering.partition_info
     if config_options.executor.quent_context is not None:
         assert local_quent_context is not None
-        logical_plan_id = ir.get_stable_plan_id()
+        logical_plan_id = optimized.get_stable_plan_id()
         plan, ops, ports, logical_op_by_id = build_plan(
-            ir,
+            optimized,
             config_options,
             query=local_quent_context.context.query,
             plan_id=logical_plan_id,
@@ -752,10 +758,6 @@ def evaluate_on_rank(
                 local_quent_context.logger, plan, ops, ports
             )
 
-    ir, partition_info, node_map = lower_ir_graph_with_node_map(
-        ir, config_options, stats, rank=comm.rank, nranks=comm.nranks
-    )
-
     if comm.rank == 0:
         log_query_plan(ir, config_options)
 

python/cudf_polars/cudf_polars/engine/options.py

@@ -26,6 +26,7 @@
     from cudf_polars.quent import QuentContext
     from cudf_polars.utils.config import (
         DynamicPlanningOptions,
+        JoinFilterPushdownOptions,
         ParquetOptions,
     )
 
@@ -248,6 +249,14 @@ class StreamingOptions:
         Env: ``CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING``.
         Default: enabled.
         Category: executor.
+    join_filter_pushdown
+        Config for join filter pushdown optimizations, dict or
+        :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions`. ``None``
+        disables the rewrite.
+        Env: ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN`` and
+        ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__*``.
+        Default: enabled.
+        Category: executor.
     sink_to_directory
         Whether multi-partition sink operations should write to a directory
         rather than a single file. The ``spmd``/``ray``/``dask`` engines
@@ -346,6 +355,9 @@ class StreamingOptions:
     dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = (
         _opt("executor")
     )
+    join_filter_pushdown: (
+        dict[str, Any] | JoinFilterPushdownOptions | None | Unspecified
+    ) = _opt("executor")
     sink_to_directory: bool | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean
     )

python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py

@@ -38,8 +38,9 @@
 import cudf_polars.dsl.tracing
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction
-from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select
+from cudf_polars.dsl.ir import Filter, GroupBy, HStack, Join, Projection, Select
 from cudf_polars.dsl.tracing import Scope
+from cudf_polars.dsl.utils.column_domain import column_domain_bindings
 from cudf_polars.dsl.utils.naming import names_to_indices
 from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager
 from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk
@@ -446,9 +447,8 @@ def _derived_ordering(
 
 def _select_column_targets(select: Select) -> dict[str, dict[str, None]]:
     old_to_new_names: defaultdict[str, dict[str, None]] = defaultdict(dict)
-    for ne in select.exprs:
-        if isinstance(ne.value, Col):
-            old_to_new_names[ne.value.name][ne.name] = None
+    for output_name, source in column_domain_bindings(select).items():
+        old_to_new_names[source.name][output_name] = None
     return dict(old_to_new_names)
 
 
@@ -603,7 +603,7 @@ def maybe_remap_partitioning(
             ),
             local=_remap_scheme_simple(ir, partitioning.local, ir.children[0]),
         )
-    if isinstance(ir, (Cache, Join, Projection, Filter)):
+    if isinstance(ir, (Join, Projection, Filter)):
         child = child_ir if child_ir is not None else ir.children[0]
         return Partitioning(
             inter_rank=_remap_scheme_simple(ir, partitioning.inter_rank, child),

python/cudf_polars/cudf_polars/streaming/explain.py

@@ -123,8 +123,10 @@ def explain_query(
     if physical:
         with cm:
             stats = collect_statistics(ir, config, executor)
-        lowered_ir, partition_info = lower_ir_graph(ir, config, stats)
-        return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config)
+        lowered = lower_ir_graph(ir, config, stats)
+        return _repr_ir_tree(
+            lowered.lowered, lowered.partition_info, stats=stats, config=config
+        )
     else:
         if config.executor.name == "streaming":
             # Include row-count statistics for the logical plan
@@ -150,7 +152,9 @@ def collect_partition_plan(
 
     with concurrent.futures.ThreadPoolExecutor() as executor:
         stats = collect_statistics(ir, config, executor)
-    lowered_ir, partition_info = lower_ir_graph(ir, config, stats)
+    lowered = lower_ir_graph(ir, config, stats)
+    lowered_ir = lowered.lowered
+    partition_info = lowered.partition_info
 
     seen: set[tuple] = set()
     rows: list[PartitionPlanRow] = []
@@ -755,7 +759,9 @@ def from_ir(
         if lowered:
             with cm:
                 stats = collect_statistics(ir, config_options, executor)
-            ir, partition_info_d = lower_ir_graph(ir, config_options, stats)
+            lowering = lower_ir_graph(ir, config_options, stats)
+            ir = lowering.lowered
+            partition_info_d = lowering.partition_info
             partition_info_dict = {}
 
         nodes: dict[str, SerializableIRNode] = {}

python/cudf_polars/cudf_polars/streaming/join.py

@@ -75,7 +75,7 @@ def _make_hash_join(
         partition_info,
         output_count,
     )
-    # Always reconstruct in case children contain Cache nodes
+    # Reconstruct with the lowered and possibly shuffled children.
     ir = ir.reconstruct([left, right])
 
     # Record new partitioning info