feat: expose residual_expression and post_join_filter on join builders (#225)

Author: nielspardon

src/substrait/builders/plan.py

@@ -19,7 +19,11 @@
     resolve_expression,
 )
 from substrait.extension_registry import ExtensionRegistry
-from substrait.type_inference import infer_plan_schema, join_output_names
+from substrait.type_inference import (
+    _join_output_struct,
+    infer_plan_schema,
+    join_output_names,
+)
 from substrait.utils import (
     merge_extension_declarations,
     merge_extension_urns,
@@ -487,6 +491,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         left_ns = infer_plan_schema(bound_left, registry=registry)
         right_ns = infer_plan_schema(bound_right, registry=registry)
 
+        # The join condition binds against the combined left+right schema.
         ns = stt.NamedStruct(
             struct=stt.Type.Struct(
                 types=list(left_ns.struct.types) + list(right_ns.struct.types),
@@ -497,11 +502,28 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_expression: stee.ExtendedExpression = resolve_expression(
             expression, ns, registry
         )
-        bound_post = (
-            resolve_expression(post_join_filter, ns, registry)
-            if post_join_filter is not None
-            else None
-        )
+
+        # The output names must match the columns the join type actually emits
+        # (semi/anti drop a side, mark appends a boolean).
+        type_name = stalg.JoinRel.JoinType.Name(type)
+        out_names = join_output_names(type_name, left_ns.names, right_ns.names)
+
+        # post_join_filter is applied to each output record after
+        # join-type-specific output formation (semantically a FilterRel above the
+        # join), so it resolves against the output schema -- which for semi/anti
+        # joins is a single side, not the combined schema.
+        bound_post = None
+        if post_join_filter is not None:
+            output_ns = stt.NamedStruct(
+                names=out_names,
+                struct=_join_output_struct(
+                    type_name,
+                    bound_left.relations[-1].root.input,
+                    bound_right.relations[-1].root.input,
+                    registry=registry,
+                ),
+            )
+            bound_post = resolve_expression(post_join_filter, output_ns, registry)
 
         rel = stalg.Rel(
             join=stalg.JoinRel(
@@ -516,12 +538,6 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             )
         )
 
-        # The join condition resolves against the combined left+right schema, but
-        # the output names must match the columns the join type actually emits
-        # (semi/anti drop a side, mark appends a boolean).
-        out_names = join_output_names(
-            stalg.JoinRel.JoinType.Name(type), left_ns.names, right_ns.names
-        )
         return stp.Plan(
             version=default_version,
             relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=out_names))],
@@ -1023,6 +1039,9 @@ def builder(
         left_keys: Iterable[Union[str, int]],
         right_keys: Iterable[Union[str, int]],
         type,
+        *,
+        post_join_filter: Optional[ExtendedExpressionOrUnbound] = None,
+        residual_expression: Optional[ExtendedExpressionOrUnbound] = None,
         extension: Optional[AdvancedExtension] = None,
     ) -> UnboundPlan:
         def resolve(registry: ExtensionRegistry) -> stp.Plan:
@@ -1033,24 +1052,65 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             keys = _comparison_join_keys(
                 list(left_keys), list(right_keys), left_ns, right_ns, registry
             )
-            names = join_output_names(
-                rel_cls.JoinType.Name(type), left_ns.names, right_ns.names
-            )
+            type_name = rel_cls.JoinType.Name(type)
+            names = join_output_names(type_name, left_ns.names, right_ns.names)
+
+            # post_join_filter is applied to each output record after
+            # join-type-specific output formation (semantically a FilterRel above
+            # the join), so it resolves against the output schema -- which for
+            # semi/anti joins is a single side. residual_expression is evaluated
+            # on each candidate key-match (both rows present), so it resolves
+            # against the combined left+right schema. Each is built only when the
+            # corresponding predicate is supplied.
+            bound_post = None
+            if post_join_filter is not None:
+                output_ns = stt.NamedStruct(
+                    names=names,
+                    struct=_join_output_struct(
+                        type_name,
+                        bound_left.relations[-1].root.input,
+                        bound_right.relations[-1].root.input,
+                        registry=registry,
+                    ),
+                )
+                bound_post = resolve_expression(post_join_filter, output_ns, registry)
+
+            bound_residual = None
+            if residual_expression is not None:
+                combined_ns = stt.NamedStruct(
+                    struct=stt.Type.Struct(
+                        types=list(left_ns.struct.types) + list(right_ns.struct.types),
+                        nullability=stt.Type.Nullability.NULLABILITY_REQUIRED,
+                    ),
+                    names=list(left_ns.names) + list(right_ns.names),
+                )
+                bound_residual = resolve_expression(
+                    residual_expression, combined_ns, registry
+                )
+
             rel = stalg.Rel(
                 **{
                     rel_name: rel_cls(
                         left=bound_left.relations[-1].root.input,
                         right=bound_right.relations[-1].root.input,
                         keys=keys,
                         type=type,
+                        post_join_filter=bound_post.referred_expr[0].expression
+                        if bound_post
+                        else None,
+                        residual_expression=bound_residual.referred_expr[0].expression
+                        if bound_residual
+                        else None,
                         advanced_extension=extension,
                     )
                 }
             )
             return stp.Plan(
                 version=default_version,
                 relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-                **_merge_plan_metadata(bound_left, bound_right),
+                **_merge_plan_metadata(
+                    bound_left, bound_right, bound_post, bound_residual
+                ),
             )
 
         return resolve

src/substrait/dataframe/frame.py

@@ -406,7 +406,17 @@ def nested_loop_join(
             )
         )
 
-    def _equi_join(self, builder, rel_cls, other, left_on, right_on, how):
+    def _equi_join(
+        self,
+        builder,
+        rel_cls,
+        other,
+        left_on,
+        right_on,
+        how,
+        post_filter,
+        residual,
+    ):
         if how not in _JOIN_TYPES:
             raise ValueError(
                 f"unknown join type {how!r}; expected one of {sorted(_JOIN_TYPES)}"
@@ -420,7 +430,19 @@ def _equi_join(self, builder, rel_cls, other, left_on, right_on, how):
                 [right_on] if isinstance(right_on, (str, int)) else list(right_on)
             )
         return self._next(
-            builder(self._plan, other._plan, left_keys, right_keys, join_type)
+            builder(
+                self._plan,
+                other._plan,
+                left_keys,
+                right_keys,
+                join_type,
+                post_join_filter=(
+                    _unbound(post_filter) if post_filter is not None else None
+                ),
+                residual_expression=(
+                    _unbound(residual) if residual is not None else None
+                ),
+            )
         )
 
     def hash_join(
@@ -429,14 +451,27 @@ def hash_join(
         left_on: Union[str, int, Iterable[Union[str, int]]],
         right_on: Union[str, int, Iterable[Union[str, int]], None] = None,
         how: str = "inner",
+        *,
+        post_filter: Union[Expr, Any, None] = None,
+        residual: Union[Expr, Any, None] = None,
     ) -> "DataFrame":
         """Physical hash equi-join on key columns.
 
         ``left_on``/``right_on`` are column names/indices; ``right_on`` defaults
         to ``left_on``. ``how`` accepts the same values as :meth:`join`.
+        ``post_filter`` is an optional predicate applied to the join output;
+        ``residual`` is an optional non-equi condition evaluated alongside the
+        key equalities. Both bind against the concatenated left+right schema.
         """
         return self._equi_join(
-            _plan.hash_join, stalg.HashJoinRel, other, left_on, right_on, how
+            _plan.hash_join,
+            stalg.HashJoinRel,
+            other,
+            left_on,
+            right_on,
+            how,
+            post_filter,
+            residual,
         )
 
     def merge_join(
@@ -445,10 +480,23 @@ def merge_join(
         left_on: Union[str, int, Iterable[Union[str, int]]],
         right_on: Union[str, int, Iterable[Union[str, int]], None] = None,
         how: str = "inner",
+        *,
+        post_filter: Union[Expr, Any, None] = None,
+        residual: Union[Expr, Any, None] = None,
     ) -> "DataFrame":
-        """Physical sort-merge equi-join on key columns (inputs assumed sorted)."""
+        """Physical sort-merge equi-join on key columns (inputs assumed sorted).
+
+        ``post_filter`` and ``residual`` behave as in :meth:`hash_join`.
+        """
         return self._equi_join(
-            _plan.merge_join, stalg.MergeJoinRel, other, left_on, right_on, how
+            _plan.merge_join,
+            stalg.MergeJoinRel,
+            other,
+            left_on,
+            right_on,
+            how,
+            post_filter,
+            residual,
         )
 
     def repartition(self, n: int = 0) -> "DataFrame":

tests/builders/plan/test_join.py

@@ -3,7 +3,7 @@
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
 
-from substrait.builders.extended_expression import literal
+from substrait.builders.extended_expression import column, literal
 from substrait.builders.plan import default_version, join, read_named_table
 from substrait.builders.type import boolean, i64, string
 from substrait.extension_registry import ExtensionRegistry
@@ -69,3 +69,62 @@ def test_join():
     )
 
     assert actual == expected
+
+
+def _post_field(plan):
+    ref = plan.relations[-1].root.input.join.post_join_filter.selection
+    return ref.direct_reference.struct_field.field
+
+
+def test_join_post_join_filter_binds_output_schema():
+    # post_join_filter is applied to the join output (semantically a FilterRel
+    # above the join), so it resolves against the output schema. For an inner join
+    # the output is the combined schema [id, is_applicable, fk_id, flag], so a
+    # filter on the right-side boolean `flag` binds to index 3; for a right-semi
+    # join the output is the right side only [fk_id, flag], so it binds to index 1.
+    left = read_named_table("l", named_struct)
+    right = read_named_table(
+        "r",
+        stt.NamedStruct(
+            names=["fk_id", "flag"],
+            struct=stt.Type.Struct(
+                types=[i64(nullable=False), boolean()],
+                nullability=stt.Type.NULLABILITY_REQUIRED,
+            ),
+        ),
+    )
+
+    inner = join(
+        left,
+        right,
+        literal(True, boolean()),
+        stalg.JoinRel.JOIN_TYPE_INNER,
+        post_join_filter=column("flag"),
+    )(registry)
+    assert _post_field(inner) == 3
+
+    right_semi = join(
+        left,
+        right,
+        literal(True, boolean()),
+        stalg.JoinRel.JOIN_TYPE_RIGHT_SEMI,
+        post_join_filter=column("flag"),
+    )(registry)
+    assert list(right_semi.relations[-1].root.names) == ["fk_id", "flag"]
+    assert _post_field(right_semi) == 1
+
+
+def test_join_post_join_filter_on_dropped_side_raises():
+    # A right-semi join drops the left side from its output, so a post_join_filter
+    # on a left column cannot resolve -- it fails fast rather than emitting a
+    # dangling field reference against the combined schema.
+    left = read_named_table("l", named_struct)
+    right = read_named_table("r", named_struct_2)
+    with pytest.raises(ValueError, match="not in list"):
+        join(
+            left,
+            right,
+            literal(True, boolean()),
+            stalg.JoinRel.JOIN_TYPE_RIGHT_SEMI,
+            post_join_filter=column("id"),  # left-only column, absent from output
+        )(registry)