From 4d87e95139a1ca83b610c68231c06b8fb2fe7d81 Mon Sep 17 00:00:00 2001
From: Niels Pardon <par@zurich.ibm.com>
Date: Fri, 24 Jul 2026 13:55:02 +0200
Subject: [PATCH] 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
---
 src/substrait/builders/extended_expression.py |  14 +-
 src/substrait/builders/plan.py                | 690 ++++++++++--------
 src/substrait/dataframe/frame.py              |  24 +
 src/substrait/type_inference.py               | 219 +++++-
 src/substrait/utils/__init__.py               | 104 +++
 tests/builders/plan/test_reference.py         | 104 +++
 tests/conftest.py                             |  19 +
 tests/dataframe/test_cache.py                 | 114 +++
 tests/test_type_inference.py                  |  26 +
 9 files changed, 963 insertions(+), 351 deletions(-)
 create mode 100644 tests/builders/plan/test_reference.py
 create mode 100644 tests/conftest.py
 create mode 100644 tests/dataframe/test_cache.py

diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py
index a4d7e87..7a33aa3 100644
--- a/src/substrait/builders/extended_expression.py
+++ b/src/substrait/builders/extended_expression.py
@@ -13,8 +13,10 @@
 from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_extended_expression_schema, outer_schemas
 from substrait.utils import (
+    inline_reference_rels,
     merge_extension_declarations,
     merge_extension_urns,
+    plan_subtrees,
     type_num_names,
 )
 
@@ -1041,7 +1043,17 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema):
         plan = query(registry) if callable(query) else query
     finally:
         outer_schemas.reset(token)
-    return plan, plan.relations[-1].root.input
+    rel = plan.relations[-1].root.input
+    # An Expression.Subquery embeds only a bare Rel, but a ReferenceRel is
+    # plan-global -- it cannot resolve once lifted out of its plan. So if the
+    # subquery's plan carries shared subtrees (e.g. it uses a cached frame),
+    # inline them into the subquery Rel, making it self-contained. A cached frame
+    # used inside a subquery is thus inlined there rather than shared across the
+    # subquery boundary.
+    subtrees = plan_subtrees(plan)
+    if subtrees:
+        rel = inline_reference_rels(rel, subtrees)
+    return plan, rel
 
 
 def scalar_subquery(query, alias: Union[str, None] = None):
diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py
index 8885ca7..3e1cece 100644
--- a/src/substrait/builders/plan.py
+++ b/src/substrait/builders/plan.py
@@ -23,6 +23,8 @@
 from substrait.utils import (
     merge_extension_declarations,
     merge_extension_urns,
+    plan_subtrees,
+    rebase_reference_ordinals,
 )
 from substrait.version import substrait_version
 
@@ -67,6 +69,80 @@ def _merge_plan_metadata(*objs):
     return metadata
 
 
+def _is_identity(remap: dict) -> bool:
+    return all(old == new for old, new in remap.items())
+
+
+def _merge_input_subtrees(bound_inputs):
+    """Combine the leading shared subtrees carried in-band by resolved input plans.
+
+    Returns ``(subtree_planrels, rebased_root_inputs)``: the deduplicated combined
+    subtrees as leading ``PlanRel(rel=...)`` entries, and, per input, its root's
+    input Rel with ReferenceRel ordinals rebased into the combined list. Mirrors how
+    ``_merge_plan_metadata`` carries extension declarations upward.
+
+    Structurally-identical subtrees (byte-equal serialized ``Rel``) collapse to a
+    single ordinal, so a cached frame reused across branches that later meet at a
+    multi-input relation is emitted once and referenced many times. Inputs that
+    carry no subtrees pass their root input through untouched, keeping output for
+    the (overwhelmingly common) no-subtree case byte-identical.
+    """
+    combined: "list[stalg.Rel]" = []
+    key_to_ordinal: dict = {}  # serialized subtree bytes -> ordinal in `combined`
+    rebased_root_inputs = []
+    for plan in bound_inputs:
+        remap: dict = {}
+        for old_ordinal, subtree in enumerate(plan_subtrees(plan)):
+            # Rebase the subtree's own references (to earlier subtrees in this same
+            # input) before deduping, so structurally-equal subtrees compare equal.
+            rebased = (
+                subtree
+                if _is_identity(remap)
+                else rebase_reference_ordinals(subtree, remap)
+            )
+            key = rebased.SerializeToString(deterministic=True)
+            new_ordinal = key_to_ordinal.get(key)
+            if new_ordinal is None:
+                new_ordinal = len(combined)
+                key_to_ordinal[key] = new_ordinal
+                combined.append(rebased)
+            remap[old_ordinal] = new_ordinal
+        root_input = plan.relations[-1].root.input
+        rebased_root_inputs.append(
+            root_input
+            if _is_identity(remap)
+            else rebase_reference_ordinals(root_input, remap)
+        )
+    subtree_planrels = [stp.PlanRel(rel=r) for r in combined]
+    return subtree_planrels, rebased_root_inputs
+
+
+def _plan_from(
+    bound_inputs, make_rel, names, metadata_sources, *, include_version=True
+):
+    """Assemble a relational builder's output Plan.
+
+    Merges the shared subtrees carried by ``bound_inputs`` (deduping and rebasing
+    ordinals), builds the output ``Rel`` by calling ``make_rel`` with the list of
+    rebased input rels (one per bound input, in order), and prepends the combined
+    subtrees as leading ``rel`` entries ahead of the query root. Metadata (extension
+    declarations / execution behavior) is merged from ``metadata_sources`` (input
+    plans and bound expressions). This is the single place the CTE subtree
+    propagation and Plan assembly live, so every relational builder is one call.
+    """
+    subtree_planrels, input_rels = _merge_input_subtrees(bound_inputs)
+    root = stp.PlanRel(
+        root=stalg.RelRoot(input=make_rel(input_rels), names=list(names))
+    )
+    kwargs = {
+        "relations": [*subtree_planrels, root],
+        **_merge_plan_metadata(*metadata_sources),
+    }
+    if include_version:
+        kwargs["version"] = default_version
+    return stp.Plan(**kwargs)
+
+
 def with_execution_behavior(
     plan: PlanOrUnbound,
     variable_eval_mode: "stp.ExecutionBehavior.VariableEvaluationMode.ValueType",
@@ -251,20 +327,21 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             e.output_names[0] for ee in bound_expressions for e in ee.referred_expr
         ]
 
-        rel = stalg.Rel(
-            project=stalg.ProjectRel(
-                input=_plan.relations[-1].root.input,
-                expressions=[
-                    e.expression for ee in bound_expressions for e in ee.referred_expr
-                ],
-                advanced_extension=extension,
-            )
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(_plan, *bound_expressions),
+        return _plan_from(
+            [_plan],
+            lambda inp: stalg.Rel(
+                project=stalg.ProjectRel(
+                    input=inp[0],
+                    expressions=[
+                        e.expression
+                        for ee in bound_expressions
+                        for e in ee.referred_expr
+                    ],
+                    advanced_extension=extension,
+                )
+            ),
+            names,
+            (_plan, *bound_expressions),
         )
 
     return resolve
@@ -302,25 +379,26 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             e.output_names[0] for ee in bound_expressions for e in ee.referred_expr
         ]
 
-        rel = stalg.Rel(
-            project=stalg.ProjectRel(
-                common=stalg.RelCommon(
-                    emit=stalg.RelCommon.Emit(
-                        output_mapping=[i + start_index for i in range(len(names))]
-                    )
-                ),
-                input=_plan.relations[-1].root.input,
-                expressions=[
-                    e.expression for ee in bound_expressions for e in ee.referred_expr
-                ],
-                advanced_extension=extension,
-            )
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(_plan, *bound_expressions),
+        return _plan_from(
+            [_plan],
+            lambda inp: stalg.Rel(
+                project=stalg.ProjectRel(
+                    common=stalg.RelCommon(
+                        emit=stalg.RelCommon.Emit(
+                            output_mapping=[i + start_index for i in range(len(names))]
+                        )
+                    ),
+                    input=inp[0],
+                    expressions=[
+                        e.expression
+                        for ee in bound_expressions
+                        for e in ee.referred_expr
+                    ],
+                    advanced_extension=extension,
+                )
+            ),
+            names,
+            (_plan, *bound_expressions),
         )
 
     return resolve
@@ -338,20 +416,17 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             expression, ns, registry
         )
 
-        rel = stalg.Rel(
-            filter=stalg.FilterRel(
-                input=bound_plan.relations[-1].root.input,
-                condition=bound_expression.referred_expr[0].expression,
-                advanced_extension=extension,
-            )
-        )
-
-        names = ns.names
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(bound_plan, bound_expression),
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                filter=stalg.FilterRel(
+                    input=inp[0],
+                    condition=bound_expression.referred_expr[0].expression,
+                    advanced_extension=extension,
+                )
+            ),
+            ns.names,
+            (bound_plan, bound_expression),
         )
 
     return resolve
@@ -381,24 +456,23 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             (resolve_expression(e[0], ns, registry), e[1]) for e in bound_expressions
         ]
 
-        rel = stalg.Rel(
-            sort=stalg.SortRel(
-                input=bound_plan.relations[-1].root.input,
-                sorts=[
-                    stalg.SortField(
-                        expr=e[0].referred_expr[0].expression,
-                        direction=e[1],
-                    )
-                    for e in bound_expressions
-                ],
-                advanced_extension=extension,
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                sort=stalg.SortRel(
+                    input=inp[0],
+                    sorts=[
+                        stalg.SortField(
+                            expr=e[0].referred_expr[0].expression,
+                            direction=e[1],
+                        )
+                        for e in bound_expressions
+                    ],
+                    advanced_extension=extension,
+                ),
             ),
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
-            **_merge_plan_metadata(bound_plan, *[e[0] for e in bound_expressions]),
+            ns.names,
+            (bound_plan, *[e[0] for e in bound_expressions]),
         )
 
     return resolve
@@ -407,22 +481,45 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
 def set(inputs: Iterable[PlanOrUnbound], op: stalg.SetRel.SetOp) -> UnboundPlan:
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_inputs = [i if isinstance(i, stp.Plan) else i(registry) for i in inputs]
-        rel = stalg.Rel(
-            set=stalg.SetRel(
-                inputs=[plan.relations[-1].root.input for plan in bound_inputs], op=op
-            )
+        return _plan_from(
+            bound_inputs,
+            lambda inp: stalg.Rel(set=stalg.SetRel(inputs=inp, op=op)),
+            bound_inputs[0].relations[-1].root.names,
+            tuple(bound_inputs),
         )
 
+    return resolve
+
+
+def reference(plan: PlanOrUnbound) -> UnboundPlan:
+    """Promote a plan to a shared subtree (a CTE) and reference it.
+
+    Returns a plan whose query root is a ``ReferenceRel`` pointing at ``plan``'s
+    root, which is carried along as a leading shared subtree (a ``PlanRel(rel=...)``
+    entry). Any subtrees ``plan`` already carries are preserved ahead of it, so
+    their ordinals stay valid; the promoted root takes the next ordinal.
+
+    Every downstream use of the returned plan carries the subtree upward; when two
+    branches that share the subtree later meet at a multi-input builder, the copies
+    collapse to one (see :func:`_merge_input_subtrees`). This is the building block
+    for :meth:`substrait.dataframe.DataFrame.cache`.
+    """
+
+    def resolve(registry: ExtensionRegistry) -> stp.Plan:
+        bound = plan if isinstance(plan, stp.Plan) else plan(registry)
+        nested = [stp.PlanRel(rel=s) for s in plan_subtrees(bound)]
+        ordinal = len(nested)  # the promoted root sits after the plan's own subtrees
+        promoted = stp.PlanRel(rel=bound.relations[-1].root.input)
+        names = list(bound.relations[-1].root.names)
+        ref = stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=ordinal))
         return stp.Plan(
             version=default_version,
             relations=[
-                stp.PlanRel(
-                    root=stalg.RelRoot(
-                        input=rel, names=bound_inputs[0].relations[-1].root.names
-                    )
-                )
+                *nested,
+                promoted,
+                stp.PlanRel(root=stalg.RelRoot(input=ref, names=names)),
             ],
-            **_merge_plan_metadata(*bound_inputs),
+            **_merge_plan_metadata(bound),
         )
 
     return resolve
@@ -444,29 +541,22 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             resolve_expression(count, ns, registry) if count is not None else None
         )
 
-        rel = stalg.Rel(
-            fetch=stalg.FetchRel(
-                input=bound_plan.relations[-1].root.input,
-                offset_expr=bound_offset.referred_expr[0].expression
-                if bound_offset
-                else None,
-                count_expr=bound_count.referred_expr[0].expression
-                if bound_count
-                else None,
-                advanced_extension=extension,
-            )
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[
-                stp.PlanRel(
-                    root=stalg.RelRoot(
-                        input=rel, names=bound_plan.relations[-1].root.names
-                    )
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                fetch=stalg.FetchRel(
+                    input=inp[0],
+                    offset_expr=bound_offset.referred_expr[0].expression
+                    if bound_offset
+                    else None,
+                    count_expr=bound_count.referred_expr[0].expression
+                    if bound_count
+                    else None,
+                    advanced_extension=extension,
                 )
-            ],
-            **_merge_plan_metadata(bound_plan, bound_offset, bound_count),
+            ),
+            bound_plan.relations[-1].root.names,
+            (bound_plan, bound_offset, bound_count),
         )
 
     return resolve
@@ -503,31 +593,28 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             else None
         )
 
-        rel = stalg.Rel(
-            join=stalg.JoinRel(
-                left=bound_left.relations[-1].root.input,
-                right=bound_right.relations[-1].root.input,
-                expression=bound_expression.referred_expr[0].expression,
-                post_join_filter=bound_post.referred_expr[0].expression
-                if bound_post
-                else None,
-                type=type,
-                advanced_extension=extension,
-            )
-        )
-
         # The join condition resolves against the combined left+right schema, but
         # the output names must match the columns the join type actually emits
         # (semi/anti drop a 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 _plan_from(
+            [bound_left, bound_right],
+            lambda inp: stalg.Rel(
+                join=stalg.JoinRel(
+                    left=inp[0],
+                    right=inp[1],
+                    expression=bound_expression.referred_expr[0].expression,
+                    post_join_filter=bound_post.referred_expr[0].expression
+                    if bound_post
+                    else None,
+                    type=type,
+                    advanced_extension=extension,
+                )
             ),
+            out_names,
+            (bound_left, bound_right, bound_expression, bound_post),
         )
 
     return resolve
@@ -552,18 +639,17 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             names=list(left_ns.names) + list(right_ns.names),
         )
 
-        rel = stalg.Rel(
-            cross=stalg.CrossRel(
-                left=bound_left.relations[-1].root.input,
-                right=bound_right.relations[-1].root.input,
-                advanced_extension=extension,
-            )
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=ns.names))],
-            **_merge_plan_metadata(bound_left, bound_right),
+        return _plan_from(
+            [bound_left, bound_right],
+            lambda inp: stalg.Rel(
+                cross=stalg.CrossRel(
+                    left=inp[0],
+                    right=inp[1],
+                    advanced_extension=extension,
+                )
+            ),
+            ns.names,
+            (bound_left, bound_right),
         )
 
     return resolve
@@ -611,35 +697,35 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             else [list(range(len(bound_grouping_expressions)))]
         )
 
-        rel = stalg.Rel(
-            aggregate=stalg.AggregateRel(
-                input=bound_input.relations[-1].root.input,
-                grouping_expressions=[
-                    e.referred_expr[0].expression for e in bound_grouping_expressions
-                ],
-                groupings=[
-                    stalg.AggregateRel.Grouping(expression_references=refs)
-                    for refs in sets
-                ],
-                measures=[
-                    stalg.AggregateRel.Measure(
-                        measure=m.referred_expr[0].measure,
-                        filter=bf.referred_expr[0].expression if bf else None,
-                    )
-                    for m, bf in zip(bound_measures, bound_filters)
-                ],
-                advanced_extension=extension,
-            )
-        )
-
         names = [
             e.referred_expr[0].output_names[0] for e in bound_grouping_expressions
         ] + [e.referred_expr[0].output_names[0] for e in bound_measures]
 
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(
+        return _plan_from(
+            [bound_input],
+            lambda inp: stalg.Rel(
+                aggregate=stalg.AggregateRel(
+                    input=inp[0],
+                    grouping_expressions=[
+                        e.referred_expr[0].expression
+                        for e in bound_grouping_expressions
+                    ],
+                    groupings=[
+                        stalg.AggregateRel.Grouping(expression_references=refs)
+                        for refs in sets
+                    ],
+                    measures=[
+                        stalg.AggregateRel.Measure(
+                            measure=m.referred_expr[0].measure,
+                            filter=bf.referred_expr[0].expression if bf else None,
+                        )
+                        for m, bf in zip(bound_measures, bound_filters)
+                    ],
+                    advanced_extension=extension,
+                )
+            ),
+            names,
+            (
                 bound_input,
                 *bound_grouping_expressions,
                 *bound_measures,
@@ -664,23 +750,23 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         _create_mode = create_mode or stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS
         _op = op if op is not None else stalg.WriteRel.WRITE_OP_CTAS
 
-        write_rel = stalg.Rel(
-            write=stalg.WriteRel(
-                input=bound_input.relations[-1].root.input,
-                table_schema=ns,
-                op=_op,
-                create_mode=_create_mode,
-                output=output_mode
-                if output_mode is not None
-                else stalg.WriteRel.OUTPUT_MODE_UNSPECIFIED,
-                named_table=stalg.NamedObjectWrite(names=_table_names),
-            )
-        )
-        return stp.Plan(
-            relations=[
-                stp.PlanRel(root=stalg.RelRoot(input=write_rel, names=ns.names))
-            ],
-            **_merge_plan_metadata(bound_input),
+        return _plan_from(
+            [bound_input],
+            lambda inp: stalg.Rel(
+                write=stalg.WriteRel(
+                    input=inp[0],
+                    table_schema=ns,
+                    op=_op,
+                    create_mode=_create_mode,
+                    output=output_mode
+                    if output_mode is not None
+                    else stalg.WriteRel.OUTPUT_MODE_UNSPECIFIED,
+                    named_table=stalg.NamedObjectWrite(names=_table_names),
+                )
+            ),
+            ns.names,
+            (bound_input,),
+            include_version=False,
         )
 
     return resolve
@@ -703,7 +789,7 @@ def ddl(
 
     def resolve(registry: ExtensionRegistry) -> stp.Plan:
         merge_sources = []
-        view_rel = None
+        bound_inputs = []
         schema = table_schema
         if view_definition is not None:
             view_plan = (
@@ -711,26 +797,26 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                 if isinstance(view_definition, stp.Plan)
                 else view_definition(registry)
             )
-            view_rel = view_plan.relations[-1].root.input
+            bound_inputs = [view_plan]
             merge_sources.append(view_plan)
             if schema is None:
                 schema = infer_plan_schema(view_plan, registry=registry)
 
-        ddl_rel = stalg.Rel(
-            ddl=stalg.DdlRel(
-                named_object=stalg.NamedObjectWrite(names=_names),
-                table_schema=schema,
-                object=object_type,
-                op=op,
-                view_definition=view_rel,
-                advanced_extension=extension,
-            )
-        )
         out_names = list(schema.names) if schema is not None else []
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=ddl_rel, names=out_names))],
-            **_merge_plan_metadata(*merge_sources),
+        return _plan_from(
+            bound_inputs,
+            lambda inp: stalg.Rel(
+                ddl=stalg.DdlRel(
+                    named_object=stalg.NamedObjectWrite(names=_names),
+                    table_schema=schema,
+                    object=object_type,
+                    op=op,
+                    view_definition=inp[0] if inp else None,
+                    advanced_extension=extension,
+                )
+            ),
+            out_names,
+            tuple(merge_sources),
         )
 
     return resolve
@@ -853,28 +939,27 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             for i, wf_ee in enumerate(bound_window_fns)
         ]
 
-        rel = stalg.Rel(
-            window=stalg.ConsistentPartitionWindowRel(
-                input=bound_plan.relations[-1].root.input,
-                window_functions=window_rel_functions,
-                partition_expressions=[
-                    e.referred_expr[0].expression for e in bound_partitions
-                ],
-                sorts=[
-                    stalg.SortField(
-                        expr=e[0].referred_expr[0].expression,
-                        direction=e[1],
-                    )
-                    for e in bound_sorts
-                ],
-                advanced_extension=extension,
-            )
-        )
-
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                window=stalg.ConsistentPartitionWindowRel(
+                    input=inp[0],
+                    window_functions=window_rel_functions,
+                    partition_expressions=[
+                        e.referred_expr[0].expression for e in bound_partitions
+                    ],
+                    sorts=[
+                        stalg.SortField(
+                            expr=e[0].referred_expr[0].expression,
+                            direction=e[1],
+                        )
+                        for e in bound_sorts
+                    ],
+                    advanced_extension=extension,
+                )
+            ),
+            names,
+            (
                 bound_plan,
                 *bound_partitions,
                 *[e[0] for e in bound_sorts],
@@ -924,16 +1009,16 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                     )
                 )
 
-        rel = stalg.Rel(
-            expand=stalg.ExpandRel(
-                input=bound_input.relations[-1].root.input,
-                fields=expand_fields,
-            )
-        )
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=list(names)))],
-            **_merge_plan_metadata(*merge_sources),
+        return _plan_from(
+            [bound_input],
+            lambda inp: stalg.Rel(
+                expand=stalg.ExpandRel(
+                    input=inp[0],
+                    fields=expand_fields,
+                )
+            ),
+            list(names),
+            tuple(merge_sources),
         )
 
     return resolve
@@ -963,24 +1048,24 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
         bound_expression = resolve_expression(expression, ns, registry)
 
-        rel = stalg.Rel(
-            nested_loop_join=stalg.NestedLoopJoinRel(
-                left=bound_left.relations[-1].root.input,
-                right=bound_right.relations[-1].root.input,
-                expression=bound_expression.referred_expr[0].expression,
-                type=type,
-                advanced_extension=extension,
-            )
-        )
         out_names = join_output_names(
             stalg.NestedLoopJoinRel.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),
+        return _plan_from(
+            [bound_left, bound_right],
+            lambda inp: stalg.Rel(
+                nested_loop_join=stalg.NestedLoopJoinRel(
+                    left=inp[0],
+                    right=inp[1],
+                    expression=bound_expression.referred_expr[0].expression,
+                    type=type,
+                    advanced_extension=extension,
+                )
+            ),
+            out_names,
+            (bound_left, bound_right, bound_expression),
         )
 
     return resolve
@@ -1036,21 +1121,21 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             names = join_output_names(
                 rel_cls.JoinType.Name(type), left_ns.names, right_ns.names
             )
-            rel = stalg.Rel(
-                **{
-                    rel_name: rel_cls(
-                        left=bound_left.relations[-1].root.input,
-                        right=bound_right.relations[-1].root.input,
-                        keys=keys,
-                        type=type,
-                        advanced_extension=extension,
-                    )
-                }
-            )
-            return stp.Plan(
-                version=default_version,
-                relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-                **_merge_plan_metadata(bound_left, bound_right),
+            return _plan_from(
+                [bound_left, bound_right],
+                lambda inp: stalg.Rel(
+                    **{
+                        rel_name: rel_cls(
+                            left=inp[0],
+                            right=inp[1],
+                            keys=keys,
+                            type=type,
+                            advanced_extension=extension,
+                        )
+                    }
+                ),
+                names,
+                (bound_left, bound_right),
             )
 
         return resolve
@@ -1106,15 +1191,15 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             names = list(detail.derive_schema(input_struct).names)
         else:
             names = list(bound_plan.relations[-1].root.names)
-        rel = stalg.Rel(
-            extension_single=stalg.ExtensionSingleRel(
-                input=bound_plan.relations[-1].root.input, detail=_detail_any(detail)
-            )
-        )
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(bound_plan),
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                extension_single=stalg.ExtensionSingleRel(
+                    input=inp[0], detail=_detail_any(detail)
+                )
+            ),
+            names,
+            (bound_plan,),
         )
 
     return resolve
@@ -1129,16 +1214,16 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             infer_plan_schema(b, registry=registry).struct for b in bound_inputs
         ]
         names = list(detail.derive_schema(input_structs).names)
-        rel = stalg.Rel(
-            extension_multi=stalg.ExtensionMultiRel(
-                inputs=[b.relations[-1].root.input for b in bound_inputs],
-                detail=_detail_any(detail),
-            )
-        )
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=names))],
-            **_merge_plan_metadata(*bound_inputs),
+        return _plan_from(
+            bound_inputs,
+            lambda inp: stalg.Rel(
+                extension_multi=stalg.ExtensionMultiRel(
+                    inputs=inp,
+                    detail=_detail_any(detail),
+                )
+            ),
+            names,
+            tuple(bound_inputs),
         )
 
     return resolve
@@ -1162,23 +1247,17 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             if broadcast
             else {"round_robin": stalg.ExchangeRel.RoundRobin()}
         )
-        rel = stalg.Rel(
-            exchange=stalg.ExchangeRel(
-                input=bound_plan.relations[-1].root.input,
-                partition_count=partition_count,
-                **kind,
-            )
-        )
-        return stp.Plan(
-            version=default_version,
-            relations=[
-                stp.PlanRel(
-                    root=stalg.RelRoot(
-                        input=rel, names=bound_plan.relations[-1].root.names
-                    )
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                exchange=stalg.ExchangeRel(
+                    input=inp[0],
+                    partition_count=partition_count,
+                    **kind,
                 )
-            ],
-            **_merge_plan_metadata(bound_plan),
+            ),
+            bound_plan.relations[-1].root.names,
+            (bound_plan,),
         )
 
     return resolve
@@ -1211,40 +1290,29 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             resolve_expression(offset, ns, registry) if offset is not None else None
         )
 
-        rel = stalg.Rel(
-            top_n=stalg.TopNRel(
-                input=bound_plan.relations[-1].root.input,
-                sorts=[
-                    stalg.SortField(
-                        expr=s.referred_expr[0].expression, direction=direction
-                    )
-                    for s, direction in bound_sorts
-                ],
-                count=bound_count.referred_expr[0].expression,
-                offset=bound_offset.referred_expr[0].expression
-                if bound_offset
-                else None,
-                mode=stalg.FetchMode.FETCH_MODE_WITH_TIES
-                if with_ties
-                else stalg.FetchMode.FETCH_MODE_ROWS_ONLY,
-                advanced_extension=extension,
-            )
-        )
-        return stp.Plan(
-            version=default_version,
-            relations=[
-                stp.PlanRel(
-                    root=stalg.RelRoot(
-                        input=rel, names=bound_plan.relations[-1].root.names
-                    )
+        return _plan_from(
+            [bound_plan],
+            lambda inp: stalg.Rel(
+                top_n=stalg.TopNRel(
+                    input=inp[0],
+                    sorts=[
+                        stalg.SortField(
+                            expr=s.referred_expr[0].expression, direction=direction
+                        )
+                        for s, direction in bound_sorts
+                    ],
+                    count=bound_count.referred_expr[0].expression,
+                    offset=bound_offset.referred_expr[0].expression
+                    if bound_offset
+                    else None,
+                    mode=stalg.FetchMode.FETCH_MODE_WITH_TIES
+                    if with_ties
+                    else stalg.FetchMode.FETCH_MODE_ROWS_ONLY,
+                    advanced_extension=extension,
                 )
-            ],
-            **_merge_plan_metadata(
-                bound_plan,
-                *[s for s, _ in bound_sorts],
-                bound_count,
-                bound_offset,
             ),
+            bound_plan.relations[-1].root.names,
+            (bound_plan, *[s for s, _ in bound_sorts], bound_count, bound_offset),
         )
 
     return resolve
diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py
index 0517a8b..82c3451 100644
--- a/src/substrait/dataframe/frame.py
+++ b/src/substrait/dataframe/frame.py
@@ -500,6 +500,15 @@ def resolve(registry: ExtensionRegistry):
             bound = inner(registry)
             rel = bound.relations[-1].root.input
             rel_inner = getattr(rel, rel.WhichOneof("rel_type"))
+            # A few relations (a ReferenceRel from .cache(), an UpdateRel) carry no
+            # RelCommon and so cannot hold a hint -- fail with a clear message
+            # rather than an opaque AttributeError on `.common`.
+            if "common" not in rel_inner.DESCRIPTOR.fields_by_name:
+                raise TypeError(
+                    f"cannot attach a hint to a {rel_inner.DESCRIPTOR.name} "
+                    "(e.g. a cached/reference relation); apply .hint(...) before "
+                    ".cache()"
+                )
             common = rel_inner.common
             if row_count is not None:
                 common.hint.stats.row_count = row_count
@@ -514,6 +523,21 @@ def resolve(registry: ExtensionRegistry):
 
         return self._next(resolve)
 
+    def cache(self) -> "DataFrame":
+        """Mark this DataFrame as a reusable common subplan (a CTE).
+
+        The returned frame is emitted as a shared subtree (a leading ``rel`` entry
+        in the plan) and referenced via a ``ReferenceRel`` wherever it is used.
+        When the same cached frame feeds two branches that later meet at a
+        multi-input relation (e.g. ``base.filter(...).union(base.filter(...))``),
+        the shared subtree is emitted **once** and referenced from both branches
+        instead of being inlined twice.
+
+        Because the subtree carries no ``RelCommon``, apply :meth:`hint` (and any
+        node-level annotation) *before* ``cache()``, not after.
+        """
+        return self._next(_plan.reference(self._plan))
+
     def with_execution_behavior(self, variable_eval_mode: str) -> "DataFrame":
         """Set how often execution context variables are evaluated (plan-level).
 
diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py
index 71e27dd..393a8ee 100644
--- a/src/substrait/type_inference.py
+++ b/src/substrait/type_inference.py
@@ -5,6 +5,54 @@
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
 
+from substrait.utils import plan_subtrees
+
+
+class _SubtreeScope:
+    """The shared-subtree list a ``ReferenceRel`` resolves against, with per-ordinal
+    schema memoization and cycle detection.
+
+    A ``ReferenceRel``'s schema is that of ``subtrees[subtree_ordinal]``, and a
+    subtree may itself reference an earlier one, so resolution recurses. Memoizing
+    each ordinal's schema avoids re-inferring a subtree referenced from many places
+    (otherwise exponential for chained cached self-joins), and the in-progress set
+    turns a cyclic/self reference in a malformed plan into a clear error rather than
+    a ``RecursionError``. Behaves as a read-only sequence (``len`` / indexing) so
+    callers that pass a plain list of subtrees still work unchanged.
+    """
+
+    __slots__ = ("_rels", "_schemas", "_resolving")
+
+    def __init__(self, rels):
+        self._rels = list(rels)
+        self._schemas: dict = {}
+        self._resolving: set = set()
+
+    def __len__(self):
+        return len(self._rels)
+
+    def __getitem__(self, ordinal):
+        return self._rels[ordinal]
+
+    def schema_of(self, ordinal, registry):
+        if ordinal in self._schemas:
+            return self._schemas[ordinal]
+        if ordinal in self._resolving:
+            raise Exception(
+                f"ReferenceRel subtree_ordinal {ordinal} forms a cycle; a shared "
+                "subtree cannot reference itself directly or transitively"
+            )
+        self._resolving.add(ordinal)
+        try:
+            schema = infer_rel_schema(
+                self._rels[ordinal], registry=registry, subtrees=self
+            )
+        finally:
+            self._resolving.discard(ordinal)
+        self._schemas[ordinal] = schema
+        return schema
+
+
 # Stack of enclosing-query schemas (NamedStruct) for correlated subqueries, so a
 # field reference with an OuterReference root resolves against the right level.
 # Pushed by the subquery builders while resolving their inner plan.
@@ -159,7 +207,7 @@ def infer_literal_type(literal: stalg.Expression.Literal) -> stt.Type:
 
 
 def infer_nested_type(
-    nested: stalg.Expression.Nested, parent_schema, *, registry=None
+    nested: stalg.Expression.Nested, parent_schema, *, registry=None, subtrees=()
 ) -> stt.Type:
     nested_type = nested.WhichOneof("nested_type")
 
@@ -173,7 +221,9 @@ def infer_nested_type(
         return stt.Type(
             struct=stt.Type.Struct(
                 types=[
-                    infer_expression_type(f, parent_schema, registry=registry)
+                    infer_expression_type(
+                        f, parent_schema, registry=registry, subtrees=subtrees
+                    )
                     for f in nested.struct.fields
                 ],
                 nullability=nullability,
@@ -183,7 +233,10 @@ def infer_nested_type(
         return stt.Type(
             list=stt.Type.List(
                 type=infer_expression_type(
-                    nested.list.values[0], parent_schema, registry=registry
+                    nested.list.values[0],
+                    parent_schema,
+                    registry=registry,
+                    subtrees=subtrees,
                 ),
                 nullability=nullability,
             )
@@ -192,10 +245,16 @@ def infer_nested_type(
         return stt.Type(
             map=stt.Type.Map(
                 key=infer_expression_type(
-                    nested.map.key_values[0].key, parent_schema, registry=registry
+                    nested.map.key_values[0].key,
+                    parent_schema,
+                    registry=registry,
+                    subtrees=subtrees,
                 ),
                 value=infer_expression_type(
-                    nested.map.key_values[0].value, parent_schema, registry=registry
+                    nested.map.key_values[0].value,
+                    parent_schema,
+                    registry=registry,
+                    subtrees=subtrees,
                 ),
                 nullability=nullability,
             )
@@ -205,7 +264,11 @@ def infer_nested_type(
 
 
 def infer_expression_type(
-    expression: stalg.Expression, parent_schema: stt.Type.Struct, *, registry=None
+    expression: stalg.Expression,
+    parent_schema: stt.Type.Struct,
+    *,
+    registry=None,
+    subtrees=(),
 ) -> stt.Type:
     rex_type = expression.WhichOneof("rex_type")
     if rex_type == "selection":
@@ -254,11 +317,17 @@ def infer_expression_type(
         return expression.window_function.output_type
     elif rex_type == "if_then":
         return infer_expression_type(
-            expression.if_then.ifs[0].then, parent_schema, registry=registry
+            expression.if_then.ifs[0].then,
+            parent_schema,
+            registry=registry,
+            subtrees=subtrees,
         )
     elif rex_type == "switch_expression":
         return infer_expression_type(
-            expression.switch_expression.ifs[0].then, parent_schema, registry=registry
+            expression.switch_expression.ifs[0].then,
+            parent_schema,
+            registry=registry,
+            subtrees=subtrees,
         )
     elif rex_type == "cast":
         return expression.cast.type
@@ -267,12 +336,16 @@ def infer_expression_type(
             bool=stt.Type.Boolean(nullability=stt.Type.Nullability.NULLABILITY_NULLABLE)
         )
     elif rex_type == "nested":
-        return infer_nested_type(expression.nested, parent_schema, registry=registry)
+        return infer_nested_type(
+            expression.nested, parent_schema, registry=registry, subtrees=subtrees
+        )
     elif rex_type == "lambda":
         # A lambda's type is func<param_types -> body_type>; the body's parameter
         # references resolve against the lambda's own parameter struct.
         lam = getattr(expression, "lambda")
-        body_type = infer_expression_type(lam.body, lam.parameters, registry=registry)
+        body_type = infer_expression_type(
+            lam.body, lam.parameters, registry=registry, subtrees=subtrees
+        )
         return stt.Type(
             func=stt.Type.Func(
                 parameter_types=list(lam.parameters.types),
@@ -294,7 +367,9 @@ def infer_expression_type(
         try:
             if subquery_type == "scalar":
                 scalar_rel = infer_rel_schema(
-                    expression.subquery.scalar.input, registry=registry
+                    expression.subquery.scalar.input,
+                    registry=registry,
+                    subtrees=subtrees,
                 )
                 return scalar_rel.types[0]
             elif (
@@ -378,12 +453,12 @@ def join_output_names(type_name: str, left_names, right_names) -> list:
 
 
 def _join_output_struct(
-    type_name: str, left_rel, right_rel, *, registry=None
+    type_name: str, left_rel, right_rel, *, registry=None, subtrees=()
 ) -> stt.Type.Struct:
     """Join output column types by join-type NAME (shared across all join
     relations, whose enum integer values differ)."""
-    left = infer_rel_schema(left_rel, registry=registry)
-    right = infer_rel_schema(right_rel, registry=registry)
+    left = infer_rel_schema(left_rel, registry=registry, subtrees=subtrees)
+    right = infer_rel_schema(right_rel, registry=registry, subtrees=subtrees)
     required = stt.Type.Nullability.NULLABILITY_REQUIRED
     shape = _join_column_shape(type_name)
     if shape == "left":
@@ -466,7 +541,14 @@ def _set_output_struct(op_name: str, inputs: list) -> stt.Type.Struct:
     return stt.Type.Struct(types=types, nullability=primary.nullability)
 
 
-def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
+def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type.Struct:
+    """Infer a relation's output struct.
+
+    ``subtrees`` is the plan's list of shared subtree ``Rel``s (the leading ``rel``
+    entries of a ``Plan``, extracted by :func:`infer_plan_schema`); a
+    ``ReferenceRel`` resolves its schema against ``subtrees[subtree_ordinal]``. It
+    defaults to ``()`` so plans without shared subtrees behave exactly as before.
+    """
     rel_type = rel.WhichOneof("rel_type")
 
     if rel_type == "read":
@@ -474,17 +556,21 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
     elif rel_type == "filter":
         (common, struct) = (
             rel.filter.common,
-            infer_rel_schema(rel.filter.input, registry=registry),
+            infer_rel_schema(rel.filter.input, registry=registry, subtrees=subtrees),
         )
     elif rel_type == "fetch":
         (common, struct) = (
             rel.fetch.common,
-            infer_rel_schema(rel.fetch.input, registry=registry),
+            infer_rel_schema(rel.fetch.input, registry=registry, subtrees=subtrees),
         )
     elif rel_type == "aggregate":
-        parent_schema = infer_rel_schema(rel.aggregate.input, registry=registry)
+        parent_schema = infer_rel_schema(
+            rel.aggregate.input, registry=registry, subtrees=subtrees
+        )
         grouping_types = [
-            infer_expression_type(g, parent_schema, registry=registry)
+            infer_expression_type(
+                g, parent_schema, registry=registry, subtrees=subtrees
+            )
             for g in rel.aggregate.grouping_expressions
         ]
         measure_types = [m.measure.output_type for m in rel.aggregate.measures]
@@ -504,12 +590,16 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
     elif rel_type == "sort":
         (common, struct) = (
             rel.sort.common,
-            infer_rel_schema(rel.sort.input, registry=registry),
+            infer_rel_schema(rel.sort.input, registry=registry, subtrees=subtrees),
         )
     elif rel_type == "project":
-        parent_schema = infer_rel_schema(rel.project.input, registry=registry)
+        parent_schema = infer_rel_schema(
+            rel.project.input, registry=registry, subtrees=subtrees
+        )
         expression_types = [
-            infer_expression_type(e, parent_schema, registry=registry)
+            infer_expression_type(
+                e, parent_schema, registry=registry, subtrees=subtrees
+            )
             for e in rel.project.expressions
         ]
         raw_schema = stt.Type.Struct(
@@ -519,14 +609,21 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
 
         (common, struct) = (rel.project.common, raw_schema)
     elif rel_type == "set":
-        input_structs = [infer_rel_schema(i, registry=registry) for i in rel.set.inputs]
+        input_structs = [
+            infer_rel_schema(i, registry=registry, subtrees=subtrees)
+            for i in rel.set.inputs
+        ]
         (common, struct) = (
             rel.set.common,
             _set_output_struct(stalg.SetRel.SetOp.Name(rel.set.op), input_structs),
         )
     elif rel_type == "cross":
-        left_schema = infer_rel_schema(rel.cross.left, registry=registry)
-        right_schema = infer_rel_schema(rel.cross.right, registry=registry)
+        left_schema = infer_rel_schema(
+            rel.cross.left, registry=registry, subtrees=subtrees
+        )
+        right_schema = infer_rel_schema(
+            rel.cross.right, registry=registry, subtrees=subtrees
+        )
 
         raw_schema = stt.Type.Struct(
             types=list(left_schema.types) + list(right_schema.types),
@@ -540,10 +637,13 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
             rel.join.left,
             rel.join.right,
             registry=registry,
+            subtrees=subtrees,
         )
         (common, struct) = (rel.join.common, raw_schema)
     elif rel_type == "window":
-        parent_schema = infer_rel_schema(rel.window.input, registry=registry)
+        parent_schema = infer_rel_schema(
+            rel.window.input, registry=registry, subtrees=subtrees
+        )
         window_output_types = [wf.output_type for wf in rel.window.window_functions]
         raw_schema = stt.Type.Struct(
             types=list(parent_schema.types) + window_output_types,
@@ -551,13 +651,18 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
         )
         (common, struct) = (rel.window.common, raw_schema)
     elif rel_type == "expand":
-        parent_schema = infer_rel_schema(rel.expand.input, registry=registry)
+        parent_schema = infer_rel_schema(
+            rel.expand.input, registry=registry, subtrees=subtrees
+        )
         field_types = []
         for field in rel.expand.fields:
             if field.HasField("consistent_field"):
                 field_types.append(
                     infer_expression_type(
-                        field.consistent_field, parent_schema, registry=registry
+                        field.consistent_field,
+                        parent_schema,
+                        registry=registry,
+                        subtrees=subtrees,
                     )
                 )
             else:
@@ -571,7 +676,10 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
                 # determines the output column type.
                 field_types.append(
                     infer_expression_type(
-                        duplicates[0], parent_schema, registry=registry
+                        duplicates[0],
+                        parent_schema,
+                        registry=registry,
+                        subtrees=subtrees,
                     )
                 )
         # Expand appends an i32 column with the index of the duplicate the row
@@ -590,32 +698,56 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
             rel.nested_loop_join.left,
             rel.nested_loop_join.right,
             registry=registry,
+            subtrees=subtrees,
         )
         (common, struct) = (rel.nested_loop_join.common, raw_schema)
     elif rel_type == "hash_join":
         name = stalg.HashJoinRel.JoinType.Name(rel.hash_join.type)
         raw_schema = _join_output_struct(
-            name, rel.hash_join.left, rel.hash_join.right, registry=registry
+            name,
+            rel.hash_join.left,
+            rel.hash_join.right,
+            registry=registry,
+            subtrees=subtrees,
         )
         (common, struct) = (rel.hash_join.common, raw_schema)
     elif rel_type == "merge_join":
         name = stalg.MergeJoinRel.JoinType.Name(rel.merge_join.type)
         raw_schema = _join_output_struct(
-            name, rel.merge_join.left, rel.merge_join.right, registry=registry
+            name,
+            rel.merge_join.left,
+            rel.merge_join.right,
+            registry=registry,
+            subtrees=subtrees,
         )
         (common, struct) = (rel.merge_join.common, raw_schema)
     elif rel_type == "exchange":
         # Exchange redistributes rows without changing the schema.
         (common, struct) = (
             rel.exchange.common,
-            infer_rel_schema(rel.exchange.input, registry=registry),
+            infer_rel_schema(rel.exchange.input, registry=registry, subtrees=subtrees),
         )
     elif rel_type == "top_n":
         # TopN is a fused sort+fetch; the schema is unchanged from the input.
         (common, struct) = (
             rel.top_n.common,
-            infer_rel_schema(rel.top_n.input, registry=registry),
-        )
+            infer_rel_schema(rel.top_n.input, registry=registry, subtrees=subtrees),
+        )
+    elif rel_type == "reference":
+        # A ReferenceRel has no RelCommon/emit; its schema is that of the shared
+        # subtree its subtree_ordinal indexes into. A _SubtreeScope memoizes that
+        # resolution and guards against reference cycles; a plain sequence (direct
+        # callers/tests) is resolved inline against the full subtree list so a
+        # subtree may itself reference an earlier one.
+        ordinal = rel.reference.subtree_ordinal
+        if not 0 <= ordinal < len(subtrees):
+            raise Exception(
+                f"ReferenceRel subtree_ordinal {ordinal} is out of range "
+                f"({len(subtrees)} shared subtree(s) in scope)"
+            )
+        if isinstance(subtrees, _SubtreeScope):
+            return subtrees.schema_of(ordinal, registry)
+        return infer_rel_schema(subtrees[ordinal], registry=registry, subtrees=subtrees)
     elif rel_type == "extension_leaf":
         derived = _derive_extension_schema(rel.extension_leaf.detail, None, registry)
         if derived is None:
@@ -625,7 +757,9 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
             )
         (common, struct) = (rel.extension_leaf.common, derived.struct)
     elif rel_type == "extension_single":
-        input_struct = infer_rel_schema(rel.extension_single.input, registry=registry)
+        input_struct = infer_rel_schema(
+            rel.extension_single.input, registry=registry, subtrees=subtrees
+        )
         derived = _derive_extension_schema(
             rel.extension_single.detail, input_struct, registry
         )
@@ -636,7 +770,8 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
         )
     elif rel_type == "extension_multi":
         input_structs = [
-            infer_rel_schema(i, registry=registry) for i in rel.extension_multi.inputs
+            infer_rel_schema(i, registry=registry, subtrees=subtrees)
+            for i in rel.extension_multi.inputs
         ]
         derived = _derive_extension_schema(
             rel.extension_multi.detail, input_structs, registry
@@ -662,6 +797,12 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None) -> stt.Type.Struct:
 
 
 def infer_plan_schema(plan: stp.Plan, *, registry=None) -> stt.NamedStruct:
-    schema = infer_rel_schema(plan.relations[-1].root.input, registry=registry)
-
-    return stt.NamedStruct(names=plan.relations[-1].root.names, struct=schema)
+    # A Plan carries its shared subtrees in-band as the leading ``rel`` entries of
+    # ``relations`` (the query root is the trailing ``root`` entry), so a
+    # ReferenceRel anywhere in the tree resolves against them by ordinal. Wrap them
+    # in a _SubtreeScope so repeated references are memoized and cycles are caught.
+    subtrees = _SubtreeScope(plan_subtrees(plan))
+    root = plan.relations[-1].root
+    schema = infer_rel_schema(root.input, registry=registry, subtrees=subtrees)
+
+    return stt.NamedStruct(names=root.names, struct=schema)
diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py
index 79026fd..8dd9790 100644
--- a/src/substrait/utils/__init__.py
+++ b/src/substrait/utils/__init__.py
@@ -4,7 +4,9 @@
 
 from typing import Iterable
 
+import substrait.algebra_pb2 as stalg
 import substrait.extensions.extensions_pb2 as ste
+import substrait.plan_pb2 as stplan
 import substrait.type_pb2 as stp
 
 
@@ -72,6 +74,108 @@ def merge_extension_declarations(
     return ret
 
 
+def plan_subtrees(plan: stplan.Plan):
+    """The leading shared-subtree relations of a Plan.
+
+    A ``Plan`` carries its shared subtrees (the common subplans a ``ReferenceRel``
+    points at) in-band as the leading ``rel`` entries of ``relations``; the query
+    root is the trailing ``root`` entry. ``subtree_ordinal`` indexes into this
+    ordered list. This is the single definition of "what is a shared subtree of a
+    plan", reused by the builders and by schema inference.
+    """
+    return [pr.rel for pr in plan.relations if pr.WhichOneof("rel_type") == "rel"]
+
+
+def _child_rel_fields(node):
+    """The names of ``node``'s child-``Rel`` fields, discovered from the protobuf
+    descriptor (any field whose message type is ``substrait.Rel``).
+
+    Derived rather than hand-listed so the set stays correct as the schema evolves
+    -- e.g. ``FilterRel.input``, ``JoinRel.{left,right}``, ``SetRel.inputs``,
+    ``DdlRel.view_definition``. Relations with no child ``Rel`` (``ReadRel``,
+    ``ReferenceRel``, ``UpdateRel``, ...) yield nothing. Note this covers only
+    *direct* child relations; relations embedded inside ``Expression`` subqueries
+    are self-contained (see :func:`inline_reference_rels`) and are not walked here.
+    """
+    return [
+        f
+        for f in node.DESCRIPTOR.fields
+        if f.message_type is not None and f.message_type.full_name == "substrait.Rel"
+    ]
+
+
+def _iter_child_rels(rel: stalg.Rel):
+    """Yield ``(container, index_or_None)`` for each child ``Rel`` of ``rel``.
+
+    ``index`` is the position for a repeated field (``set.inputs``) or ``None`` for
+    a singular field; ``container`` is the repeated field or the parent message so
+    callers can read/replace the child in place.
+    """
+    rel_type = rel.WhichOneof("rel_type")
+    if rel_type is None:
+        return
+    node = getattr(rel, rel_type)
+    for field in _child_rel_fields(node):
+        if field.label == field.LABEL_REPEATED:
+            children = getattr(node, field.name)
+            for i in range(len(children)):
+                yield children, i
+        elif node.HasField(field.name):
+            yield node, field.name
+
+
+def _child_rel(container, key):
+    return container[key] if isinstance(key, int) else getattr(container, key)
+
+
+def rebase_reference_ordinals(rel: stalg.Rel, remap: dict) -> stalg.Rel:
+    """A copy of ``rel`` with every nested ``ReferenceRel.subtree_ordinal`` remapped
+    (old -> new) per ``remap``. Recurses through direct child relations only."""
+    out = stalg.Rel()
+    out.CopyFrom(rel)
+    _rebase_reference_ordinals_in_place(out, remap)
+    return out
+
+
+def _rebase_reference_ordinals_in_place(rel: stalg.Rel, remap: dict) -> None:
+    if rel.WhichOneof("rel_type") == "reference":
+        old = rel.reference.subtree_ordinal
+        rel.reference.subtree_ordinal = remap.get(old, old)
+        return
+    for container, key in _iter_child_rels(rel):
+        _rebase_reference_ordinals_in_place(_child_rel(container, key), remap)
+
+
+def inline_reference_rels(rel: stalg.Rel, subtrees) -> stalg.Rel:
+    """A copy of ``rel`` with every ``ReferenceRel`` replaced by (a recursive inline
+    of) the shared subtree it points at, yielding a self-contained relation.
+
+    Used to embed a subquery whose plan carries shared subtrees: a ``ReferenceRel``
+    is plan-global, so it cannot survive inside an ``Expression.Subquery`` (which
+    holds only a bare ``Rel``, with no place for the plan's subtree list). Inlining
+    de-references the subquery so the emitted plan is valid. ``subtrees`` must be
+    acyclic (builder-produced subtrees only reference earlier ones)."""
+    if rel.WhichOneof("rel_type") == "reference":
+        return inline_reference_rels(subtrees[rel.reference.subtree_ordinal], subtrees)
+    out = stalg.Rel()
+    out.CopyFrom(rel)
+    _inline_reference_rels_in_place(out, subtrees)
+    return out
+
+
+def _inline_reference_rels_in_place(rel: stalg.Rel, subtrees) -> None:
+    for container, key in _iter_child_rels(rel):
+        child = _child_rel(container, key)
+        if child.WhichOneof("rel_type") == "reference":
+            child.CopyFrom(
+                inline_reference_rels(
+                    subtrees[child.reference.subtree_ordinal], subtrees
+                )
+            )
+        else:
+            _inline_reference_rels_in_place(child, subtrees)
+
+
 def merge_extensions_into(target, *sources):
     """Merge the extension URNs and declarations of ``sources`` into ``target`` in place.
 
diff --git a/tests/builders/plan/test_reference.py b/tests/builders/plan/test_reference.py
new file mode 100644
index 0000000..dd02f7e
--- /dev/null
+++ b/tests/builders/plan/test_reference.py
@@ -0,0 +1,104 @@
+import pytest
+import substrait.algebra_pb2 as stalg
+import substrait.plan_pb2 as stp
+import substrait.type_pb2 as stt
+
+from substrait.builders.extended_expression import literal
+from substrait.builders.plan import (
+    default_version,
+    fetch,
+    read_named_table,
+    reference,
+    set,
+)
+from substrait.builders.type import i64
+from substrait.extension_registry import ExtensionRegistry
+from substrait.type_inference import infer_plan_schema
+
+registry = ExtensionRegistry(load_default_extensions=False)
+
+struct = stt.Type.Struct(
+    types=[
+        stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)),
+        stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)),
+    ],
+    nullability=stt.Type.NULLABILITY_REQUIRED,
+)
+named_struct = stt.NamedStruct(names=["id", "v"], struct=struct)
+
+
+def test_reference_emits_shared_subtree_and_ref_root():
+    table = read_named_table("t", named_struct)
+    plan = reference(table)(registry)
+
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    # The promoted subtree is the source read; the root references it by ordinal 0.
+    assert plan.relations[0].rel.HasField("read")
+    root = plan.relations[-1].root
+    assert root.input.WhichOneof("rel_type") == "reference"
+    assert root.input.reference.subtree_ordinal == 0
+    assert list(root.names) == ["id", "v"]
+
+
+def test_reference_schema_inference_resolves_through_ref():
+    # A plan rooted at a ReferenceRel infers its schema from the shared subtree.
+    plan = reference(read_named_table("t", named_struct))(registry)
+    ns = infer_plan_schema(plan, registry=registry)
+    assert ns == named_struct
+
+
+def test_reference_shared_subtree_deduped_across_inputs(find_reference):
+    # The same cached subtree feeding two branches that meet at a set collapses to
+    # a single shared subtree, referenced from both inputs.
+    base = reference(read_named_table("t", named_struct))
+    plan = set([base, base], stalg.SetRel.SET_OP_UNION_ALL)(registry)
+
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    set_inputs = plan.relations[-1].root.input.set.inputs
+    assert [find_reference(i) for i in set_inputs] == [0, 0]
+
+
+def test_reference_distinct_subtrees_are_rebased(find_reference):
+    # Two *different* cached subtrees merged at a set keep distinct ordinals; the
+    # second input's ReferenceRel is rebased from 0 to 1.
+    a = reference(read_named_table("a", named_struct))
+    b = reference(read_named_table("b", named_struct))
+    plan = set([a, b], stalg.SetRel.SET_OP_UNION_ALL)(registry)
+
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "rel", "root"]
+    assert list(plan.relations[0].rel.read.named_table.names) == ["a"]
+    assert list(plan.relations[1].rel.read.named_table.names) == ["b"]
+    set_inputs = plan.relations[-1].root.input.set.inputs
+    assert [find_reference(i) for i in set_inputs] == [0, 1]
+
+
+def test_reference_carried_through_single_input_builder():
+    # A downstream single-input builder keeps the leading subtree (and its ordinal),
+    # wrapping the ReferenceRel rather than inlining the source.
+    base = reference(read_named_table("t", named_struct))
+    plan = fetch(base, offset=literal(0, i64()), count=literal(5, i64()))(registry)
+
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    assert plan.relations[0].rel.HasField("read")
+    fetch_rel = plan.relations[-1].root.input
+    assert fetch_rel.WhichOneof("rel_type") == "fetch"
+    assert fetch_rel.fetch.input.reference.subtree_ordinal == 0
+    # Schema still resolves through the wrapped reference.
+    assert infer_plan_schema(plan, registry=registry) == named_struct
+
+
+def test_reference_out_of_range_ordinal_raises():
+    # A bare ReferenceRel with no subtrees in scope cannot be inferred.
+    ref_plan = stp.Plan(
+        version=default_version,
+        relations=[
+            stp.PlanRel(
+                root=stalg.RelRoot(
+                    input=stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=0)),
+                    names=["id", "v"],
+                )
+            )
+        ],
+    )
+    with pytest.raises(Exception, match="out of range"):
+        infer_plan_schema(ref_plan, registry=registry)
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..173ddac
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,19 @@
+import pytest
+
+
+@pytest.fixture
+def find_reference():
+    """Return a helper that finds the ``subtree_ordinal`` of the first
+    ``ReferenceRel`` reachable through single-input wrapper relations, or None.
+
+    Shared by the ReferenceRel / CTE tests (builders and DataFrame layers)."""
+
+    def _find_reference(rel):
+        kind = rel.WhichOneof("rel_type")
+        if kind == "reference":
+            return rel.reference.subtree_ordinal
+        if kind in ("filter", "project", "fetch", "sort"):
+            return _find_reference(getattr(rel, kind).input)
+        return None
+
+    return _find_reference
diff --git a/tests/dataframe/test_cache.py b/tests/dataframe/test_cache.py
new file mode 100644
index 0000000..bdea408
--- /dev/null
+++ b/tests/dataframe/test_cache.py
@@ -0,0 +1,114 @@
+"""Tests for DataFrame.cache() -- shared subplans / CTEs (ReferenceRel).
+
+A cached frame is emitted once as a leading shared subtree (a ``PlanRel(rel=...)``)
+and referenced via a ``ReferenceRel`` wherever it is used; repeated uses that meet
+at a multi-input relation collapse to a single shared subtree.
+"""
+
+import pytest
+
+import substrait.dataframe as sub
+
+
+def test_cache_emits_shared_subtree_referenced_twice(find_reference):
+    base = sub.read_named_table("t", {"id": sub.i64, "v": sub.i64}).cache()
+    plan = (
+        base.filter(sub.col("id") < 10)
+        .union(base.filter(sub.col("id") >= 10))
+        .to_plan()
+    )
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    assert plan.relations[0].rel.HasField("read")
+    set_inputs = plan.relations[-1].root.input.set.inputs
+    assert [find_reference(i) for i in set_inputs] == [0, 0]
+
+
+def test_no_cache_inlines_single_relation(find_reference):
+    a = sub.read_named_table("t", {"id": sub.i64})
+    plan = a.filter(sub.col("id") > 0).union(a.filter(sub.col("id") < 0)).to_plan()
+    # Without cache the source is inlined into each union input -- no shared subtree.
+    assert len(plan.relations) == 1
+    for i in plan.relations[-1].root.input.set.inputs:
+        assert find_reference(i) is None
+
+
+def test_cache_schema_inference_through_reference(find_reference):
+    # Filtering/selecting the cached frame requires inferring its schema through the
+    # ReferenceRel.
+    base = sub.read_named_table("t", {"id": sub.i64, "v": sub.i64}).cache()
+    plan = base.filter(sub.col("v") > 0).select("id").to_plan()
+    # The plan must actually carry the shared subtree and reach it via a reference
+    # (a no-op cache would inline the read and drop both, so these guard the path).
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    assert find_reference(plan.relations[-1].root.input) == 0
+    assert plan.relations[-1].root.input.project.HasField("common")
+    assert list(plan.relations[-1].root.names) == ["id"]
+
+
+def test_cache_merges_subtree_extensions():
+    base = sub.read_named_table("t", {"id": sub.i64}).filter(sub.col("id") > 5).cache()
+    plan = base.union(base).to_plan()
+    # The comparison function used inside the cached subtree must be declared on the
+    # merged plan; the shared subtree itself must be emitted (a no-op cache would
+    # inline it into both union inputs, giving a single relation).
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "root"]
+    assert plan.relations[0].rel.filter.HasField("condition")
+    urns = {u.urn for u in plan.extension_urns}
+    assert "extension:io.substrait:functions_comparison" in urns
+
+
+def test_distinct_caches_rebased_across_union(find_reference):
+    # Two *different* cached frames unioned keep distinct subtree ordinals (0, 1),
+    # rebasing the second input's ReferenceRel.
+    a = sub.read_named_table("a", {"id": sub.i64}).cache()
+    b = sub.read_named_table("b", {"id": sub.i64}).cache()
+    plan = a.union(b).to_plan()
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "rel", "root"]
+    set_inputs = plan.relations[-1].root.input.set.inputs
+    assert [find_reference(i) for i in set_inputs] == [0, 1]
+
+
+def test_nested_cache_references_earlier_subtree(find_reference):
+    # A cache built on top of another cache produces a subtree that itself
+    # references the earlier subtree.
+    a = sub.read_named_table("t", {"id": sub.i64, "v": sub.i64}).cache()
+    b = a.filter(sub.col("v") > 0).cache()
+    plan = b.union(b).to_plan()
+    assert [r.WhichOneof("rel_type") for r in plan.relations] == ["rel", "rel", "root"]
+    # subtree 0 is the read; subtree 1 is b's filter, referencing subtree 0.
+    assert plan.relations[0].rel.HasField("read")
+    assert find_reference(plan.relations[1].rel) == 0
+    # Both set inputs reference the promoted b subtree (ordinal 1).
+    set_inputs = plan.relations[-1].root.input.set.inputs
+    assert [find_reference(i) for i in set_inputs] == [1, 1]
+
+
+def test_cache_inside_subquery_is_inlined(find_reference):
+    # A cached frame consumed inside a subquery cannot carry a plan-global
+    # ReferenceRel across the subquery boundary, so it is inlined into the subquery
+    # (self-contained) and the outer plan carries no dangling shared subtree.
+    from substrait.dataframe.expr import exists
+
+    a = sub.read_named_table("a", {"x": sub.i64}).cache()
+    b = sub.read_named_table("b", {"y": sub.i64})
+    plan = b.filter(exists(a)).to_plan()
+    assert all(r.WhichOneof("rel_type") != "rel" for r in plan.relations)
+    tuples = plan.relations[
+        -1
+    ].root.input.filter.condition.subquery.set_predicate.tuples
+    assert find_reference(tuples) is None
+    assert tuples.WhichOneof("rel_type") == "read"
+
+
+def test_hint_after_cache_raises():
+    base = sub.read_named_table("t", {"id": sub.i64}).cache()
+    with pytest.raises(TypeError, match="cannot attach a hint"):
+        base.hint(alias="x").to_plan()
+
+
+def test_hint_before_cache_is_fine():
+    # Hinting the node *before* caching it is allowed (the hint lands on the read).
+    plan = (
+        sub.read_named_table("t", {"id": sub.i64}).hint(alias="src").cache().to_plan()
+    )
+    assert plan.relations[0].rel.read.common.hint.alias == "src"
diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py
index e67cfe6..ca603bc 100644
--- a/tests/test_type_inference.py
+++ b/tests/test_type_inference.py
@@ -600,3 +600,29 @@ def _read(struct):
     # UNION -> nullable if nullable in any input; types (decimal 10/2, varchar 5,
     # list<string>) and the required inner string element are preserved.
     assert infer_rel_schema(rel) == _struct(_NULL, _NULL, _NULL)
+
+
+def test_inference_reference_resolves_against_subtree():
+    # A ReferenceRel's schema is the schema of the shared subtree its
+    # subtree_ordinal indexes into, taken from the `subtrees` list.
+    ref = stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=1))
+    assert infer_rel_schema(ref, subtrees=[right_read_rel, read_rel]) == struct
+
+
+def test_inference_reference_through_wrapping_relation():
+    # A reference nested under another relation resolves too (subtrees thread down).
+    filt = stalg.Rel(
+        filter=stalg.FilterRel(
+            input=stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=0))
+        )
+    )
+    assert infer_rel_schema(filt, subtrees=[read_rel]) == struct
+
+
+def test_inference_reference_out_of_range_raises():
+    ref = stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=3))
+    with pytest.raises(Exception, match="out of range"):
+        infer_rel_schema(ref, subtrees=[read_rel])
+    # No subtrees in scope at all is also out of range.
+    with pytest.raises(Exception, match="out of range"):
+        infer_rel_schema(ref)