refactor(api): drop cache()/reference from this PR (track as #211)

Remove the cache() CTE feature and its ReferenceRel plumbing to keep this PR
focused. Per review, shared subplans should be implemented at both the builder
and dataframe layers with subtrees carried in-band in the Plan (like extensions)
rather than via a contextvar; tracked in #211.

- frame.py: remove DataFrame.cache(), _CteContext, the _cte_context contextvar,
  and _materialize's subtree prepend (to_plan/to_substrait resolve directly);
  drop the now-unreachable hint() guard for common-less relations.
- type_inference.py: remove the reference_subtrees contextvar and the
  ReferenceRel schema-inference case.
- builders/plan.py: remove the reference() builder added in fb1c49b.
- tests: remove the cache/reference and hint-after-cache tests.

Author: nielspardon

src/substrait/builders/plan.py

@@ -389,36 +389,6 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
     return resolve
 
 
-def reference(
-    subtree_ordinal: int,
-    names: Iterable[str],
-    *extension_sources: PlanOrUnbound,
-) -> UnboundPlan:
-    """A plan whose root references a shared subtree by ordinal (a ReferenceRel).
-
-    ``subtree_ordinal`` indexes into the plan's leading shared relations (the
-    common subplans a ``ReferenceRel`` points at); ``names`` labels the output.
-    ``extension_sources`` are the already-resolved subtree plans whose extension
-    declarations are propagated, so a builder merging this plan upward keeps the
-    subtree's extensions.
-
-    Inferring the resulting ``ReferenceRel``'s schema needs the subtree list to
-    be in scope via :data:`substrait.type_inference.reference_subtrees`; this
-    builder is meant for shared-subplan / CTE construction (see
-    ``substrait.dataframe.DataFrame.cache``), which sets that up.
-    """
-
-    def resolve(registry: ExtensionRegistry) -> stp.Plan:
-        rel = stalg.Rel(reference=stalg.ReferenceRel(subtree_ordinal=subtree_ordinal))
-        return stp.Plan(
-            version=default_version,
-            relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=list(names)))],
-            **_merge_extensions(*extension_sources),
-        )
-
-    return resolve
-
-
 def fetch(
     plan: PlanOrUnbound,
     offset: ExtendedExpressionOrUnbound,

src/substrait/dataframe/frame.py

@@ -28,19 +28,17 @@
 
 from __future__ import annotations
 
-import contextvars
 from itertools import combinations
 from typing import Any, Iterable, Optional, Union
 
 import substrait.algebra_pb2 as stalg
-import substrait.plan_pb2 as stpl
 import substrait.type_pb2 as stp
 
 from substrait.builders import plan as _plan
 from substrait.builders import type as _type
 from substrait.dataframe.expr import Expr, Measure, col, lit
 from substrait.extension_registry import ExtensionRegistry
-from substrait.type_inference import infer_plan_schema, reference_subtrees
+from substrait.type_inference import infer_plan_schema
 
 # All 13 JoinRel.JoinType variants (SET_OP_UNSPECIFIED excluded). "single"
 # returns at most one right match per left row (runtime error on multiple);
@@ -126,34 +124,6 @@ def _split_measure(m: Union[Expr, Measure]):
     return _unbound(m), None
 
 
-class _CteContext:
-    """Collects shared subtrees while a plan with ``cache()`` is being built."""
-
-    __slots__ = ("subtrees", "names", "plans", "ordinal_by_token")
-
-    def __init__(self):
-        self.subtrees: "list[stalg.Rel]" = []  # indexed by subtree_ordinal
-        self.names: "list[list[str]]" = []
-        self.plans: "list[stpl.Plan]" = []  # the resolved subtree plans (extensions)
-        self.ordinal_by_token: dict = {}
-
-
-# Active while a DataFrame is being materialized; None otherwise.
-#
-# ``cache()`` needs one accumulator shared by every nested ``resolve()`` in a
-# single ``to_plan()`` -- to dedupe repeated uses of a cached frame by identity
-# and hand out ``subtree_ordinal``s. The builder contract is
-# ``UnboundPlan = Callable[[ExtensionRegistry], Plan]``, which has no slot to
-# thread that accumulator, and schema inference (which resolves ReferenceRels
-# against the subtree list) runs *inside* the builder layer. Rather than widen
-# that contract across every builder, ``cache()`` publishes the accumulator here
-# and ``type_inference.reference_subtrees`` exposes its ``.subtrees`` list to
-# inference; ``_materialize`` sets both for the duration of the build.
-_cte_context: contextvars.ContextVar = contextvars.ContextVar(
-    "cte_context", default=None
-)
-
-
 _default_registry: Optional[ExtensionRegistry] = None
 
 
@@ -532,15 +502,6 @@ def resolve(registry: ExtensionRegistry):
             bound = inner(registry)
             rel = bound.relations[-1].root.input
             rel_inner = getattr(rel, rel.WhichOneof("rel_type"))
-            # A few relations (ReferenceRel from .cache(), 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
@@ -666,62 +627,13 @@ def write_named_table(
             )
         )
 
-    def cache(self) -> "DataFrame":
-        """Mark this DataFrame as a reusable common subplan (a CTE).
-
-        Every use of the returned frame in the same ``to_plan()`` emits the
-        subplan once as a shared subtree and references it (``ReferenceRel``),
-        instead of inlining a fresh copy each time.
-        """
-        inner = self._plan
-        token = object()  # identity for this cached node
-
-        def resolve(registry: ExtensionRegistry) -> stpl.Plan:
-            ctx = _cte_context.get(None)
-            if ctx is None:  # built without a context -> just inline
-                return inner(registry)
-            ordinal = ctx.ordinal_by_token.get(token)
-            if ordinal is None:
-                subplan = inner(registry)
-                ordinal = len(ctx.subtrees)
-                ctx.ordinal_by_token[token] = ordinal
-                ctx.subtrees.append(subplan.relations[-1].root.input)
-                ctx.names.append(list(subplan.relations[-1].root.names))
-                ctx.plans.append(subplan)
-            # Emit a ReferenceRel to the shared subtree, propagating the
-            # subtree's extensions so builder merges carry them up.
-            return _plan.reference(ordinal, ctx.names[ordinal], ctx.plans[ordinal])(
-                registry
-            )
-
-        return self._next(resolve)
-
-    def _materialize(self, registry: ExtensionRegistry) -> stpl.Plan:
-        ctx = _CteContext()
-        cte_token = _cte_context.set(ctx)
-        ref_token = reference_subtrees.set(ctx.subtrees)
-        try:
-            plan = self._plan(registry)
-        finally:
-            _cte_context.reset(cte_token)
-            reference_subtrees.reset(ref_token)
-        if not ctx.subtrees:
-            return plan
-        # Prepend the shared subtrees; ReferenceRel ordinals index into them.
-        subtree_rels = [stpl.PlanRel(rel=s) for s in ctx.subtrees]
-        return stpl.Plan(
-            version=_plan.default_version,
-            relations=[*subtree_rels, plan.relations[-1]],
-            **_plan._merge_extensions(plan, *ctx.plans),
-        )
-
     def to_plan(self):
         """Materialize to a ``substrait.proto.Plan``."""
-        return self._materialize(self._registry)
+        return self._plan(self._registry)
 
     # Kept for parity with the substrait.narwhals (Narwhals) wrapper's API.
     def to_substrait(self, registry: Optional[ExtensionRegistry] = None):
-        return self._materialize(registry or self._registry)
+        return self._plan(registry or self._registry)
 
 
 class GroupBy:

src/substrait/type_inference.py

@@ -5,12 +5,6 @@
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
 
-# The shared subtrees (Rels) a ReferenceRel's subtree_ordinal indexes into, set
-# for the duration of a plan build so a `reference` relation's schema resolves.
-reference_subtrees: contextvars.ContextVar = contextvars.ContextVar(
-    "reference_subtrees", default=None
-)
-
 # 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.
@@ -535,12 +529,6 @@ def infer_rel_schema(rel: stalg.Rel) -> stt.Type.Struct:
                 f"{rel.extension_multi.detail.type_url!r}"
             )
         (common, struct) = (rel.extension_multi.common, derived.struct)
-    elif rel_type == "reference":
-        subtrees = reference_subtrees.get()
-        if subtrees is None:
-            raise Exception("cannot infer a ReferenceRel's schema outside a plan build")
-        # ReferenceRel has no common/emit; its schema is the subtree's schema.
-        return infer_rel_schema(subtrees[rel.reference.subtree_ordinal])
     else:
         raise Exception(f"Unhandled rel_type {rel_type}")
 

tests/dataframe/test_frame.py

@@ -794,57 +794,6 @@ def test_unpivot_requires_on():
         _wide_df().unpivot([], index="region")
 
 
-# -- Phase (core-ext): References / CTEs -----------------------------------
-
-
-def _find_reference(rel):
-    kind = rel.WhichOneof("rel_type")
-    if kind == "reference":
-        return rel.reference.subtree_ordinal
-    for field in ("filter", "project", "fetch", "sort"):
-        if kind == field:
-            return _find_reference(getattr(rel, field).input)
-    return None
-
-
-def test_cache_emits_shared_subtree_referenced_twice():
-    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():
-    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.
-    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():
-    # 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()
-    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()
-    urns = {u.urn for u in plan.extension_urns}
-    assert "extension:io.substrait:functions_comparison" in urns
-
-
 # -- Phase (core-ext): physical joins + exchange --------------------------
 
 
@@ -1209,14 +1158,6 @@ def test_mark_join_output_names_match_types():
     assert ns.struct.types[-1].WhichOneof("kind") == "bool"
 
 
-def test_hint_after_cache_raises_clear_error():
-    # Regression: a ReferenceRel (from .cache()) has no RelCommon, so hint() must
-    # fail with a clear message rather than an opaque AttributeError.
-    df = sub.read_named_table("t", {"a": sub.i64})
-    with pytest.raises(TypeError, match="cannot attach a hint"):
-        df.cache().hint(row_count=100).to_plan()
-
-
 def test_default_registry_is_reused():
     assert sub.default_registry() is sub.default_registry()