diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index 7a33aa3..f692a74 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -1,4 +1,6 @@ import calendar +import contextlib +import contextvars import itertools import uuid as uuid_module from datetime import date, datetime, time, timedelta, timezone @@ -20,6 +22,33 @@ type_num_names, ) +# Monotonic source of unique RelCommon.rel_anchor values within a single build. +# Not reset between builds by default; reset at each top-level materialization +# (see fresh_rel_anchors) so a plan built the same way twice numbers alike. +_rel_anchor_counter: contextvars.ContextVar = contextvars.ContextVar( + "_rel_anchor_counter", default=0 +) + + +def next_rel_anchor() -> int: + """Allocate a fresh RelCommon.rel_anchor, unique within the current build.""" + n = _rel_anchor_counter.get() + 1 + _rel_anchor_counter.set(n) + return n + + +@contextlib.contextmanager +def fresh_rel_anchors(): + """Number rel_anchors from 1 within this block, restoring the prior counter + afterwards. Wrap a top-level materialization so repeated builds of the same + plan assign identical anchors.""" + token = _rel_anchor_counter.set(0) + try: + yield + finally: + _rel_anchor_counter.reset(token) + + UnboundExtendedExpression = Callable[ [stp.NamedStruct, ExtensionRegistry], stee.ExtendedExpression ] @@ -404,6 +433,52 @@ def resolve( return resolve +class LateralInput: + """A handle to a lateral join's left input, passed to the ``right`` builder + by :func:`~substrait.builders.plan.lateral_join`. + + Its :meth:`column` references resolve against the left row via an id-based + ``OuterReference`` (``rel_reference`` naming the join's + ``RelCommon.rel_anchor``), so the right (dependent) input can correlate on + the current left row by capturing the handle -- no nesting-depth bookkeeping. + """ + + def __init__(self, rel_anchor: int, schema: stp.NamedStruct): + self._rel_anchor = rel_anchor + self._schema = schema + + def column(self, field: Union[str, int]): + """A correlated reference to the left row's ``field`` (name or index).""" + rel_anchor = self._rel_anchor + schema = self._schema + + def resolve( + base_schema: stp.NamedStruct, registry: ExtensionRegistry + ) -> stee.ExtendedExpression: + # Resolve the column against the left schema, then re-root it as an + # id-based outer reference (keeping the resolved struct-field segment). + resolved = column(field)(schema, registry).referred_expr[0] + segment = resolved.expression.selection.direct_reference + expr = stalg.Expression( + selection=stalg.Expression.FieldReference( + outer_reference=stalg.Expression.FieldReference.OuterReference( + rel_reference=rel_anchor + ), + direct_reference=segment, + ) + ) + return stee.ExtendedExpression( + referred_expr=[ + stee.ExpressionReference( + expression=expr, output_names=resolved.output_names + ) + ], + base_schema=base_schema, + ) + + return resolve + + def column(field: Union[str, int], alias: Union[Iterable[str], str, None] = None): """Builds a resolver for ExtendedExpression containing a FieldReference expression diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 91efbd8..dff7d70 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -16,11 +16,15 @@ from substrait.builders.extended_expression import ( ExtendedExpressionOrUnbound, + LateralInput, + next_rel_anchor, resolve_expression, ) from substrait.extension_registry import ExtensionRegistry from substrait.type_inference import ( _join_output_struct, + _join_struct_from_schemas, + _outer_anchor_binding, infer_plan_schema, join_output_names, ) @@ -636,6 +640,107 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: return resolve +def lateral_join( + left: PlanOrUnbound, + right: Callable[[LateralInput], PlanOrUnbound], + type: stalg.JoinRel.JoinType, + *, + expression: Optional[ExtendedExpressionOrUnbound] = None, + post_join_filter: Optional[ExtendedExpressionOrUnbound] = None, + extension: Optional[AdvancedExtension] = None, +) -> UnboundPlan: + """A LateralJoinRel: the right (dependent) input is evaluated once per left + row and may reference the current left row. + + ``right`` is a function of a :class:`~substrait.builders.extended_expression.LateralInput` + handle to the left; use ``handle.column(...)`` inside it to correlate on the + current left row (an id-based ``OuterReference`` to this relation's + ``rel_anchor``). Capturing the handle avoids counting nesting levels: an + inner lateral join can reference an outer one by using its handle directly. + + Only INNER and left-oriented join types are valid for lateral joins: INNER, + LEFT, LEFT_SEMI, LEFT_ANTI, LEFT_SINGLE, LEFT_MARK. ``expression`` is an + optional match condition over the combined left+right schema. + """ + + def resolve(registry: ExtensionRegistry) -> stp.Plan: + bound_left = left if isinstance(left, stp.Plan) else left(registry) + left_ns = infer_plan_schema(bound_left, registry=registry) + + anchor = next_rel_anchor() + handle = LateralInput(anchor, left_ns) + # Bind the left schema to `anchor` while the right input is built and its + # schema inferred, so id-based references (handle.column(...)) resolve to + # the current left row during that inference. + with _outer_anchor_binding(anchor, left_ns.struct): + unbound_right = right(handle) + bound_right = ( + unbound_right + if isinstance(unbound_right, stp.Plan) + else unbound_right(registry) + ) + right_ns = infer_plan_schema(bound_right, registry=registry) + + # The join condition binds against the combined left+right input row. + 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_expression = ( + resolve_expression(expression, ns, registry) + if expression is not None + 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). + 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. + bound_post = None + if post_join_filter is not None: + output_ns = stt.NamedStruct( + names=out_names, + struct=_join_struct_from_schemas( + type_name, left_ns.struct, right_ns.struct + ), + ) + bound_post = resolve_expression(post_join_filter, output_ns, registry) + + return _plan_from( + [bound_left, bound_right], + lambda inp: stalg.Rel( + lateral_join=stalg.LateralJoinRel( + common=stalg.RelCommon(rel_anchor=anchor), + left=inp[0], + right=inp[1], + expression=( + bound_expression.referred_expr[0].expression + if bound_expression + else None + ), + post_join_filter=( + bound_post.referred_expr[0].expression if bound_post else None + ), + type=type, + advanced_extension=extension, + ) + ), + out_names, + (bound_left, bound_right, bound_expression, bound_post), + ) + + return resolve + + def cross( left: PlanOrUnbound, right: PlanOrUnbound, diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index f687b36..54f4058 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -29,7 +29,7 @@ from __future__ import annotations from itertools import combinations -from typing import Any, Iterable, Optional, Union +from typing import Any, Callable, Iterable, Optional, Union import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stplan @@ -37,6 +37,7 @@ from substrait.builders import plan as _plan from substrait.builders import type as _type +from substrait.builders.extended_expression import LateralInput, fresh_rel_anchors 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 @@ -60,6 +61,20 @@ "right_mark": stalg.JoinRel.JOIN_TYPE_RIGHT_MARK, } +# Lateral joins evaluate the right input per left row, so only INNER and +# left-oriented join types are valid (RIGHT-oriented and OUTER have no meaning). +_LATERAL_JOIN_TYPES = { + how: _JOIN_TYPES[how] + for how in ( + "inner", + "left", + "left_semi", + "left_anti", + "left_single", + "left_mark", + ) +} + # Write create-modes: what to do when the target table already exists. _CREATE_MODES = { "error": stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS, @@ -159,6 +174,26 @@ def _unbound(value: Any): return value # assume already an unbound expression callable +class LateralLeft: + """Handle to a lateral join's left input, passed to the ``right`` builder of + :meth:`DataFrame.lateral_join`. + + Its columns are correlated references to the current left row (an id-based + ``OuterReference``), so the right frame can be built as a function of the + left without counting nesting levels. + """ + + def __init__(self, handle: LateralInput): + self._handle = handle + + def col(self, name: Union[str, int]) -> Expr: + """A correlated reference to the left row's column ``name`` (or index).""" + return Expr(self._handle.column(name)) + + def __getitem__(self, name: Union[str, int]) -> Expr: + return self.col(name) + + class DataFrame: """The Substrait-native fluent DataFrame. @@ -389,6 +424,53 @@ def cross_join(self, other: "DataFrame") -> "DataFrame": right row).""" return self._next(_plan.cross(self._plan, other._plan)) + def lateral_join( + self, + right: "Callable[[LateralLeft], DataFrame]", + how: str = "inner", + *, + on: Union[Expr, Any, None] = None, + post_filter: Union[Expr, Any, None] = None, + ) -> "DataFrame": + """Lateral join: evaluate the right frame once per row of this frame. + + ``right`` is a function of a :class:`LateralLeft` handle to this frame; + use ``left.col(...)`` inside it to correlate on the current left row:: + + left.lateral_join(lambda lat: inner.filter(sub.col("k") == lat.col("k"))) + + Capturing the handle avoids counting nesting levels -- an inner lateral + join can reference an outer one via its own handle. + + ``on`` is an optional match condition over the combined left+right + schema. Only ``inner`` and left-oriented join types are valid for + lateral joins: ``inner``, ``left``, ``left_semi``, ``left_anti``, + ``left_single``, ``left_mark``. ``post_filter`` is an optional predicate + applied to the join output. + """ + try: + join_type = _LATERAL_JOIN_TYPES[how] + except KeyError: + raise ValueError( + f"unknown lateral join type {how!r}; expected one of " + f"{sorted(_LATERAL_JOIN_TYPES)}" + ) from None + + def build_right(handle: LateralInput): + return right(LateralLeft(handle))._plan + + return self._next( + _plan.lateral_join( + self._plan, + build_right, + type=join_type, + expression=_unbound(on) if on is not None else None, + post_join_filter=( + _unbound(post_filter) if post_filter is not None else None + ), + ) + ) + def nested_loop_join( self, other: "DataFrame", on: Union[Expr, Any], how: str = "inner" ) -> "DataFrame": @@ -718,9 +800,14 @@ def write_named_table( 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)) + correlated (outer) references in the id-based form (``rel_reference``): a + lateral join's builder assigns them directly, and any offset-based + ``steps_out`` a correlated subquery produced is rewritten here. Building + under ``fresh_rel_anchors`` numbers lateral-join anchors from 1 per + materialization, so building the same frame twice yields identical plans.""" + with fresh_rel_anchors(): + plan = self._plan(registry) + return to_id_based_outer_references(plan) def to_plan(self) -> stplan.Plan: """Materialize to a ``substrait.proto.Plan``.""" diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index f99c150..8c431a6 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -1,3 +1,4 @@ +import contextlib import contextvars import substrait.algebra_pb2 as stalg @@ -54,31 +55,43 @@ def schema_of(self, ordinal, registry): 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. + """Map of ``RelCommon.rel_anchor`` -> the relation carrying it, with lazy + per-anchor 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 + and resolves against that relation's schema: a ``LateralJoinRel`` anchor denotes + the *current left row* (its left input's schema), any other anchor its output + schema. Resolution is lazy (only referenced anchors are inferred) and memoized; + 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 on a shared subtree still resolves. + + ``register`` pre-binds an anchor to an already-known schema, for a correlated + sub-tree whose anchoring relation is not yet assembled (a lateral join's right + input, at build or inference time). ``parent`` chains to an enclosing scope so + nested correlations still resolve outer anchors. """ - __slots__ = ("_rels", "_subtrees", "_schemas", "_resolving") + __slots__ = ("_rels", "_subtrees", "_schemas", "_resolving", "_parent") - def __init__(self, rels: dict, subtrees): + def __init__(self, rels: dict, subtrees, *, parent=None): self._rels = rels self._subtrees = subtrees self._schemas: dict = {} self._resolving: set = set() + self._parent = parent + + def register(self, anchor, struct: stt.Type.Struct) -> None: + """Pre-bind ``anchor`` to an already-known schema.""" + self._schemas[anchor] = struct def schema_of(self, anchor, registry) -> stt.Type.Struct: if anchor in self._schemas: return self._schemas[anchor] if anchor not in self._rels: + if self._parent is not None: + return self._parent.schema_of(anchor, registry) raise Exception(f"outer reference to unknown rel_anchor {anchor}") if anchor in self._resolving: raise Exception( @@ -86,8 +99,16 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: ) self._resolving.add(anchor) try: + rel = self._rels[anchor] + # A lateral-join anchor denotes the current left row, not the join's + # own output; every other anchor resolves against its output schema. + target = ( + rel.lateral_join.left + if rel.WhichOneof("rel_type") == "lateral_join" + else rel + ) struct = infer_rel_schema( - self._rels[anchor], registry=registry, subtrees=self._subtrees + target, registry=registry, subtrees=self._subtrees ) finally: self._resolving.discard(anchor) @@ -96,22 +117,43 @@ def schema_of(self, anchor, registry) -> stt.Type.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. Used for -# offset-based (steps_out) outer references; id-based (rel_reference) ones resolve -# against ``anchor_scope`` instead. +# field reference with an OuterReference.steps_out root resolves against the right +# level (offset-based, for tree-shaped plans). Pushed by the subquery builders +# while resolving their inner plan. Id-based (rel_reference) outer references +# 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. +# The anchor scope (an ``_AnchorScope``) for the plan / correlated sub-tree +# currently being inferred, or None outside inference. Set for the duration of a +# whole-plan inference (``infer_plan_schema``) so a ``rel_reference`` anywhere in +# the tree resolves, and temporarily narrowed by ``_outer_anchor_binding`` while a +# lateral join's right input is inferred. anchor_scope: contextvars.ContextVar = contextvars.ContextVar( "anchor_scope", default=None ) +@contextlib.contextmanager +def _outer_anchor_binding(anchor, struct): + """Bind ``anchor`` -> ``struct`` for id-based outer references resolved within + the block, chaining to any enclosing anchor scope. + + Used while a ``LateralJoinRel``'s right (dependent) input is built or inferred: + references to the join's ``rel_anchor`` resolve to the current left row even + though the join relation is not yet in an anchor index. Nested lateral joins + compose via the parent chain. + """ + scope = _AnchorScope({}, (), parent=anchor_scope.get()) + scope.register(anchor, struct) + token = anchor_scope.set(scope) + try: + yield + finally: + anchor_scope.reset(token) + + def _derive_extension_schema(detail, inputs, registry): """Derive a registered extension relation's output NamedStruct, or None. @@ -324,15 +366,16 @@ def infer_expression_type( rex_type = expression.WhichOneof("rex_type") if rex_type == "selection": root_type = expression.selection.WhichOneof("root_type") - # An OuterReference resolves against an enclosing query's schema (from - # the correlated-subquery stack); a lambda parameter reference against - # the lambda's parameter struct; otherwise against the input row. + # An OuterReference resolves against an enclosing query's schema; a lambda + # parameter reference against the lambda's parameter struct; otherwise + # against the input row. if root_type == "outer_reference": 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. + # of the relation carrying that rel_anchor, via the anchor scope -- a + # LateralJoinRel anchor gives the current left row) 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: @@ -517,13 +560,12 @@ def join_output_names(type_name: str, left_names, right_names) -> list: return list(left_names) + list(right_names) -def _join_output_struct( - type_name: str, left_rel, right_rel, *, registry=None, subtrees=() +def _join_struct_from_schemas( + type_name: str, left: stt.Type.Struct, right: stt.Type.Struct ) -> stt.Type.Struct: - """Join output column types by join-type NAME (shared across all join - relations, whose enum integer values differ).""" - left = infer_rel_schema(left_rel, registry=registry, subtrees=subtrees) - right = infer_rel_schema(right_rel, registry=registry, subtrees=subtrees) + """Combine already-inferred left/right schemas into a join's output struct + by join-type NAME (shared across all join relations, whose enum integer + values differ).""" required = stt.Type.Nullability.NULLABILITY_REQUIRED shape = _join_column_shape(type_name) if shape == "left": @@ -545,6 +587,42 @@ def _join_output_struct( return stt.Type.Struct(types=types, nullability=required) +def _join_output_struct( + type_name: str, left_rel, right_rel, *, registry=None, subtrees=() +) -> stt.Type.Struct: + """Join output column types by join-type NAME (shared across all join + relations, whose enum integer values differ).""" + left = infer_rel_schema(left_rel, registry=registry, subtrees=subtrees) + right = infer_rel_schema(right_rel, registry=registry, subtrees=subtrees) + return _join_struct_from_schemas(type_name, left, right) + + +def _lateral_join_output_struct( + type_name: str, lateral_join: stalg.LateralJoinRel, *, registry=None, subtrees=() +) -> stt.Type.Struct: + """Lateral-join output column types. + + A lateral join forms output like a regular join (same per-join-type column + shapes), except its right (dependent) input is evaluated once per left row and + may reference the current left row via ``OuterReference.rel_reference`` pointing + to this relation's ``RelCommon.rel_anchor``. The left schema is bound to that + anchor (via the anchor scope) while the right schema is inferred, so those + id-based references resolve. + """ + left = infer_rel_schema(lateral_join.left, registry=registry, subtrees=subtrees) + common = lateral_join.common + if common.HasField("rel_anchor"): + with _outer_anchor_binding(common.rel_anchor, left): + right = infer_rel_schema( + lateral_join.right, registry=registry, subtrees=subtrees + ) + else: + right = infer_rel_schema( + lateral_join.right, registry=registry, subtrees=subtrees + ) + return _join_struct_from_schemas(type_name, left, right) + + def _field_nullability(t: stt.Type): """The nullability of a (concrete) field type, or UNSPECIFIED if it has none.""" kind = t.WhichOneof("kind") @@ -705,6 +783,14 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. subtrees=subtrees, ) (common, struct) = (rel.join.common, raw_schema) + elif rel_type == "lateral_join": + raw_schema = _lateral_join_output_struct( + stalg.JoinRel.JoinType.Name(rel.lateral_join.type), + rel.lateral_join, + registry=registry, + subtrees=subtrees, + ) + (common, struct) = (rel.lateral_join.common, raw_schema) elif rel_type == "window": parent_schema = infer_rel_schema( rel.window.input, registry=registry, subtrees=subtrees @@ -873,7 +959,10 @@ def infer_plan_schema(plan: stp.Plan, *, registry=None) -> stt.NamedStruct: 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)) + # Chain to any enclosing anchor scope (e.g. a lateral join binding its left + # schema while its right input -- a separate plan being inferred here -- is + # built) so references to an outer anchor still resolve. + token = anchor_scope.set(_AnchorScope(anchors, subtrees, parent=anchor_scope.get())) try: root = plan.relations[-1].root schema = infer_rel_schema(root.input, registry=registry, subtrees=subtrees) diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 9fedf2e..ebffd4d 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -388,6 +388,11 @@ def to_id_based_outer_references(plan: stplan.Plan) -> stplan.Plan: 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. + * a ``LateralJoinRel``'s ``rel_anchor`` is reserved (per the Substrait spec) for + its right input's reference to the current *left* row, so it does **not** name + the join's output row. A ``steps_out`` correlation into a lateral join's output + therefore cannot be anchored on it and is left offset-based, rather than + aliasing (and corrupting) the left-row anchor. Raises only if a resolvable binding carries no ``RelCommon`` at all (e.g. an ``UpdateRel``), which no correlated-subquery shape produces. @@ -429,6 +434,15 @@ def anchor_for(binding: stalg.Rel) -> int: anchor_by_id[key] = counter return counter + def _binding_is_lateral_join(binding: stalg.Rel) -> bool: + # A LateralJoinRel's rel_anchor is reserved (per the Substrait spec) for its + # right input's reference to the current *left* row, so it does not denote + # the join's output row and must not be reused to anchor a correlation into + # that output. Unwrap a ReferenceRel to the subtree it points at first. + while binding.WhichOneof("rel_type") == "reference": + binding = subtrees[binding.reference.subtree_ordinal] + return binding.WhichOneof("rel_type") == "lateral_join" + def convert_expr(expr, scope, binding): rex = expr.WhichOneof("rex_type") if rex == "selection": @@ -444,8 +458,11 @@ def convert_expr(expr, scope, binding): ) 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: + # (a reducing join's condition). A lateral join's rel_anchor is + # reserved for its right input's left-row reference, so it cannot + # double as the output-row anchor a correlation here would need. + # Both are left offset-based (spec-valid, read by inference). + if target is not None and not _binding_is_lateral_join(target): oref.rel_reference = anchor_for(target) elif rex == "subquery": for inner in _iter_subquery_rels(expr): diff --git a/tests/builders/plan/test_lateral_join.py b/tests/builders/plan/test_lateral_join.py new file mode 100644 index 0000000..4428ee6 --- /dev/null +++ b/tests/builders/plan/test_lateral_join.py @@ -0,0 +1,196 @@ +import pytest +import substrait.algebra_pb2 as stalg + +from substrait.builders.extended_expression import ( + column, + fresh_rel_anchors, + literal, +) +from substrait.builders.plan import lateral_join, project, read_named_table +from substrait.builders.type import boolean, i64, named_struct, string, struct +from substrait.extension_registry import ExtensionRegistry +from substrait.type_inference import infer_plan_schema + +registry = ExtensionRegistry(load_default_extensions=False) + +left_ns = named_struct( + ["k", "v"], struct([i64(nullable=False), string()], nullable=False) +) +right_ns = named_struct(["w"], struct([i64(nullable=False)], nullable=False)) + + +def _left(): + return read_named_table("left", left_ns) + + +def _right(): + return read_named_table("right", right_ns) + + +def _correlated_right(left): + # A right input that projects the current left row's first column ("k") via + # the left handle, on top of its own column. + return project(_right(), expressions=[left.column("k")]) + + +def test_lateral_join_optional_args_are_keyword_only(): + with pytest.raises(TypeError): + lateral_join( + _left(), + _correlated_right, + stalg.JoinRel.JOIN_TYPE_INNER, + None, # would have bound to expression positionally + ) + + +def test_lateral_join_sets_anchor_matching_right_reference(): + with fresh_rel_anchors(): + plan = lateral_join( + _left(), _correlated_right, type=stalg.JoinRel.JOIN_TYPE_INNER + )(registry) + + rel = plan.relations[-1].root.input + assert rel.WhichOneof("rel_type") == "lateral_join" + lj = rel.lateral_join + # The join assigns a rel_anchor and the right's handle reference names it. + assert lj.common.HasField("rel_anchor") + ref = lj.right.project.expressions[0].selection + assert ref.WhichOneof("root_type") == "outer_reference" + assert ref.outer_reference.WhichOneof("outer_reference_type") == "rel_reference" + assert ref.outer_reference.rel_reference == lj.common.rel_anchor + # "k" is the left's first column. + assert ref.direct_reference.struct_field.field == 0 + + +def test_lateral_join_inner_output_and_inference(): + with fresh_rel_anchors(): + plan = lateral_join( + _left(), _correlated_right, type=stalg.JoinRel.JOIN_TYPE_INNER + )(registry) + + # Output names: left + right (right = its own column + the correlated one). + assert list(plan.relations[-1].root.names) == ["k", "v", "w", "k"] + ns = infer_plan_schema(plan, registry=registry) + assert [t.WhichOneof("kind") for t in ns.struct.types] == [ + "i64", # left.k + "string", # left.v + "i64", # right.w + "i64", # right's correlated reference to left.k + ] + + +def test_lateral_join_left_semi_drops_right(): + with fresh_rel_anchors(): + plan = lateral_join( + _left(), _correlated_right, type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI + )(registry) + + assert list(plan.relations[-1].root.names) == ["k", "v"] + ns = infer_plan_schema(plan, registry=registry) + assert [t.WhichOneof("kind") for t in ns.struct.types] == ["i64", "string"] + + +def test_lateral_join_left_mark_appends_boolean(): + with fresh_rel_anchors(): + plan = lateral_join( + _left(), _correlated_right, type=stalg.JoinRel.JOIN_TYPE_LEFT_MARK + )(registry) + + ns = infer_plan_schema(plan, registry=registry) + assert list(ns.names)[-1] == "mark" + assert ns.struct.types[-1].WhichOneof("kind") == "bool" + assert len(ns.names) == len(ns.struct.types) + + +def test_lateral_join_nested_handles_reference_distinct_anchors(): + # An inner lateral join can reference the outer left via its captured handle, + # with no depth bookkeeping; each join gets a distinct anchor. + def inner_of(outer): + def middle(_middle_left): + return project(_right(), expressions=[outer.column("k")]) + + return lateral_join(_right(), middle, type=stalg.JoinRel.JOIN_TYPE_INNER) + + with fresh_rel_anchors(): + plan = lateral_join(_left(), inner_of, type=stalg.JoinRel.JOIN_TYPE_INNER)( + registry + ) + + outer = plan.relations[-1].root.input.lateral_join + inner = outer.right.lateral_join + assert outer.common.rel_anchor != inner.common.rel_anchor + # The innermost projection references the OUTER join's anchor. + ref = inner.right.project.expressions[0].selection + assert ref.outer_reference.rel_reference == outer.common.rel_anchor + infer_plan_schema(plan, registry=registry) + + +def _uncorrelated_right(_left): + # A right input that ignores the left row (no correlation), for exercising + # the match condition / post-filter arguments in isolation. + return _right() + + +def test_lateral_join_expression_condition(): + # An optional match condition binds against the combined left+right inputs and + # is emitted on LateralJoinRel.expression. + with fresh_rel_anchors(): + plan = lateral_join( + _left(), + _uncorrelated_right, + type=stalg.JoinRel.JOIN_TYPE_INNER, + expression=literal(True, boolean()), + )(registry) + + lj = plan.relations[-1].root.input.lateral_join + assert lj.HasField("expression") + assert lj.expression.literal.boolean is True + infer_plan_schema(plan, registry=registry) + + +def _post_field(plan): + ref = plan.relations[-1].root.input.lateral_join.post_join_filter.selection + return ref.direct_reference.struct_field.field + + +def test_lateral_join_post_join_filter_binds_output_schema(): + # post_join_filter is applied to the join output (semantically a FilterRel + # 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. + with fresh_rel_anchors(): + inner = lateral_join( + _left(), + _uncorrelated_right, + type=stalg.JoinRel.JOIN_TYPE_INNER, + post_join_filter=column("w"), + )(registry) + assert list(inner.relations[-1].root.names) == ["k", "v", "w"] + assert _post_field(inner) == 2 + + with fresh_rel_anchors(): + mark = lateral_join( + _left(), + _uncorrelated_right, + 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 + infer_plan_schema(mark, registry=registry) + + +def test_lateral_join_post_join_filter_on_dropped_side_raises(): + # A left-semi lateral join drops the right side from its output, so a + # post_join_filter on a right column cannot resolve -- it fails fast rather + # than emitting a dangling reference against the combined inputs. + with pytest.raises(ValueError, match="not in list"): + with fresh_rel_anchors(): + lateral_join( + _left(), + _uncorrelated_right, + type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, + post_join_filter=column("w"), # right-only, absent from output + )(registry) diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 590b3d7..76d5f9a 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -1251,6 +1251,154 @@ def test_outer_steps_out_below_one_raises(): outer.filter(sub.exists(corr)).to_plan() +# -- Lateral joins (handle-based id correlation) -------------------------- + + +def test_lateral_join_correlated_filter(): + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"k": sub.i64, "v": sub.i64}) + inner = sub.read_named_table("r", {"k": sub.i64, "w": sub.i64}) + # The right filters on the current left row via the left handle l.col("k"). + plan = left.lateral_join( + lambda lat: inner.filter(sub.col("k") == lat.col("k")), how="inner" + ).to_plan() + + lj = plan.relations[-1].root.input.lateral_join + assert lj.common.HasField("rel_anchor") + rhs = lj.right.filter.condition.scalar_function.arguments[1].value.selection + assert rhs.WhichOneof("root_type") == "outer_reference" + assert rhs.outer_reference.rel_reference == lj.common.rel_anchor + # Inner + full right schema. + assert list(infer_plan_schema(plan).names) == ["k", "v", "k", "w"] + + +def test_lateral_join_nested_handles_no_depth(): + # An inner lateral join references the outer left via its captured handle; + # the innermost predicate correlates on both levels with no depth argument. + from substrait.type_inference import infer_plan_schema + + outer = sub.read_named_table("o", {"j": sub.i64}) + mid = sub.read_named_table("m", {"k": sub.i64}) + inner = sub.read_named_table("i", {"k": sub.i64, "j": sub.i64}) + plan = outer.lateral_join( + lambda o: mid.lateral_join( + lambda m: inner.filter( + (sub.col("k") == m.col("k")) & (sub.col("j") == o.col("j")) + ), + how="inner", + ), + how="inner", + ).to_plan() + + outer_lj = plan.relations[-1].root.input.lateral_join + inner_lj = outer_lj.right.lateral_join + assert outer_lj.common.rel_anchor != inner_lj.common.rel_anchor + infer_plan_schema(plan) # both id references resolve + + +def test_lateral_join_is_deterministic(): + left = sub.read_named_table("l", {"k": sub.i64}) + inner = sub.read_named_table("r", {"k": sub.i64}) + + def build(): + return left.lateral_join( + lambda lat: inner.filter(sub.col("k") == lat.col("k")), how="inner" + ) + + # Building the same frame twice assigns identical anchors -> equal plans, + # and both materialization entry points reset anchor numbering alike. + assert build().to_plan() == build().to_plan() + assert build().to_plan() == build().to_substrait() + + +def test_lateral_join_left_semi_drops_right(): + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"k": sub.i64, "v": sub.i64}) + inner = sub.read_named_table("r", {"k": sub.i64}) + plan = left.lateral_join( + lambda lat: inner.filter(sub.col("k") == lat.col("k")), how="left_semi" + ).to_plan() + assert list(infer_plan_schema(plan).names) == ["k", "v"] + + +def test_lateral_join_unknown_how_raises(): + left = sub.read_named_table("l", {"k": sub.i64}) + inner = sub.read_named_table("r", {"k": sub.i64}) + with pytest.raises(ValueError, match="unknown lateral join type 'right'"): + left.lateral_join(lambda lat: inner, how="right") + + +def test_lateral_join_on_condition(): + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"k": sub.i64, "v": sub.i64}) + inner = sub.read_named_table("r", {"w": sub.i64}) + # `on` is a match condition over the combined left+right schema. + plan = left.lateral_join( + lambda lat: inner, how="inner", on=sub.col("v") == sub.col("w") + ).to_plan() + + lj = plan.relations[-1].root.input.lateral_join + assert lj.HasField("expression") + assert lj.expression.WhichOneof("rex_type") == "scalar_function" + assert list(infer_plan_schema(plan).names) == ["k", "v", "w"] + + +def test_lateral_join_post_filter_binds_output_schema(): + from substrait.type_inference import infer_plan_schema + + left = sub.read_named_table("l", {"k": sub.i64, "v": sub.i64}) + inner = sub.read_named_table("r", {"w": sub.i64}) + # A left-mark join appends a `mark` column to the output; post_filter resolves + # against that output schema, so it can reference `mark` -- which the combined + # left+right inputs do not carry. + plan = left.lateral_join( + lambda lat: inner, how="left_mark", post_filter=sub.col("mark") + ).to_plan() + + 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"] + + +def test_correlated_exists_above_lateral_join_stays_steps_out(): + # Regression: a correlated subquery stacked ABOVE a lateral join references the + # join's OUTPUT row. A lateral join's rel_anchor is reserved (per the Substrait + # spec) for its right input's reference to the current LEFT row, so it must NOT + # be reused to anchor this correlation -- doing so aliases the left-row anchor + # and corrupts any reference beyond the left columns (here the right-side `w`). + # The reference is left offset-based (steps_out) instead. + from substrait.type_inference import infer_plan_schema + + outer = sub.read_named_table("outer", {"k": sub.i64, "v": sub.i64}) + inner = sub.read_named_table("inner", {"k": sub.i64, "w": sub.i64}) + subq = sub.read_named_table("subq", {"k": sub.i64}) + + lj = outer.lateral_join( + lambda lat: inner.filter(sub.col("k") == lat.col("k")), how="inner" + ) + # The EXISTS correlates on the lateral join's OUTPUT column `w` (right-side, + # index 3 in the output [k, v, k, w]). + plan = lj.filter( + sub.exists(subq.filter(sub.col("k") == sub.outer("w", steps_out=1))) + ).to_plan() + + top = plan.relations[-1].root.input + lat_anchor = top.filter.input.lateral_join.common.rel_anchor + oref = top.filter.condition.subquery.set_predicate.tuples.filter.condition.scalar_function.arguments[ + 1 + ].value.selection.outer_reference + # Offset-based, NOT aliasing the lateral join's (left-row) anchor. + assert oref.WhichOneof("outer_reference_type") == "steps_out" + assert oref.steps_out == 1 + assert lat_anchor >= 1 # the lateral join still carries its own left-row anchor + infer_plan_schema(plan) # resolves without corruption + + def test_correlated_subquery_projecting_outer_column_then_chaining(): # Regression: a correlated subquery whose *output* is the outer column forces # the enclosing plan's schema inference to resolve the OuterReference. This diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 0a042dc..b5b9dd0 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -345,6 +345,169 @@ def test_inference_join_left_mark(): assert infer_rel_schema(rel) == expected +def test_inference_lateral_join_inner(): + # A lateral join emits the same columns as the equivalent JoinRel; only the + # right input's evaluation semantics differ. + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + left=read_rel, + right=right_read_rel, + type=stalg.JoinRel.JOIN_TYPE_INNER, + ) + ) + + expected = stt.Type.Struct( + types=[ + 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)), + ], + nullability=stt.Type.Nullability.NULLABILITY_REQUIRED, + ) + + assert infer_rel_schema(rel) == expected + + +def test_inference_lateral_join_left_semi(): + # Left-oriented semi/anti joins drop the right side, just like JoinRel. + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + left=read_rel, + right=right_read_rel, + type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, + ) + ) + + expected = stt.Type.Struct( + types=[ + 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)), + ], + nullability=stt.Type.Nullability.NULLABILITY_REQUIRED, + ) + + assert infer_rel_schema(rel) == expected + + +def test_inference_lateral_join_left_mark(): + # Left-mark joins append a nullable boolean marker column. + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + left=read_rel, + right=right_read_rel, + type=stalg.JoinRel.JOIN_TYPE_LEFT_MARK, + ) + ) + + expected = stt.Type.Struct( + types=[ + 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, + ) + + assert infer_rel_schema(rel) == expected + + +def _outer_rel_reference(anchor: int, field: int) -> stalg.Expression: + """An OuterReference resolved by id: rel_reference -> the given rel_anchor, + selecting the struct field at ``field``.""" + return stalg.Expression( + selection=stalg.Expression.FieldReference( + outer_reference=stalg.Expression.FieldReference.OuterReference( + rel_reference=anchor + ), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField(field=field) + ), + ) + ) + + +def test_inference_lateral_join_correlated_rel_reference(): + # The right (dependent) input references the current left row via an + # OuterReference.rel_reference pointing to the lateral join's rel_anchor. + # Inference must resolve that against the left schema registered under the + # anchor. Here the right input projects the left's first column (i64) on top + # of its own columns. + anchor = 7 + correlated_right = stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, + expressions=[_outer_rel_reference(anchor, 0)], + ) + ) + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + common=stalg.RelCommon(rel_anchor=anchor), + left=read_rel, + right=correlated_right, + type=stalg.JoinRel.JOIN_TYPE_INNER, + ) + ) + + expected = stt.Type.Struct( + types=[ + # left columns + 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)), + # right columns + stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), + stt.Type(bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_NULLABLE)), + # right's projected OuterReference to the left's first column (i64) + stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), + ], + nullability=stt.Type.Nullability.NULLABILITY_REQUIRED, + ) + + assert infer_rel_schema(rel) == expected + + +def test_inference_lateral_join_unknown_rel_anchor_raises(): + # A rel_reference that does not match the (only) enclosing lateral join's + # rel_anchor cannot be resolved. + correlated_right = stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, + expressions=[_outer_rel_reference(99, 0)], + ) + ) + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + common=stalg.RelCommon(rel_anchor=7), + left=read_rel, + right=correlated_right, + type=stalg.JoinRel.JOIN_TYPE_INNER, + ) + ) + + with pytest.raises(Exception, match="unknown rel_anchor 99"): + infer_rel_schema(rel) + + +def test_infer_expression_type_rel_reference_resolves_against_anchor(): + # infer_expression_type resolves an id-based OuterReference against the schema + # bound to the matching rel_anchor in the current anchor scope. + from substrait.type_inference import _outer_anchor_binding + + # rel_anchor 5 -> `struct` ([i64, string, fp32]); field 1 is the string. + with _outer_anchor_binding(5, struct): + result = infer_expression_type(_outer_rel_reference(5, 1), right_struct) + + assert result == stt.Type( + string=stt.Type.String(nullability=stt.Type.NULLABILITY_NULLABLE) + ) + + def test_infer_expression_type_literal(): """Test infer_expression_type with a literal expression.""" expr = stalg.Expression(literal=stalg.Expression.Literal(i64=42, nullable=False)) diff --git a/tests/test_utils.py b/tests/test_utils.py index aef1f6c..7b221fd 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -424,6 +424,44 @@ def test_convert_reducing_join_condition_left_as_steps_out(): assert ref.steps_out == 1 +def _lateral_join( + left: stalg.Rel, + right: stalg.Rel, + *, + rel_anchor: int, + type=stalg.JoinRel.JOIN_TYPE_INNER, +) -> stalg.Rel: + return stalg.Rel( + lateral_join=stalg.LateralJoinRel( + common=stalg.RelCommon(rel_anchor=rel_anchor), + left=left, + right=right, + type=type, + ) + ) + + +def test_convert_correlation_above_lateral_join_left_as_steps_out(): + # A LateralJoinRel's rel_anchor is reserved (per the Substrait spec) for its + # right input's reference to the current *left* row, so it does not name the + # join's output row. A correlation stacked above the lateral join (into its + # output) must not reuse that anchor -- doing so would alias the left-row anchor + # and corrupt any reference beyond the left columns. Such a reference is left + # offset-based (still spec-valid), like a reducing join's condition. + lj = _lateral_join(_read("l", ncols=2), _read("r", ncols=2), rel_anchor=5) + plan = _plan(_filter(lj, _exists(_filter(_read("i"), _outer(1, field=3))))) + out = to_id_based_outer_references(plan) + + top = out.relations[-1].root.input + ref = top.filter.condition.subquery.set_predicate.tuples.filter.condition.selection.outer_reference + assert ref.WhichOneof("outer_reference_type") == "steps_out" + assert ref.steps_out == 1 + # The lateral join keeps its own (left-row) anchor; no new anchor is minted for + # the un-rewritable correlation. + assert rel_anchor_of(top.filter.input) == 5 + assert {a for r in iter_plan_rels(out) if (a := rel_anchor_of(r))} == {5} + + 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