Summary
Add support for CTEs / shared subplans — ReferenceRel + subtree_ordinal — to both the substrait.builders.plan layer and the substrait.dataframe native DataFrame API (as DataFrame.cache()).
This was prototyped in #204 but pulled out to keep that PR focused, and because the first cut used a contextvars-based mechanism that circumvented the builder layer. This issue tracks doing it properly at both layers. See discussion: #204 (comment)
Design (agreed in review)
Carry the shared subtrees in-band inside the Plan — in Plan.relations (the leading rel entries), exactly the way a Plan already carries its extension declarations — rather than out-of-band via a contextvar.
Concretely:
- type_inference: thread the subtree list through
infer_rel_schema / infer_plan_schema as a parameter (extracted from plan.relations), and add a reference case that resolves a ReferenceRel's schema against subtrees[subtree_ordinal]. No contextvar.
- builders (
substrait.builders.plan):
- a
reference(...) builder that emits a ReferenceRel-rooted plan;
- every builder propagates its input plans' shared subtrees into its output plan (mirroring how
_merge_extensions already carries extensions upward);
- multi-input builders (
join, set, …) rebase subtree_ordinals when merging inputs that each carry their own subtrees, so ordinals stay unique and correct.
- dataframe (
substrait.dataframe): DataFrame.cache() on top of the builder support — dedupe repeated uses of a cached frame and reference the shared subtree; to_plan() emits the subtrees as the plan's leading relations.
Because the subtrees then live inside the resolved Plan, the builder-internal infer_plan_schema(bound_plan) calls have the subtree list available directly — which is what makes the contextvar unnecessary.
Notes
Summary
Add support for CTEs / shared subplans —
ReferenceRel+subtree_ordinal— to both thesubstrait.builders.planlayer and thesubstrait.dataframenative DataFrame API (asDataFrame.cache()).This was prototyped in #204 but pulled out to keep that PR focused, and because the first cut used a
contextvars-based mechanism that circumvented the builder layer. This issue tracks doing it properly at both layers. See discussion: #204 (comment)Design (agreed in review)
Carry the shared subtrees in-band inside the
Plan— inPlan.relations(the leadingrelentries), exactly the way aPlanalready carries its extension declarations — rather than out-of-band via a contextvar.Concretely:
infer_rel_schema/infer_plan_schemaas a parameter (extracted fromplan.relations), and add areferencecase that resolves aReferenceRel's schema againstsubtrees[subtree_ordinal]. No contextvar.substrait.builders.plan):reference(...)builder that emits aReferenceRel-rooted plan;_merge_extensionsalready carries extensions upward);join,set, …) rebasesubtree_ordinals when merging inputs that each carry their own subtrees, so ordinals stay unique and correct.substrait.dataframe):DataFrame.cache()on top of the builder support — dedupe repeated uses of a cached frame and reference the shared subtree;to_plan()emits the subtrees as the plan's leading relations.Because the subtrees then live inside the resolved
Plan, the builder-internalinfer_plan_schema(bound_plan)calls have the subtree list available directly — which is what makes the contextvar unnecessary.Notes
hint()guard for relations withoutRelCommon(e.g.ReferenceRel) was removed alongside the prototype in feat: ergonomic native DataFrame/Expr API (substrait.api) #204; it should return with this feature.substrait-javacross-check onsubtree_ordinalindexing semantics while implementing.