feat: infer LateralJoinRel schema with id-based outer references

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
 

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,

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
 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: