perf: infer each relation's schema once per build, not once per level - #250
perf: infer each relation's schema once per build, not once per level#250nielspardon wants to merge 4 commits into
Conversation
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 substrait-io#207
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.
…ntion 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.
bvolpato
left a comment
There was a problem hiding this comment.
Other than the inline comment, this looks good to me.
|
|
||
| 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) |
There was a problem hiding this comment.
Could this be a concern for chains whose builders never infer their inputs? exchange, reference, set, and with_execution_behavior can keep adding _pending entries without reaching struct_of, so _release_inputs_of never runs. In a 32-level exchange chain over 10,000 columns, I saw peak RSS of 93,664 KiB versus 39,544 KiB with recording disabled. Would it make sense to collapse or release superseded unresolved entries here?
There was a problem hiding this comment.
Good catch — reproduced it (121 MB against 55 MB peak RSS for 32 stacked exchangees over 10k columns) and fixed. Recording now also drops the unresolved entries it supersedes, and the sweep covers the root input as well, which is where with_execution_behavior and reference record. Retention is flat at any depth now, and the retention test covers a chain that resolves nothing.
An entry is dropped once nothing can reach the relation it keys on, but until now that was only noticed at resolution time -- and a chain of builders that never ask their input for a schema (`set`, `exchange`, `reference`, `with_execution_behavior`) resolves nothing, so it held one live intermediate plan per level: 121 MB against 55 MB of peak RSS for 32 stacked exchanges over a 10,000-column table. Recording now also drops the unresolved entries it supersedes, which are exactly the ones only reachable by walking into the plan the new entry names. Resolved entries stay, since they are what makes that walk unnecessary everywhere else and they key on submessages of a plan the new entry pins anyway. Walking such a run directly turns out to be cheaper than hopping the boundaries recorded through it, so this also lowers inference counts where the two meet: 93 to 78 calls for 16 exchanges under 16 projections. Emitted plans are unchanged.
Completes #207. PR #245 removed the extension-merging half; this is the schema half.
Problem
Every verb resolves its input's schema, and plans are built as nested
resolve()closures, so an N-verb chain re-walked the whole subtree beneath it at every level.
Profiling a 40-verb
projectchain (45 ms total) shows the cost is not only wherethe issue points:
rel_anchorindex ininfer_plan_schema(iter_plan_rels)infer_rel_schemarecursion — the half #207 names_plan_from)Indexing anchors walks every relation and every expression of the plan, to reach
the relations embedded in subqueries. Doing that per level made it the larger term,
so fixing only the recursion would have left two thirds of the cost in place.
Approach
Assigning an input's root
Relinto the output relation copies it, so the schema justinferred for that input is unreachable from the copy by identity — a plain identity
cache never hits across levels. The copy is reachable as it is made, though, and
protobuf wrappers are identity-stable, so
_plan_fromrecords that each copy carriesits input plan's output schema;
infer_rel_schemaconsults those records beforedispatching and stops there instead of recursing through.
referenceandwith_execution_behaviorassemble aPlandirectly and record their own root; theread builders and
updateare already leaves.What is recorded is the plan, not the schema, resolved on first lookup —
set,referenceandexchangenever look at their input's schema, and inference can failwhere building does not, so resolving eagerly would reject plans that build today.
Only builders write to the memo: a relation's output struct can depend on ambient
correlation context (
outer_schemas,anchor_scope), and anything crossing intoanother context is copied on the way, so a record is only ever read back under the
context it was made in. Caching inference results wholesale would not have that
property.
build_scopescopes the memo alongside the build'sExtensionCollector, soinference used directly as a library function is untouched.
An entry keys on a live submessage, and a submessage keeps its whole plan's arena, so
entries left to accumulate would hold every intermediate plan of the build — 26 MB
against 10 MB over a 32-verb chain on a 2000-column table. Each is released as soon as
nothing can reach the relation it keys on, which happens at two points: when a lookup
has resolved through it, and when a new entry is recorded over an unresolved one. The
second is the only release a run of verbs that never asks its input for a schema
(
set,exchange,reference,with_execution_behavior) ever gets; without it, 32stacked
exchangees over a 10,000-column table peaked at 121 MB against 55 MB. Walkingsuch a run directly is also cheaper than hopping the boundaries recorded through it, so
inference counts fall where the two meet — 93 → 78 calls for 16 exchanges under 16
projections.
DataFrame.rename,dropandhintbuilt their resolvers as plain closures ratherthan via
build_scoped, so no build scope covered them and they were the one publicpath the memo could not reach. They are wrapped like every other verb.
Two supporting details:
infer_plan_schemapasses its anchor index as a factory that_AnchorScopecalls on the first lookup needing one, and the pairing between a recordand its relation is positional — the i-th child
Relin declaration order is the i-thbound input, which is how every builder places them. The count is checked, and a
parametrized test over every multi-input builder pins the order, since a silent swap
would hand the level above a join its two sides' schemas the wrong way round.
Result
infer_rel_schemacalls per build, and relations visited while indexing anchors:mainN(N+1)/2 becomes 4N-4, and the anchor index is never built for a plan with no
id-based outer reference. Every shape follows: at 16 verbs,
join+selectgoes 800 →187,
lateral_join+select1088 → 251,group_by/agg528 → 124, andrename—which no build scope used to cover — 272 → 91. Peak memory stays flat in chain length,
matching
main(10.2 MB against 9.9 MB at 32 verbs over a 2000-column table).Build time, best of 7 through the DataFrame API:
main@ 8 / 16 verbswith_columnsfilterovercache()join+selectlateral_join+selectgroup_by/aggWhat remains per level is the protobuf copy in
_plan_from, which is proportional tothe plan's byte size and inherent to assembling nested immutable messages — so
building is still quadratic in plan bytes, with a constant roughly twenty times
smaller than the term removed here.
One behavior change
join,hash_joinandmerge_joinderived theirpost_join_filteroutput schema byre-inferring from the input relations, passing neither input's shared-subtree list — so a
reference()-promoted (i.e..cache()d) input, whose root is a plan-globalReferenceRel, could not resolve. Onmainall three raiseReferenceRel subtree_ordinal 0 is out of range; the memo silently answered theReferenceRelandmasked it, which would have left the memo load-bearing for correctness on an untested
path. They now combine the schemas already inferred a line above, as
lateral_joinalways did, so the case builds independently of the memo — and one redundant walk of both
input subtrees per join goes away with it.
tests/builders/plan/test_reference.pycoversall three; it passes with the memo disabled, which is what shows the hidden dependence is
gone.
Verification
Emitted plans are unchanged, which is the property that matters most here, so it was
checked rather than assumed: 42 plans covering every builder — including the
multi-input,
cache(),lateral_joinand subquery shapes — are byte-identical to amainworktree, and every example prints output identical to it — which CI does notcheck, since it runs four of them and only asserts they exit cleanly
(
duckdb_examplestill executes its plan to the same result set).Both new cost tests fail with the memo lookup removed (528 inferences against a bound
of 216), and all six pairing tests fail with the pairing reversed.
That comparison also turned up two pre-existing coverage gaps, closed here: nothing
inferred a
SortRel's schema, since sort is always terminal in the suite and had nodirect 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.
Closes #207
🤖 Generated with AI