From 4cd3462bc63b40bf0284d30b4bd5cb564f7ede1e Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Mon, 7 Sep 2026 23:24:59 +0300 Subject: [PATCH] fix!: derive mark join output from the selected input `LEFT_MARK` and `RIGHT_MARK` currently include both inputs in their inferred schema and output names. Consumers expect only the selected input followed by the marker, so generated plans can fail schema validation. Match [Substrait spec v0.99.0](https://github.com/substrait-io/substrait/blob/v0.99.0/site/docs/relations/logical_relations.md#join-operation): keep the left input for `LEFT_MARK` and the right input for `RIGHT_MARK`, then append a nullable boolean marker. Use that output for names, post-join filters, and emit across logical, lateral, and physical joins. Join conditions and residual expressions still use both inputs. Keep correlations into that combined scope offset-based, since the mark join output cannot represent it. Closes #263 BREAKING CHANGE: Mark joins no longer expose columns from the other input. Output field indices change, and post-join filters must use only the selected input and the marker. --- src/substrait/builders/plan.py | 27 ++-- src/substrait/dataframe/frame.py | 3 +- src/substrait/type_inference.py | 34 +++-- src/substrait/utils/__init__.py | 10 +- tests/builders/plan/test_lateral_join.py | 9 +- tests/builders/plan/test_mark_join.py | 150 +++++++++++++++++++++++ tests/dataframe/test_frame.py | 14 ++- tests/test_type_inference.py | 4 - tests/test_utils.py | 30 ++++- 9 files changed, 224 insertions(+), 57 deletions(-) create mode 100644 tests/builders/plan/test_mark_join.py diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 519217f..d5aabde 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -636,15 +636,15 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: expression, ns, registry ) - # The output names must match the columns the join type actually emits - # (semi/anti drop a side, mark appends a boolean). + # Output names use the selected side for semi/anti/mark joins, + # with a nullable boolean marker appended for mark joins. 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. + # join), so it resolves against the output schema. Semi/anti joins expose + # the selected side; mark joins expose that side plus the marker. bound_post = None if post_join_filter is not None: output_ns = stt.NamedStruct( @@ -730,16 +730,15 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: else None ) - # Output names/columns follow the same per-join-type shape as a - # regular join (semi/anti drop the right side, mark appends a boolean). + # Output names use the selected side for semi/anti/mark joins, + # with a nullable boolean marker appended for mark joins. 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 and for a mark join carries the - # appended marker column -- not the combined input row. + # the join), so it resolves against the output schema. Semi/anti joins + # expose the selected side; mark joins expose that side plus the marker. bound_post = None if post_join_filter is not None: output_ns = stt.NamedStruct( @@ -1295,11 +1294,11 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: # 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. + # the join), so it resolves against the output schema. Semi/anti joins + # expose the selected side; mark joins expose that side plus the marker. + # residual_expression resolves against the combined left+right schema + # because each candidate key-match contains both rows. Each schema is + # built only when its corresponding predicate is supplied. bound_post = None if post_join_filter is not None: output_ns = stt.NamedStruct( diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index 54f4058..d6ab4a1 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -544,7 +544,8 @@ def hash_join( 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. + key equalities. ``post_filter`` binds against the join output schema; + ``residual`` binds against the concatenated left+right schema. """ return self._equi_join( _plan.hash_join, diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 8c431a6..1fcd4e1 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -535,14 +535,16 @@ def infer_extended_expression_schema( def _join_column_shape(type_name: str) -> str: """Which columns a join emits, by join-type NAME (shared across all join relations, whose enum integer values differ): ``left`` / ``right`` only for - semi/anti, ``both+mark`` for mark joins, ``both`` otherwise. Single source of - truth for both the inferred type list and the RelRoot names.""" + semi/anti, ``left+mark`` / ``right+mark`` for mark joins, ``both`` otherwise. + Single source of truth for the inferred type list and the RelRoot names.""" if type_name in ("JOIN_TYPE_LEFT_SEMI", "JOIN_TYPE_LEFT_ANTI"): return "left" if type_name in ("JOIN_TYPE_RIGHT_SEMI", "JOIN_TYPE_RIGHT_ANTI"): return "right" - if type_name in ("JOIN_TYPE_LEFT_MARK", "JOIN_TYPE_RIGHT_MARK"): - return "both+mark" + if type_name == "JOIN_TYPE_LEFT_MARK": + return "left+mark" + if type_name == "JOIN_TYPE_RIGHT_MARK": + return "right+mark" return "both" # inner / outer / left / right / single @@ -555,8 +557,10 @@ def join_output_names(type_name: str, left_names, right_names) -> list: return list(left_names) if shape == "right": return list(right_names) - if shape == "both+mark": - return list(left_names) + list(right_names) + [JOIN_MARK_COLUMN_NAME] + if shape == "left+mark": + return list(left_names) + [JOIN_MARK_COLUMN_NAME] + if shape == "right+mark": + return list(right_names) + [JOIN_MARK_COLUMN_NAME] return list(left_names) + list(right_names) @@ -568,22 +572,16 @@ def _join_struct_from_schemas( values differ).""" required = stt.Type.Nullability.NULLABILITY_REQUIRED shape = _join_column_shape(type_name) - if shape == "left": + if shape in ("left", "left+mark"): types = list(left.types) - elif shape == "right": + elif shape in ("right", "right+mark"): types = list(right.types) - elif shape == "both+mark": - types = ( - list(left.types) - + list(right.types) - + [ - stt.Type( - bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE) - ) - ] - ) else: types = list(left.types) + list(right.types) + if shape in ("left+mark", "right+mark"): + types.append( + stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)) + ) return stt.Type.Struct(types=types, nullability=required) diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index c41782a..af3c497 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -471,19 +471,21 @@ def _plan_has_steps_out(plan: stplan.Plan) -> bool: # against the join *output* (semantically a Filter above the join), the condition # and ``residual_expression`` against the *combined* left+right inputs. The join # relation's own output equals that combined row for every non-reducing join, but a -# reducing join (semi/anti) emits a single side -- so a correlation into its +# reducing join (semi/anti/mark) drops a side -- so a correlation into its # condition scope names columns the output drops and has no anchorable relation. _JOIN_COMBINED_SCOPED_FIELDS = frozenset({"expression", "residual_expression"}) def _is_reducing_join(node) -> bool: - """Whether a join relation-variant ``node`` emits only one side (semi/anti), so + """Whether a join relation-variant ``node`` drops one side (semi/anti/mark), so its output row differs from its combined left+right condition scope.""" field = node.DESCRIPTOR.fields_by_name.get("type") if field is None or field.enum_type is None: return False name = field.enum_type.values_by_number.get(node.type) - return name is not None and ("SEMI" in name.name or "ANTI" in name.name) + return name is not None and ( + "SEMI" in name.name or "ANTI" in name.name or "MARK" in name.name + ) def to_id_based_outer_references(plan: stplan.Plan) -> stplan.Plan: @@ -512,7 +514,7 @@ def to_id_based_outer_references(plan: stplan.Plan) -> stplan.Plan: output row, so the host is anchored. * a join *condition* / ``residual_expression`` exposes the **combined** left+right row; the join's own output equals that row for a non-reducing join, so the join - is anchored. For a *reducing* join (semi/anti) the two differ and no relation + is anchored. For a *reducing* join (semi/anti/mark) the two differ and no relation carries that row -- such a reference is left offset-based (still spec-valid, and read by inference), rather than mis-anchored. * a ``LateralJoinRel``'s ``rel_anchor`` is reserved (per the Substrait spec) for diff --git a/tests/builders/plan/test_lateral_join.py b/tests/builders/plan/test_lateral_join.py index 4428ee6..19a7ed5 100644 --- a/tests/builders/plan/test_lateral_join.py +++ b/tests/builders/plan/test_lateral_join.py @@ -97,7 +97,7 @@ def test_lateral_join_left_mark_appends_boolean(): )(registry) ns = infer_plan_schema(plan, registry=registry) - assert list(ns.names)[-1] == "mark" + assert list(ns.names) == ["k", "v", "mark"] assert ns.struct.types[-1].WhichOneof("kind") == "bool" assert len(ns.names) == len(ns.struct.types) @@ -158,8 +158,7 @@ def test_lateral_join_post_join_filter_binds_output_schema(): # above the lateral join), so it resolves against the *output* schema, not the # combined left+right inputs. For an inner join the output is [k, v, w], so a # filter on the right column `w` binds to index 2; for a left-mark join the - # output appends a `mark` column absent from the combined inputs, binding to - # index 3. + # output is [k, v, mark], with the marker at index 2. with fresh_rel_anchors(): inner = lateral_join( _left(), @@ -177,8 +176,8 @@ def test_lateral_join_post_join_filter_binds_output_schema(): type=stalg.JoinRel.JOIN_TYPE_LEFT_MARK, post_join_filter=column("mark"), )(registry) - assert list(mark.relations[-1].root.names) == ["k", "v", "w", "mark"] - assert _post_field(mark) == 3 + assert list(mark.relations[-1].root.names) == ["k", "v", "mark"] + assert _post_field(mark) == 2 infer_plan_schema(mark, registry=registry) diff --git a/tests/builders/plan/test_mark_join.py b/tests/builders/plan/test_mark_join.py new file mode 100644 index 0000000..24832cd --- /dev/null +++ b/tests/builders/plan/test_mark_join.py @@ -0,0 +1,150 @@ +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 ( + hash_join, + join, + lateral_join, + merge_join, + nested_loop_join, + read_named_table, +) +from substrait.builders.type import boolean, i64, string +from substrait.extension_registry import ExtensionRegistry +from substrait.type_inference import infer_plan_schema, infer_rel_schema + +REGISTRY = ExtensionRegistry() +LEFT = stt.NamedStruct( + names=["l_id", "payload", "text", "l_flag"], + struct=stt.Type.Struct( + types=[ + i64(nullable=False), + stt.Type( + struct=stt.Type.Struct( + types=[string()], nullability=stt.Type.NULLABILITY_REQUIRED + ) + ), + boolean(nullable=False), + ], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), +) +RIGHT = stt.NamedStruct( + names=["r_id", "r_flag"], + struct=stt.Type.Struct( + types=[i64(), boolean()], nullability=stt.Type.NULLABILITY_REQUIRED + ), +) +JOIN_CASES = [ + pytest.param(builder, field, cls, side, id=f"{field}-{side.lower()}") + for builder, field, cls in [ + (join, "join", stalg.JoinRel), + (lateral_join, "lateral_join", stalg.JoinRel), + (nested_loop_join, "nested_loop_join", stalg.NestedLoopJoinRel), + (hash_join, "hash_join", stalg.HashJoinRel), + (merge_join, "merge_join", stalg.MergeJoinRel), + ] + for side in ("LEFT", "RIGHT") + if builder is not lateral_join or side == "LEFT" +] + + +def _mark_plan(builder, cls, side, post_field="mark"): + left = read_named_table("l", LEFT) + right = read_named_table("r", RIGHT) + condition = scalar_function( + "extension:io.substrait:functions_comparison", + "equal", + expressions=[column("l_id"), column("r_id")], + ) + join_type = getattr(cls, f"JOIN_TYPE_{side}_MARK") + if builder is nested_loop_join: + plan = builder(left, right, condition, join_type) + elif builder is lateral_join: + plan = builder( + left, + lambda _: right, + join_type, + expression=condition, + post_join_filter=column(post_field), + ) + elif builder is join: + plan = builder( + left, right, condition, join_type, post_join_filter=column(post_field) + ) + else: + plan = builder( + left, + right, + ["l_id"], + ["r_id"], + join_type, + post_join_filter=column(post_field), + residual_expression=condition, + ) + return plan(REGISTRY) + + +@pytest.mark.parametrize("builder,field,cls,side", JOIN_CASES) +def test_mark_join_output_and_condition_bindings(builder, field, cls, side): + plan = _mark_plan(builder, cls, side) + root = plan.relations[-1].root + rel = getattr(root.input, field) + condition = ( + rel.residual_expression + if builder in (hash_join, merge_join) + else rel.expression + ) + # Conditions use both inputs, counting top-level fields rather than DFS names. + assert [ + arg.value.selection.direct_reference.struct_field.field + for arg in condition.scalar_function.arguments + ] == [0, 3] + + kept = LEFT if side == "LEFT" else RIGHT + expected = stt.NamedStruct( + names=[*kept.names, "mark"], + struct=stt.Type.Struct( + types=[*kept.struct.types, boolean()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + assert list(root.names) == list(expected.names) + assert infer_rel_schema(root.input, registry=REGISTRY) == expected.struct + assert infer_plan_schema(plan, registry=REGISTRY) == expected + if builder is not nested_loop_join: + ref = rel.post_join_filter.selection.direct_reference.struct_field + assert ref.field == len(kept.struct.types) + + +@pytest.mark.parametrize( + "builder,field,cls,side", + [case for case in JOIN_CASES if case.values[0] is not nested_loop_join], +) +def test_mark_join_post_filter_cannot_reference_dropped_side(builder, field, cls, side): + dropped_field = "r_flag" if side == "LEFT" else "l_flag" + with pytest.raises(ValueError, match=dropped_field): + _mark_plan(builder, cls, side, dropped_field) + + +@pytest.mark.parametrize("builder,field,cls,side", JOIN_CASES) +def test_mark_join_emit_indexes_selected_side_and_marker(builder, field, cls, side): + plan = _mark_plan(builder, cls, side) + kept = LEFT if side == "LEFT" else RIGHT + root = plan.relations[-1].root + rel = getattr(root.input, field) + rel.common.emit.output_mapping[:] = [ + len(kept.struct.types), + 0, + len(kept.struct.types), + ] + root.names[:] = ["first_mark", "id", "second_mark"] + assert infer_plan_schema(plan, registry=REGISTRY) == stt.NamedStruct( + names=list(root.names), + struct=stt.Type.Struct( + types=[boolean(), kept.struct.types[0], boolean()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index a46fbb5..49a85a8 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -1387,8 +1387,8 @@ def test_lateral_join_post_filter_binds_output_schema(): lj = plan.relations[-1].root.input.lateral_join assert lj.HasField("post_join_filter") field = lj.post_join_filter.selection.direct_reference.struct_field.field - assert field == 3 # output is [k, v, w, mark] - assert list(infer_plan_schema(plan).names) == ["k", "v", "w", "mark"] + assert field == 2 # output is [k, v, mark] + assert list(infer_plan_schema(plan).names) == ["k", "v", "mark"] def test_correlated_exists_above_lateral_join_stays_steps_out(): @@ -1478,14 +1478,16 @@ def test_semi_join_output_names_match_types(): assert len(ns.names) == len(ns.struct.types) -def test_mark_join_output_names_match_types(): +@pytest.mark.parametrize( + "how,names", [("left_mark", ["x", "y", "mark"]), ("right_mark", ["w", "z", "mark"])] +) +def test_mark_join_output_names_match_types(how, names): from substrait.type_inference import infer_plan_schema left, right = _ab() - plan = left.hash_join(right, "x", "w", how="left_mark").to_plan() + plan = left.hash_join(right, "x", "w", how=how).to_plan() ns = infer_plan_schema(plan) - # left + right + a trailing boolean mark column. - assert list(ns.names) == ["x", "y", "w", "z", "mark"] + assert list(ns.names) == names assert len(ns.names) == len(ns.struct.types) assert ns.struct.types[-1].WhichOneof("kind") == "bool" diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index b5b9dd0..0f7a13c 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -335,8 +335,6 @@ def test_inference_join_left_mark(): stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), stt.Type(string=stt.Type.String(nullability=stt.Type.NULLABILITY_NULLABLE)), stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE)), - stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), - stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)), stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)), ], nullability=stt.Type.Nullability.NULLABILITY_REQUIRED, @@ -407,8 +405,6 @@ def test_inference_lateral_join_left_mark(): stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), stt.Type(string=stt.Type.String(nullability=stt.Type.NULLABILITY_NULLABLE)), stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE)), - stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), - stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)), stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)), ], nullability=stt.Type.Nullability.NULLABILITY_REQUIRED, diff --git a/tests/test_utils.py b/tests/test_utils.py index f56e21b..1261ff7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -662,12 +662,21 @@ def test_convert_multi_input_join_condition_anchors_the_join(): assert ref.rel_reference == 1 -def test_convert_post_join_filter_anchors_the_join(): +@pytest.mark.parametrize( + "join_type", + [ + stalg.JoinRel.JOIN_TYPE_INNER, + stalg.JoinRel.JOIN_TYPE_LEFT_MARK, + stalg.JoinRel.JOIN_TYPE_RIGHT_MARK, + ], +) +def test_convert_post_join_filter_anchors_the_join(join_type): # A correlation in a join's post_join_filter resolves against the join output, # i.e. the join itself -- anchored the same way. join = _join( _read("l"), _read("r"), + type=join_type, post_join_filter=_exists(_filter(_read("i"), _outer(1))), ) out = to_id_based_outer_references(_plan(join)) @@ -678,15 +687,26 @@ def test_convert_post_join_filter_anchors_the_join(): assert ref.rel_reference == 1 -def test_convert_reducing_join_condition_left_as_steps_out(): - # A reducing join (semi/anti) emits only one side, so its output row differs - # from the combined condition scope the reference sees. No relation carries +@pytest.mark.parametrize( + "join_type", + [ + stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, + stalg.JoinRel.JOIN_TYPE_RIGHT_SEMI, + stalg.JoinRel.JOIN_TYPE_LEFT_ANTI, + stalg.JoinRel.JOIN_TYPE_RIGHT_ANTI, + stalg.JoinRel.JOIN_TYPE_LEFT_MARK, + stalg.JoinRel.JOIN_TYPE_RIGHT_MARK, + ], +) +def test_convert_reducing_join_condition_left_as_steps_out(join_type): + # A semi/anti/mark join emits one side (plus a marker for mark joins), so its + # output differs from the combined condition scope. No relation carries # that row, so the reference stays offset-based (still spec-valid) rather than # being mis-anchored to the join's narrower output. join = _join( _read("l"), _read("r"), - type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, + type=join_type, expression=_exists(_filter(_read("i"), _outer(1))), ) out = to_id_based_outer_references(_plan(join))