planner: preserve outer refs in window subqueries#66869
planner: preserve outer refs in window subqueries#66869hawkingrei wants to merge 12 commits intopingcap:masterfrom
Conversation
ti-chi-bot
bot
added
the
release-note
|
❌ Failed to Create Project We encountered an error while creating your project. Please try again or contact support if the issue persists. |
ti-chi-bot
bot
added
sig/planner
|
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds test data and tests plus planner changes to allow subqueries inside window functions to reference outer-query columns by extracting auxiliary fields, preserving outer-correlated columns during window rewrite, and propagating context into window resolution. Also updates planner docs and increases test shard_count. Changes
Sequence Diagram(s)sequenceDiagram
rect rgba(200,220,255,0.5)
participant Client as Client
participant Parser as Parser
participant PlanBuilder as PlanBuilder
participant Extractor as subqueryExprExtractor
participant LogicalPlan as LogicalPlan
end
Client->>Parser: submit SQL with window + subquery
Parser->>PlanBuilder: AST (SelectStmt)
PlanBuilder->>Extractor: traverse window args, WINDOW specs, ORDER BY
Extractor-->>PlanBuilder: collected subqueries & outer refs
PlanBuilder->>LogicalPlan: append auxiliary select fields for outer-correlated cols
PlanBuilder->>LogicalPlan: build projection/schema aligned with rewritten window args
LogicalPlan-->>Client: planned query (explain/result used by tests)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Follow-up for the review items from #66823:
Validation:
Not verified locally:
|
ti-chi-bot
bot
added
size/XXL
hawkingrei
force-pushed
the
fix-window-subquery-56532
branch
from
1d28639 to
2e6b187
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@pkg/planner/core/logical_plan_builder.go`:
- Around line 2529-2555: The helper findColumnNameByUniqueID currently only
inspects the current plan node and a top-level *logicalop.LogicalJoin; update it
to walk down wrapper/plan layers (e.g., LogicalSelection, LogicalProjection,
LogicalAggregation, LogicalSort, LogicalLimit, LogicalWindow, etc.) to find an
inner *logicalop.LogicalJoin that may hold the FullSchema/FullNames for the
requested uniqueID; once you locate such a join, perform the same
FullSchema/FullNames lookup as the existing join branch and return the
ast.ColumnName if found, otherwise fall back to the existing return nil
behavior. Ensure you keep references to the function name
findColumnNameByUniqueID and the type *logicalop.LogicalJoin so reviewers can
locate the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 49bafa21-0eb5-4b73-ad5f-2cec8237b377
📒 Files selected for processing (5)
pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_out.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_xut.jsonpkg/planner/core/casetest/windows/widow_with_exist_subquery_test.gopkg/planner/core/logical_plan_builder.go
hawkingrei
force-pushed
the
fix-window-subquery-56532
branch
from
2e6b187 to
359d687
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #66869 +/- ##
================================================
- Coverage 77.6754% 77.0993% -0.5762%
================================================
Files 2009 1930 -79
Lines 550314 538215 -12099
================================================
- Hits 427459 414960 -12499
- Misses 121147 123248 +2101
+ Partials 1708 7 -1701
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
♻️ Duplicate comments (1)
pkg/planner/core/logical_plan_builder.go (1)
2559-2585:⚠️ Potential issue | 🟠 MajorWalk through pass-through wrappers before falling back from
FullSchema.
findColumnNameByUniqueIDstill only inspects the current node and a top-level*logicalop.LogicalJoin. A shape likeLogicalSelection -> LogicalJoin(from an innerJOIN ... ON ...) can leave redundant USING/NATURAL columns only in the child join’sFullSchema, so this returnsniland the new window-subquery flow still skips the auxiliary field. Mirroring the wrapper walk infindColFromNaturalUsingJoinshould close that hole.💡 Suggested change
func findColumnNameByUniqueID(p base.LogicalPlan, uniqueID int64) *ast.ColumnName { for idx, pCol := range p.Schema().Columns { if uniqueID != pCol.UniqueID { continue } pName := p.OutputNames()[idx] return &ast.ColumnName{ Schema: pName.DBName, Table: pName.TblName, Name: pName.ColName, } } - // USING/NATURAL JOIN can keep table-qualified outer references only in FullSchema/FullNames. - if join, ok := p.(*logicalop.LogicalJoin); ok && join.FullSchema != nil && len(join.FullNames) != 0 { - for idx, pCol := range join.FullSchema.Columns { + switch x := p.(type) { + case *logicalop.LogicalLimit, *logicalop.LogicalSelection, *logicalop.LogicalTopN, *logicalop.LogicalSort, *logicalop.LogicalMaxOneRow: + return findColumnNameByUniqueID(x.Children()[0], uniqueID) + case *logicalop.LogicalJoin: + if x.FullSchema == nil || len(x.FullNames) == 0 { + return nil + } + // USING/NATURAL JOIN can keep table-qualified outer references only in FullSchema/FullNames. + for idx, pCol := range x.FullSchema.Columns { if uniqueID != pCol.UniqueID { continue } - pName := join.FullNames[idx] + pName := x.FullNames[idx] return &ast.ColumnName{ Schema: pName.DBName, Table: pName.TblName, Name: pName.ColName, } } } return nil }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/planner/core/logical_plan_builder.go` around lines 2559 - 2585, findColumnNameByUniqueID currently only checks the current plan node and a top-level *logicalop.LogicalJoin, so it misses cases where a join with FullSchema/FullNames is wrapped by pass-through operators (e.g., LogicalSelection) as handled by findColFromNaturalUsingJoin; update findColumnNameByUniqueID to walk through pass-through wrappers (mirror the wrapper traversal logic in findColFromNaturalUsingJoin) until you locate the underlying *logicalop.LogicalJoin with FullSchema/FullNames, then inspect join.FullSchema.Columns and join.FullNames for the matching uniqueID and return the ast.ColumnName as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@pkg/planner/core/logical_plan_builder.go`:
- Around line 2559-2585: findColumnNameByUniqueID currently only checks the
current plan node and a top-level *logicalop.LogicalJoin, so it misses cases
where a join with FullSchema/FullNames is wrapped by pass-through operators
(e.g., LogicalSelection) as handled by findColFromNaturalUsingJoin; update
findColumnNameByUniqueID to walk through pass-through wrappers (mirror the
wrapper traversal logic in findColFromNaturalUsingJoin) until you locate the
underlying *logicalop.LogicalJoin with FullSchema/FullNames, then inspect
join.FullSchema.Columns and join.FullNames for the matching uniqueID and return
the ast.ColumnName as before.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f8224853-c786-4e46-bba7-b4104537fea2
📒 Files selected for processing (8)
docs/note/planner/README.mddocs/note/planner/window_subquery_notes.mdpkg/planner/core/casetest/windows/BUILD.bazelpkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_out.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_xut.jsonpkg/planner/core/casetest/windows/widow_with_exist_subquery_test.gopkg/planner/core/logical_plan_builder.go
🚧 Files skipped from review as they are similar to previous changes (2)
- pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.json
- pkg/planner/core/casetest/windows/testdata/window_push_down_suite_out.json
|
/retest |
|
@pantheon-bot review |
|
❌ Failed to Create Project We encountered an error while creating your project. Please try again or contact support if the issue persists. |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.json (1)
45-54: Add one regression for correlated subqueries in windowORDER BY.These cases cover window arguments and
PARTITION BY, but the PR also changes outer-ref preservation for windowORDER BYitems. A focusedOVER (ORDER BY <correlated-subquery>)case would keep that path protected too.Example addition
"cases": [ "select count((select /* issue:56532 */ ref_1.c2 from t1 limit 1)) over (partition by ref_1.c1) from t2 as ref_1 order by ref_1.c1, ref_1.c2", "select coalesce(count(*) over (partition by ref_1.c1), (select /* issue:56532 */ ref_1.c2 from t1)) from t2 as ref_1 order by ref_1.c1, ref_1.c2", "select count(exists(select /* issue:56532 */ 1 from t1 where ref_1.c1 is not null)) over (partition by ref_1.c1) from t2 as ref_1 order by ref_1.c1, ref_1.c2", "select count((select /* issue:56532 */ 1 from t1 where ref_1.c1 is not null limit 1)) over (partition by ref_1.c1) from t2 as ref_1 join t3 using (c1) order by ref_1.c1, ref_1.c2", "select count(*) over w from t2 as ref_1 join t3 using (c1) window w as (partition by exists(select /* issue:56532 */ 1 from t1 where ref_1.c1 is not null)) order by ref_1.c1, ref_1.c2", + "select count(*) over (order by (select /* issue:56532 */ ref_1.c2 from t1 limit 1)) from t2 as ref_1 order by ref_1.c1, ref_1.c2", "select count((select /* issue:56532 */ 1 from t1 where t3.c1 is not null limit 1)) over (partition by ref_1.c1) from t2 as ref_1 join t3 using (c1) where ref_1.c2 is not null order by ref_1.c1, ref_1.c2" ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.json` around lines 45 - 54, Add a focused regression case to the "cases" array of TestWindowSubqueryOuterRef that exercises an OVER (ORDER BY <correlated-subquery>) expression so the change to preserve outer refs for window ORDER BY items is covered; update the "cases" array (in the same JSON object under "name": "TestWindowSubqueryOuterRef") by appending a SQL test that uses a correlated subquery inside an OVER (ORDER BY ...) clause (refer to the existing cases for naming and ordering style).pkg/planner/core/logical_plan_builder.go (1)
2741-2769: Consider batching these auxiliary-field scans.
appendAuxiliaryFieldsForSubqueriesalready accepts variadic nodes, so the three loops here repeat AST walks, rewrites, and dedup scans over the samesel.Fields.Fields. Collecting the window field exprs, specs, and windowedORDER BYexprs first and making one call would keep the behavior while trimming the extra pre-build work.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/planner/core/logical_plan_builder.go` around lines 2741 - 2769, The three loops each call appendAuxiliaryFieldsForSubqueries separately causing repeated AST walks; instead gather all window-related nodes into a single slice and call appendAuxiliaryFieldsForSubqueries once. Concretely, iterate over sel.Fields.Fields, sel.WindowSpecs, and sel.OrderBy.Items, collect expressions/nodes that satisfy ast.HasWindowFlag (and the WindowSpec pointers) into one []ast.Node (or the expected variadic type), then replace the three separate calls with a single call: sel.Fields.Fields, err = b.appendAuxiliaryFieldsForSubqueries(ctx, p, sel.Fields.Fields, collectedNodes...), handling error as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.json`:
- Around line 45-54: Add a focused regression case to the "cases" array of
TestWindowSubqueryOuterRef that exercises an OVER (ORDER BY
<correlated-subquery>) expression so the change to preserve outer refs for
window ORDER BY items is covered; update the "cases" array (in the same JSON
object under "name": "TestWindowSubqueryOuterRef") by appending a SQL test that
uses a correlated subquery inside an OVER (ORDER BY ...) clause (refer to the
existing cases for naming and ordering style).
In `@pkg/planner/core/logical_plan_builder.go`:
- Around line 2741-2769: The three loops each call
appendAuxiliaryFieldsForSubqueries separately causing repeated AST walks;
instead gather all window-related nodes into a single slice and call
appendAuxiliaryFieldsForSubqueries once. Concretely, iterate over
sel.Fields.Fields, sel.WindowSpecs, and sel.OrderBy.Items, collect
expressions/nodes that satisfy ast.HasWindowFlag (and the WindowSpec pointers)
into one []ast.Node (or the expected variadic type), then replace the three
separate calls with a single call: sel.Fields.Fields, err =
b.appendAuxiliaryFieldsForSubqueries(ctx, p, sel.Fields.Fields,
collectedNodes...), handling error as before.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: ec0a8d12-d204-4800-8421-53cb0e095151
📒 Files selected for processing (4)
pkg/planner/core/casetest/windows/testdata/window_push_down_suite_in.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_out.jsonpkg/planner/core/casetest/windows/testdata/window_push_down_suite_xut.jsonpkg/planner/core/logical_plan_builder.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/planner/core/casetest/windows/testdata/window_push_down_suite_out.json
qw4990
left a comment
qw4990
left a comment
There was a problem hiding this comment.
Could you also add this as a test case?
This case is from #57594
func TestCorrelatedExistsWithWindowFunction(t *testing.T) {
store := testkit.CreateMockStore(t)
tk := testkit.NewTestKit(t, store)
tk.MustExec("use test")
tk.MustExec("create table t1 (c1 int)")
// Empty table should return no rows.
tk.MustQuery("select case when (EXISTS (select 1 where ref_0.c1)) then count(*) over () else 1 end from t1 as ref_0").Check(testkit.Rows())
tk.MustExec("insert into t1 values (1), (null), (0), (2)")
// c1=1 → EXISTS true → count(*) over () = 4
// c1=null → EXISTS false → 1
// c1=0 → EXISTS false → 1
// c1=2 → EXISTS true → count(*) over () = 4
tk.MustQuery("select c1, case when (EXISTS (select 1 where ref_0.c1)) then count(*) over () else 1 end as res from t1 as ref_0 order by c1").Check(testkit.Rows(
"<nil> 1",
"0 1",
"1 4",
"2 4",
))
// Correlated EXISTS with window function using PARTITION BY.
tk.MustQuery("select c1, case when (EXISTS (select 1 where ref_0.c1 > 0)) then sum(c1) over () else 0 end as res from t1 as ref_0 order by c1").Check(testkit.Rows(
"<nil> 0",
"0 0",
"1 3",
"2 3",
))
// Correlated EXISTS with a text column (original bug reproduction from issue).
tk.MustExec("drop table t1")
tk.MustExec("create table t1 (c1 text)")
tk.MustQuery("select case when (EXISTS (select 1 where ref_0.c1)) then count(*) over (partition by 1) else 1 end from t1 as ref_0").Check(testkit.Rows())
}
|
@pantheon-bot review |
|
@Mini256 I've received your request and will start reviewing the pull request. I'll conduct a thorough review covering code quality, potential issues, and implementation details. ⏳ This process typically takes 10-30 minutes depending on the complexity of the changes. ℹ️ Learn more details on Pantheon AI. |
![pantheon-ai[bot]](https://avatars.githubusercontent.com/in/2265341?s=60&v=4){kind=small}
|
|
||
| Validation commands: | ||
| - `make failpoint-enable` | ||
| - `go test -run 'TestWindowWithCorrelatedSubQuery|TestWindowSubqueryOuterRef|TestWindowNonScalarSubqueryInWindowExpr' -tags=intest,deadlock ./pkg/planner/core/casetest/windows` |
There was a problem hiding this comment.
[P2] Stale validation command references non-existent test functions
Why: The documented go test -run command references TestWindowWithCorrelatedSubQuery and TestWindowNonScalarSubqueryInWindowExpr, neither of which exists in the package. Only TestWindowSubqueryOuterRef matches an actual test. The intended non-scalar regression coverage lives inside TestWindowSubqueryRewrite (see pkg/planner/core/casetest/windows/widow_with_exist_subquery_test.go:86-87). This means a developer following the checklist validation command would silently miss running the non-scalar subquery regression tests.
Evidence: docs/note/planner/window_subquery_notes.md:32 and docs/note/planner/window_subquery_notes.md:38 — the -run regex TestWindowWithCorrelatedSubQuery|TestWindowSubqueryOuterRef|TestWindowNonScalarSubqueryInWindowExpr only matches TestWindowSubqueryOuterRef in the target package.
Suggestion: Replace with -run 'TestWindowSubqueryRewrite|TestWindowSubqueryOuterRef' (or rename/split the test to match the documented name).
There was a problem hiding this comment.
Good catch. The note was still using the old test names after I merged the hand-written regressions into TestWindowSubqueryRewrite.
Updated docs/note/planner/window_subquery_notes.md to keep the documented coverage and validation command aligned with the current package tests:
TestWindowSubqueryRewriteTestWindowSubqueryOuterRef
The command is now:
go test -run 'TestWindowSubqueryRewrite|TestWindowSubqueryOuterRef' -tags=intest,deadlock ./pkg/planner/core/casetest/windows
Fixed in commit e4799aa.
There was a problem hiding this comment.
Thanks for the quick fix! Confirmed — the updated command correctly maps to the two existing test functions. LGTM.
ℹ️ Learn more details on Pantheon AI.
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: AilinKid, qw4990 The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
ti-chi-bot
bot
added
lgtm
and removed
needs-1-more-lgtm
|
@hawkingrei: The following tests failed, say
Full PR test history. Your PR dashboard. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here. |
4 participants
What problem does this PR solve?
Issue Number: close #56532
Problem Summary:
Window function expressions and peer expressions that contain correlated scalar subqueries can lose access to outer columns during the deferred window rewrite path. This makes queries that work in MySQL fail in TiDB with unknown-column errors.
Review also found that the auxiliary-field collection path could rewrite non-scalar subqueries inside window expressions through the scalar-subquery path, incorrectly introducing
MaxOneRowfor validIN/ANYcases.What changed and how does it work?
EXISTS,IN, and quantified subquery expressions as complete expressions during auxiliary-field extraction so non-scalar subqueries do not get rewritten as scalar subqueriesCheck List
Tests
make failpoint-enablego test -run 'TestWindowWithCorrelatedSubQuery|TestWindowSubqueryOuterRef|TestWindowNonScalarSubqueryInWindowExpr' -tags=intest,deadlock ./pkg/planner/core/casetest/windowsmake failpoint-disablemake lintSide effects
Documentation
Release note
Summary by CodeRabbit
Bug Fixes
Tests
Documentation
Chores