From 279cfb86a95b50eadd4e944509b0e3c627a8a38a Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 10 Aug 2026 14:08:56 +0200 Subject: [PATCH 1/3] perf: infer each relation's schema once per build, not once per level Every verb resolves its input's schema and plans are built as nested resolvers, so an N-verb chain re-walked the whole subtree beneath it at every level. Profiling a 40-verb chain put 41 of 45 ms in `infer_plan_schema`, split between the anchor index (30 ms) and the inference recursion (9 ms); assembling the protobuf was 2 ms. An input's root `Rel` is copied when it is assigned into the output relation, so the schema just inferred for it is unreachable from the copy by identity. Message wrappers are identity-stable, though, so `_plan_from` names the copies as it makes them and records that each has its input plan's output schema; `infer_rel_schema` stops at that boundary instead of recursing through it. The record is the plan, not the schema, so a builder that never needs its input's schema still never causes one to be inferred. The memo lives in the build scope, next to the ExtensionCollector. `infer_plan_schema` also builds its rel_anchor index -- a walk of every relation and expression in the plan -- only when an id-based outer reference asks for one. Closes #207 --- src/substrait/builders/plan.py | 45 +++- src/substrait/extension_registry/collector.py | 11 +- src/substrait/type_inference.py | 153 ++++++++++++- src/substrait/utils/__init__.py | 12 + tests/builders/plan/test_schema_memo.py | 210 ++++++++++++++++++ tests/test_type_inference.py | 51 +++++ tests/test_utils.py | 26 +++ 7 files changed, 497 insertions(+), 11 deletions(-) create mode 100644 tests/builders/plan/test_schema_memo.py diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 519217f..ba9ba0a 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -31,8 +31,10 @@ _outer_anchor_binding, infer_plan_schema, join_output_names, + remember_rel_output_schema, ) from substrait.utils import ( + child_rels, plan_subtrees, rebase_reference_ordinals, remap_function_references, @@ -148,6 +150,33 @@ def _merge_input_subtrees(bound_inputs): return subtree_planrels, rebased_root_inputs +def _remember_input_schemas(plan: stp.Plan, bound_inputs) -> None: + """Record each embedded input relation's output schema for the build in progress. + + Assigning an input's root ``Rel`` into the output relation copies it, so the + schema just inferred for that input is unreachable from the copy by identity and + every enclosing level would re-walk the whole subtree below it. Naming the copies + here is what keeps that walk from happening (see + ``type_inference.remember_rel_output_schema``); the schemas themselves are not + inferred, only pointed at, so a builder that never needs its input's schema still + never causes one to be inferred. + + The pairing is positional: the i-th child ``Rel`` of the assembled relation, in + field declaration order, is the i-th bound input. Every builder here hands + ``make_rel``'s ``inp`` to its relation in that order (``left=inp[0], + right=inp[1]``, ``inputs=inp``); the count is checked below, and the order is + pinned by a test over every builder that embeds more than one input, since a + silent swap would hand a level its sides' schemas the wrong way round. + """ + children = list(child_rels(plan.relations[-1].root.input)) + assert len(children) == len(bound_inputs), ( + f"assembled relation embeds {len(children)} input relation(s) but the builder " + f"bound {len(bound_inputs)}" + ) + for child, bound_input in zip(children, bound_inputs): + remember_rel_output_schema(child, bound_input) + + def _plan_from( bound_inputs, make_rel, names, metadata_sources, *, include_version=True ): @@ -172,7 +201,9 @@ def _plan_from( } if include_version: kwargs["version"] = default_version - return stp.Plan(**kwargs) + plan = stp.Plan(**kwargs) + _remember_input_schemas(plan, bound_inputs) + return plan def with_execution_behavior( @@ -210,6 +241,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: result.ClearField("extension_urns") result.ClearField("extensions") result.execution_behavior.variable_eval_mode = variable_eval_mode + # The copy carries the input's relations verbatim, so its root has the input's + # output schema -- recorded against the copy, which is a different object. + remember_rel_output_schema(result.relations[-1].root.input, bound_plan) return result return build_scoped(resolve) @@ -559,7 +593,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: 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( + result = stp.Plan( version=default_version, relations=[ *nested, @@ -568,6 +602,13 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ], **_merge_plan_metadata(bound), ) + # Both the ReferenceRel and the subtree it points at have the promoted plan's + # output schema. Recording them means resolving a reference costs a lookup + # rather than a walk of the shared subtree -- which every downstream verb of a + # cached frame would otherwise repeat. + for rel in (result.relations[-1].root.input, result.relations[ordinal].rel): + remember_rel_output_schema(rel, bound) + return result return build_scoped(resolve) diff --git a/src/substrait/extension_registry/collector.py b/src/substrait/extension_registry/collector.py index 332e3d9..640e310 100644 --- a/src/substrait/extension_registry/collector.py +++ b/src/substrait/extension_registry/collector.py @@ -28,6 +28,8 @@ import substrait.extensions.extensions_pb2 as ste import substrait.plan_pb2 as stplan +from substrait.type_inference import schema_memo_scope + # Identity of a function as declared in a plan: (extension URN, function name). # The name is the compound form carried by SimpleExtensionDeclaration (e.g. # "add:i64_i64"), which is what makes an identity resolvable without the catalog. @@ -217,6 +219,12 @@ def build_scope(): extensions onto its output. Nested resolvers get the same collector and write nothing, which is what lets a build accumulate extensions once instead of re-merging them at every level. + + The outermost resolver also opens the build's schema memo + (:func:`substrait.type_inference.schema_memo_scope`), which spares the builders + the matching re-derivation on the schema side: both are state derived from the + plan being assembled and meaningless once it is finished, so both live and die + with this scope. """ collector = _collector.get() if collector is not None: @@ -225,7 +233,8 @@ def build_scope(): collector = ExtensionCollector() token = _collector.set(collector) try: - yield collector, True + with schema_memo_scope(): + yield collector, True finally: _collector.reset(token) diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 8c431a6..6ace83d 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -1,5 +1,6 @@ import contextlib import contextvars +from typing import Optional import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee @@ -71,12 +72,23 @@ class _AnchorScope: sub-tree whose anchoring relation is not yet assembled (a lateral join's right input, at build or inference time). ``parent`` chains to an enclosing scope so nested correlations still resolve outer anchors. + + The index itself may be supplied as a zero-argument factory instead of a dict, + and is then built only if an id-based reference actually asks for an anchor. + Indexing a plan means walking every relation *and* every expression in it to + find the relations embedded in subqueries, which is whole-plan work; plans + carrying an id-based ``OuterReference`` at all are the exception, and a builder + re-infers its input's schema at every level (see :class:`_SchemaMemo`), so + building the index eagerly made that walk the single largest cost of assembling + a long pipeline. """ - __slots__ = ("_rels", "_subtrees", "_schemas", "_resolving", "_parent") + __slots__ = ("_rels", "_index", "_subtrees", "_schemas", "_resolving", "_parent") - def __init__(self, rels: dict, subtrees, *, parent=None): - self._rels = rels + def __init__(self, rels, subtrees, *, parent=None): + # Either the index or a factory for it; whichever it is not stays None. + self._rels: Optional[dict] = rels if isinstance(rels, dict) else None + self._index = None if isinstance(rels, dict) else rels self._subtrees = subtrees self._schemas: dict = {} self._resolving: set = set() @@ -86,10 +98,17 @@ def register(self, anchor, struct: stt.Type.Struct) -> None: """Pre-bind ``anchor`` to an already-known schema.""" self._schemas[anchor] = struct + def _anchors(self) -> dict: + """The anchor index, built on first use if it was supplied as a factory.""" + if self._rels is None: + self._rels = self._index() + self._index = None + return self._rels + def schema_of(self, anchor, registry) -> stt.Type.Struct: if anchor in self._schemas: return self._schemas[anchor] - if anchor not in self._rels: + if anchor not in self._anchors(): if self._parent is not None: return self._parent.schema_of(anchor, registry) raise Exception(f"outer reference to unknown rel_anchor {anchor}") @@ -135,6 +154,108 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: ) +class _SchemaMemo: + """Output structs already known for particular ``Rel`` *objects*, for the + duration of one build. + + A builder assembles its output relation by assigning its input's root ``Rel`` + into a fresh message, which protobuf copies. The schema it inferred for that + input one level down is therefore unreachable by object identity from the copy, + so every level re-walks the whole subtree beneath it and an N-verb chain does + O(N^2) inference (#207). The copy *is* reachable at the moment it is made, + though, and protobuf message wrappers are identity-stable, so the builder + records "this relation's output schema is that plan's output schema" and + :func:`infer_rel_schema` stops at the boundary instead of recursing through it. + + What is recorded is the ``Plan`` the schema comes from, not the schema, resolved + on first lookup and then kept. Recording it eagerly would mean inferring schemas + no one asked for, and inference can legitimately fail where building does not -- + ``set``, ``reference`` and ``exchange`` never look at their input's schema today, + so a plan they accept (an extension relation with no registered deriver, say) + has to keep building. The cost of deferring is that an unresolved entry keeps its + input plan alive until the build finishes rather than until the level above it + returns; the entries a chain of verbs produces are consumed one level up, and the + rest is bounded by the plan under construction. + + Only builders write here; inference never memoizes on its own. A relation's + output struct can depend on ambient correlation context (``outer_schemas``, + ``anchor_scope``) -- a projection of a correlated column is one -- and every + relation that crosses into another context is copied on the way, so an entry can + only ever be read back under the context it was recorded in. Caching inference + results wholesale would not have that property. + """ + + __slots__ = ("_structs", "_pending", "_resolving") + + def __init__(self) -> None: + # id(Rel) -> (Rel, struct) resolved, and id(Rel) -> (Rel, Plan) not yet. The + # Rel is held in the value so its id stays valid -- and stays *that* Rel's -- + # for the lifetime of the memo. + self._structs: dict = {} + self._pending: dict = {} + self._resolving: set = set() + + def remember_plan_output(self, rel: stalg.Rel, plan: stp.Plan) -> None: + """Record that ``rel``'s output schema is the output schema of ``plan``.""" + self._pending[id(rel)] = (rel, plan) + + def struct_of(self, rel: stalg.Rel, registry) -> Optional[stt.Type.Struct]: + """``rel``'s remembered output struct, or None if nothing was recorded.""" + key = id(rel) + known = self._structs.get(key) + if known is not None: + return known[1] + pending = self._pending.get(key) + if pending is None: + return None + if key in self._resolving: + raise Exception( + "remembered schema resolves to itself; a relation cannot be its own " + "input" + ) + self._resolving.add(key) + try: + struct = infer_plan_schema(pending[1], registry=registry).struct + finally: + self._resolving.discard(key) + self._structs[key] = (rel, struct) + # The plan was held only to answer this; the struct replaces it. + del self._pending[key] + return struct + + +# The schema memo (a ``_SchemaMemo``) for the build in progress, or None outside a +# build -- so inference used directly as a library function memoizes nothing and +# behaves exactly as before. Entered by ``extension_registry.build_scope`` alongside +# the build's ExtensionCollector, the two being the same kind of state: derived from +# the plan being assembled, and meaningless once it is finished. +schema_memo: contextvars.ContextVar = contextvars.ContextVar( + "schema_memo", default=None +) + + +@contextlib.contextmanager +def schema_memo_scope(): + """Scope a fresh :class:`_SchemaMemo` to the build in progress.""" + token = schema_memo.set(_SchemaMemo()) + try: + yield + finally: + schema_memo.reset(token) + + +def remember_rel_output_schema(rel: stalg.Rel, plan: stp.Plan) -> None: + """Record that ``rel``'s output schema is ``plan``'s, for the build in progress. + + Called by a builder for the input relations it has just embedded in the output it + assembled, so a later inference of that output stops at the boundary rather than + re-walking everything below it. A no-op outside a build. + """ + memo = schema_memo.get() + if memo is not None: + memo.remember_plan_output(rel, plan) + + @contextlib.contextmanager def _outer_anchor_binding(anchor, struct): """Bind ``anchor`` -> ``struct`` for id-based outer references resolved within @@ -692,6 +813,16 @@ def infer_rel_schema(rel: stalg.Rel, *, registry=None, subtrees=()) -> stt.Type. ``ReferenceRel`` resolves its schema against ``subtrees[subtree_ordinal]``. It defaults to ``()`` so plans without shared subtrees behave exactly as before. """ + # A builder may already have recorded this exact relation's output schema while + # embedding it (see _SchemaMemo), in which case the subtree below it needs no + # walking. The recorded struct is the post-emit output struct -- what this + # function returns -- so it is returned as-is. + memo = schema_memo.get() + if memo is not None: + remembered = memo.struct_of(rel, registry) + if remembered is not None: + return remembered + rel_type = rel.WhichOneof("rel_type") if rel_type == "read": @@ -953,12 +1084,18 @@ def infer_plan_schema(plan: stp.Plan, *, registry=None) -> stt.NamedStruct: # 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)) + # Index every RelCommon.rel_anchor in the plan (across subtrees, the root, and # subquery-embedded relations) so an id-based OuterReference (rel_reference) - # anywhere resolves against the anchored relation's output schema. - anchors = { - a: rel for rel in iter_plan_rels(plan) if (a := rel_anchor_of(rel)) is not None - } + # anywhere resolves against the anchored relation's output schema. Passed as a + # factory: indexing walks the whole plan, and most plans never ask for an anchor. + def anchors(): + return { + a: rel + for rel in iter_plan_rels(plan) + if (a := rel_anchor_of(rel)) is not None + } + # Chain to any enclosing anchor scope (e.g. a lateral join binding its left # schema while its right input -- a separate plan being inferred here -- is # built) so references to an outer anchor still resolve. diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index c41782a..e46fe02 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -138,6 +138,18 @@ def _child_rel(container, key): return container[key] if isinstance(key, int) else getattr(container, key) +def child_rels(rel: stalg.Rel): + """``rel``'s direct child ``Rel`` messages, in field declaration order. + + Yields the messages themselves rather than the ``(container, key)`` pairs + :func:`_iter_child_rels` uses for in-place rewriting, for callers that only read + them -- note they are the live submessages, so identity is meaningful. + Subquery-embedded relations are not direct children and are not yielded. + """ + for container, key in _iter_child_rels(rel): + yield _child_rel(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.""" diff --git a/tests/builders/plan/test_schema_memo.py b/tests/builders/plan/test_schema_memo.py new file mode 100644 index 0000000..43ef1bd --- /dev/null +++ b/tests/builders/plan/test_schema_memo.py @@ -0,0 +1,210 @@ +"""The per-build schema memo: what it costs to build a pipeline, and the pairing it +relies on. + +Every verb resolves its input's schema, and a plan is built as nested resolvers, so +without memoization an N-verb chain re-walks the whole subtree beneath it at every +level -- O(N^2) inference (#207). ``builders.plan._remember_input_schemas`` records +each embedded input relation's schema as the plan is assembled, and +``infer_plan_schema`` builds its ``rel_anchor`` index only when an id-based outer +reference asks for one. Both are invisible in the emitted plan, so they need tests +that observe cost, plus tests that the recorded schemas land on the right relations. +""" + +import collections + +import pytest +import substrait.algebra_pb2 as stalg +import substrait.plan_pb2 as stp +import substrait.type_pb2 as stt + +import substrait.type_inference as type_inference +from substrait.builders.extended_expression import column, literal +from substrait.builders.plan import ( + cross, + hash_join, + join, + lateral_join, + merge_join, + nested_loop_join, + project, + read_named_table, + reference, + with_execution_behavior, + write_named_table, +) +from substrait.builders.type import boolean, i64, string +from substrait.extension_registry import ExtensionRegistry +from substrait.type_inference import infer_plan_schema, schema_memo + +registry = ExtensionRegistry(load_default_extensions=False) + +named_struct = stt.NamedStruct( + names=["k", "v"], + struct=stt.Type.Struct( + types=[i64(nullable=False), i64(nullable=False)], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), +) + +# Same arity as `named_struct` but a different second column type, so pairing the +# two sides of a join the wrong way round shows up in the types alone. +right_named_struct = stt.NamedStruct( + names=["rk", "rv"], + struct=stt.Type.Struct( + types=[i64(nullable=False), string()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), +) + + +@pytest.fixture +def counts(monkeypatch): + """Counts of the two whole-subtree walks a build must not repeat per level. + + Both are patched on ``substrait.type_inference`` because that is where the + recursion and the anchor index resolve them from. + """ + counted: collections.Counter = collections.Counter() + + infer_rel_schema = type_inference.infer_rel_schema + iter_plan_rels = type_inference.iter_plan_rels + + def counting_infer_rel_schema(rel, **kwargs): + counted["infer_rel_schema"] += 1 + return infer_rel_schema(rel, **kwargs) + + def counting_iter_plan_rels(plan): + counted["iter_plan_rels"] += 1 + return iter_plan_rels(plan) + + monkeypatch.setattr(type_inference, "infer_rel_schema", counting_infer_rel_schema) + monkeypatch.setattr(type_inference, "iter_plan_rels", counting_iter_plan_rels) + return counted + + +def _project_chain(length: int): + plan = read_named_table("t", named_struct) + for _ in range(length): + plan = project(plan, expressions=[column("v")]) + return plan + + +# Deliberately well above the 4N-4 this currently does and well below the ~N^2/2 it +# did before, so the test tracks the complexity class rather than the exact count. +_CALLS_PER_VERB = 6 + + +@pytest.mark.parametrize("length", [4, 8, 16, 32]) +def test_building_a_chain_infers_each_level_a_bounded_number_of_times(counts, length): + built = _project_chain(length)(registry) + + assert len(built.relations[-1].root.names) == 2 + length + assert counts["infer_rel_schema"] <= _CALLS_PER_VERB * length + + +def test_chain_inference_grows_linearly_not_quadratically(counts): + _project_chain(8)(registry) + short = counts["infer_rel_schema"] + counts.clear() + _project_chain(32)(registry) + long = counts["infer_rel_schema"] + + # Four times the verbs, so linear allows roughly four times the inferences (with + # headroom); quadratic would be sixteen. + assert long <= 6 * short + + +def test_building_a_chain_never_indexes_rel_anchors(counts): + # Indexing walks every relation and expression in the plan. Nothing here carries + # an id-based OuterReference, so nothing should ask for the index. + _project_chain(8)(registry) + + assert counts["iter_plan_rels"] == 0 + + +def test_memo_does_not_outlive_the_build(): + # The memo keys on object identity and holds its keys alive, so leaking it past + # the build would both pin memory and answer for relations of a later one. + assert schema_memo.get() is None + built = _project_chain(2)(registry) + assert schema_memo.get() is None + + # Inference of the finished plan is unmemoized and still correct. + assert list(infer_plan_schema(built, registry=registry).names) == [ + "k", + "v", + "v", + "v", + ] + + +def _left(): + return read_named_table("left", named_struct) + + +def _right(): + return read_named_table("right", right_named_struct) + + +def _true(): + return literal(True, boolean()) + + +# Every builder that embeds more than one input relation, since those are the ones +# whose recorded schemas could be paired with the wrong side. +TWO_INPUT_BUILDERS = { + "join": lambda: join(_left(), _right(), _true(), stalg.JoinRel.JOIN_TYPE_INNER), + "cross": lambda: cross(_left(), _right()), + "nested_loop_join": lambda: nested_loop_join( + _left(), _right(), _true(), stalg.NestedLoopJoinRel.JOIN_TYPE_INNER + ), + "hash_join": lambda: hash_join( + _left(), _right(), ["k"], ["rk"], stalg.HashJoinRel.JOIN_TYPE_INNER + ), + "merge_join": lambda: merge_join( + _left(), _right(), ["k"], ["rk"], stalg.MergeJoinRel.JOIN_TYPE_INNER + ), + "lateral_join": lambda: lateral_join( + _left(), lambda handle: _right(), stalg.JoinRel.JOIN_TYPE_INNER + ), +} + + +@pytest.mark.parametrize("builder", TWO_INPUT_BUILDERS.values(), ids=TWO_INPUT_BUILDERS) +def test_two_input_builders_record_each_side_against_its_own_relation(builder): + # write_named_table emits the schema it inferred for its input, so it reports what + # the level above the join sees: left columns then right columns. Swapping the + # recorded schemas keeps the arity and the names but reorders the types. + written = write_named_table("out", builder())(registry) + + table_schema = written.relations[-1].root.input.write.table_schema + assert list(table_schema.names) == ["k", "v", "rk", "rv"] + assert list(table_schema.struct.types) == [ + i64(nullable=False), + i64(nullable=False), + i64(nullable=False), + string(), + ] + + +def test_reference_records_the_promoted_subtree_and_the_reference(): + # A ReferenceRel's schema is its subtree's, and `reference` records both so a + # downstream verb resolves it by lookup rather than by walking the subtree. + written = write_named_table("out", reference(_left()))(registry) + + table_schema = written.relations[-1].root.input.write.table_schema + assert table_schema == named_struct + + +def test_with_execution_behavior_records_the_copied_root(): + # This builder copies its input plan wholesale rather than assembling a fresh + # relation, so the copy needs its own record. + written = write_named_table( + "out", + with_execution_behavior( + _left(), stp.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD + ), + )(registry) + + table_schema = written.relations[-1].root.input.write.table_schema + assert table_schema == named_struct diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index b5b9dd0..3768754 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -903,3 +903,54 @@ def test_infer_rel_reference_anchor_zero_is_a_distinct_anchor(): + [stt.Type(fp32=stt.Type.FP32(nullability=stt.Type.NULLABILITY_NULLABLE))] ) assert infer_plan_schema(plan).struct == expected + + +def test_anchor_index_is_built_only_when_a_rel_reference_needs_it(monkeypatch): + # Indexing rel_anchors walks every relation *and* every expression of the plan to + # reach the relations embedded in subqueries, so it is whole-plan work on every + # call. Plans carrying an id-based OuterReference are the exception, and the + # builders re-infer their input's schema at every level, so the index is built on + # demand rather than for every inference. + import substrait.type_inference as type_inference + + iter_plan_rels = type_inference.iter_plan_rels + indexed = [] + + def counting_iter_plan_rels(plan): + indexed.append(plan) + return iter_plan_rels(plan) + + monkeypatch.setattr(type_inference, "iter_plan_rels", counting_iter_plan_rels) + + plain = stp.Plan( + relations=[stp.PlanRel(root=stalg.RelRoot(input=read_rel, names=["a"]))] + ) + assert infer_plan_schema(plain).struct == struct + assert indexed == [] + + # A rel_reference does need it, and still gets it. + anchored = stalg.Rel( + read=stalg.ReadRel( + base_schema=named_struct, + common=stalg.RelCommon(rel_anchor=7), + named_table=stalg.ReadRel.NamedTable(names=["shared"]), + ) + ) + correlated = stp.Plan( + relations=[ + stp.PlanRel(rel=anchored), + stp.PlanRel( + root=stalg.RelRoot( + input=stalg.Rel( + project=stalg.ProjectRel( + input=right_read_rel, + expressions=[_outer_ref(2, rel_reference=7)], + ) + ), + names=["a", "b", "c"], + ) + ), + ] + ) + assert len(infer_plan_schema(correlated).struct.types) == 3 + assert len(indexed) == 1 diff --git a/tests/test_utils.py b/tests/test_utils.py index f56e21b..0714786 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,6 +6,7 @@ import substrait.type_pb2 as stt from substrait.utils import ( + child_rels, iter_plan_rels, merge_extension_declarations, merge_extension_urns, @@ -762,3 +763,28 @@ def test_convert_is_idempotent(): once = to_id_based_outer_references(plan) twice = to_id_based_outer_references(once) assert twice == once + + +def test_child_rels_yields_the_live_submessages_in_declaration_order(): + left = _read("l") + right = _read("r") + rel = stalg.Rel( + join=stalg.JoinRel(left=left, right=right, type=stalg.JoinRel.JOIN_TYPE_INNER) + ) + + children = list(child_rels(rel)) + + # Declaration order, and the messages themselves -- callers key on identity, so + # yielding a copy would silently break them. + assert children == [left, right] + assert children[0] is rel.join.left + assert children[1] is rel.join.right + + +def test_child_rels_yields_repeated_inputs_and_nothing_for_a_leaf(): + inputs = [_read("a"), _read("b"), _read("c")] + union = stalg.Rel(set=stalg.SetRel(inputs=inputs, op=stalg.SetRel.SET_OP_UNION_ALL)) + + assert list(child_rels(union)) == inputs + assert list(child_rels(_read("a"))) == [] + assert list(child_rels(stalg.Rel())) == [] From e415a8b7d26b0ea09876816041a326651e3b265a Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 10 Aug 2026 14:33:37 +0200 Subject: [PATCH 2/3] test: cover the pass-through relations' schemas and column() by ordinal Two gaps the byte-for-byte comparison against main turned up while verifying the schema memo, both pre-existing: nothing inferred a SortRel's schema (sort is always terminal in the suite, and it had no direct unit test), and column() was only ever called with a name, never an ordinal. The sort branch sits in the function the memo now short-circuits, so leaving it uncovered would mean a change there could only be caught downstream. --- .../extended_expression/test_column.py | 21 ++++++++ tests/test_type_inference.py | 50 +++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/tests/builders/extended_expression/test_column.py b/tests/builders/extended_expression/test_column.py index 9449848..ec75228 100644 --- a/tests/builders/extended_expression/test_column.py +++ b/tests/builders/extended_expression/test_column.py @@ -103,3 +103,24 @@ def test_column_nested_struct(): ], base_schema=nested_named_struct, ) + + +def test_column_by_ordinal(): + # column() takes an index as well as a name, and the index is a top-level field + # ordinal -- so the output name is read from the same slice of base_schema.names + # that a lookup by name would have landed on. + assert column(1)(named_struct, None) == column("description")(named_struct, None) + + +def test_column_by_ordinal_over_a_nested_struct(): + # A struct field consumes several entries of base_schema.names (its own plus one + # per member), so the ordinal indexes top-level fields while the output names come + # from the flattened list: field 1 is shop_details, carrying its two members. + by_ordinal = column(1)(nested_named_struct, None) + + assert by_ordinal == column("shop_details")(nested_named_struct, None) + assert list(by_ordinal.referred_expr[0].output_names) == [ + "shop_details", + "shop_id", + "shop_total", + ] diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 3768754..4e86f5b 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -954,3 +954,53 @@ def counting_iter_plan_rels(plan): ) assert len(infer_plan_schema(correlated).struct.types) == 3 assert len(indexed) == 1 + + +# The relations that emit their input's rows unchanged. Their schema is their input's, +# so they are worth pinning together: the builders reach them only when a further verb +# resolves the schema above one, which no test happened to do for sort. +PASS_THROUGH_RELS = { + "filter": stalg.Rel(filter=stalg.FilterRel(input=read_rel)), + "fetch": stalg.Rel(fetch=stalg.FetchRel(input=read_rel)), + "sort": stalg.Rel( + sort=stalg.SortRel( + input=read_rel, + sorts=[ + stalg.SortField( + expr=stalg.Expression( + selection=stalg.Expression.FieldReference( + root_reference=stalg.Expression.FieldReference.RootReference(), + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField( + field=0 + ) + ), + ) + ), + direction=stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST, + ) + ], + ) + ), + "exchange": stalg.Rel(exchange=stalg.ExchangeRel(input=read_rel)), + "top_n": stalg.Rel(top_n=stalg.TopNRel(input=read_rel)), +} + + +@pytest.mark.parametrize("rel", PASS_THROUGH_RELS.values(), ids=PASS_THROUGH_RELS) +def test_inference_pass_through_rels_keep_the_input_schema(rel): + assert infer_rel_schema(rel) == struct + + +@pytest.mark.parametrize("rel", PASS_THROUGH_RELS.values(), ids=PASS_THROUGH_RELS) +def test_inference_pass_through_rels_apply_emit(rel): + # A pass-through relation still projects through its own emit, so the schema it + # reports is not unconditionally its input's. + emitted = stalg.Rel() + emitted.CopyFrom(rel) + node = getattr(emitted, emitted.WhichOneof("rel_type")) + node.common.emit.output_mapping.extend([2, 0]) + + assert infer_rel_schema(emitted) == stt.Type.Struct( + types=[struct.types[2], struct.types[0]], nullability=struct.nullability + ) From 8be46d5a287be160d1dec4a3833abd8da8ffbd26 Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 10 Aug 2026 15:33:53 +0200 Subject: [PATCH 3/3] fix: resolve post_join_filter over a cached input, and keep memo retention flat Review follow-ups on the schema memo. `join`, `hash_join` and `merge_join` derived their post_join_filter output schema by re-inferring from the input relations, without either input's shared-subtree list in scope -- so a `reference()`-promoted (cached) input, whose root is a plan-global ReferenceRel, could not resolve. That raises on main; the memo happened to answer the ReferenceRel and mask it. Combining the schemas already inferred one line above, as `lateral_join` did, fixes it independently of the memo and drops a redundant walk of both subtrees. `with_execution_behavior` recorded its copied root unconditionally, which broke a Plan carrying no relations -- it copies a caller-supplied Plan rather than assembling one, so it has to stay total over what it accepted. A resolved memo entry keys on a live submessage, and a submessage keeps its whole plan's arena, so entries left to accumulate held every intermediate plan: 26 MB against main's 10 MB over a 32-verb chain on a 2000-column table. Releasing the entries for a plan's own inputs once a lookup has resolved through it puts that back to 10 MB with the inference counts unchanged. `DataFrame.rename`, `drop` and `hint` built their resolvers as plain closures, so no build scope covered them and they stayed quadratic (272 inferences at 16 verbs, the same as main). Wrapping them in `build_scoped`, as every other verb is, brings them to 91. Also: the anchor index is always passed as a factory rather than sometimes a dict, so a caller cannot silently land in the wrong branch; the pairing guard raises ValueError rather than a stripped-under-O assert; and three docstring claims that measurement contradicted are corrected. --- src/substrait/builders/plan.py | 52 +++++---- src/substrait/dataframe/frame.py | 8 +- src/substrait/type_inference.py | 84 +++++++++---- src/substrait/utils/__init__.py | 12 +- tests/builders/plan/test_reference.py | 53 ++++++++- tests/builders/plan/test_schema_memo.py | 110 +++++++++++++++--- .../plan/test_with_execution_behavior.py | 30 +++++ tests/test_type_inference.py | 24 +--- 8 files changed, 284 insertions(+), 89 deletions(-) diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index ba9ba0a..ae65cc7 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -26,7 +26,6 @@ current_collector, ) from substrait.type_inference import ( - _join_output_struct, _join_struct_from_schemas, _outer_anchor_binding, infer_plan_schema, @@ -165,14 +164,21 @@ def _remember_input_schemas(plan: stp.Plan, bound_inputs) -> None: field declaration order, is the i-th bound input. Every builder here hands ``make_rel``'s ``inp`` to its relation in that order (``left=inp[0], right=inp[1]``, ``inputs=inp``); the count is checked below, and the order is - pinned by a test over every builder that embeds more than one input, since a - silent swap would hand a level its sides' schemas the wrong way round. + pinned by tests over the builders that embed more than one input -- one per + two-field builder, plus ``set`` for the repeated-field shape -- since a silent + swap would hand a level its sides' schemas the wrong way round. + + The count mismatch raises rather than asserting: ``assert`` is stripped under + ``python -O``, and with it gone ``zip`` would truncate silently, leaving a child + holding another input's schema -- a wrong schema in the emitted plan rather than a + loud failure. """ children = list(child_rels(plan.relations[-1].root.input)) - assert len(children) == len(bound_inputs), ( - f"assembled relation embeds {len(children)} input relation(s) but the builder " - f"bound {len(bound_inputs)}" - ) + if len(children) != len(bound_inputs): + raise ValueError( + f"assembled relation embeds {len(children)} input relation(s) but the " + f"builder bound {len(bound_inputs)}" + ) for child, bound_input in zip(children, bound_inputs): remember_rel_output_schema(child, bound_input) @@ -243,7 +249,12 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: result.execution_behavior.variable_eval_mode = variable_eval_mode # The copy carries the input's relations verbatim, so its root has the input's # output schema -- recorded against the copy, which is a different object. - remember_rel_output_schema(result.relations[-1].root.input, bound_plan) + # Guarded on there being a query root to record: this builder takes a raw + # caller-supplied Plan and, unlike every other one here, does not assemble the + # relations itself, so it must stay total over the Plans it accepted before -- + # including one carrying no relations at all, or none that is a root. + if result.relations and result.relations[-1].WhichOneof("rel_type") == "root": + remember_rel_output_schema(result.relations[-1].root.input, bound_plan) return result return build_scoped(resolve) @@ -685,16 +696,17 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: # post_join_filter is applied to each output record after # join-type-specific output formation (semantically a FilterRel above the # join), so it resolves against the output schema -- which for semi/anti - # joins is a single side, not the combined schema. + # joins is a single side, not the combined schema. Combined from the schemas + # already inferred above rather than re-inferred from the input relations: + # re-inference would walk both subtrees again, and it would do so without the + # inputs' shared-subtree lists in scope, so a `reference()`-promoted input + # (whose root is a plan-global ReferenceRel) could not resolve at all. bound_post = None if post_join_filter is not None: output_ns = stt.NamedStruct( names=out_names, - struct=_join_output_struct( - type_name, - bound_left.relations[-1].root.input, - bound_right.relations[-1].root.input, - registry=registry, + struct=_join_struct_from_schemas( + type_name, left_ns.struct, right_ns.struct ), ) bound_post = resolve_expression(post_join_filter, output_ns, registry) @@ -1340,16 +1352,16 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: # semi/anti joins is a single side. residual_expression is evaluated # on each candidate key-match (both rows present), so it resolves # against the combined left+right schema. Each is built only when the - # corresponding predicate is supplied. + # corresponding predicate is supplied. The output schema is combined from + # the schemas already inferred above rather than re-inferred from the + # input relations -- see `join` for why re-inference cannot resolve a + # `reference()`-promoted input. bound_post = None if post_join_filter is not None: output_ns = stt.NamedStruct( names=names, - struct=_join_output_struct( - type_name, - bound_left.relations[-1].root.input, - bound_right.relations[-1].root.input, - registry=registry, + struct=_join_struct_from_schemas( + type_name, left_ns.struct, right_ns.struct ), ) bound_post = resolve_expression(post_join_filter, output_ns, registry) diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index 54f4058..f491754 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -39,7 +39,7 @@ from substrait.builders import type as _type from substrait.builders.extended_expression import LateralInput, fresh_rel_anchors from substrait.dataframe.expr import Expr, Measure, col, lit, sort_direction -from substrait.extension_registry import ExtensionRegistry +from substrait.extension_registry import ExtensionRegistry, build_scoped from substrait.type_inference import infer_plan_schema from substrait.utils import to_id_based_outer_references @@ -260,7 +260,7 @@ def resolve(registry: ExtensionRegistry): ] return _plan.select(bound, expressions=expressions)(registry) - return self._next(resolve) + return self._next(build_scoped(resolve)) def drop(self, *columns: str) -> "DataFrame": """Drop the named columns, keeping the rest in their original order.""" @@ -278,7 +278,7 @@ def resolve(registry: ExtensionRegistry): raise ValueError("drop would remove every column") return _plan.select(bound, expressions=expressions)(registry) - return self._next(resolve) + return self._next(build_scoped(resolve)) def unpivot( self, @@ -652,7 +652,7 @@ def resolve(registry: ExtensionRegistry): common.hint.output_names.extend(output_names) return bound - return self._next(resolve) + return self._next(build_scoped(resolve)) def cache(self) -> "DataFrame": """Mark this DataFrame as a reusable common subplan (a CTE). diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index 6ace83d..dffb046 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -7,7 +7,7 @@ import substrait.plan_pb2 as stp import substrait.type_pb2 as stt -from substrait.utils import iter_plan_rels, plan_subtrees, rel_anchor_of +from substrait.utils import child_rels, iter_plan_rels, plan_subtrees, rel_anchor_of class _SubtreeScope: @@ -73,22 +73,24 @@ class _AnchorScope: input, at build or inference time). ``parent`` chains to an enclosing scope so nested correlations still resolve outer anchors. - The index itself may be supplied as a zero-argument factory instead of a dict, - and is then built only if an id-based reference actually asks for an anchor. - Indexing a plan means walking every relation *and* every expression in it to - find the relations embedded in subqueries, which is whole-plan work; plans - carrying an id-based ``OuterReference`` at all are the exception, and a builder - re-infers its input's schema at every level (see :class:`_SchemaMemo`), so - building the index eagerly made that walk the single largest cost of assembling - a long pipeline. + The index arrives as a zero-argument factory and is built only if an id-based + reference actually asks for an anchor. Indexing a plan means walking every relation + *and* every expression in it to find the relations embedded in subqueries, which is + whole-plan work; plans carrying an id-based ``OuterReference`` at all are the + exception, and a builder re-infers its input's schema at every level (see + :class:`_SchemaMemo`), so building the index eagerly made that walk the single + largest cost of assembling a long pipeline. """ __slots__ = ("_rels", "_index", "_subtrees", "_schemas", "_resolving", "_parent") - def __init__(self, rels, subtrees, *, parent=None): - # Either the index or a factory for it; whichever it is not stays None. - self._rels: Optional[dict] = rels if isinstance(rels, dict) else None - self._index = None if isinstance(rels, dict) else rels + def __init__(self, index, subtrees, *, parent=None): + # ``index`` is always a factory, never the index itself: a caller with an + # already-built dict passes ``dict`` or a lambda, which keeps this from having + # to tell the two apart -- a test that got it wrong would report the mistake as + # an uncallable-object TypeError from deep inside a later lookup. + self._rels: Optional[dict] = None + self._index = index self._subtrees = subtrees self._schemas: dict = {} self._resolving: set = set() @@ -99,7 +101,7 @@ def register(self, anchor, struct: stt.Type.Struct) -> None: self._schemas[anchor] = struct def _anchors(self) -> dict: - """The anchor index, built on first use if it was supplied as a factory.""" + """The anchor index, built on first use and kept.""" if self._rels is None: self._rels = self._index() self._index = None @@ -118,7 +120,7 @@ def schema_of(self, anchor, registry) -> stt.Type.Struct: ) self._resolving.add(anchor) try: - rel = self._rels[anchor] + rel = self._anchors()[anchor] # A lateral-join anchor denotes the current left row, not the join's # own output; every other anchor resolves against its output schema. target = ( @@ -163,19 +165,32 @@ class _SchemaMemo: input one level down is therefore unreachable by object identity from the copy, so every level re-walks the whole subtree beneath it and an N-verb chain does O(N^2) inference (#207). The copy *is* reachable at the moment it is made, - though, and protobuf message wrappers are identity-stable, so the builder - records "this relation's output schema is that plan's output schema" and - :func:`infer_rel_schema` stops at the boundary instead of recursing through it. + though, so the builder records "this relation's output schema is that plan's + output schema" and :func:`infer_rel_schema` stops at the boundary instead of + recursing through it. + + Keying that on ``id(rel)`` only works because the entry holds the ``Rel`` itself: + a protobuf message wrapper is *not* permanently cached, so under the upb + implementation a submessage whose last Python reference is dropped is re-wrapped + at a fresh -- possibly recycled -- address on the next access. Holding it pins + both the object and its id for the memo's lifetime, which is what makes a hit + provably the relation that was recorded rather than a later tenant of its id. What is recorded is the ``Plan`` the schema comes from, not the schema, resolved on first lookup and then kept. Recording it eagerly would mean inferring schemas no one asked for, and inference can legitimately fail where building does not -- - ``set``, ``reference`` and ``exchange`` never look at their input's schema today, - so a plan they accept (an extension relation with no registered deriver, say) - has to keep building. The cost of deferring is that an unresolved entry keeps its - input plan alive until the build finishes rather than until the level above it - returns; the entries a chain of verbs produces are consumed one level up, and the - rest is bounded by the plan under construction. + ``set``, ``reference`` and ``exchange`` (and, among the builders that copy rather + than assemble, ``with_execution_behavior``) never look at their input's schema + today, so a plan they accept (an extension relation with no registered deriver, + say) has to keep building. + + An entry costs more memory than a schema: keying on identity means holding a live + submessage, and a submessage keeps its whole plan's arena allocated, so each entry + holds an entire intermediate plan. Left to accumulate, that makes peak memory over + an N-verb build the sum of the intermediates rather than the couple of levels in + flight (measured at 10 MB flat versus 26 MB at 32 verbs over a 2000-column table). + :meth:`_release_inputs_of` is what keeps it flat, and where it can drop an entry + safely is the whole of the reasoning -- see there. Only builders write here; inference never memoizes on its own. A relation's output struct can depend on ambient correlation context (``outer_schemas``, @@ -199,6 +214,24 @@ def remember_plan_output(self, rel: stalg.Rel, plan: stp.Plan) -> None: """Record that ``rel``'s output schema is the output schema of ``plan``.""" self._pending[id(rel)] = (rel, plan) + def _release_inputs_of(self, plan: stp.Plan) -> None: + """Drop the entries recorded for ``plan``'s own input relations. + + Called once a lookup has resolved through ``plan``: those entries existed to + answer that resolution, and nothing above can reach the relations they key on + again, because each enclosing level embedded a *copy*. Dropping them is what + keeps retention flat -- the key of a resolved entry is a live submessage, and + a submessage keeps its whole plan's arena allocated, so holding one holds an + entire intermediate plan. Kept, they would make peak memory over an N-verb + build the sum of the intermediates instead of the two levels in flight. + """ + relations = plan.relations + if not relations or relations[-1].WhichOneof("rel_type") != "root": + return + for child in child_rels(relations[-1].root.input): + self._structs.pop(id(child), None) + self._pending.pop(id(child), None) + def struct_of(self, rel: stalg.Rel, registry) -> Optional[stt.Type.Struct]: """``rel``'s remembered output struct, or None if nothing was recorded.""" key = id(rel) @@ -218,6 +251,7 @@ def struct_of(self, rel: stalg.Rel, registry) -> Optional[stt.Type.Struct]: struct = infer_plan_schema(pending[1], registry=registry).struct finally: self._resolving.discard(key) + self._release_inputs_of(pending[1]) self._structs[key] = (rel, struct) # The plan was held only to answer this; the struct replaces it. del self._pending[key] @@ -266,7 +300,7 @@ def _outer_anchor_binding(anchor, struct): though the join relation is not yet in an anchor index. Nested lateral joins compose via the parent chain. """ - scope = _AnchorScope({}, (), parent=anchor_scope.get()) + scope = _AnchorScope(dict, (), parent=anchor_scope.get()) scope.register(anchor, struct) token = anchor_scope.set(scope) try: diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index e46fe02..2b72450 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -143,8 +143,16 @@ def child_rels(rel: stalg.Rel): Yields the messages themselves rather than the ``(container, key)`` pairs :func:`_iter_child_rels` uses for in-place rewriting, for callers that only read - them -- note they are the live submessages, so identity is meaningful. - Subquery-embedded relations are not direct children and are not yielded. + them -- note they are the live submessages, not copies, so mutating one mutates + ``rel``. Subquery-embedded relations are not direct children and are not yielded. + + A yielded message is only *identity*-stable while a reference to it is held: a + protobuf submessage wrapper is not permanently cached, so under the upb + implementation dropping the last reference lets the next access re-wrap it at a + fresh -- possibly recycled -- address. A caller keying on ``id()`` (as + ``type_inference._SchemaMemo`` does) must therefore keep the message alive + alongside the key; consuming this generator lazily and keeping only the ids will + collide. """ for container, key in _iter_child_rels(rel): yield _child_rel(container, key) diff --git a/tests/builders/plan/test_reference.py b/tests/builders/plan/test_reference.py index dd02f7e..6f9f3f7 100644 --- a/tests/builders/plan/test_reference.py +++ b/tests/builders/plan/test_reference.py @@ -3,15 +3,18 @@ import substrait.plan_pb2 as stp import substrait.type_pb2 as stt -from substrait.builders.extended_expression import literal +from substrait.builders.extended_expression import column, literal from substrait.builders.plan import ( default_version, fetch, + hash_join, + join, + merge_join, read_named_table, reference, set, ) -from substrait.builders.type import i64 +from substrait.builders.type import boolean, i64 from substrait.extension_registry import ExtensionRegistry from substrait.type_inference import infer_plan_schema @@ -102,3 +105,49 @@ def test_reference_out_of_range_ordinal_raises(): ) with pytest.raises(Exception, match="out of range"): infer_plan_schema(ref_plan, registry=registry) + + +# The joins that resolve a post_join_filter against their output schema. That schema is +# combined from the schemas already inferred for each side; deriving it by re-inferring +# from the input relations instead would do so without either input's shared-subtree +# list in scope, and a promoted input's root is a plan-global ReferenceRel that cannot +# resolve on its own. +POST_JOIN_FILTER_BUILDERS = { + "join": lambda left, right, predicate: join( + left, + right, + literal(True, boolean()), + stalg.JoinRel.JOIN_TYPE_INNER, + post_join_filter=predicate, + ), + "hash_join": lambda left, right, predicate: hash_join( + left, + right, + ["id"], + ["id"], + stalg.HashJoinRel.JOIN_TYPE_INNER, + post_join_filter=predicate, + ), + "merge_join": lambda left, right, predicate: merge_join( + left, + right, + ["id"], + ["id"], + stalg.MergeJoinRel.JOIN_TYPE_INNER, + post_join_filter=predicate, + ), +} + + +@pytest.mark.parametrize( + "builder", POST_JOIN_FILTER_BUILDERS.values(), ids=POST_JOIN_FILTER_BUILDERS +) +def test_post_join_filter_resolves_over_a_referenced_input(builder): + cached = reference(read_named_table("t", named_struct)) + plan = builder(cached, read_named_table("u", named_struct), column("id"))(registry) + + # Both sides' columns are in scope for the filter, so it binds to the first "id". + root = plan.relations[-1].root + assert list(root.names) == ["id", "v", "id", "v"] + node = getattr(root.input, root.input.WhichOneof("rel_type")) + assert node.post_join_filter.selection.direct_reference.struct_field.field == 0 diff --git a/tests/builders/plan/test_schema_memo.py b/tests/builders/plan/test_schema_memo.py index 43ef1bd..4b23f44 100644 --- a/tests/builders/plan/test_schema_memo.py +++ b/tests/builders/plan/test_schema_memo.py @@ -29,6 +29,7 @@ project, read_named_table, reference, + set, with_execution_behavior, write_named_table, ) @@ -90,7 +91,10 @@ def _project_chain(length: int): # Deliberately well above the 4N-4 this currently does and well below the ~N^2/2 it -# did before, so the test tracks the complexity class rather than the exact count. +# did before, so the test tracks the complexity class rather than the exact count. Only +# the longer chains can discriminate: unmemoized costs 10/36/136/528 for the four +# lengths, so 4 and 8 come in under any bound that 16 and 32 fail, and are here to +# exercise short chains rather than to catch the regression. _CALLS_PER_VERB = 6 @@ -102,24 +106,41 @@ def test_building_a_chain_infers_each_level_a_bounded_number_of_times(counts, le assert counts["infer_rel_schema"] <= _CALLS_PER_VERB * length -def test_chain_inference_grows_linearly_not_quadratically(counts): +def test_building_a_chain_never_indexes_rel_anchors(counts): + # Indexing walks every relation and expression in the plan. Nothing here carries + # an id-based OuterReference, so nothing should ask for the index. Asserted against + # a positive control, because a Counter reads 0 for a key nothing ever wrote: were + # the patch to stop intercepting, a bare `== 0` would keep passing. _project_chain(8)(registry) - short = counts["infer_rel_schema"] - counts.clear() - _project_chain(32)(registry) - long = counts["infer_rel_schema"] + assert counts["infer_rel_schema"] > 0 + assert counts["iter_plan_rels"] == 0 - # Four times the verbs, so linear allows roughly four times the inferences (with - # headroom); quadratic would be sixteen. - assert long <= 6 * short +def test_memo_retention_does_not_grow_with_chain_length(): + # An entry keys on a live submessage, which keeps its whole plan's arena alive, so + # entries that accumulate hold every intermediate plan of the build rather than the + # levels in flight. `_release_inputs_of` drops each entry once it has been resolved + # through; this pins that, by watching how many are ever live at once. + peaks = {} + for length in (4, 8, 16, 32): + peak = 0 + original = type_inference._SchemaMemo.remember_plan_output -def test_building_a_chain_never_indexes_rel_anchors(counts): - # Indexing walks every relation and expression in the plan. Nothing here carries - # an id-based OuterReference, so nothing should ask for the index. - _project_chain(8)(registry) + def counting_remember(self, rel, plan): + nonlocal peak + original(self, rel, plan) + peak = max(peak, len(self._structs) + len(self._pending)) - assert counts["iter_plan_rels"] == 0 + type_inference._SchemaMemo.remember_plan_output = counting_remember + try: + _project_chain(length)(registry) + finally: + type_inference._SchemaMemo.remember_plan_output = original + peaks[length] = peak + + # Bounded by the levels in flight, not by the chain: 8x the verbs must not mean + # meaningfully more live entries. + assert peaks[32] <= peaks[4] + 2, peaks def test_memo_does_not_outlive_the_build(): @@ -150,8 +171,10 @@ def _true(): return literal(True, boolean()) -# Every builder that embeds more than one input relation, since those are the ones -# whose recorded schemas could be paired with the wrong side. +# The builders that embed more than one input relation into a *pair* of fields, since +# those are the ones whose recorded schemas could be paired with the wrong side. The +# repeated-field ones (`set`, `extension_multi`) take their inputs in one list and are +# covered separately below, where order shows up in the output rather than in the types. TWO_INPUT_BUILDERS = { "join": lambda: join(_left(), _right(), _true(), stalg.JoinRel.JOIN_TYPE_INNER), "cross": lambda: cross(_left(), _right()), @@ -187,6 +210,61 @@ def test_two_input_builders_record_each_side_against_its_own_relation(builder): ] +def test_lateral_join_records_a_correlated_right_input_against_its_own_relation(): + # `lateral_join` assembles its relation *outside* the anchor binding its right input + # was built under, so the recorded schema is resolved later, with the binding + # re-established by inference rather than by the builder. A right input that + # actually correlates is what exercises that: it can only be inferred while the + # left row is bound to the join's rel_anchor. + lateral = lateral_join( + _left(), + lambda handle: project(_right(), expressions=[handle.column("k")]), + stalg.JoinRel.JOIN_TYPE_INNER, + ) + written = write_named_table("out", lateral)(registry) + + table_schema = written.relations[-1].root.input.write.table_schema + assert list(table_schema.names) == ["k", "v", "rk", "rv", "k"] + assert list(table_schema.struct.types) == [ + i64(nullable=False), + i64(nullable=False), + i64(nullable=False), + string(), + i64(nullable=False), + ] + + +def test_set_records_each_input_against_its_own_relation(): + # SetRel takes its inputs as one repeated field, so a mispairing shows up in the + # ops whose output is not symmetric in the inputs: MINUS takes each field's + # nullability from the *primary* (first) input alone, so pairing the recorded + # schemas the wrong way round reports the secondary's nullability instead. + primary = read_named_table( + "primary", + stt.NamedStruct( + names=["k"], + struct=stt.Type.Struct( + types=[i64(nullable=False)], nullability=stt.Type.NULLABILITY_REQUIRED + ), + ), + ) + secondary = read_named_table( + "secondary", + stt.NamedStruct( + names=["k"], + struct=stt.Type.Struct( + types=[i64(nullable=True)], nullability=stt.Type.NULLABILITY_REQUIRED + ), + ), + ) + written = write_named_table( + "out", set([primary, secondary], stalg.SetRel.SET_OP_MINUS_PRIMARY) + )(registry) + + table_schema = written.relations[-1].root.input.write.table_schema + assert list(table_schema.struct.types) == [i64(nullable=False)] + + def test_reference_records_the_promoted_subtree_and_the_reference(): # A ReferenceRel's schema is its subtree's, and `reference` records both so a # downstream verb resolves it by lookup rather than by walking the subtree. diff --git a/tests/builders/plan/test_with_execution_behavior.py b/tests/builders/plan/test_with_execution_behavior.py index 97b254a..3e7c307 100644 --- a/tests/builders/plan/test_with_execution_behavior.py +++ b/tests/builders/plan/test_with_execution_behavior.py @@ -90,3 +90,33 @@ def test_plans_have_no_execution_behavior_by_default(): plan = read_named_table("example_table", named_struct)(None) assert not plan.HasField("execution_behavior") + + +def test_execution_behavior_on_a_plan_carrying_no_relations(): + # This is the one builder that copies a caller-supplied Plan wholesale instead of + # assembling the relations itself, so it has to stay total over the Plans it + # accepts -- including degenerate ones it only sets a field on. Recording the + # copied root's schema for the build (see builders.plan._remember_input_schemas) + # must not turn that into an IndexError. + actual = with_execution_behavior(stp.Plan(version=default_version), PER_RECORD)( + None + ) + + assert actual.execution_behavior.variable_eval_mode == PER_RECORD + assert not actual.relations + + +def test_execution_behavior_on_a_plan_whose_last_relation_is_a_subtree(): + # Likewise for a Plan whose trailing entry is a shared subtree rather than a query + # root: there is no root schema to record, and reading one would key the build's + # memo on an unset-oneof stub that a later write could reify in place. + base = read_named_table("example_table", named_struct)(None) + subtree_only = stp.Plan( + version=default_version, + relations=[stp.PlanRel(rel=base.relations[-1].root.input)], + ) + + actual = with_execution_behavior(subtree_only, PER_RECORD)(None) + + assert actual.execution_behavior.variable_eval_mode == PER_RECORD + assert [r.WhichOneof("rel_type") for r in actual.relations] == ["rel"] diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 4e86f5b..9de8146 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -962,26 +962,10 @@ def counting_iter_plan_rels(plan): PASS_THROUGH_RELS = { "filter": stalg.Rel(filter=stalg.FilterRel(input=read_rel)), "fetch": stalg.Rel(fetch=stalg.FetchRel(input=read_rel)), - "sort": stalg.Rel( - sort=stalg.SortRel( - input=read_rel, - sorts=[ - stalg.SortField( - expr=stalg.Expression( - selection=stalg.Expression.FieldReference( - root_reference=stalg.Expression.FieldReference.RootReference(), - direct_reference=stalg.Expression.ReferenceSegment( - struct_field=stalg.Expression.ReferenceSegment.StructField( - field=0 - ) - ), - ) - ), - direction=stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST, - ) - ], - ) - ), + # Bare of sort fields, like the fetch/top_n entries are of counts: a pass-through + # relation's schema comes from its input and its own emit, and inference reads + # neither the sorts nor the counts. + "sort": stalg.Rel(sort=stalg.SortRel(input=read_rel)), "exchange": stalg.Rel(exchange=stalg.ExchangeRel(input=read_rel)), "top_n": stalg.Rel(top_n=stalg.TopNRel(input=read_rel)), }