feat: add LateralJoinRel support (inference + builders)

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
 

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,

src/substrait/dataframe/frame.py

@@ -29,14 +29,15 @@
 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
 import substrait.type_pb2 as stp
 
 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``."""