fix!: derive mark join output from the selected input

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(

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,

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)
 
 

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

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)
 
 

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,
+        ),
+    )

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"
 

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,