feat: infer LateralJoinRel schema with id-based outer references

[nielspardon]

Summary

Adds full support for LateralJoinRel, the relation introduced in Substrait v0.96.0 (#205): schema inference, the id-based outer reference resolution it depends on (introduced in v0.89.0, #188), and builders across both the low-level and DataFrame layers.

A lateral join is semantically identical to a JoinRel for output formation, 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 the join's RelCommon.rel_anchor.

Inference (type_inference.py)

• infer_rel_schema handles lateral_join. Output columns follow the same per-join-type shapes as a regular JoinRel: INNER/LEFT/LEFT_SINGLE → left+right, LEFT_SEMI/LEFT_ANTI → left only, LEFT_MARK → +nullable mark.
• Id-based outer reference resolution. OuterReference is a oneof of steps_out (offset) and rel_reference (id). infer_expression_type dispatches on it and resolves rel_reference against a new outer_anchors context map keyed by RelCommon.rel_anchor. The existing offset-based (steps_out) mechanism is unchanged, and the lateral join is deliberately not pushed onto the steps_out stack (it is not a subquery boundary), so offset counting for nested subqueries is unaffected.

Builders — handle-based correlation

The right (dependent) input is built as a function of a handle to the left, mirroring substrait-java's Function<JoinInput, Expression> idiom. Since rel_reference is id-based, capturing a handle (identity) fits it better than counting nesting levels.

# DataFrame layer
left.lateral_join(lambda lat: inner.filter(sub.col("k") == lat.col("k")), how="inner")

# nested joins need no depth bookkeeping — reference the handle you captured
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",
)

• builders.plan.lateral_join(left, right, type, *, expression, post_join_filter, extension) — right: Callable[[LateralInput], PlanOrUnbound]. The builder allocates a rel_anchor, builds the LateralInput handle, and calls right(handle) while the left schema is registered under the anchor (so id references resolve during the right's schema inference). handle.column(field) emits the id-based OuterReference.
• DataFrame.lateral_join(right, how, *, on, post_filter) — right: Callable[[LateralLeft], DataFrame]; lat.col(name) / lat[name] are the correlated references.
• Anchor allocation — next_rel_anchor() draws unique values within a build; fresh_rel_anchors() resets numbering per materialization, and both DataFrame.to_plan and to_substrait wrap in it, so repeated builds are reproducible.
• Referencing outside a lateral join is structurally impossible (the handle only exists inside the builder), rather than a runtime error.

Testing

• pixi run lint / pixi run format --check — clean.
• Full suite: 560 passed, 30 skipped.
• Inference: inner / left-semi / left-mark shapes, correlated rel_reference→rel_anchor resolution, unknown-anchor error, direct infer_expression_type id resolution.
• Builders (both layers): anchor↔reference wiring, output shapes, nested-handle correlation (no depth) with distinct anchors, determinism across to_plan/to_substrait, and invalid-join-type error path.

Closes #205. Contributes to #188.

🤖 Generated with AI