Skip to content

infer_rel_schema emits both inputs plus the mark column for mark joins #263

Description

@alexandrefimov

_join_column_shape sends both mark join types to both+mark, and the output is then built as left input, right input, mark column (type_inference.py:558 and :575; same on main). Join Operation (v0.102.0), Direct Output Order: "For mark joins, the emit order is either left or right with a "mark" column appended at the end." So LEFT_MARK emits the left input plus the mark column, RIGHT_MARK the right input plus it, and the opposite input is in neither.

Reproduced with 0.31.0 through the builders (mark below is the library's own JOIN_MARK_COLUMN_NAME):

import substrait.algebra_pb2 as stalg, substrait.type_pb2 as stt
from substrait.builders.plan import read_named_table, join
from substrait.builders.extended_expression import literal
from substrait.extension_registry import ExtensionRegistry
from substrait.type_inference import infer_plan_schema

REQ = stt.Type.NULLABILITY_REQUIRED
reg, i64 = ExtensionRegistry(), stt.Type(i64=stt.Type.I64(nullability=REQ))
ns = lambda *n: stt.NamedStruct(names=list(n), struct=stt.Type.Struct(
    types=[i64] * len(n), nullability=REQ))

p = join(read_named_table("l", ns("l0", "l1")), read_named_table("r", ns("r0", "r1")),
         literal(True, stt.Type(bool=stt.Type.Boolean(nullability=REQ))),
         stalg.JoinRel.JOIN_TYPE_LEFT_MARK)(reg)

print(list(infer_plan_schema(p, registry=reg).names))
# actual:   ['l0', 'l1', 'r0', 'r1', 'mark']
# expected: ['l0', 'l1', 'mark']

substrait-java derives left-plus-mark for that same plan and refuses to load it, which is how I ran into this: Plan.Root names count (5) must match input record type depth-first named-field count (3).

One thing I did not want to assert: the join condition still references fields of both inputs, so narrowing the output should leave condition-side numbering alone, but that is worth a second pair of eyes. Happy to send the PR if you want it.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions