planner: use local CTE counts for MPP shared CTE

[winoros]

What problem does this PR solve?

Issue Number: close #68837, ref #61241

Problem Summary:

MPP shared CTE currently fills CteSinkNum and CteSourceNum by global CTE id counts. This is unsafe for TiFlash because the runtime CTEManager is node-local. When a producer-side UNION ALL creates sink fragments on different TiFlash nodes, a source running on one node can only observe the local sinks on that node, but the plan may still tell it to wait for the global sink count.

What changed and how does it work?

This PR changes MPP shared CTE sink/source count filling to use TiFlash-address-local counts instead of global CTE-id counts. It records the fragment self tasks for each CTE sink/source plan node, computes the local sink/source instance count per TiFlash address, and writes those local counts back to the CTE operators.

If one CTE plan node spans tasks that would require different local counts, the planner now returns an error instead of sending counts that cannot be represented in the current TiPB fields.

The comments and helper naming are also updated to explain why local executor-instance counts are required.

Check List

Tests

• Unit test
• Integration test
• Manual test (add detailed scripts or steps below)
• No need to test

  • I checked and no code files have been changed.

Commands:

./tools/check/failpoint-go-test.sh pkg/planner/core/operator/physicalop -run TestFillLocalCTECountsUsesLocalTaskCounts -count=1
go test -run TestMPPSharedCTEScan -tags=intest,deadlock ./pkg/planner/core/casetest/enforcempp -count=1
GOFLAGS=-buildvcs=false make lint
git diff --check e53c50c1c5ae0bd926d75d278ed6fa958670b648..HEAD

Side effects

• Performance regression: Consumes more CPU
• Performance regression: Consumes more Memory
• Breaking backward compatibility

Documentation

• Affects user behaviors
• Contains syntax changes
• Contains variable changes
• Contains experimental features
• Changes MySQL compatibility

Release note

Please refer to Release Notes Language Style Guide to write a quality release note.

Fix a potential MPP shared CTE hang caused by using global sink and source counts instead of TiFlash-node-local counts.

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Fixed executor count calculation for Common Table Expressions (CTEs) in distributed TiFlash queries. Counts are now calculated using per-address local task information, improving accuracy of CTE query execution.

• Tests

  • Added test coverage for local CTE executor count calculation.

• Chores

  • Updated test infrastructure and documentation.

[pantheon-ai]

@winoros I've received your pull request and will start the review. 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.}

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR replaces global CTE executor counting with local per-address counting to fix a deadlock risk in MPP shared CTE execution when producers contain UNION ALL branches scheduled across different TiFlash workers. The implementation introduces grouping structures, traversal logic, counting validation, and a comprehensive unit test.

Changes

Local CTE Count Calculation for MPP Shared CTEs

Layer / File(s)	Summary
CTE grouping structures for local task tracking pkg/planner/core/operator/physicalop/fragment.go	Introduces cteSinkInMPPTasks and cteSourceInMPPTasks structures and getCTEInMPPTasks helper to group plan nodes and their associated fragment-local tasks by CTE storage ID, replacing the prior flat cteInMPPTasks layout.
Fragment traversal for CTE node collection pkg/planner/core/operator/physicalop/fragment.go	Reworks traverseFragForCTE to accept and record fragment task lists alongside each PhysicalCTESink/PhysicalCTESource node within the per-CTE grouping structure, enabling local count computation per address.
Local count aggregation and validation pkg/planner/core/operator/physicalop/fragment.go	Adds addCTELocalCount to count tasks by TiFlash address and getUniformCTELocalCount to enforce uniform per-address counts for a given sink/source node, returning an error if counts diverge across addresses within that node's task set.
Main CTE count population logic pkg/planner/core/operator/physicalop/fragment.go	fillLocalCTECounts builds the CTE grouping index across fragments, computes per-address local sink and source task counts, validates count uniformity per plan node, and writes CteSinkNum/CteSourceNum using local counts instead of global duplicated totals.
Integration into task generation pkg/planner/core/operator/physicalop/fragment.go	generateMPPTasks now invokes fillLocalCTECounts(e.frags) to populate CTE counts, replacing the prior call to fixDuplicatedTimesForCTE.
Updated documentation for CTE count fields pkg/planner/core/operator/physicalop/physical_cte.go	Inline comments on PhysicalCTESink and PhysicalCTESource fields CteSinkNum/CteSourceNum are updated to reference mppTaskGenerator.fillLocalCTECounts instead of fixDuplicatedTimesForCTE.
Unit test for local CTE count assignment pkg/planner/core/operator/physicalop/fragment_test.go	TestFillLocalCTECountsUsesLocalTaskCounts constructs a CTE fragment scenario with sinks on different TiFlash addresses and sources wrapped by exchange senders, invokes fillLocalCTECounts, and verifies that each address sees correct local counts (1 sink, 2 sources) rather than global totals.
Test build configuration pkg/planner/core/operator/physicalop/BUILD.bazel	Increase physicalop_test shard count from 3 to 4 and add //pkg/kv dependency to support construction of test kv.MPPTask instances.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 Counting on Local Tasks

The CTE sinks were split too wide,
Global counts left TiFlash to hide.
Now local tasks on each address shine,
With uniform counts, no more waiting in line! 🌟

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'planner: use local CTE counts for MPP shared CTE' directly describes the main change: replacing global CTE counts with local (TiFlash-address-local) counts.
Linked Issues check	✅ Passed	The changes fully address issue #68837 [#68837] by implementing local CTE counts: fragment.go replaces global count logic with local count computation per TiFlash address, fragment_test.go validates the new behavior, and physical_cte.go updates documentation.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to implementing local CTE counts: BUILD.bazel changes support the new test, fragment.go implements the core fix, fragment_test.go validates it, and physical_cte.go updates related comments.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

Tools execution failed with the following error:

Failed to run tools: 13 INTERNAL: Received RST_STREAM with code 2 (Internal server error)

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

pkg/planner/core/operator/physicalop/fragment_test.go (1)

56-96: ⚡ Quick win

Consider adding a negative-path test for the local-count mismatch validation.

This test only covers the uniform colocated case. Per the PR objectives, fillLocalCTECounts is expected to return an error when a CTE node's local per-address counts cannot be represented uniformly (the getUniformCTELocalCount validation path). Adding a case where sink/source self-tasks yield differing per-address counts would guard the error branch against regressions.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/planner/core/operator/physicalop/fragment_test.go` around lines 56 - 96,
Add a negative-path unit test in fragment_test.go that exercises
mppTaskGenerator.fillLocalCTECounts when per-address local counts are
non-uniform so getUniformCTELocalCount returns an error: create CTE
sinks/sources (PhysicalCTESink, PhysicalCTESource) and ExchangeSenders with
self-tasks arranged such that one address has a different number of local tasks
than another (e.g., sink0 has tasks ["tiflash0"], sink1 has tasks
["tiflash0","tiflash1"] or similar), call
(&mppTaskGenerator{}).fillLocalCTECounts(frags) and assert it returns a non-nil
error, and optionally assert that the error message indicates non-uniform local
CTE counts.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@pkg/planner/core/operator/physicalop/fragment_test.go`:
- Around line 56-96: Add a negative-path unit test in fragment_test.go that
exercises mppTaskGenerator.fillLocalCTECounts when per-address local counts are
non-uniform so getUniformCTELocalCount returns an error: create CTE
sinks/sources (PhysicalCTESink, PhysicalCTESource) and ExchangeSenders with
self-tasks arranged such that one address has a different number of local tasks
than another (e.g., sink0 has tasks ["tiflash0"], sink1 has tasks
["tiflash0","tiflash1"] or similar), call
(&mppTaskGenerator{}).fillLocalCTECounts(frags) and assert it returns a non-nil
error, and optionally assert that the error message indicates non-uniform local
CTE counts.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 26fe0d37-93c5-4051-8feb-d0bc2895654c

📥 Commits

Reviewing files that changed from the base of the PR and between a9add5c and 05ebc1e.

📒 Files selected for processing (4)

• pkg/planner/core/operator/physicalop/BUILD.bazel
• pkg/planner/core/operator/physicalop/fragment.go
• pkg/planner/core/operator/physicalop/fragment_test.go
• pkg/planner/core/operator/physicalop/physical_cte.go

[codecov]

Codecov Report

❌ Patch coverage is 0% with 75 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.9176%. Comparing base (a9add5c) to head (e899b75).
⚠️ Report is 123 commits behind head on master.

Additional details and impacted files

@@               Coverage Diff                @@
##             master     #68848        +/-   ##
================================================
- Coverage   76.3085%   73.9176%   -2.3910%     
================================================
  Files          2041       2042         +1     
  Lines        563262     597537     +34275     
================================================
+ Hits         429817     441685     +11868     
- Misses       132529     153915     +21386     
- Partials        916       1937      +1021     

Flag	Coverage Δ
integration	42.8095% <0.0000%> (+3.0310%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Components	Coverage Δ
dumpling	60.4610% <ø> (ø)
parser	∅ <ø> (∅)
br	44.6341% <ø> (-18.1969%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[winoros]

/retest

[xzhangxian1008]

lgtm

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: time-and-fate, xzhangxian1008

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

• pkg/planner/OWNERS [time-and-fate]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[ti-chi-bot]

[LGTM Timeline notifier]

Timeline:

• 2026-06-23 19:09:57.461651434 +0000 UTC m=+61081.124876957: ☑️ agreed by time-and-fate.
• 2026-06-24 09:35:58.217816166 +0000 UTC m=+113041.881041679: ☑️ agreed by xzhangxian1008.