diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index d95b119..f687b36 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -40,6 +40,7 @@ from substrait.dataframe.expr import Expr, Measure, col, lit, sort_direction from substrait.extension_registry import ExtensionRegistry from substrait.type_inference import infer_plan_schema +from substrait.utils import to_id_based_outer_references # All 13 JoinRel.JoinType variants (SET_OP_UNSPECIFIED excluded). "single" # returns at most one right match per left row (runtime error on multiple); @@ -715,13 +716,19 @@ def write_named_table( ) ) - def to_plan(self): + def _finalize(self, registry: Optional[ExtensionRegistry]) -> stplan.Plan: + """Build the plan and normalize it for output. The DataFrame layer emits + correlated (outer) references in the id-based form (``rel_reference``), so + any offset-based ``steps_out`` the builders produced is rewritten here.""" + return to_id_based_outer_references(self._plan(registry)) + + def to_plan(self) -> stplan.Plan: """Materialize to a ``substrait.proto.Plan``.""" - return self._plan(self._registry) + return self._finalize(self._registry) # Kept for parity with the substrait.narwhals (Narwhals) wrapper's API. - def to_substrait(self, registry: Optional[ExtensionRegistry] = None): - return self._plan(registry or self._registry) + def to_substrait(self, registry: Optional[ExtensionRegistry] = None) -> stplan.Plan: + return self._finalize(registry or self._registry) class GroupBy: diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 393a8ee..f99c150 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -5,7 +5,7 @@ import substrait.plan_pb2 as stp import substrait.type_pb2 as stt -from substrait.utils import plan_subtrees +from substrait.utils import iter_plan_rels, plan_subtrees, rel_anchor_of class _SubtreeScope: @@ -53,13 +53,64 @@ def schema_of(self, ordinal, registry): return schema +class _AnchorScope: + """Plan-wide map of ``RelCommon.rel_anchor`` -> the relation carrying it, with + lazy per-anchor output-schema memoization, for resolving id-based outer + references (``OuterReference.rel_reference``). + + A ``rel_reference`` names the anchor of the relation the reference is rooted on; + it resolves against that relation's output schema. Resolution is lazy (only + anchors actually referenced are inferred) and memoized. Inference of an anchored + relation may itself resolve a ``rel_reference``, so an in-progress set turns a + malformed cyclic reference into a clear error rather than a ``RecursionError``. + Anchored relations are inferred with the plan's ``_SubtreeScope`` in scope so an + anchor set on a shared subtree still resolves. + """ + + __slots__ = ("_rels", "_subtrees", "_schemas", "_resolving") + + def __init__(self, rels: dict, subtrees): + self._rels = rels + self._subtrees = subtrees + self._schemas: dict = {} + self._resolving: set = set() + + def schema_of(self, anchor, registry) -> stt.Type.Struct: + if anchor in self._schemas: + return self._schemas[anchor] + if anchor not in self._rels: + raise Exception(f"outer reference to unknown rel_anchor {anchor}") + if anchor in self._resolving: + raise Exception( + f"outer reference rel_anchor {anchor} forms a resolution cycle" + ) + self._resolving.add(anchor) + try: + struct = infer_rel_schema( + self._rels[anchor], registry=registry, subtrees=self._subtrees + ) + finally: + self._resolving.discard(anchor) + self._schemas[anchor] = struct + return struct + + # Stack of enclosing-query schemas (NamedStruct) for correlated subqueries, so a # field reference with an OuterReference root resolves against the right level. -# Pushed by the subquery builders while resolving their inner plan. +# Pushed by the subquery builders while resolving their inner plan. Used for +# offset-based (steps_out) outer references; id-based (rel_reference) ones resolve +# against ``anchor_scope`` instead. outer_schemas: contextvars.ContextVar = contextvars.ContextVar( "outer_schemas", default=() ) +# The plan-wide anchor index (an ``_AnchorScope``) for the plan currently being +# inferred, or None outside ``infer_plan_schema``. Set for the duration of a +# whole-plan inference so a ``rel_reference`` anywhere in the tree resolves. +anchor_scope: contextvars.ContextVar = contextvars.ContextVar( + "anchor_scope", default=None +) + def _derive_extension_schema(detail, inputs, registry): """Derive a registered extension relation's output NamedStruct, or None. @@ -277,20 +328,34 @@ def infer_expression_type( # the correlated-subquery stack); a lambda parameter reference against # the lambda's parameter struct; otherwise against the input row. if root_type == "outer_reference": - stack = outer_schemas.get() - # steps_out is 1-based per the Substrait spec (1 = the immediately - # enclosing query), so it indexes back from the top of the stack. - steps = expression.selection.outer_reference.steps_out - if steps < 1: - raise Exception( - f"outer reference has steps_out={steps}; Substrait requires " - "steps_out >= 1 (1 = the immediately enclosing query)" - ) - if steps > len(stack): - raise Exception( - "outer reference outside an enclosing (correlated) query" - ) - schema = stack[len(stack) - steps].struct + outer_ref = expression.selection.outer_reference + # An outer reference resolves either by id (rel_reference -> the schema + # of the relation carrying that rel_anchor, via the plan-wide anchor + # index) or by offset (steps_out -> that many levels up the + # correlated-subquery stack). These are a protobuf oneof. + if outer_ref.WhichOneof("outer_reference_type") == "rel_reference": + anchors = anchor_scope.get() + if anchors is None: + raise Exception( + "rel_reference outer reference requires whole-plan context; " + "infer via infer_plan_schema" + ) + schema = anchors.schema_of(outer_ref.rel_reference, registry) + else: + stack = outer_schemas.get() + # steps_out is 1-based per the Substrait spec (1 = the immediately + # enclosing query), so it indexes back from the top of the stack. + steps = outer_ref.steps_out + if steps < 1: + raise Exception( + f"outer reference has steps_out={steps}; Substrait requires " + "steps_out >= 1 (1 = the immediately enclosing query)" + ) + if steps > len(stack): + raise Exception( + "outer reference outside an enclosing (correlated) query" + ) + schema = stack[len(stack) - steps].struct else: assert root_type in ("root_reference", "lambda_parameter_reference") schema = parent_schema @@ -802,7 +867,17 @@ def infer_plan_schema(plan: stp.Plan, *, registry=None) -> stt.NamedStruct: # ReferenceRel anywhere in the tree resolves against them by ordinal. Wrap them # in a _SubtreeScope so repeated references are memoized and cycles are caught. subtrees = _SubtreeScope(plan_subtrees(plan)) - root = plan.relations[-1].root - schema = infer_rel_schema(root.input, registry=registry, subtrees=subtrees) + # Index every RelCommon.rel_anchor in the plan (across subtrees, the root, and + # subquery-embedded relations) so an id-based OuterReference (rel_reference) + # anywhere resolves against the anchored relation's output schema. + anchors = { + a: rel for rel in iter_plan_rels(plan) if (a := rel_anchor_of(rel)) is not None + } + token = anchor_scope.set(_AnchorScope(anchors, subtrees)) + try: + root = plan.relations[-1].root + schema = infer_rel_schema(root.input, registry=registry, subtrees=subtrees) + finally: + anchor_scope.reset(token) return stt.NamedStruct(names=root.names, struct=schema) diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 8dd9790..9fedf2e 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -116,7 +116,7 @@ def _iter_child_rels(rel: stalg.Rel): return node = getattr(rel, rel_type) for field in _child_rel_fields(node): - if field.label == field.LABEL_REPEATED: + if field.is_repeated: children = getattr(node, field.name) for i in range(len(children)): yield children, i @@ -176,6 +176,314 @@ def _inline_reference_rels_in_place(rel: stalg.Rel, subtrees) -> None: _inline_reference_rels_in_place(child, subtrees) +def _iter_direct_subexpressions(msg): + """Yield the immediate ``Expression`` messages owned by ``msg``. + + Recurses through sub-messages that are neither ``Expression`` nor ``Rel`` (e.g. + ``FunctionArgument``, ``IfClause``, the ``Subquery`` wrappers), yields each + ``Expression``-typed field without descending into it, and stops at ``Rel`` + fields (child relations / subquery inputs, handled separately). Discovered from + the protobuf descriptor so it stays correct as the schema evolves -- a shallow + scan of top-level ``Expression`` fields would miss e.g. an aggregate measure's + arguments (``measures[].measure.arguments[].value``) or a sort key + (``sorts[].expr``). + """ + for field in msg.DESCRIPTOR.fields: + if field.message_type is None: + continue + if field.message_type.GetOptions().map_entry: + continue + full = field.message_type.full_name + if field.is_repeated: + values = getattr(msg, field.name) + if full == "substrait.Expression": + yield from values + elif full != "substrait.Rel": + for value in values: + yield from _iter_direct_subexpressions(value) + elif msg.HasField(field.name): + if full == "substrait.Expression": + yield getattr(msg, field.name) + elif full != "substrait.Rel": + yield from _iter_direct_subexpressions(getattr(msg, field.name)) + + +def _iter_named_direct_expressions(node): + """Like :func:`_iter_direct_subexpressions` but yields ``(field_name, expr)``, + tagging each expression with the name of the top-level field on ``node`` it was + reached through. Used to tell a join's ``post_join_filter`` (output-scoped) from + its condition / ``residual_expression`` (combined-inputs-scoped).""" + for field in node.DESCRIPTOR.fields: + if field.message_type is None or field.message_type.GetOptions().map_entry: + continue + full = field.message_type.full_name + if full == "substrait.Rel": + continue + if field.is_repeated: + values = getattr(node, field.name) + elif node.HasField(field.name): + values = [getattr(node, field.name)] + else: + continue + for value in values: + if full == "substrait.Expression": + yield field.name, value + else: + for expr in _iter_direct_subexpressions(value): + yield field.name, expr + + +def _iter_rel_expressions(rel: stalg.Rel): + """Yield the root ``Expression`` messages a relation owns (its own scope): the + filter condition, project expressions, join condition, aggregate/sort/expand + expressions, etc. Expressions inside child relations belong to those relations.""" + rel_type = rel.WhichOneof("rel_type") + if rel_type is None: + return + yield from _iter_direct_subexpressions(getattr(rel, rel_type)) + + +def _iter_subquery_rels(expr: stalg.Expression): + """Yield the input ``Rel``(s) of a subquery expression (``scalar.input``, + ``set_predicate.tuples``, ``set_comparison.right``, ``in_predicate.haystack``), + discovered from the descriptor. Empty for a non-subquery expression.""" + if expr.WhichOneof("rex_type") != "subquery": + return + variant = expr.subquery.WhichOneof("subquery_type") + if variant is None: + return + inner = getattr(expr.subquery, variant) + for field in _child_rel_fields(inner): + if field.is_repeated: + yield from getattr(inner, field.name) + elif inner.HasField(field.name): + yield getattr(inner, field.name) + + +def _iter_subquery_rels_in_expr(expr: stalg.Expression): + """Yield every subquery input ``Rel`` reachable from ``expr`` in its own scope + (i.e. not descending into those inner relations).""" + yield from _iter_subquery_rels(expr) + for sub in _iter_direct_subexpressions(expr): + yield from _iter_subquery_rels_in_expr(sub) + + +def _walk_rel(rel: stalg.Rel): + yield rel + for container, key in _iter_child_rels(rel): + yield from _walk_rel(_child_rel(container, key)) + for expr in _iter_rel_expressions(rel): + for inner in _iter_subquery_rels_in_expr(expr): + yield from _walk_rel(inner) + + +def iter_plan_rels(plan: stplan.Plan): + """Yield every ``Rel`` in a Plan: the shared subtrees, the query root, all direct + child relations, and relations embedded inside ``Expression`` subqueries.""" + for pr in plan.relations: + kind = pr.WhichOneof("rel_type") + if kind == "rel": + yield from _walk_rel(pr.rel) + elif kind == "root": + yield from _walk_rel(pr.root.input) + + +def _rel_node_with_common(rel: stalg.Rel): + """The active relation-variant submessage of ``rel`` if it carries a + ``RelCommon``, else ``None`` (a ``ReferenceRel`` has no ``common``).""" + rel_type = rel.WhichOneof("rel_type") + if rel_type is None: + return None + node = getattr(rel, rel_type) + if node.DESCRIPTOR.fields_by_name.get("common") is None: + return None + return node + + +def rel_anchor_of(rel: stalg.Rel): + """The ``RelCommon.rel_anchor`` of ``rel`` if set, else ``None``.""" + node = _rel_node_with_common(rel) + if node is None or not node.common.HasField("rel_anchor"): + return None + return node.common.rel_anchor + + +def _all_subexpressions(expr: stalg.Expression): + """Yield every ``Expression`` transitively owned by ``expr`` in its own scope + (through operators and subquery-wrapping expressions, but not into subquery + input relations).""" + for sub in _iter_direct_subexpressions(expr): + yield sub + yield from _all_subexpressions(sub) + + +def _is_steps_out_ref(expr: stalg.Expression) -> bool: + if expr.WhichOneof("rex_type") != "selection": + return False + sel = expr.selection + return ( + sel.WhichOneof("root_type") == "outer_reference" + and sel.outer_reference.WhichOneof("outer_reference_type") == "steps_out" + ) + + +def _plan_has_steps_out(plan: stplan.Plan) -> bool: + """Whether any ``OuterReference`` in ``plan`` is still offset-based (``steps_out``) + -- i.e. whether :func:`to_id_based_outer_references` has anything to rewrite. A + cheap pre-scan so the common (non-correlated) plan skips the copy-and-walk.""" + for rel in iter_plan_rels(plan): + for expr in _iter_rel_expressions(rel): + if _is_steps_out_ref(expr) or any( + _is_steps_out_ref(sub) for sub in _all_subexpressions(expr) + ): + return True + return False + + +# A join's expression fields bind against two different rows: ``post_join_filter`` +# 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 +# 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 + 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) + + +def to_id_based_outer_references(plan: stplan.Plan) -> stplan.Plan: + """A copy of ``plan`` with every offset-based ``OuterReference`` (``steps_out``) + rewritten to the id-based form (``rel_reference`` naming a + ``RelCommon.rel_anchor``). + + The binding relation an ``OuterReference`` resolves against is stamped with a + plan-wide-unique ``rel_anchor`` (``>= 1``, per the Substrait spec), and the + reference is rewritten to name it; several references to the same scope share one + anchor, and an anchor already present is reused. References already id-based are + left unchanged, so the pass is idempotent and tolerates a partially-converted + input. A plan with nothing to rewrite is returned unchanged (no copy). + + A ``steps_out`` reference resolves against an enclosing query's *row*, which is + the output of a specific relation, so the relation whose output is that row is + anchored: + + * a single-input host (``Filter`` / ``Project`` / ...) exposes its **input**'s + row. If that input is a ``ReferenceRel`` (a ``cache()``-shared subtree, which + carries no ``RelCommon``), the shared subtree it points at is anchored instead + -- the ``ReferenceRel`` has no ``emit``, so its output is exactly the subtree's. + This is the shared-subtree / DAG case that offset-based ``steps_out`` cannot + address unambiguously. + * a ``post_join_filter``, or a leaf host's own filter, exposes the **host's** + 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 + carries that row -- such a reference is left offset-based (still spec-valid, and + read by inference), rather than mis-anchored. + + Raises only if a resolvable binding carries no ``RelCommon`` at all (e.g. an + ``UpdateRel``), which no correlated-subquery shape produces. + """ + if not _plan_has_steps_out(plan): + return plan + + out = stplan.Plan() + out.CopyFrom(plan) + + subtrees = plan_subtrees(out) + existing = [ + a for a in (rel_anchor_of(r) for r in iter_plan_rels(out)) if a is not None + ] + counter = max(existing) if existing else 0 + anchor_by_id: dict = {} + + def anchor_for(binding: stalg.Rel) -> int: + nonlocal counter + # A ReferenceRel is a plan-global pointer with no RelCommon of its own; its + # output is the shared subtree's, so anchor the subtree instead. + while binding.WhichOneof("rel_type") == "reference": + binding = subtrees[binding.reference.subtree_ordinal] + found = rel_anchor_of(binding) + if found is not None: + return found + key = id(binding) + if key in anchor_by_id: + return anchor_by_id[key] + node = _rel_node_with_common(binding) + if node is None: + raise Exception( + "cannot resolve an outer reference into a " + f"{binding.WhichOneof('rel_type')!r} relation to an id-based " + "rel_reference: it carries no RelCommon to hold a rel_anchor" + ) + counter += 1 + node.common.rel_anchor = counter + anchor_by_id[key] = counter + return counter + + def convert_expr(expr, scope, binding): + rex = expr.WhichOneof("rex_type") + if rex == "selection": + sel = expr.selection + if sel.WhichOneof("root_type") == "outer_reference": + oref = sel.outer_reference + if oref.WhichOneof("outer_reference_type") == "steps_out": + steps = oref.steps_out + if not 1 <= steps <= len(scope): + raise Exception( + f"outer reference steps_out={steps} escapes its " + f"{len(scope)} enclosing query scope(s)" + ) + target = scope[-steps] + # None marks a combined-inputs scope with no anchorable relation + # (a reducing join's condition); leave such a reference as-is. + if target is not None: + oref.rel_reference = anchor_for(target) + elif rex == "subquery": + for inner in _iter_subquery_rels(expr): + convert_rel(inner, scope + [binding]) + for sub in _iter_direct_subexpressions(expr): + convert_expr(sub, scope, binding) + + def convert_rel(rel, scope): + children = list(_iter_child_rels(rel)) + rel_type = rel.WhichOneof("rel_type") + node = getattr(rel, rel_type) if rel_type is not None else None + if node is not None: + # The relation whose output row a subquery here would see one level up: + # a single-input host exposes its input; a leaf or multi-input host its + # own output -- except a reducing join's combined-inputs-scoped fields, + # whose scope no relation's output carries (binding None -> left as-is). + single_input = _child_rel(*children[0]) if len(children) == 1 else None + reducing = single_input is None and _is_reducing_join(node) + for name, expr in _iter_named_direct_expressions(node): + if single_input is not None: + binding = single_input + elif reducing and name in _JOIN_COMBINED_SCOPED_FIELDS: + binding = None + else: + binding = rel + convert_expr(expr, scope, binding) + for container, key in children: + convert_rel(_child_rel(container, key), scope) + + for pr in out.relations: + kind = pr.WhichOneof("rel_type") + if kind == "rel": + convert_rel(pr.rel, []) + elif kind == "root": + convert_rel(pr.root.input, []) + return out + + def merge_extensions_into(target, *sources): """Merge the extension URNs and declarations of ``sources`` into ``target`` in place. diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 03048b4..590b3d7 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -1123,13 +1123,18 @@ def test_correlated_exists(): corr = inner.filter(sub.col("k") == sub.outer("k")) # inner.k == outer.k plan = outer.filter(sub.exists(corr)).to_plan() - inner_cond = plan.relations[ - -1 - ].root.input.filter.condition.subquery.set_predicate.tuples.filter.condition + root_input = plan.relations[-1].root.input + inner_cond = ( + root_input.filter.condition.subquery.set_predicate.tuples.filter.condition + ) lhs, rhs = (a.value.selection for a in inner_cond.scalar_function.arguments) assert lhs.WhichOneof("root_type") == "root_reference" # inner.k assert rhs.WhichOneof("root_type") == "outer_reference" # outer.k - assert rhs.outer_reference.steps_out == 1 # 1 = the immediately enclosing query + # The DataFrame emits id-based correlations: the outer reference names the + # rel_anchor stamped on the binding relation (the enclosing filter's input). + anchor = root_input.filter.input.read.common.rel_anchor + assert rhs.outer_reference.WhichOneof("outer_reference_type") == "rel_reference" + assert rhs.outer_reference.rel_reference == anchor assert rhs.direct_reference.struct_field.field == 0 # "k" in the outer schema @@ -1141,7 +1146,7 @@ def test_correlated_scalar_subquery_chains(): assert plan.relations[-1].root.input.HasField("filter") -def test_nested_correlation_steps_out(): +def test_nested_correlation_id_based(): outer = sub.read_named_table("o", {"k": sub.i64}) mid = sub.read_named_table("m", {"k": sub.i64}) inner = sub.read_named_table("i", {"k": sub.i64}) @@ -1150,11 +1155,85 @@ def test_nested_correlation_steps_out(): mid_corr = mid.filter(sub.exists(inner_corr)) plan = outer.filter(sub.exists(mid_corr)).to_plan() - inner_cond = plan.relations[ - -1 - ].root.input.filter.condition.subquery.set_predicate.tuples.filter.condition.subquery.set_predicate.tuples.filter.condition # mid # inner + root_input = plan.relations[-1].root.input + inner_cond = root_input.filter.condition.subquery.set_predicate.tuples.filter.condition.subquery.set_predicate.tuples.filter.condition # mid # inner rhs = inner_cond.scalar_function.arguments[1].value.selection - assert rhs.outer_reference.steps_out == 2 # two levels out -> the outermost + # steps_out=2 selects the outermost query; it is emitted id-based as a + # rel_reference to the rel_anchor stamped on that outermost relation. + outermost_read = root_input.filter.input.read + assert rhs.outer_reference.WhichOneof("outer_reference_type") == "rel_reference" + assert rhs.outer_reference.rel_reference == outermost_read.common.rel_anchor + + +def test_correlated_exists_over_cached_outer_is_id_based(): + # Regression: the correlation's enclosing scope is a cache()d frame, so its + # binding relation is a ReferenceRel (which carries no RelCommon). The anchor + # lands on the shared subtree the reference points at, and the whole plan still + # infers. Previously to_plan() raised. + from substrait.type_inference import infer_plan_schema + + outer = sub.read_named_table("o", {"k": sub.i64}).cache() + inner = sub.read_named_table("i", {"k": sub.i64}) + corr = inner.filter(sub.col("k") == sub.outer("k")) + plan = outer.filter(sub.exists(corr)).to_plan() + + root_input = plan.relations[-1].root.input + # The binding is a ReferenceRel; the anchor is on the shared subtree. + assert root_input.filter.input.WhichOneof("rel_type") == "reference" + anchor = plan.relations[0].rel.read.common.rel_anchor + rhs = root_input.filter.condition.subquery.set_predicate.tuples.filter.condition.scalar_function.arguments[ + 1 + ].value.selection.outer_reference + assert rhs.WhichOneof("outer_reference_type") == "rel_reference" + assert rhs.rel_reference == anchor + infer_plan_schema(plan) + + +def test_correlated_exists_in_join_post_filter_is_id_based(): + # Regression: a correlated subquery in a join's post_join_filter binds against + # the join output, i.e. the join itself -- so the join is anchored. Previously + # to_plan() raised on the multi-input host. + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"a": sub.i64}) + right = sub.read_named_table("r", {"b": sub.i64}) + inner = sub.read_named_table("i", {"k": sub.i64}) + corr = inner.filter(sub.col("k") == sub.outer("a")) + plan = left.join( + right, sub.col("a") == sub.col("b"), how="inner", post_filter=sub.exists(corr) + ).to_plan() + + join = plan.relations[-1].root.input + assert join.WhichOneof("rel_type") == "join" + anchor = join.join.common.rel_anchor + rhs = join.join.post_join_filter.subquery.set_predicate.tuples.filter.condition.scalar_function.arguments[ + 1 + ].value.selection.outer_reference + assert rhs.WhichOneof("outer_reference_type") == "rel_reference" + assert rhs.rel_reference == anchor + infer_plan_schema(plan) + + +def test_correlated_exists_in_reducing_join_condition_stays_steps_out(): + # A reducing join (semi) emits one side, so its output row can't carry the + # combined condition scope the correlation sees; the reference is left + # offset-based (steps_out) -- still spec-valid and read by inference -- rather + # than mis-anchored. to_plan() must still succeed. + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"a": sub.i64}) + right = sub.read_named_table("r", {"b": sub.i64}) + inner = sub.read_named_table("i", {"k": sub.i64}) + corr = inner.filter(sub.col("k") == sub.outer("b")) # references the right side + plan = left.join(right, sub.exists(corr), how="left_semi").to_plan() + + join = plan.relations[-1].root.input + assert not join.join.common.HasField("rel_anchor") + rhs = join.join.expression.subquery.set_predicate.tuples.filter.condition.scalar_function.arguments[ + 1 + ].value.selection.outer_reference + assert rhs.WhichOneof("outer_reference_type") == "steps_out" + infer_plan_schema(plan) def test_outer_outside_subquery_raises(): diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index ca603bc..0a042dc 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -1,10 +1,12 @@ import pytest import substrait.algebra_pb2 as stalg +import substrait.plan_pb2 as stp import substrait.type_pb2 as stt from substrait.type_inference import ( infer_expression_type, infer_nested_type, + infer_plan_schema, infer_rel_schema, ) @@ -626,3 +628,115 @@ def test_inference_reference_out_of_range_raises(): # No subtrees in scope at all is also out of range. with pytest.raises(Exception, match="out of range"): infer_rel_schema(ref) + + +def _outer_ref(field, *, rel_reference=None, steps_out=None): + outer = stalg.Expression.FieldReference.OuterReference() + if rel_reference is not None: + outer.rel_reference = rel_reference + else: + outer.steps_out = steps_out + return stalg.Expression( + selection=stalg.Expression.FieldReference( + outer_reference=outer, + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=field) + ), + ) + ) + + +def test_infer_rel_reference_resolves_against_anchored_subtree(): + # A rel_reference resolves against the output schema of whatever relation in the + # plan carries the matching rel_anchor -- here a shared subtree -- which the + # offset-based steps_out could not address. The project appends order_total + # (field 2, fp32 nullable) pulled from the anchored subtree. + anchored = stalg.Rel( + read=stalg.ReadRel( + base_schema=named_struct, + common=stalg.RelCommon(rel_anchor=7), + named_table=stalg.ReadRel.NamedTable(names=["shared"]), + ) + ) + root_input = stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, + expressions=[_outer_ref(2, rel_reference=7)], + ) + ) + plan = stp.Plan( + relations=[ + stp.PlanRel(rel=anchored), + stp.PlanRel( + root=stalg.RelRoot( + input=root_input, + names=["order_id", "is_refundable", "order_total"], + ) + ), + ] + ) + + expected = stt.Type.Struct( + types=list(right_struct.types) + + [stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE))] + ) + assert infer_plan_schema(plan).struct == expected + + +def test_infer_rel_reference_unknown_anchor_raises(): + root_input = stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, expressions=[_outer_ref(0, rel_reference=99)] + ) + ) + plan = stp.Plan( + relations=[ + stp.PlanRel(root=stalg.RelRoot(input=root_input, names=["a", "b", "c"])) + ] + ) + with pytest.raises(Exception, match="unknown rel_anchor 99"): + infer_plan_schema(plan) + + +def test_infer_expression_rel_reference_without_plan_context_raises(): + # Resolving a rel_reference needs the plan-wide anchor index; a bare + # infer_expression_type call (no infer_plan_schema) has no index in scope. + with pytest.raises(Exception, match="whole-plan context"): + infer_expression_type(_outer_ref(0, rel_reference=1), struct) + + +def test_infer_rel_reference_anchor_zero_is_a_distinct_anchor(): + # rel_anchor has explicit field presence, so 0 is a set, valid anchor distinct + # from "absent". The anchor index must key on presence, not truthiness, or a + # legitimate anchor 0 would be silently dropped. (The builder-side converter + # never emits 0, but an externally-produced or #228 lateral-join plan can.) + anchored = stalg.Rel( + read=stalg.ReadRel( + base_schema=named_struct, + common=stalg.RelCommon(rel_anchor=0), + named_table=stalg.ReadRel.NamedTable(names=["shared"]), + ) + ) + root_input = stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, + expressions=[_outer_ref(2, rel_reference=0)], + ) + ) + plan = stp.Plan( + relations=[ + stp.PlanRel(rel=anchored), + stp.PlanRel( + root=stalg.RelRoot( + input=root_input, + names=["order_id", "is_refundable", "order_total"], + ) + ), + ] + ) + + expected = stt.Type.Struct( + types=list(right_struct.types) + + [stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE))] + ) + assert infer_plan_schema(plan).struct == expected diff --git a/tests/test_utils.py b/tests/test_utils.py index aa66ff0..aef1f6c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,12 +1,17 @@ import pytest +import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee import substrait.extensions.extensions_pb2 as ste +import substrait.plan_pb2 as stplan import substrait.type_pb2 as stt from substrait.utils import ( + iter_plan_rels, merge_extension_declarations, merge_extension_urns, merge_extensions_into, + rel_anchor_of, + to_id_based_outer_references, type_num_names, ) @@ -177,3 +182,271 @@ def test_merge_extension_declarations_rejects_non_function_mapping(): with pytest.raises(NotImplementedError, match="extension_type"): merge_extension_declarations([declaration]) + + +# --- to_id_based_outer_references ---------------------------------------------- +# +# Compact hand-built plans exercising the steps_out -> rel_reference conversion. +# A filter condition here is a bare (outer) reference rather than a realistic +# boolean predicate -- the converter only walks expressions to find and rewrite +# OuterReferences, so the surrounding operator shape is what matters. + + +def _read(name: str, ncols: int = 1) -> stalg.Rel: + return stalg.Rel( + read=stalg.ReadRel( + base_schema=stt.NamedStruct( + names=[f"c{i}" for i in range(ncols)], + struct=stt.Type.Struct( + types=[stt.Type(i64=stt.Type.I64()) for _ in range(ncols)] + ), + ), + named_table=stalg.ReadRel.NamedTable(names=[name]), + ) + ) + + +def _outer(steps_out: int, field: int = 0) -> stalg.Expression: + return stalg.Expression( + selection=stalg.Expression.FieldReference( + outer_reference=stalg.Expression.FieldReference.OuterReference( + steps_out=steps_out + ), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=field) + ), + ) + ) + + +def _exists(inner: stalg.Rel) -> stalg.Expression: + return stalg.Expression( + subquery=stalg.Expression.Subquery( + set_predicate=stalg.Expression.Subquery.SetPredicate( + predicate_op=stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_EXISTS, + tuples=inner, + ) + ) + ) + + +def _filter(input: stalg.Rel, condition: stalg.Expression) -> stalg.Rel: + return stalg.Rel(filter=stalg.FilterRel(input=input, condition=condition)) + + +def _plan(root_input: stalg.Rel, *subtrees: stalg.Rel) -> stplan.Plan: + relations = [stplan.PlanRel(rel=s) for s in subtrees] + relations.append(stplan.PlanRel(root=stalg.RelRoot(input=root_input, names=["c0"]))) + return stplan.Plan(relations=relations) + + +def _outer_refs(plan: stplan.Plan): + """Every OuterReference in a plan (across subquery-embedded relations).""" + return [ + rel.filter.condition + for rel in iter_plan_rels(plan) + if rel.WhichOneof("rel_type") == "filter" + and rel.filter.condition.WhichOneof("rex_type") == "selection" + and rel.filter.condition.selection.WhichOneof("root_type") == "outer_reference" + ] + + +def test_convert_correlated_exists_stamps_anchor_and_rewrites(): + plan = _plan(_filter(_read("o"), _exists(_filter(_read("i"), _outer(1))))) + out = to_id_based_outer_references(plan) + + binding = out.relations[-1].root.input.filter.input # the outer read + assert rel_anchor_of(binding) == 1 + + ref = out.relations[ + -1 + ].root.input.filter.condition.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.rel_reference == 1 + # The input plan is untouched (conversion works on a copy). + assert rel_anchor_of(plan.relations[-1].root.input.filter.input) is None + + +def test_convert_dedup_same_scope_shares_one_anchor(): + # Two correlated columns from the same enclosing scope share a single anchor. + inner = _filter( + _read("i"), + stalg.Expression( + nested=stalg.Expression.Nested( + struct=stalg.Expression.Nested.Struct( + fields=[_outer(1, 0), _outer(1, 1)] + ) + ) + ), + ) + plan = _plan(_filter(_read("o", ncols=2), _exists(inner))) + out = to_id_based_outer_references(plan) + + anchors = {a for r in iter_plan_rels(out) if (a := rel_anchor_of(r))} + assert anchors == {1} + refs = out.relations[ + -1 + ].root.input.filter.condition.subquery.set_predicate.tuples.filter.condition.nested.struct.fields + assert [f.selection.outer_reference.rel_reference for f in refs] == [1, 1] + + +def test_convert_distinct_scopes_get_distinct_anchors(): + # An inner subquery referencing two different enclosing scopes (steps_out 1 and + # 2) anchors each binding relation separately. + inner = _filter( + _read("i"), + stalg.Expression( + nested=stalg.Expression.Nested( + struct=stalg.Expression.Nested.Struct(fields=[_outer(1), _outer(2)]) + ) + ), + ) + mid = _filter(_read("m"), _exists(inner)) + plan = _plan(_filter(_read("o"), _exists(mid))) + out = to_id_based_outer_references(plan) + + root = out.relations[-1].root.input + outer_read = root.filter.input + mid_read = root.filter.condition.subquery.set_predicate.tuples.filter.input + assert rel_anchor_of(mid_read) != rel_anchor_of(outer_read) + + fields = root.filter.condition.subquery.set_predicate.tuples.filter.condition.subquery.set_predicate.tuples.filter.condition.nested.struct.fields + assert fields[0].selection.outer_reference.rel_reference == rel_anchor_of(mid_read) + assert fields[1].selection.outer_reference.rel_reference == rel_anchor_of( + outer_read + ) + + +def test_convert_preexisting_anchor_reused_and_counter_continues(): + # An anchor already present in the plan is left alone; freshly allocated + # anchors continue past the maximum existing one. + marked = _read("i") + marked.read.common.rel_anchor = 5 + plan = _plan(_filter(_read("o"), _exists(_filter(marked, _outer(1))))) + out = to_id_based_outer_references(plan) + + assert rel_anchor_of(out.relations[-1].root.input.filter.input) == 6 + + +def _join( + left: stalg.Rel, + right: stalg.Rel, + *, + type=stalg.JoinRel.JOIN_TYPE_INNER, + expression: stalg.Expression = None, + post_join_filter: stalg.Expression = None, +) -> stalg.Rel: + return stalg.Rel( + join=stalg.JoinRel( + left=left, + right=right, + type=type, + expression=expression, + post_join_filter=post_join_filter, + ) + ) + + +def test_convert_reference_rel_binding_anchors_the_shared_subtree(): + # A correlation whose enclosing host input is a ReferenceRel (a cache()d frame) + # has no RelCommon on the reference itself; the anchor lands on the shared + # subtree it points at, whose output is exactly the reference's. This is the + # DAG case that offset-based steps_out cannot address unambiguously. + ref_rel = stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=0)) + plan = _plan( + _filter(ref_rel, _exists(_filter(_read("i"), _outer(1)))), + _read("s"), # subtree at ordinal 0 + ) + out = to_id_based_outer_references(plan) + + subtree = out.relations[0].rel + assert rel_anchor_of(subtree) == 1 + # The ReferenceRel binding is left as-is (it carries no RelCommon). + assert rel_anchor_of(out.relations[-1].root.input.filter.input) is None + ref = out.relations[ + -1 + ].root.input.filter.condition.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.rel_reference == 1 + + +def test_convert_multi_input_join_condition_anchors_the_join(): + # A correlation into a (non-reducing) join's condition resolves against the + # combined left+right row, which equals the join's own output -- so the join + # relation is anchored and the reference names it. + join = _join( + _read("l"), + _read("r"), + expression=_exists(_filter(_read("i"), _outer(1))), + ) + out = to_id_based_outer_references(_plan(join)) + + out_join = out.relations[-1].root.input + assert rel_anchor_of(out_join) == 1 + ref = out_join.join.expression.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.rel_reference == 1 + + +def test_convert_post_join_filter_anchors_the_join(): + # 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"), + post_join_filter=_exists(_filter(_read("i"), _outer(1))), + ) + out = to_id_based_outer_references(_plan(join)) + + out_join = out.relations[-1].root.input + assert rel_anchor_of(out_join) == 1 + ref = out_join.join.post_join_filter.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + 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 + # 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, + expression=_exists(_filter(_read("i"), _outer(1))), + ) + out = to_id_based_outer_references(_plan(join)) + + out_join = out.relations[-1].root.input + assert rel_anchor_of(out_join) is None + ref = out_join.join.expression.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "steps_out" + assert ref.steps_out == 1 + + +def test_convert_binding_without_rel_common_raises(): + # A binding relation that carries no RelCommon at all (an UpdateRel) cannot hold + # an anchor. No correlated-subquery shape produces this, but the converter + # guards it rather than silently dropping the reference. + update = stalg.Rel( + update=stalg.UpdateRel( + named_table=stalg.NamedTable(names=["t"]), + ) + ) + plan = _plan(_filter(update, _exists(_filter(_read("i"), _outer(1))))) + with pytest.raises(Exception, match="no RelCommon"): + to_id_based_outer_references(plan) + + +def test_convert_noncorrelated_plan_unchanged(): + plan = _plan( + _filter(_read("o"), stalg.Expression(literal=stalg.Expression.Literal(i64=1))) + ) + assert to_id_based_outer_references(plan) == plan + + +def test_convert_is_idempotent(): + plan = _plan(_filter(_read("o"), _exists(_filter(_read("i"), _outer(1))))) + once = to_id_based_outer_references(plan) + twice = to_id_based_outer_references(once) + assert twice == once