Skip to content

Commit 5dd720d

Browse files
committed
feat: shared subplans (ReferenceRel) / DataFrame.cache() (CTEs)
Add CTE / shared-subplan support at both the builders.plan and native DataFrame layers, carrying the shared subtrees in-band inside the resolved Plan (the leading `rel` entries of Plan.relations, the way a Plan already carries its extension declarations) rather than out-of-band via a contextvar. type_inference: thread an optional `subtrees` list through infer_rel_schema / infer_expression_type and 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 (so a subtree referenced from many places is inferred once, not exponentially) and detects reference cycles (a clear error instead of RecursionError). builders: 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 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. 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. 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. Closes #211
1 parent 154ebbe commit 5dd720d

9 files changed

Lines changed: 968 additions & 358 deletions

File tree

src/substrait/builders/extended_expression.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
from substrait.extension_registry import ExtensionRegistry
1414
from substrait.type_inference import infer_extended_expression_schema, outer_schemas
1515
from substrait.utils import (
16+
inline_reference_rels,
1617
merge_extension_declarations,
1718
merge_extension_urns,
19+
plan_subtrees,
1820
type_num_names,
1921
)
2022

@@ -1041,7 +1043,17 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema):
10411043
plan = query(registry) if callable(query) else query
10421044
finally:
10431045
outer_schemas.reset(token)
1044-
return plan, plan.relations[-1].root.input
1046+
rel = plan.relations[-1].root.input
1047+
# An Expression.Subquery embeds only a bare Rel, but a ReferenceRel is
1048+
# plan-global -- it cannot resolve once lifted out of its plan. So if the
1049+
# subquery's plan carries shared subtrees (e.g. it uses a cached frame),
1050+
# inline them into the subquery Rel, making it self-contained. A cached frame
1051+
# used inside a subquery is thus inlined there rather than shared across the
1052+
# subquery boundary.
1053+
subtrees = plan_subtrees(plan)
1054+
if subtrees:
1055+
rel = inline_reference_rels(rel, subtrees)
1056+
return plan, rel
10451057

10461058

10471059
def scalar_subquery(query, alias: Union[str, None] = None):

0 commit comments

Comments
 (0)