diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index d99c5c0..ac808d1 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 @@ -18,6 +20,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 ] @@ -376,6 +405,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 8885ca7..0756b7c 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -16,10 +16,12 @@ from substrait.builders.extended_expression import ( ExtendedExpressionOrUnbound, + LateralInput, + next_rel_anchor, resolve_expression, ) from substrait.extension_registry import ExtensionRegistry -from substrait.type_inference import infer_plan_schema, join_output_names +from substrait.type_inference import infer_plan_schema, join_output_names, outer_anchors from substrait.utils import ( merge_extension_declarations, merge_extension_urns, @@ -533,6 +535,103 @@ 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) + # Register the left schema under `anchor` while the right input is built + # and its schema inferred, so id-based references (handle.column(...)) + # resolve during that inference. + anchors = outer_anchors.get() + anchor_token = outer_anchors.set({**anchors, anchor: left_ns}) + try: + 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) + + 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 + ) + bound_post = ( + resolve_expression(post_join_filter, ns, registry) + if post_join_filter is not None + else None + ) + finally: + outer_anchors.reset(anchor_token) + + rel = stalg.Rel( + lateral_join=stalg.LateralJoinRel( + common=stalg.RelCommon(rel_anchor=anchor), + left=bound_left.relations[-1].root.input, + right=bound_right.relations[-1].root.input, + 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, + ) + ) + + # Output names follow the same per-join-type shape as a regular join + # (semi/anti drop the right side, mark appends a boolean). + out_names = join_output_names( + stalg.JoinRel.JoinType.Name(type), left_ns.names, right_ns.names + ) + return stp.Plan( + version=default_version, + relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=out_names))], + **_merge_plan_metadata( + 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 2b7cd6f..561cec2 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 from substrait.extension_registry import ExtensionRegistry from substrait.type_inference import infer_plan_schema @@ -59,6 +60,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, @@ -166,6 +181,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. @@ -397,6 +432,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": @@ -654,11 +736,15 @@ def write_named_table( def to_plan(self): """Materialize to a ``substrait.proto.Plan``.""" - return self._plan(self._registry) + # Number lateral-join rel_anchors from 1 per materialization so building + # the same frame twice yields identical plans. + with fresh_rel_anchors(): + return self._plan(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) + with fresh_rel_anchors(): + return self._plan(registry or self._registry) class GroupBy: diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 941b2ed..d3bfc4f 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -6,12 +6,22 @@ import substrait.type_pb2 as stt # 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. +# 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. outer_schemas: contextvars.ContextVar = contextvars.ContextVar( "outer_schemas", default=() ) +# Map of RelCommon.rel_anchor -> enclosing NamedStruct, for id-based outer +# reference resolution (OuterReference.rel_reference). Used where an offset +# (steps_out) would be ambiguous -- notably a LateralJoinRel, whose right input +# references the current left row via its rel_anchor. Registered by the relation +# that owns the anchor while resolving the sub-tree that references it. +outer_anchors: contextvars.ContextVar = contextvars.ContextVar( + "outer_anchors", default={} +) + def _derive_extension_schema(detail, inputs, registry): """Derive a registered extension relation's output NamedStruct, or None. @@ -210,17 +220,27 @@ 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. An OuterReference resolves either by id + # (rel_reference -> the schema registered under that rel_anchor) or by + # offset (steps_out -> that many levels up the correlated-subquery stack). if root_type == "outer_reference": - stack = outer_schemas.get() - steps = expression.selection.outer_reference.steps_out - if steps >= len(stack): - raise Exception( - "outer reference outside an enclosing (correlated) query" - ) - schema = stack[len(stack) - 1 - steps].struct + outer_ref = expression.selection.outer_reference + if outer_ref.WhichOneof("outer_reference_type") == "rel_reference": + anchors = outer_anchors.get() + anchor = outer_ref.rel_reference + if anchor not in anchors: + raise Exception(f"outer reference to unknown rel_anchor {anchor}") + schema = anchors[anchor].struct + else: + stack = outer_schemas.get() + steps = outer_ref.steps_out + if steps >= len(stack): + raise Exception( + "outer reference outside an enclosing (correlated) query" + ) + schema = stack[len(stack) - 1 - steps].struct else: assert root_type in ("root_reference", "lambda_parameter_reference") schema = parent_schema @@ -370,13 +390,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 +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) - right = infer_rel_schema(right_rel, registry=registry) + """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": @@ -398,6 +417,44 @@ 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 +) -> 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) + right = infer_rel_schema(right_rel, registry=registry) + return _join_struct_from_schemas(type_name, left, right) + + +def _lateral_join_output_struct( + type_name: str, lateral_join: stalg.LateralJoinRel, *, registry=None +) -> stt.Type.Struct: + """Lateral-join output column types. + + A lateral join is semantically identical to a regular join for output + formation (same per-join-type column shapes), except the right (dependent) + input is evaluated once per left row and may reference fields of the current + left row via ``OuterReference.rel_reference`` pointing to this relation's + ``RelCommon.rel_anchor``. The left schema is registered under that anchor + while the right schema is inferred so those id-based references resolve. + """ + left = infer_rel_schema(lateral_join.left, registry=registry) + common = lateral_join.common + if common.HasField("rel_anchor"): + anchors = outer_anchors.get() + token = outer_anchors.set( + {**anchors, common.rel_anchor: stt.NamedStruct(struct=left)} + ) + try: + right = infer_rel_schema(lateral_join.right, registry=registry) + finally: + outer_anchors.reset(token) + else: + right = infer_rel_schema(lateral_join.right, registry=registry) + 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") @@ -535,6 +592,13 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct: registry=registry, ) (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, + ) + (common, struct) = (rel.lateral_join.common, raw_schema) elif rel_type == "window": parent_schema = infer_rel_schema(rel.window.input, registry=registry) window_output_types = [wf.output_type for wf in rel.window.window_functions] diff --git a/tests/builders/plan/test_lateral_join.py b/tests/builders/plan/test_lateral_join.py new file mode 100644 index 0000000..ac3004d --- /dev/null +++ b/tests/builders/plan/test_lateral_join.py @@ -0,0 +1,121 @@ +import pytest +import substrait.algebra_pb2 as stalg + +from substrait.builders.extended_expression import fresh_rel_anchors +from substrait.builders.plan import lateral_join, project, read_named_table +from substrait.builders.type import 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) diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index ee77a99..f74aec8 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -1094,6 +1094,85 @@ def test_outer_outside_subquery_raises(): df.filter(sub.col("x") == sub.outer("x")).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_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 e67cfe6..b995d69 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -343,6 +343,172 @@ 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 registered under the matching rel_anchor. + from substrait.type_inference import outer_anchors + + token = outer_anchors.set({5: named_struct}) + try: + # rel_anchor 5 -> `struct` ([i64, string, fp32]); field 1 is the string. + result = infer_expression_type(_outer_rel_reference(5, 1), right_struct) + finally: + outer_anchors.reset(token) + + 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))