fix(planner): remove probe state-pollution in heap top-K rewrite (#347, #335) - #349
Open
temporaryfix wants to merge 4 commits into
Open
fix(planner): remove probe state-pollution in heap top-K rewrite (#347, #335)#349temporaryfix wants to merge 4 commits into
temporaryfix wants to merge 4 commits into
Conversation
The two formulas that map LogicalExpression to a column name were
duplicated across the LPG and RDF planners:
- alias.unwrap_or_else(|| expression_to_string(...)) [4 sites]
- format!("{v}_{p}") / format!("__expr_{expr:?}") [11 sites]
Producers (Aggregate/Sort augmenting projections, Return/Project items)
and the consumer (`resolve_expression_to_column`) all re-derived the
same strings by hand. Any drift between them silently breaks resolution.
Extract two helpers in `query/planner/common.rs`:
- `output_column_name(alias, expr)` — for Return/Project item names.
- `resolved_column_name(expr)` — symmetric with the lookup performed by
`resolve_expression_to_column`, which now calls it internally.
Replace every duplicated call site with these helpers. Pure refactor:
no behavior change.
…rafeoDB#347, GrafeoDB#335) `try_heap_topk_rewrite` planned the input subtree as a speculative probe, then fell through to the unfused path when sort-key resolution failed. The probe mutated `Planner::scalar_columns` (and `edge_columns`) via the side effects in `plan_return_projection`: each output column gets registered as scalar so an enclosing Apply doesn't try to re-resolve it as a NodeId. When resolution failed, the caller re-planned the same Return through `plan_sort`'s unfused path, but now with the polluted state — and `plan_return_projection`'s `Variable(name)` arm flips from `NodeResolve` to a raw `Column` passthrough when it sees the name already in `scalar_columns`. Result: `MATCH (n) RETURN n ORDER BY <key> LIMIT k` returned raw `Int64` NodeIds instead of resolved maps for any sort key whose resolver column name was not in the Return's output. Issue GrafeoDB#335 covered Property keys; issue GrafeoDB#347 covered FunctionCall (text_score) keys. The class extends to Case, Binary, and IndexAccess sort keys — none of which had reports, but all produce the same `Value::Int64(0)` output (confirmed via the new regression tests, which fail without this commit on a Variable Return item across all four key shapes). PR GrafeoDB#337 patched the predicate `sort_needs_augmenting_projection` to return `true` for Property keys, gating the probe off for that specific shape. That fix was a partial mitigation: it widened the gate without removing the state-pollution mechanism, so any sort key the predicate considered "safe" continued to leak. The predicate-as-gate approach also has the inverse footgun of save/restore — every new expression form requires another carve-out, with no compiler enforcement. This change removes the probe entirely. `try_heap_topk_rewrite` now: 1. Predicts Return's output columns from `ret.items` via the shared `output_column_name` helper (the same one `plan_return_projection` uses, so prediction can't drift from reality). 2. Resolves sort keys against the predicted columns — pure function, no planner state touched. The resolved `SortKey`s are held and reused after planning, so resolution runs once. 3. Plans for real only when resolution succeeds. A `debug_assert_eq!` between predicted and actual columns guards against future drift if a new shape is added to `predict_subtree_columns`. `RETURN *` is explicitly skipped because its column count depends on the input subtree, which can only be known by planning — mirrors the same `Variable("*")` check in `plan_return_projection`. `sort_needs_augmenting_projection` is no longer consulted by the rewrite. It is retained for `plan_sort`'s pre-Return-vs-post-Return augmenting decision, where its variable-membership rule answers the correct question. This also obviates PR GrafeoDB#337: the probe is gone, so the predicate-as-gate that GrafeoDB#337 hardened is no longer reachable from the bug class.
Adds four tests, one per sort-key shape that triggered the bug class fixed in the previous commit. Each asserts that `RETURN n ... LIMIT k` yields a `Value::Map` for the entity column, never a raw `Int64` NodeId — the symptom of the pollution failure mode. - order_by_limit_property_key_yields_map (issue GrafeoDB#335) - order_by_limit_text_score_key_yields_map (issue GrafeoDB#347, behind text-index) - order_by_limit_case_key_yields_map (Case expression) - order_by_limit_binary_key_yields_map (Binary expression) Verified to fail on the parent commit reverted without the fix: row 0: expected Map, got Int64(0)
`RETURN v.p ORDER BY v.p LIMIT k` was getting routed through the
unfused path even though the heap top-K rewrite could safely handle
it. The cause: the resolver builds the column-name to look up via
`resolved_column_name` ("v_p"), but the Return outputs that column
under `output_column_name`'s `expression_to_string` fallback
("v.p"). The two formulas don't match, so the rewrite's predictive
resolve fails and we fall through.
`plan_sort` already handles this in its prelude: for each
`Property{v,p}` Return item, it registers an alias entry mapping
"v_p" -> the index of column "v.p". That prelude logic is now
extracted as `register_return_property_sort_aliases` and called from
both `plan_sort` (unchanged behavior) and `try_heap_topk_rewrite`
(new — closes the coverage gap).
`resolve_logical_to_physical_keys` now takes a pre-built
`&HashMap<String, usize>` instead of building one from a `&[String]`,
since callers need to inject the aliases before resolving.
Pure coverage win: every shape the rewrite already handled keeps
producing identical output; queries of the form `RETURN v.p ORDER BY
v.p LIMIT k` now fire TopK instead of full sort + limit (existing
`cypher_order_by_limit_uses_topk` exercises this shape — it still
passes, and now via the fused path). The §2.9 e2e benchmark in
`benches/topk_e2e.rs` is the silent-fall-through canary if this
regresses.
3 tasks
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Contributor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #347 and the broader bug class also surfaced in #335.
try_heap_topk_rewriteplanned the input subtree as a speculative probe, then fell through to the unfused path when sort-key resolution failed. The probe mutatedPlanner::scalar_columnsvia the side effects inplan_return_projection, and the unfused re-planning of the sameReturnthen choseProjectExpr::Column(raw passthrough) instead ofProjectExpr::NodeResolve—MATCH (n) RETURN n ORDER BY <key> LIMIT kreturned a rawInt64NodeId instead of a resolved map for any sort key whose resolver column name wasn't materialised by Return.Relation to the prior #335 fix
The earlier fix for #335 (PR #337, closed; the predicate change was cherry-picked into
release/0.5.43as 530edc5) patched the gate predicatesort_needs_augmenting_projectionto returntrueforPropertysort keys. That kept the probe from running for that one shape, but the underlying mechanism — the probe mutating planner state on fall-through — was untouched, and the same root cause re-fires for any other sort-key shape the predicate doesn't gate. #347 surfaced it withtext_score(aFunctionCall); the audit in this PR found the same break forCaseandBinarysort keys, neither of which has a GH issue but both reproduce onmaintoday.maindoes not currently contain therelease/0.5.43predicate cherry-pick, so this PR is what fixes #335 and #347 onmain. The cherry-picked release-branch fix and this PR are not in conflict — this one removes the architectural condition that made the predicate-as-gate necessary in the first place, so future sort-key shapes don't need new carve-outs.Changes (commit-by-commit)
refactor(planner): centralize column-naming formulas— extracts two helpers inquery/planner/common.rs:output_column_name(alias, expr)for Return/Project item names.resolved_column_name(expr)symmetric withresolve_expression_to_column's lookup.Replaces 15 duplicated
alias.unwrap_or_else(|| expression_to_string(...))andformat!("{v}_{p}")/format!("__expr_{expr:?}")callsites acrosslpg/{aggregate,project}.rs,rdf/mod.rs, andcommon.rs. Pure refactor — no behavior change.fix(planner): remove probe state-pollution in try_heap_topk_rewrite (#347, #335)— replaces the probe with predictive column resolution.predict_subtree_columnsderives Return's output columns fromret.itemsvia the sharedoutput_column_namehelper (no planner mutation), and sort keys resolve against the prediction before any planning happens. Adebug_assert_eq!between predicted and actual columns guards against future drift if a new shape is added.RETURN *is explicitly skipped (input columns aren't knowable without planning) — mirrors the same check inplan_return_projection.test(topk_rewrite): regression coverage for probe state-pollution— four tests, one per sort-key shape that triggered the bug class:order_by_limit_property_key_yields_map(issue [Bug]: ORDER BY + LIMIT changes returned result #335)order_by_limit_text_score_key_yields_map(issue [Bug]: Order + Limit + text_score changes returned format #347, behindtext-index)order_by_limit_case_key_yields_map(Case expression — no GH issue, found via audit)order_by_limit_binary_key_yields_map(Binary expression — no GH issue, found via audit)Each asserts the entity Return item resolves to
Value::Map. Verified to fail on the parent commit (without the fix) withexpected Map, got Int64(0).fix(planner): restore TopK coverage for property-alias-mismatch case—RETURN v.p ORDER BY v.p LIMIT kwas getting routed through the unfused path because Return outputs the column under"v.p"(viaexpression_to_stringfallback) but the resolver looks up"v_p"(viaresolved_column_name).plan_sortalready handles this in its prelude; that prelude logic is now extracted asregister_return_property_sort_aliasesand called from bothplan_sort(unchanged behavior) andtry_heap_topk_rewrite(new — closes the coverage gap).Test plan
cargo test -p grafeo-engine --no-default-features --features "lpg gql ai parallel" --test topk_rewrite— 11 tests pass (7 existing + 4 new regression).planner_lpg_coverage(26),planner_coverage(34),project_coverage(18),hybrid_query(34),spec_compliance(106),mutation_planning(39),factorized_aggregation_test,seam_aggregates_expressions(25),null_and_coercion(27),push_pipeline_integration(42),expression_and_projection(65) — all green.Value::Int64(0)symptom — proving the bug class extends beyond what was reported in [Bug]: ORDER BY + LIMIT changes returned result #335 and [Bug]: Order + Limit + text_score changes returned format #347.cypher_order_by_limit_uses_topk(RETURN n.r ORDER BY n.r DESC LIMIT 5) still passes — this is the shape that now fires TopK after the alias-mismatch fix.upstream/main(merge_rollback,query_correctness,regression_externalcompile errors;gql_spec_complianceLIKE tests) confirmed independent of this branch.Notes for the reviewer
query/planner/. Tests assert correctness via the existingsession.execute()path.RETURN. Submitted separately because it's a user-visible behavior change in column naming and deserves its own review and CHANGELOG note.