fine grained shuffle occupy too much memory when max_threads is large

[guo-shaoge]

Enhancement

On the dev machine (72C), When running the following query, memory usage is 12G+, but when change max_threads to 8, the memory usage is smaller.

Most memory is occupied by FineGrainedShuffleWriter, we need to optimize the memory usage

Env is as follow:

query: explain analyze select * from (select *, row_number() over (partition by ps_partkey order by ps_partkey) as rn from partsupp) d where rn = 1;
workload: tpch-50G
cluster: 1 tidb, 3 tikv, 3 tiflash, 1 pd
version:

mysql> select tidb_version();
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| tidb_version()                                                                                                                                                                                                                                                                             |
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Release Version: v9.0.0-beta.2.pre-275-g3409448
Edition: Community
Git Commit Hash: 3409448d3b9c405df37755f047c28bb0195ee891
Git Branch: HEAD
UTC Build Time: 2025-08-06 05:59:46
GoVersion: go1.23.11
Race Enabled: false
Check Table Before Drop: false
Store: tikv
Kernel Type: Classic |
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
1 row in set (0.00 sec)

[guo-shaoge]

Analysis: 72 * 3 * 72
There are 72 ExchangeOp operators, each containing a FineGrainedShuffleWriter.
Each writer has a stream_count of 72, and needs to send data to 3 TiFlash nodes,
which means each writer needs to split every column into 216 parts.

This query has 5 columns in total:

• 3 × Int64
• 1 × Decimal(15, x)
• 1 × String

Each split column will reserve 1024 rows.

Final memory usage:
(16384 * 5 + 131072) * 72 * 72 * 3 = 3312451584
—approximately at that order of magnitude.

Where:

16384: Memory usage for each Int64 column after reserving 1024 rows.
131072: Memory usage for ColumnString::chars after reserving 1024 rows.