feat: shared subplans (ReferenceRel) / DataFrame.cache() (CTEs)

[nielspardon]

Summary

Adds CTEs / shared subplans — ReferenceRel + subtree_ordinal — to both the substrait.builders.plan layer and the native substrait.dataframe API (as DataFrame.cache()), closing #211.

Per the design agreed in the #204 review, the shared subtrees are carried in-band inside the resolved Plan — as the leading rel entries of Plan.relations, the same way a Plan already carries its extension declarations — rather than out-of-band via a contextvar. Because the subtrees then live inside the resolved Plan, the builder-internal infer_plan_schema(bound_plan) calls have the subtree list in hand directly, so no contextvar is needed.

What's in it

type_inference

• Thread an optional subtrees list through infer_rel_schema / infer_expression_type; add a reference case that resolves a ReferenceRel's schema against subtrees[subtree_ordinal].
• infer_plan_schema extracts the leading subtrees and wraps them in a _SubtreeScope that memoizes each subtree's schema (a subtree referenced from many places is inferred once, not exponentially) and detects reference cycles (a clear error instead of RecursionError).

builders.plan

• A reference(plan) builder promotes a plan to a shared subtree and returns a ReferenceRel-rooted plan.
• Every relational builder propagates its inputs' shared subtrees upward and rebases subtree_ordinals when merging inputs, structurally deduping identical subtrees so a cached frame reused across branches that later meet at a join/union is emitted once and referenced many times.
• A shared _plan_from helper owns subtree propagation + Plan assembly, so each builder is a single call.

dataframe

• DataFrame.cache() marks a frame as a reusable common subplan.
• The hint() guard for RelCommon-less relations (e.g. a ReferenceRel) is restored (removed alongside the prototype in feat: ergonomic native DataFrame/Expr API (substrait.api) #204).

Subqueries

• A cached frame consumed inside a subquery is inlined into that subquery — a plan-global ReferenceRel cannot cross the subquery boundary — keeping the emitted plan valid.

Internals

• Rel-walking primitives (plan_subtrees, rebase_reference_ordinals, inline_reference_rels) live in substrait.utils and discover child-Rel fields from the protobuf descriptor rather than a hand-maintained table.

Notes

• subtree_ordinal indexing: subtrees are the leading rel entries and the query root is last, so the ordinal indexes both relations and the ordered subtree list identically. The proto / substrait-java are underspecified here; this matches the agreed design and the earlier prototype.
• Structural dedup collapses byte-equal serialized Rels, which are semantically identical — this is the CTE "compute once, reference many" behavior.

Testing

• New tests for the reference builder (dedup / rebasing / schema inference), DataFrame.cache() (shared/nested/distinct caches, extension merge, subquery inlining, hint()-after-cache()), and the reference inference case (cycle + out-of-range).
• Full suite green (572 passed), pixi run lint and pixi run format clean. Existing byte-exact builder comparisons still pass, confirming the no-subtree path is unchanged.

🤖 Generated with AI