From ee597918e088b5353309e5ac0bb54e407218ddca Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Fri, 17 Jul 2026 12:05:02 +0200 Subject: [PATCH] feat: expose residual_expression and post_join_filter on join builders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Surface the two Expression fields spec #1044 added to HashJoinRel and MergeJoinRel — residual_expression and post_join_filter — on the ergonomic physical-join builders (hash_join / merge_join) and the DataFrame hash_join / merge_join methods, completing the remaining work for #190. Bind each expression against the schema the spec prescribes: - residual_expression is evaluated on each candidate key-match (both rows present), so it resolves against the combined left+right schema. - post_join_filter is applied to each output record after join-type-specific output formation ("semantically a FilterRel directly above this join"), so it resolves against the join output schema, reusing the existing _join_output_struct / join_output_names helpers. The same fix is applied to the logical join builder, which previously bound post_join_filter against the combined schema. For inner/outer/left/right/ single/mark joins the output equals (or prefixes) the combined schema so behavior is unchanged; for right_semi / right_anti (output is the right side only) a filter on a right column now uses the correct output-relative field index instead of a dangling combined-schema index, and a filter on a dropped side fails fast instead of emitting an invalid reference. The optional predicate/extension arguments are keyword-only, matching the logical join builder's convention. --- src/substrait/builders/plan.py | 92 ++++++++++++++---- src/substrait/dataframe/frame.py | 58 +++++++++++- tests/builders/plan/test_join.py | 61 +++++++++++- tests/builders/plan/test_physical.py | 137 ++++++++++++++++++++++++++- tests/dataframe/test_frame.py | 69 ++++++++++++++ 5 files changed, 393 insertions(+), 24 deletions(-) diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 8885ca72..f3cb5c76 100644 --- a/src/substrait/builders/plan.py +++ b/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,9 +1052,42 @@ 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( @@ -1043,6 +1095,12 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: 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, ) } @@ -1050,7 +1108,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: 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 diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index 2b7cd6ff..7b13831c 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -415,7 +415,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)}" @@ -429,7 +439,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( @@ -438,14 +460,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( @@ -454,10 +489,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": diff --git a/tests/builders/plan/test_join.py b/tests/builders/plan/test_join.py index 2fb46377..a044d171 100644 --- a/tests/builders/plan/test_join.py +++ b/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) diff --git a/tests/builders/plan/test_physical.py b/tests/builders/plan/test_physical.py index 102b3818..e4bbc09c 100644 --- a/tests/builders/plan/test_physical.py +++ b/tests/builders/plan/test_physical.py @@ -1,8 +1,15 @@ +import pytest import substrait.algebra_pb2 as stalg import substrait.type_pb2 as stt -from substrait.builders.extended_expression import column, scalar_function -from substrait.builders.plan import exchange, nested_loop_join, read_named_table +from substrait.builders.extended_expression import column, literal, scalar_function +from substrait.builders.plan import ( + exchange, + hash_join, + merge_join, + nested_loop_join, + read_named_table, +) from substrait.builders.type import fp64, i64, string from substrait.extension_registry import ExtensionRegistry from substrait.type_inference import infer_plan_schema @@ -84,3 +91,129 @@ def test_exchange_broadcast(): plan.relations[-1].root.input.exchange.WhichOneof("exchange_kind") == "broadcast" ) + + +@pytest.mark.parametrize( + "builder, rel_field, rel_cls", + [ + (hash_join, "hash_join", stalg.HashJoinRel), + (merge_join, "merge_join", stalg.MergeJoinRel), + ], +) +def test_equi_join_inner_residual_and_post_filter(builder, rel_field, rel_cls): + # For an inner join the output schema equals the combined schema [x, y, w, z], + # so post_join_filter on the right-only column z binds to index 3. + # residual_expression references one column from each side (x on the left, w on + # the right), confirming it binds against the combined left+right schema. + plan = builder( + _left(), + _right(), + ["x"], + ["w"], + rel_cls.JOIN_TYPE_INNER, + post_join_filter=scalar_function( + COMPARISON, "gt", expressions=[column("z"), literal(1.0, fp64())] + ), + residual_expression=scalar_function( + COMPARISON, "gt", expressions=[column("x"), column("w")] + ), + )(registry) + + rel = getattr(plan.relations[-1].root.input, rel_field) + assert rel.HasField("post_join_filter") + assert rel.HasField("residual_expression") + # z is the 4th field (index 3) of the inner-join output (== combined schema). + post_ref = rel.post_join_filter.scalar_function.arguments[0].value.selection + assert post_ref.direct_reference.struct_field.field == 3 + # residual: x -> 0 (left input), w -> 2 (right input). + res_args = rel.residual_expression.scalar_function.arguments + assert res_args[0].value.selection.direct_reference.struct_field.field == 0 + assert res_args[1].value.selection.direct_reference.struct_field.field == 2 + # The comparison extension used by the predicates is declared on the plan. + assert len(plan.extension_urns) == 1 + assert len(plan.extensions) == 1 + + +@pytest.mark.parametrize( + "builder, rel_field, rel_cls", + [ + (hash_join, "hash_join", stalg.HashJoinRel), + (merge_join, "merge_join", stalg.MergeJoinRel), + ], +) +def test_equi_join_post_filter_binds_output_schema(builder, rel_field, rel_cls): + # post_join_filter applies to the join OUTPUT. A right-semi join emits the + # right side only ([w, z]), so a filter on the right column z must bind to the + # output-relative index 1 -- not the combined-schema index 3. residual_expression + # still binds against the combined schema (x -> 0, w -> 2). + plan = builder( + _left(), + _right(), + ["x"], + ["w"], + rel_cls.JOIN_TYPE_RIGHT_SEMI, + post_join_filter=scalar_function( + COMPARISON, "gt", expressions=[column("z"), literal(1.0, fp64())] + ), + residual_expression=scalar_function( + COMPARISON, "gt", expressions=[column("x"), column("w")] + ), + )(registry) + + assert list(plan.relations[-1].root.names) == ["w", "z"] + rel = getattr(plan.relations[-1].root.input, rel_field) + post_ref = rel.post_join_filter.scalar_function.arguments[0].value.selection + assert post_ref.direct_reference.struct_field.field == 1 + res_args = rel.residual_expression.scalar_function.arguments + assert res_args[0].value.selection.direct_reference.struct_field.field == 0 + assert res_args[1].value.selection.direct_reference.struct_field.field == 2 + + +@pytest.mark.parametrize( + "builder, rel_cls", + [(hash_join, stalg.HashJoinRel), (merge_join, stalg.MergeJoinRel)], +) +def test_equi_join_post_filter_on_dropped_side_raises(builder, rel_cls): + # A right-semi join drops the left side from its output, so a post_join_filter + # referencing a left column cannot resolve -- it fails fast rather than emitting + # a dangling field reference. + with pytest.raises(ValueError, match="not in list"): + builder( + _left(), + _right(), + ["x"], + ["w"], + rel_cls.JOIN_TYPE_RIGHT_SEMI, + post_join_filter=scalar_function( + COMPARISON, "gt", expressions=[column("x"), literal(1, i64())] + ), + )(registry) + + +@pytest.mark.parametrize( + "builder, rel_field, rel_cls", + [ + (hash_join, "hash_join", stalg.HashJoinRel), + (merge_join, "merge_join", stalg.MergeJoinRel), + ], +) +def test_equi_join_predicates_unset_by_default(builder, rel_field, rel_cls): + plan = builder(_left(), _right(), ["x"], ["w"], rel_cls.JOIN_TYPE_INNER)(registry) + rel = getattr(plan.relations[-1].root.input, rel_field) + assert not rel.HasField("post_join_filter") + assert not rel.HasField("residual_expression") + + +@pytest.mark.parametrize("builder, rel_cls", [(hash_join, stalg.HashJoinRel)]) +def test_equi_join_optional_args_are_keyword_only(builder, rel_cls): + # post_join_filter / residual_expression / extension are keyword-only so a + # positional arg after `type` cannot silently rebind one of them. + with pytest.raises(TypeError): + builder( + _left(), + _right(), + ["x"], + ["w"], + rel_cls.JOIN_TYPE_INNER, + None, # would have bound to post_join_filter positionally + ) diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index ee77a994..f1891f58 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -20,8 +20,10 @@ from substrait.builders.plan import extension_table as b_extension_table from substrait.builders.plan import fetch as b_fetch from substrait.builders.plan import filter as b_filter +from substrait.builders.plan import hash_join as b_hash_join from substrait.builders.plan import join as b_join from substrait.builders.plan import local_files as b_local_files +from substrait.builders.plan import merge_join as b_merge_join from substrait.builders.plan import read_named_table as b_read from substrait.builders.plan import select as b_select from substrait.builders.plan import set as b_set @@ -986,6 +988,73 @@ def test_merge_join_default_right_on(): assert len(mj.keys) == 1 +def _equi_join_tables(): + left_ns = named_struct( + names=["cust_id", "name"], + struct=struct(types=[i64(), string()], nullable=False), + ) + right_ns = named_struct( + names=["order_id", "cust_ref", "amount"], + struct=struct(types=[i64(), i64(), fp64()], nullable=False), + ) + left = sub.read_named_table("customers", {"cust_id": sub.i64, "name": sub.string}) + right = sub.read_named_table( + "orders", {"order_id": sub.i64, "cust_ref": sub.i64, "amount": sub.fp64} + ) + return left, right, left_ns, right_ns + + +# "amount" is a right column; for right_semi the fluent and raw pipelines route +# through the same builder, so this checks the DataFrame layer forwards +# post_filter/residual correctly regardless of join type. +@pytest.mark.parametrize("how", ["inner", "right_semi"]) +@pytest.mark.parametrize( + "method, b_builder, rel_cls", + [ + ("hash_join", b_hash_join, stalg.HashJoinRel), + ("merge_join", b_merge_join, stalg.MergeJoinRel), + ], +) +def test_equi_join_post_filter_and_residual_match_builder( + method, b_builder, rel_cls, how +): + left, right, left_ns, right_ns = _equi_join_tables() + fluent = getattr(left, method)( + right, + "cust_id", + "cust_ref", + how=how, + post_filter=sub.col("amount") > 100.0, + residual=sub.col("amount") > 50.0, + ).to_plan() + raw = b_builder( + b_read("customers", left_ns), + b_read("orders", right_ns), + ["cust_id"], + ["cust_ref"], + getattr(rel_cls, "JOIN_TYPE_" + how.upper()), + post_join_filter=scalar_function( + COMPARISON, "gt", expressions=[column("amount"), literal(100.0, fp64())] + ), + residual_expression=scalar_function( + COMPARISON, "gt", expressions=[column("amount"), literal(50.0, fp64())] + ), + )(registry) + assert fluent.SerializeToString() == raw.SerializeToString() + + +def test_hash_join_without_predicates_leaves_them_unset(): + left, right, _, _ = _equi_join_tables() + hj = ( + left.hash_join(right, "cust_id", "cust_ref") + .to_plan() + .relations[-1] + .root.input.hash_join + ) + assert not hj.HasField("post_join_filter") + assert not hj.HasField("residual_expression") + + def test_dynamic_parameter(): from substrait.type_inference import infer_plan_schema