diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 519217f..ae65cc7 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -26,13 +26,14 @@ current_collector, ) from substrait.type_inference import ( - _join_output_struct, _join_struct_from_schemas, _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 +149,40 @@ 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 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)) + 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) + + def _plan_from( bound_inputs, make_rel, names, metadata_sources, *, include_version=True ): @@ -172,7 +207,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 +247,14 @@ 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. + # 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) @@ -559,7 +604,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 +613,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) @@ -644,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) @@ -1299,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/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..dffb046 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -1,12 +1,13 @@ import contextlib import contextvars +from typing import Optional import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee 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: @@ -71,12 +72,25 @@ 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 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", "_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, 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() @@ -86,10 +100,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 and kept.""" + 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}") @@ -99,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 = ( @@ -135,6 +156,140 @@ 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, 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`` (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``, + ``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 _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) + 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._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] + 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 @@ -145,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: @@ -692,6 +847,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 +1118,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..2b72450 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -138,6 +138,26 @@ 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, 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) + + 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/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/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 new file mode 100644 index 0000000..4b23f44 --- /dev/null +++ b/tests/builders/plan/test_schema_memo.py @@ -0,0 +1,288 @@ +"""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, + set, + 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. 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 + + +@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_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) + assert counts["infer_rel_schema"] > 0 + assert counts["iter_plan_rels"] == 0 + + +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 counting_remember(self, rel, plan): + nonlocal peak + original(self, rel, plan) + peak = max(peak, len(self._structs) + len(self._pending)) + + 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(): + # 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()) + + +# 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()), + "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_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. + 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/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 b5b9dd0..9de8146 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -903,3 +903,88 @@ 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 + + +# 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)), + # 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)), +} + + +@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 + ) 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())) == []