From 7b20e92b01671d9f8029f61a469133e145b58c4c Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Sun, 26 Jul 2026 21:42:02 +0800
Subject: [PATCH 01/18] Add running local TopN MSLM design

---
 .../2026-07-26-running-local-topn-for-mslm.md | 527 ++++++++++++++++++
 1 file changed, 527 insertions(+)
 create mode 100644 docs/design/2026-07-26-running-local-topn-for-mslm.md

diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
new file mode 100644
index 00000000000..9a42dd93fd8
--- /dev/null
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -0,0 +1,527 @@
+# TiFlash MSLM Running Local TopN 设计
+
+- Author: TBD
+- Discussion PR: TBD
+- Tracking Issue: TBD
+
+## 目录
+
+- [背景](#背景)
+- [目标](#目标)
+- [非目标](#非目标)
+- [MVP 范围](#mvp-范围)
+- [核心思路](#核心思路)
+- [正确性](#正确性)
+- [执行流程](#执行流程)
+- [Running Local TopN 状态](#running-local-topn-状态)
+- [Filter 坐标](#filter-坐标)
+- [启用规则](#启用规则)
+- [实现计划](#实现计划)
+- [测试方案](#测试方案)
+- [性能评估](#性能评估)
+- [风险](#风险)
+- [备选方案](#备选方案)
+- [未解决问题](#未解决问题)
+
+## 背景
+
+TiFlash 现有 multi-stage late materialization (MSLM) 通过把 TableScan 上方的 residual Selection 下沉到 DeltaMerge scan 内部，先读取较少的 filter columns，再延迟读取最终输出所需的 final rest columns。
+
+对于下面这类 SQL，MSLM 仍然存在进一步优化空间：
+
+```sql
+SELECT payload_cols
+FROM t
+WHERE selection_pred
+ORDER BY order_key
+LIMIT k;
+```
+
+如果 `payload_cols` 很宽，而 `ORDER BY ... LIMIT` 只需要少量 rows，理想情况下 TiFlash 可以先读取 selection columns 和 order-by columns，筛出 TopN 候选 rows，再读取这些候选 rows 的 payload columns。
+
+更彻底的方案是让 storage 支持按 candidate row id 回读 rest columns。但当前 DeltaMerge storage 还没有提供这种能力。短期内可以在现有 MSLM 的 block lock-step 模型里做一个更保守的方案：在每个 MSLM stream 内维护 running local TopN，用它生成当前 block 的 candidate filter，只减少 final rest columns 的读取量。
+
+## 目标
+
+本设计希望在不引入 candidate row id 回读能力的前提下，减少 `TopN + Selection + TableScan` 场景中 final rest columns 的读取。
+
+目标：
+
+- 复用现有 MSLM 框架，不改变 DeltaMerge storage 的 block-level read/skip/readWithFilter 能力。
+- 在 MSLM Stage 1 中读取 residual filter columns 和 order-by columns。
+- 在每个 MSLM stream 内维护 running local TopN 状态。
+- 使用 running local TopN 生成当前 block 的 candidate filter。
+- 将 candidate filter 合并到 final rest columns 的读取 filter 中。
+- 保留上层原有 TopN executor，由上层 TopN 负责最终全局结果。
+
+## 非目标
+
+第一版不解决以下问题：
+
+- 不支持按 candidate row id 随机回读 rest columns。
+- 不移除或替换上层 TopN executor。
+- 不做跨 stream 的 storage 内全局 TopN merge。
+- 不支持复杂 ORDER BY expression。
+- 不支持 join、aggregation、window 或其他会改变 row cardinality 的算子下的 TopN 下推。
+- 不对旧的非 pipeline DAG executor 路径做支持。
+- 不追求和未优化路径在 ORDER BY ties 下选择完全相同的物理 rows。
+
+## MVP 范围
+
+第一版只支持下面范围：
+
+- 只支持 `ORDER BY` plain columns。
+- 只支持 constant `LIMIT` 和 constant `OFFSET`。
+- 只支持单表 `Selection + TableScan + TopN` 或 `TableScan + TopN`。
+- 不移除上层 TopN。
+- 只在 MSLM 已经满足启用条件时启用。
+- TopN filter 只用于减少 final rest columns 的读取。
+
+其中本地 TopN 使用的 K 为：
+
+```text
+K = LIMIT + OFFSET
+```
+
+如果没有 OFFSET，则：
+
+```text
+K = LIMIT
+```
+
+## 核心思路
+
+现有 MSLM 的逻辑可以简化为：
+
+```text
+Stage 0:
+  读取 pushed filter columns
+  执行 pushed filters
+  得到 stage0_filter
+
+Stage 1:
+  根据 stage0_filter 读取 residual filter columns
+  执行 residual filters
+  得到 residual_filter
+
+Final:
+  根据 stage0_filter/residual_filter 读取 final rest columns
+```
+
+加入 running local TopN 后，逻辑变为：
+
+```text
+Stage 0:
+  读取 pushed filter columns
+  执行 pushed filters
+  得到 stage0_filter
+
+Stage 1:
+  根据 stage0_filter 读取 residual filter columns + order-by columns
+  执行 residual filters
+  执行 running local TopN
+  得到 topn_candidate_filter
+
+Final:
+  根据 stage0_filter && residual_filter && topn_candidate_filter
+  读取 final rest columns
+```
+
+`topn_candidate_filter` 只用于减少 final rest columns 的读取。上层 TopN executor 仍然保留，因此 storage 内输出可以是全局 TopN 的 superset。
+
+## 正确性
+
+### Local TopK 覆盖 Global TopK
+
+对于任意一个 stream，如果某一行 `r` 是全局 TopK 结果中的一行，那么 `r` 必然属于它所在 stream 的 local TopK。
+
+证明：
+
+如果 `r` 不属于所在 stream 的 local TopK，说明这个 stream 中至少有 K 行按照相同 comparator 排在 `r` 前面。这 K 行在全局也排在 `r` 前面，因此 `r` 不可能是全局 TopK，矛盾。
+
+所以：
+
+```text
+global TopK subset union(each stream local TopK)
+```
+
+第一版保留上层全局 TopN executor，因此只要 storage 内输出包含全局 TopK 的 superset，最终结果就是正确的。
+
+### Running Local TopN 的 Streaming Superset
+
+MSLM stream 是 streaming 输出，不能撤回已经输出的 rows。Running local TopN 处理后续 blocks 时，可能会把之前已经输出的 candidate rows 挤出当前 local TopK。
+
+这不会影响正确性，因为这些 rows 只是额外候选。上层 TopN 会再次排序并裁剪。
+
+因此 running local TopN 输出的是：
+
+```text
+final local TopK 的 superset
+```
+
+而不是严格等于 final local TopK。
+
+### 被跳过 rows 不会重新进入 TopK
+
+当 local TopN heap 已满时，如果当前 row 按 comparator 不优于 heap worst row，则该 row 可以被跳过。原因是后续扫描只会增加竞争者，不会让这个 row 的排序位置变得更靠前。
+
+对于和 heap worst row 相等的 ties，第一版采用保守策略：相等 rows 作为 candidate 输出，但不一定放入 bounded heap。这会降低 ties 较多场景下的剪枝效果，但可以避免因为 ties 选择不同物理 rows 引起不必要的行为差异。
+
+## 执行流程
+
+### 计划识别
+
+TiFlash 在构建 TableScan pipeline 时识别下面模式：
+
+```text
+TopN
+  |
+Selection (optional)
+  |
+TableScan
+```
+
+识别成功后，生成一个 MSLM TopN 描述：
+
+```text
+struct MSLMTopNDescription
+{
+    SortDescription sort_description;
+    UInt64 limit;
+    UInt64 offset;
+    UInt64 topk; // limit + offset
+    ColumnDefines order_by_columns;
+};
+```
+
+第一版 `order_by_columns` 只允许 plain columns，不构造复杂 expression actions。
+
+### Stage 1 列集合
+
+开启 TopN-enhanced MSLM 后，Stage 1 需要读取：
+
+```text
+stage1_columns = residual_filter_columns union order_by_columns
+```
+
+如果某个 order-by column 已经属于 residual filter columns，则去重后只读一次。
+
+如果 order-by column 也属于最终输出 columns，它仍然可以作为 Stage 1 column 提前读取，final rest columns 中需要排除该列，避免重复读取。
+
+### Running Local TopN 处理
+
+每个 `MultiStageLateMaterializationBlockInputStream` 持有一个 local TopN heap。
+
+处理一个 block 时：
+
+1. Stage 0 读取 pushed filter columns，得到 `stage0_filter`。
+2. Stage 1 根据 `stage0_filter` 读取 `stage1_columns`。
+3. 执行 residual filters，得到 `residual_filter`。
+4. 对 residual passed rows 计算 order key。
+5. 使用 order key 更新 local TopN heap。
+6. 为当前 block 生成 `topn_candidate_filter`。
+7. 使用 combined filter 读取 final rest columns。
+8. hstack Stage 1 columns 和 final rest columns。
+9. 上层 TopN executor 继续执行全局 TopN。
+
+伪代码：
+
+```text
+for each stage0 block:
+    stage0_filter = execute_stage0_filter(block)
+    if stage0_filter passes no rows:
+        skip stage1 stream
+        skip final_rest stream
+        continue
+
+    stage1_block = read stage1 columns with stage0_filter
+    residual_filter = execute_residual_filter(stage1_block)
+    if residual_filter passes no rows:
+        skip final_rest stream
+        continue
+
+    topn_candidate_filter = running_topn.update(stage1_block, residual_filter)
+    if topn_candidate_filter passes no rows:
+        skip final_rest stream
+        continue
+
+    combined_filter = compose(stage0_filter, residual_filter, topn_candidate_filter)
+    final_rest_block = read final_rest columns with combined_filter
+    stage1_block = filter stage1_block by residual_filter && topn_candidate_filter
+    return hstack(stage1_block, final_rest_block)
+```
+
+### Current Block 内 Eviction
+
+如果当前 block 的某个 row 先进入 heap，但随后在同一个 block 内被更优 row 挤出，第一版可以把它从当前 block 的 candidate filter 中清除。
+
+如果被挤出的 row 来自之前已经输出的 block，则不能撤回。该 row 会作为额外候选留给上层 TopN 过滤。
+
+Heap entry 可以记录：
+
+```text
+struct HeapEntry
+{
+    SortKey key;
+    UInt64 stream_sequence;
+    UInt64 block_sequence;
+    UInt32 row_index_in_stage1_block;
+    bool current_block_output_candidate;
+};
+```
+
+`block_sequence` 用于判断被 eviction 的 entry 是否属于当前 block。
+
+## Running Local TopN 状态
+
+每个 MSLM stream 维护独立状态：
+
+```text
+class RunningLocalTopN
+{
+public:
+    Filter update(const Block & stage1_block, const FilterPtr & residual_filter);
+
+private:
+    UInt64 topk;
+    SortDescription sort_description;
+    PriorityQueue heap; // heap top is current worst candidate
+};
+```
+
+Heap 大小最多为 `topk`。如果 `topk` 超过启发式阈值，则禁用该优化。
+
+比较器必须和上层 TopN 使用的 comparator 保持一致。第一版只支持 plain order-by columns，减少 comparator 语义不一致风险。
+
+## Filter 坐标
+
+MSLM 中存在多套 filter 坐标：
+
+- `stage0_filter`: 原始 Stage 0 block rows 坐标。
+- `residual_filter`: Stage 1 block rows 坐标，即 Stage 0 passed rows 坐标。
+- `topn_candidate_filter`: Stage 1 block rows 坐标。
+- `combined_filter`: 原始 Stage 0 block rows 坐标，用于读取 final rest columns。
+
+组合规则：
+
+```text
+stage1_final_filter = residual_filter && topn_candidate_filter
+combined_filter = compose(stage0_filter, stage1_final_filter)
+```
+
+当 `stage0_filter` 为 nullptr，即 Stage 0 全通过时，`combined_filter` 可以直接基于 `stage1_final_filter` 构造。
+
+当 `stage1_final_filter` 全通过时，可以退化为现有 MSLM 路径。
+
+当 `stage1_final_filter` 全过滤时，需要 skip final rest stream。
+
+## 启用规则
+
+第一版使用保守 rule-based 判断。
+
+必须满足：
+
+- `dt_enable_multi_stage_late_materialization` 为 true。
+- 当前查询已经满足 MSLM 启用条件。
+- 查询包含 TopN，且 TopN 位于单表 Selection/TableScan 之上。
+- TopN 上方仍保留全局 TopN executor。
+- TopN 的 `LIMIT` 和 `OFFSET` 是常量。
+- `topk = limit + offset` 未超过阈值。
+- TopN 之前的 Selection predicates 全部在 MSLM Stage 0 或 Stage 1 内执行。
+- ORDER BY 只包含 plain columns。
+- ORDER BY columns 可加入 Stage 1 column set。
+- ORDER BY columns 数量不超过阈值。
+- final rest columns 数量满足 MSLM 的宽列启发式要求。
+
+建议第一版阈值：
+
+```text
+topk <= 4096
+order_by_column_count <= 4
+```
+
+实际阈值可以根据 benchmark 调整。
+
+禁用场景：
+
+- ORDER BY expression 不是 plain column。
+- ORDER BY 依赖 generated column。
+- TopN 包含非常大的 offset。
+- TopN 之前仍有无法进入 MSLM 的 Selection。
+- 查询要求 keep order 且和 MSLM 现有限制冲突。
+- TopN 位于 join、aggregation、window 等算子之上。
+
+## 实现计划
+
+### Step 1: 文档和测试用例设计
+
+- 明确 MVP 范围。
+- 增加基于 DAG request 的单测设计。
+- 明确 ties、offset、NULL ordering、collation 的测试覆盖。
+
+### Step 2: Plan 识别和元数据传递
+
+- 在 TableScan pipeline 构建阶段识别 `TopN -> Selection -> TableScan` 模式。
+- 构造 `MSLMTopNDescription`。
+- 将 order-by columns 合并到 MSLM Stage 1 column set。
+- 保留上层 TopN executor，不修改最终 DAG 语义。
+
+### Step 3: RunningLocalTopN 状态
+
+- 实现 per-stream bounded TopK heap。
+- 支持 plain column sort key 提取。
+- 支持 ASC/DESC、NULL ordering 和 TiFlash 现有 comparator 语义。
+- 生成当前 block 的 `topn_candidate_filter`。
+
+### Step 4: 接入 MSLM final rest 读取
+
+- 在 residual filter 之后执行 running local TopN。
+- 将 `residual_filter` 和 `topn_candidate_filter` 合并。
+- 使用 combined filter 读取 final rest columns。
+- 当 TopN filter 无剪枝效果时，允许退化到现有路径。
+
+### Step 5: Runtime Stats 和可观测性
+
+- 保留现有 TableScan/Selection actRows 语义。
+- 可选增加 debug log，输出 TopN candidate rows、heap size、filtered rows。
+- 可选增加 scan context counters，用于统计 TopN filter 节省的 final rest rows。
+
+## 测试方案
+
+### 功能测试
+
+- `ORDER BY c ASC LIMIT k`。
+- `ORDER BY c DESC LIMIT k`。
+- `ORDER BY c ASC LIMIT offset, count`。
+- `WHERE + ORDER BY + LIMIT`，其中 WHERE 同时包含 Stage 0 pushed filter 和 Stage 1 residual filter。
+- ORDER BY column 不在最终 projection 中。
+- ORDER BY column 同时属于 residual filter columns。
+- 多 block 数据，确保 running local TopN 可以跨 block 收紧 threshold。
+- 多 stream 并发，确保上层 TopN 保证最终结果。
+
+### 正确性测试
+
+- 和未启用优化的执行结果比较。
+- LIMIT/OFFSET 场景比较。
+- NULL ordering 场景比较。
+- ties 场景比较，验证结果满足 SQL order semantics。
+- Stage 0 全过滤。
+- residual filter 全过滤。
+- TopN candidate filter 全过滤。
+- TopN candidate filter 全通过。
+
+### 回归测试
+
+- MSLM 原有测试必须全部通过。
+- TopN 上方的 projection 仍然正确。
+- EXPLAIN ANALYZE actRows 不出现明显异常。
+- 禁用条件下必须走原路径。
+
+## 性能评估
+
+重点 benchmark 场景：
+
+- 宽表，payload columns 多且冷数据。
+- `LIMIT` 小，`topk` 远小于 block rows。
+- `ORDER BY` columns 少且窄。
+- 数据分布有利于 running local TopN threshold 快速收紧。
+
+需要对比：
+
+- MSLM disabled。
+- MSLM enabled but without running local TopN。
+- MSLM enabled with running local TopN。
+
+关键指标：
+
+- Query latency。
+- final rest columns read rows。
+- DMFile read bytes。
+- CPU time。
+- memory usage。
+- candidate rows output to upper TopN。
+
+## 风险
+
+### Comparator 语义不一致
+
+如果 storage 内 running local TopN 的 comparator 和上层 TopN 不一致，可能错误过滤掉本应保留的 rows。
+
+缓解：
+
+- MVP 只支持 plain columns。
+- 复用现有排序比较逻辑。
+- 对 ASC/DESC、NULL、collation 做单测。
+
+### Selection 顺序错误
+
+如果 TopN filter 在某些 Selection predicates 执行之前生效，可能过滤掉最终应该进入 TopN 的 rows。
+
+缓解：
+
+- 只在 TopN 之前的 Selection 已全部进入 Stage 0/Stage 1 时启用。
+- 无法确认时禁用。
+
+### 收益不足或性能回退
+
+如果 `topk` 大、order-by columns 宽、ties 很多，running local TopN 可能减少不了 final rest 读取，还增加 Stage 1 读取和 heap 维护成本。
+
+缓解：
+
+- 使用保守启发式启用。
+- `topk` 和 order-by column count 设置上限。
+- 对 ties 多的场景 benchmark。
+
+### Streaming Superset 过大
+
+Running local TopN 不能撤回之前已经输出的 rows。极端数据分布下，每个 block 都可能输出一批候选 rows。
+
+缓解：
+
+- 上层 TopN 保留，保证正确性。
+- 通过 stats 观察 candidate rows 数量。
+- 后续考虑更强的 candidate row id materialization。
+
+## 备选方案
+
+### Block-local TopN
+
+每个 block 独立做 TopK，不维护跨 block state。
+
+优点：
+
+- 实现更简单。
+- 正确性证明直接。
+- 不需要处理 previous block candidate eviction。
+
+缺点：
+
+- 最多输出 `K * block_count` 个候选 rows。
+- 无法利用前面 blocks 已经得到的 threshold 剪枝后续 blocks。
+
+Running local TopN 是 block-local TopN 的增强版，复杂度增加有限，但剪枝能力更好。
+
+### Candidate Row ID 回读
+
+先读取 selection/order-by columns，得到全局或局部 TopN row ids，再按 row ids 回读 payload columns。
+
+优点：
+
+- 最接近理想 late materialization。
+- 可以只 materialize 很少 rows。
+
+缺点：
+
+- 需要 DeltaMerge storage 支持按 candidate row id 高效回读 columns。
+- 当前不具备该能力。
+- 设计和实现复杂度明显更高。
+
+本设计选择 running local TopN 作为短期方案。
+
+## 未解决问题
+
+- `topk` 阈值应该如何设置。
+- ties 场景是否需要为了结果稳定性保留所有 equal-worst rows。
+- 是否需要支持 ORDER BY expression，以及 expression action 如何复用。
+- 是否需要在 EXPLAIN ANALYZE 中暴露 TopN-enhanced MSLM 的 candidate rows。
+- 是否需要在 cost model 中引入列宽和 TopN selectivity。

From 7d17c0688aa5597e335a67fed112503116f51834 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Sun, 26 Jul 2026 21:51:04 +0800
Subject: [PATCH 02/18] Clarify running local TopN heap memory model

---
 .../2026-07-26-running-local-topn-for-mslm.md | 61 +++++++++++++++++--
 1 file changed, 56 insertions(+), 5 deletions(-)

diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index 9a42dd93fd8..b9c248d029f 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -257,20 +257,55 @@ for each stage0 block:
 
 如果被挤出的 row 来自之前已经输出的 block，则不能撤回。该 row 会作为额外候选留给上层 TopN 过滤。
 
-Heap entry 可以记录：
+Heap entry 只能记录轻量信息，不能持有 `Block` 或 column 引用：
 
 ```text
 struct HeapEntry
 {
-    SortKey key;
+    OwnedSortKey key;
     UInt64 stream_sequence;
     UInt64 block_sequence;
     UInt32 row_index_in_stage1_block;
-    bool current_block_output_candidate;
 };
 ```
 
-`block_sequence` 用于判断被 eviction 的 entry 是否属于当前 block。
+`block_sequence` 用于判断被 eviction 的 entry 是否属于当前 block。`row_index_in_stage1_block` 只在 entry 属于当前 block 时有效。
+
+如果 eviction 的 entry 来自当前 block，可以将 `topn_candidate_filter[row_index_in_stage1_block]` 置为 false，避免对该 row 读取 final rest columns。
+
+如果 eviction 的 entry 来自历史 block，则不做任何回撤。历史 block 已经输出给上层 TopN，该 row 只是额外候选。
+
+### Heap 内存模型
+
+Running local TopN 不能保存 `{Block, row_id}`。如果 heap entry 持有 `Block` 或 column memory 引用，会导致历史 blocks 被 retain，内存随扫描推进累积。
+
+Heap 中只保存排序需要的 owned key：
+
+```text
+heap memory = O(topk * order_key_size)
+current block filter memory = O(current_block_rows)
+```
+
+不保存：
+
+- 历史 block。
+- 历史 row payload。
+- 历史 column 引用。
+- 用于回读的 row locator。
+
+历史 candidate rows 一旦输出给上层 TopN，MSLM stream 不再负责保存它们。Heap eviction 只释放 owned sort key。
+
+对于 fixed-size order-by 类型，owned key 可以直接拷贝值和 NULL flag。对于 variable-length 类型，例如 String，owned key 可能带来较大内存开销。第一版可以先禁用 variable-length order-by columns，或者增加 per-stream memory guard。
+
+如果采用 memory guard，触发后应退化为普通 MSLM：
+
+```text
+后续 blocks 不再生成 TopN candidate filter。
+final rest columns 按现有 MSLM 逻辑读取。
+上层 TopN 保留，保证最终结果正确。
+```
+
+Ties 也需要避免破坏内存上界。第一版 heap 仍最多保存 `topk` 个 entries。与 heap worst 相等的 rows 可以作为当前 block candidate 输出，但不插入 heap。这样 ties 较多时剪枝效果会下降，但 heap 内存保持 bounded。
 
 ## Running Local TopN 状态
 
@@ -285,7 +320,8 @@ public:
 private:
     UInt64 topk;
     SortDescription sort_description;
-    PriorityQueue heap; // heap top is current worst candidate
+    PriorityQueue<HeapEntry> heap; // heap top is current worst candidate
+    UInt64 current_block_sequence;
 };
 ```
 
@@ -293,6 +329,8 @@ Heap 大小最多为 `topk`。如果 `topk` 超过启发式阈值，则禁用该
 
 比较器必须和上层 TopN 使用的 comparator 保持一致。第一版只支持 plain order-by columns，减少 comparator 语义不一致风险。
 
+第一版不要求 heap entry 能定位历史 block。Heap entry 的 locator 只服务于 current block eviction。历史 entry 被 eviction 时只释放 key，不访问历史 block。
+
 ## Filter 坐标
 
 MSLM 中存在多套 filter 坐标：
@@ -472,6 +510,18 @@ order_by_column_count <= 4
 - `topk` 和 order-by column count 设置上限。
 - 对 ties 多的场景 benchmark。
 
+### Heap 持有历史 Block 导致内存累积
+
+如果 heap entry 设计成 `{Block, row_id}`，或者持有历史 column memory 引用，会导致已经处理过的 blocks 无法释放，内存随扫描推进累积。
+
+缓解：
+
+- Heap entry 只保存 owned sort key。
+- Heap entry 不持有 `Block`、column 引用或历史 row payload。
+- Row locator 只用于 current block 内 eviction。
+- 历史 block entry 被 eviction 时只释放 key，不访问历史 block。
+- 对 variable-length order-by keys 使用禁用策略或 memory guard。
+
 ### Streaming Superset 过大
 
 Running local TopN 不能撤回之前已经输出的 rows。极端数据分布下，每个 block 都可能输出一批候选 rows。
@@ -522,6 +572,7 @@ Running local TopN 是 block-local TopN 的增强版，复杂度增加有限，
 
 - `topk` 阈值应该如何设置。
 - ties 场景是否需要为了结果稳定性保留所有 equal-worst rows。
+- 第一版是否禁用 variable-length order-by columns，还是实现 per-stream memory guard。
 - 是否需要支持 ORDER BY expression，以及 expression action 如何复用。
 - 是否需要在 EXPLAIN ANALYZE 中暴露 TopN-enhanced MSLM 的 candidate rows。
 - 是否需要在 cost model 中引入列宽和 TopN selectivity。

From bdd4d0eafea029a561c88b2d70f958673a3f308c Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Mon, 27 Jul 2026 18:26:36 +0800
Subject: [PATCH 03/18] Update MSLM running local TopN design

---
 .../2026-07-26-running-local-topn-for-mslm.md | 125 ++++++++++++++++--
 1 file changed, 111 insertions(+), 14 deletions(-)

diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index b9c248d029f..f3c4a22eed1 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -13,6 +13,7 @@
 - [核心思路](#核心思路)
 - [正确性](#正确性)
 - [执行流程](#执行流程)
+- [ORDER BY 列对齐](#order-by-列对齐)
 - [Running Local TopN 状态](#running-local-topn-状态)
 - [Filter 坐标](#filter-坐标)
 - [启用规则](#启用规则)
@@ -63,6 +64,7 @@ LIMIT k;
 - 不做跨 stream 的 storage 内全局 TopN merge。
 - 不支持复杂 ORDER BY expression。
 - 不支持 join、aggregation、window 或其他会改变 row cardinality 的算子下的 TopN 下推。
+- 第一版不支持没有进入现有 MSLM 路径的 `TopN + TableScan`。例如没有 TableScan pushed filter，或者没有 residual Selection 的查询，后续单独扩展。
 - 不对旧的非 pipeline DAG executor 路径做支持。
 - 不追求和未优化路径在 ORDER BY ties 下选择完全相同的物理 rows。
 
@@ -70,13 +72,19 @@ LIMIT k;
 
 第一版只支持下面范围：
 
-- 只支持 `ORDER BY` plain columns。
+- 只支持已经满足现有 MSLM 启用条件的查询。也就是说，TableScan 必须已有 pushed-down filters 作为 Stage 0，且 TableScan 上方必须有 residual Selection 作为 Stage 1。
+- 只支持 `ORDER BY` plain columns，即 TiPB `ByItem.expr` 必须直接是 `ColumnRef`。
 - 只支持 constant `LIMIT` 和 constant `OFFSET`。
-- 只支持单表 `Selection + TableScan + TopN` 或 `TableScan + TopN`。
+- 只支持单表 `TopN -> Selection -> TableScan`。在 TiFlash physical plan 中，`Selection` 通常已经被 push 到 `PhysicalTableScan::filter_conditions`，所以实际识别形态可以是 `TopN -> PhysicalTableScan(with filter_conditions)`。
 - 不移除上层 TopN。
 - 只在 MSLM 已经满足启用条件时启用。
 - TopN filter 只用于减少 final rest columns 的读取。
 
+后续可以分阶段扩展：
+
+- 阶段二支持有 TableScan pushed filter 但没有 residual Selection 的 `TopN -> TableScan`。这要求把 residual filter 变成 optional，Stage 1 只读取 order-by columns。
+- 阶段三支持完全没有 filter 的 `TopN -> TableScan`。这需要把 MSLM 从当前的 filter-driven 模型泛化为 order-by-column-driven 模型，Stage 1 order-by stream 需要成为 block driver，工程风险更高。
+
 其中本地 TopN 使用的 K 为：
 
 ```text
@@ -129,6 +137,12 @@ Final:
 
 `topn_candidate_filter` 只用于减少 final rest columns 的读取。上层 TopN executor 仍然保留，因此 storage 内输出可以是全局 TopN 的 superset。
 
+当前 MSLM 已经不再使用 adaptive mode。TopN-enhanced MSLM 因此不需要根据 residual 过滤率在 direct/late 两种策略之间做启发式选择，而是只保留下面三个自然分支：
+
+- `residual_filter && topn_candidate_filter` 全过滤时，skip final rest stream。
+- `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。
+- 其他部分通过场景，把 Stage 1 坐标的 filter compose 回 Stage 0 原始坐标，再用 combined filter 读取 final rest columns。
+
 ## 正确性
 
 ### Local TopK 覆盖 Global TopK
@@ -176,7 +190,7 @@ TiFlash 在构建 TableScan pipeline 时识别下面模式：
 ```text
 TopN
   |
-Selection (optional)
+Selection
   |
 TableScan
 ```
@@ -186,9 +200,7 @@ TableScan
 ```text
 struct MSLMTopNDescription
 {
-    SortDescription sort_description;
-    UInt64 limit;
-    UInt64 offset;
+    SortDescription storage_sort_description;
     UInt64 topk; // limit + offset
     ColumnDefines order_by_columns;
 };
@@ -196,6 +208,61 @@ struct MSLMTopNDescription
 
 第一版 `order_by_columns` 只允许 plain columns，不构造复杂 expression actions。
 
+Plan 识别建议在 `PhysicalTopN` 和 `PhysicalTableScan` 之间完成。`PhysicalTopN` 只在 child 是 `PhysicalTableScan` 时尝试生成 MSLM TopN 描述；如果 child 是 Projection、Join、Aggregation、Window 等其他节点，则不启用该优化。上层 `PhysicalTopN` executor 仍然保留，storage 内描述只用于减少 final rest columns 的读取。
+
+因为 TiFlash planner 会尽量把单表 `Selection` push 到 `PhysicalTableScan::filter_conditions`，所以第一版实际需要识别的 physical plan 形态通常是：
+
+```text
+PhysicalTopN
+  |
+PhysicalTableScan(with filter_conditions and pushed_down_filters)
+```
+
+这里的 `filter_conditions` 是 TiFlash physical planner 内部保存的 residual Selection。对于 TiPB DAG 中形态为 `Selection -> TableScan` 的单个 Selection，TiFlash 构建 physical plan 时会把它放入 `PhysicalTableScan::filter_conditions`，而不是保留一个单独的 `PhysicalFilter`。如果 Selection 的 child 不是 TableScan，或者同一个 TableScan 已经设置过 `filter_conditions`，则会保留普通 `PhysicalFilter`，第一版不处理这种情况。
+
+需要注意，`PhysicalTableScan::filter_conditions` 不等于 TiPB `TableScan.pushed_down_filter_conditions`：
+
+- `TableScan.pushed_down_filter_conditions` 是现有 MSLM 的 Stage 0 pushed filter。
+- `PhysicalTableScan::filter_conditions` 是 Stage 1 residual Selection。
+- 第一版 TopN-enhanced MSLM 要求这两者都存在。
+
+### ORDER BY 列对齐
+
+即使第一版只支持 plain column，也需要显式处理 TopN order-by column 和 DeltaMerge storage column 的对齐。
+
+不要直接使用 `PhysicalTopN` 现有 `SortDescription.column_name` 去匹配 storage column。这个名字来自 TopN child schema，而 DeltaMerge Stage 1 block 使用的是 `ColumnDefine.name`。TableScan 后还有 schema projection、特殊列名和 cast-after-TS 等逻辑，直接按 name 对齐容易和 storage header 不一致。
+
+第一版采用基于 column id 的对齐流程：
+
+```text
+tipb::ByItem.expr() 必须是 ColumnRef
+  -> decode ColumnRef index
+  -> table_scan.getColumns()[index].id
+  -> 在 columns_to_read 或 table column defines 中按 column id 找到 ColumnDefine
+  -> 将该 ColumnDefine 加入 Stage 1 columns
+  -> 使用 ColumnDefine.name 构造 storage-side SortDescription
+```
+
+`MSLMTopNDescription::storage_sort_description` 中的 column name 必须是 Stage 1 block 中真实存在的 storage column name，而不是 TopN child schema name。
+
+Stage 1 column union 和 final rest column subtraction 都必须按 `ColumnID` 去重：
+
+```text
+stage1_columns = residual_filter_columns union order_by_columns
+final_rest_columns = final_columns_to_read - stage1_columns
+```
+
+如果 order-by column 已经属于 residual filter columns，则 Stage 1 只读一次。如果 order-by column 也是最终输出列，则 final rest columns 中必须排除该列，避免重复读取。
+
+第一版不支持没有 residual Selection 的 `TopN -> TableScan`。这种场景没有 Selection executor 可以承载 candidate rows 的 `EXPLAIN ANALYZE` 展示，也需要让 residual filter 变成 optional，后续单独设计。
+
+plain column 还需要保证 storage 内比较语义和上层 TopN 一致。第一版建议保守禁用下面场景：
+
+- order-by expr 不是 direct `ColumnRef`。
+- order-by column 是 generated column 或 virtual `_tidb_tid`。
+- order-by column 需要 TableScan 后的 extra cast 才能得到上层 TopN 的比较值。
+- order-by column 是 variable-length string 且暂未确认 collator 和 owned key 内存上界。
+
 ### Stage 1 列集合
 
 开启 TopN-enhanced MSLM 后，Stage 1 需要读取：
@@ -361,12 +428,15 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 
 - `dt_enable_multi_stage_late_materialization` 为 true。
 - 当前查询已经满足 MSLM 启用条件。
-- 查询包含 TopN，且 TopN 位于单表 Selection/TableScan 之上。
+- 查询包含 TopN，且 TopN 的 child 是单表 `PhysicalTableScan`。
+- `PhysicalTableScan` 已经包含 residual `filter_conditions`，并且 TableScan 已经包含 pushed-down filters。
 - TopN 上方仍保留全局 TopN executor。
 - TopN 的 `LIMIT` 和 `OFFSET` 是常量。
 - `topk = limit + offset` 未超过阈值。
 - TopN 之前的 Selection predicates 全部在 MSLM Stage 0 或 Stage 1 内执行。
-- ORDER BY 只包含 plain columns。
+- ORDER BY 只包含 direct `ColumnRef` plain columns。
+- ORDER BY columns 能通过 `ColumnRef` index 映射到 `table_scan.getColumns()` 中的 column id，并能进一步映射到 DeltaMerge `ColumnDefine`。
+- ORDER BY columns 不需要 TableScan 后 extra cast。
 - ORDER BY columns 可加入 Stage 1 column set。
 - ORDER BY columns 数量不超过阈值。
 - final rest columns 数量满足 MSLM 的宽列启发式要求。
@@ -383,9 +453,12 @@ order_by_column_count <= 4
 禁用场景：
 
 - ORDER BY expression 不是 plain column。
+- ORDER BY plain column 无法按 column id 映射到 storage column。
+- ORDER BY column 需要 TableScan 后 extra cast。
 - ORDER BY 依赖 generated column。
 - TopN 包含非常大的 offset。
 - TopN 之前仍有无法进入 MSLM 的 Selection。
+- 查询没有 residual Selection，或者没有 TableScan pushed-down filters。
 - 查询要求 keep order 且和 MSLM 现有限制冲突。
 - TopN 位于 join、aggregation、window 等算子之上。
 
@@ -399,9 +472,12 @@ order_by_column_count <= 4
 
 ### Step 2: Plan 识别和元数据传递
 
-- 在 TableScan pipeline 构建阶段识别 `TopN -> Selection -> TableScan` 模式。
+- 在 `PhysicalTopN` 构建阶段识别 child 是 `PhysicalTableScan` 的模式。
+- 只在 `PhysicalTableScan` 同时包含 residual `filter_conditions` 和 TableScan pushed-down filters 时启用。
 - 构造 `MSLMTopNDescription`。
+- 按 `ColumnRef` index -> `table_scan.getColumns()[index].id` -> `ColumnDefine` 的顺序对齐 order-by columns。
 - 将 order-by columns 合并到 MSLM Stage 1 column set。
+- 使用 Stage 1 storage column names 构造 storage-side `SortDescription`。
 - 保留上层 TopN executor，不修改最终 DAG 语义。
 
 ### Step 3: RunningLocalTopN 状态
@@ -416,13 +492,18 @@ order_by_column_count <= 4
 - 在 residual filter 之后执行 running local TopN。
 - 将 `residual_filter` 和 `topn_candidate_filter` 合并。
 - 使用 combined filter 读取 final rest columns。
-- 当 TopN filter 无剪枝效果时，允许退化到现有路径。
+- 当 `residual_filter && topn_candidate_filter` 全过滤时，skip final rest stream。
+- 当 `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。
+- 当 `residual_filter && topn_candidate_filter` 部分通过时，compose 回 Stage 0 原始坐标并读取 final rest columns。
 
 ### Step 5: Runtime Stats 和可观测性
 
-- 保留现有 TableScan/Selection actRows 语义。
-- 可选增加 debug log，输出 TopN candidate rows、heap size、filtered rows。
-- 可选增加 scan context counters，用于统计 TopN filter 节省的 final rest rows。
+- 保留 `stage1_output_rows` internal counter，表示 residual Selection 通过的真实行数。
+- 增加 `topn_candidate_rows` internal counter，表示 `residual_filter && topn_candidate_filter` 之后实际读取 final rest columns 的行数。
+- 在 TopN-enhanced MSLM 启用时，可以用 `topn_candidate_rows` overwrite residual Selection executor 的 `actRows`，让 `EXPLAIN ANALYZE` 直接展示 final rest columns materialized rows。
+- 该 overwrite 会改变 Selection `actRows` 的展示语义：它不再表示纯 residual Selection 后的逻辑行数，而是表示 storage 内输出给上层 TopN 的 candidate rows。
+- TableScan executor 的 `actRows` 仍然使用 `stage0_output_rows`。
+- 可选增加 debug log，输出 residual passed rows、TopN candidate rows、heap size、filtered rows。
 
 ## 测试方案
 
@@ -434,6 +515,7 @@ order_by_column_count <= 4
 - `WHERE + ORDER BY + LIMIT`，其中 WHERE 同时包含 Stage 0 pushed filter 和 Stage 1 residual filter。
 - ORDER BY column 不在最终 projection 中。
 - ORDER BY column 同时属于 residual filter columns。
+- ORDER BY column 同时属于最终输出列，验证 final rest columns 不重复读取该列。
 - 多 block 数据，确保 running local TopN 可以跨 block 收紧 threshold。
 - 多 stream 并发，确保上层 TopN 保证最终结果。
 
@@ -447,12 +529,16 @@ order_by_column_count <= 4
 - residual filter 全过滤。
 - TopN candidate filter 全过滤。
 - TopN candidate filter 全通过。
+- ORDER BY plain column 按 column id 对齐到 storage column，避免依赖 child schema name。
+- ORDER BY column 需要 TableScan 后 extra cast 时禁用优化。
+- 没有 residual Selection 时禁用优化。
+- 没有 TableScan pushed-down filters 时禁用优化。
 
 ### 回归测试
 
 - MSLM 原有测试必须全部通过。
 - TopN 上方的 projection 仍然正确。
-- EXPLAIN ANALYZE actRows 不出现明显异常。
+- TopN-enhanced MSLM 启用时，EXPLAIN ANALYZE 中 residual Selection `actRows` 展示 TopN candidate rows，即 final rest columns materialized rows。
 - 禁用条件下必须走原路径。
 
 ## 性能评估
@@ -491,6 +577,17 @@ order_by_column_count <= 4
 - 复用现有排序比较逻辑。
 - 对 ASC/DESC、NULL、collation 做单测。
 
+### ORDER BY 列对齐错误
+
+如果直接用 TopN child schema name 匹配 DeltaMerge storage column name，可能因为 TableScan projection、特殊列名或 cast-after-TS 造成错配。
+
+缓解：
+
+- 只支持 direct `ColumnRef`。
+- 使用 `ColumnRef` index 映射到 `table_scan.getColumns()[index].id`。
+- 后续所有 Stage 1 column union、final rest column subtraction 和 storage-side sort description 都基于 `ColumnID` / `ColumnDefine` 构造。
+- 禁用无法按 column id 映射到 storage column 的场景。
+
 ### Selection 顺序错误
 
 如果 TopN filter 在某些 Selection predicates 执行之前生效，可能过滤掉最终应该进入 TopN 的 rows。

From 639f1eabb16686c08d08080315fb9ea9bfccaaaa Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Mon, 27 Jul 2026 21:29:28 +0800
Subject: [PATCH 04/18] docs: refine MSLM running local TopN plan

---
 .../2026-07-26-running-local-topn-for-mslm.md | 174 ++++++++++++++++--
 1 file changed, 157 insertions(+), 17 deletions(-)

diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index f3c4a22eed1..f8c7027ceb2 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -75,7 +75,8 @@ LIMIT k;
 - 只支持已经满足现有 MSLM 启用条件的查询。也就是说，TableScan 必须已有 pushed-down filters 作为 Stage 0，且 TableScan 上方必须有 residual Selection 作为 Stage 1。
 - 只支持 `ORDER BY` plain columns，即 TiPB `ByItem.expr` 必须直接是 `ColumnRef`。
 - 只支持 constant `LIMIT` 和 constant `OFFSET`。
-- 只支持单表 `TopN -> Selection -> TableScan`。在 TiFlash physical plan 中，`Selection` 通常已经被 push 到 `PhysicalTableScan::filter_conditions`，所以实际识别形态可以是 `TopN -> PhysicalTableScan(with filter_conditions)`。
+- 只支持 TiDB 下发的单表 `TopN -> Selection -> TableScan`。在 TiFlash physical plan 中，这一层 `Selection` 必须已经被合并到 `PhysicalTableScan::filter_conditions`，所以第一版实际只识别 `PhysicalTopN -> PhysicalTableScan(with filter_conditions and pushed_down_filters)`。
+- 第一版不穿透 `PhysicalFilter`。如果 TiFlash physical plan 中仍然存在独立的 `PhysicalFilter -> PhysicalTableScan`，说明 residual Selection 没有进入当前 MSLM Stage 1 识别范围，TopN-enhanced MSLM 直接禁用。
 - 不移除上层 TopN。
 - 只在 MSLM 已经满足启用条件时启用。
 - TopN filter 只用于减少 final rest columns 的读取。
@@ -208,9 +209,9 @@ struct MSLMTopNDescription
 
 第一版 `order_by_columns` 只允许 plain columns，不构造复杂 expression actions。
 
-Plan 识别建议在 `PhysicalTopN` 和 `PhysicalTableScan` 之间完成。`PhysicalTopN` 只在 child 是 `PhysicalTableScan` 时尝试生成 MSLM TopN 描述；如果 child 是 Projection、Join、Aggregation、Window 等其他节点，则不启用该优化。上层 `PhysicalTopN` executor 仍然保留，storage 内描述只用于减少 final rest columns 的读取。
+Plan 识别建议在 `PhysicalTopN` 和 `PhysicalTableScan` 之间完成。`PhysicalTopN` 只在 child 直接是 `PhysicalTableScan` 时尝试生成 MSLM TopN 描述；如果 child 是 `PhysicalFilter`、Projection、Join、Aggregation、Window 等其他节点，则不启用该优化。上层 `PhysicalTopN` executor 仍然保留，storage 内描述只用于减少 final rest columns 的读取。
 
-因为 TiFlash planner 会尽量把单表 `Selection` push 到 `PhysicalTableScan::filter_conditions`，所以第一版实际需要识别的 physical plan 形态通常是：
+对于 TiPB DAG 中的单层 `TopN -> Selection -> TableScan`，TiFlash physical planner 会先尝试把 `Selection` 合并到 `PhysicalTableScan::filter_conditions`。因此第一版实际需要识别的 physical plan 形态是：
 
 ```text
 PhysicalTopN
@@ -262,6 +263,21 @@ plain column 还需要保证 storage 内比较语义和上层 TopN 一致。第
 - order-by column 是 generated column 或 virtual `_tidb_tid`。
 - order-by column 需要 TableScan 后的 extra cast 才能得到上层 TopN 的比较值。
 - order-by column 是 variable-length string 且暂未确认 collator 和 owned key 内存上界。
+- order-by column 不在第一版支持的类型白名单内。
+
+第一版 order-by column 类型白名单为：
+
+- integer numeric types。
+- decimal types。
+- date/datetime types。
+- float types。
+
+第一版明确不支持：
+
+- `String` / `FixedString` / `Bytes`，因为需要处理 collation 和 owned key 内存上界。
+- `Timestamp`，因为它通常需要 TableScan 后的 timezone extra cast，storage 内底层值不一定等于上层 TopN 的比较值。
+- `Time` / `Duration`，除非后续明确确认它和上层 TopN 使用完全一致的底层比较表示。
+- JSON、Enum/Set、Bit、Vector、Array、Tuple、Map 等其他类型。
 
 ### Stage 1 列集合
 
@@ -362,7 +378,7 @@ current block filter memory = O(current_block_rows)
 
 历史 candidate rows 一旦输出给上层 TopN，MSLM stream 不再负责保存它们。Heap eviction 只释放 owned sort key。
 
-对于 fixed-size order-by 类型，owned key 可以直接拷贝值和 NULL flag。对于 variable-length 类型，例如 String，owned key 可能带来较大内存开销。第一版可以先禁用 variable-length order-by columns，或者增加 per-stream memory guard。
+对于第一版支持的 order-by 类型，owned key 只拷贝 fixed-size value 和 NULL flag。对于 variable-length 类型，例如 String，owned key 可能带来较大内存开销，并且需要处理 collation。第一版禁用 variable-length order-by columns；后续如果要支持，可以再增加 per-stream memory guard 或 owned collation key。
 
 如果采用 memory guard，触发后应退化为普通 MSLM：
 
@@ -398,6 +414,102 @@ Heap 大小最多为 `topk`。如果 `topk` 超过启发式阈值，则禁用该
 
 第一版不要求 heap entry 能定位历史 block。Heap entry 的 locator 只服务于 current block eviction。历史 entry 被 eviction 时只释放 key，不访问历史 block。
 
+### Sort Key 表示
+
+第一版不使用 order-preserving bytes key。TiFlash 现有 sort 路径本身也是基于 `IColumn::compareAt()` / typed compare 直接比较 column values，而不是先构造可按 bytes 比较的 sort key。引入 bytes key 需要重新设计 signed integer、float、NaN、DESC、NULL、decimal 等类型的 order-preserving encoding，正确性风险较高。
+
+第一版也不直接使用通用 `Field` 作为 heap key。`Field` 虽然已经是 owned variant，但它支持 String/Array/Tuple 等当前不需要的类型，大小和通用抽取路径都偏重；float 的 NaN 比较语义也需要额外绕开。为了降低 hot path 开销，第一版使用一个更窄的 typed sort key：
+
+```text
+static constexpr size_t max_sort_key_columns = 4;
+
+enum class SortKeyKind
+{
+    Int64,
+    UInt64,
+    Float32,
+    Float64,
+    Decimal32,
+    Decimal64,
+    Decimal128,
+    Decimal256,
+    Date,
+    DateTime,
+};
+
+struct SortKeyField
+{
+    bool is_null;
+    SortKeyKind kind;
+    TypedFixedSizeValue value;
+};
+
+struct OwnedSortKey
+{
+    std::array<SortKeyField, max_sort_key_columns> fields;
+    size_t size;
+};
+```
+
+`TypedFixedSizeValue` 是 conceptual storage，实际实现可以使用 union 或其他 fixed-size typed storage。它只需要覆盖第一版白名单内的 fixed-size values，不支持 String/Array/Tuple。
+
+`RunningLocalTopN` 本身不模板化。每个 order-by column 在初始化时生成一个 descriptor，descriptor 中保存 typed extractor 和 comparator：
+
+```text
+struct SortKeyColumnDesc
+{
+    String column_name;
+    size_t column_pos;
+    SortKeyKind kind;
+    int direction;
+
+    extract_owned(column, row, SortKeyField &);
+    compare_owned(SortKeyField, SortKeyField);
+    compare_column_with_owned(column, row, SortKeyField);
+};
+```
+
+模板只用于生成具体类型的 extractor/comparator。这样一个 `OwnedSortKey` 可以同时包含多个不同类型的 order-by columns，而不需要把整个 `RunningLocalTopN` 或 `SortKeyField` 做成模板。
+
+### Sort Key 比较语义
+
+NULL 在所有支持类型中统一处理：
+
+```text
+NULL is minimum before applying direction.
+ASC  -> NULL first
+DESC -> NULL last
+```
+
+非 NULL values 的比较规则：
+
+- integer numeric、date、datetime 直接比较底层整数值。
+- decimal 直接比较底层 decimal value。因为 lhs/rhs 来自同一个 order-by column，decimal type 和 scale 一致。
+- float 使用 TiFlash 现有 `CompareHelper<Float32/Float64>::compare(..., nan_direction_hint=-1)` 语义，避免 NaN 行为和上层 sort 不一致。
+
+每个字段先得到不带 direction 的比较结果，再乘以 `SortKeyColumnDesc::direction`。多列 ORDER BY 按顺序比较，直到第一个非 0 结果。
+
+### Lazy Owned Key Materialization
+
+第一版不对 residual passed 的每一行都构造 `OwnedSortKey`。处理当前 `stage1_block` 时，可以直接使用 order-by columns 参与比较；只有当前 row 真的需要进入 heap 时，才 lazy materialize owned key。
+
+处理一个 block 前，先从 `stage1_block` 中收集 order-by column 指针：
+
+```text
+sort_columns[i] = stage1_block.getByPosition(desc[i].column_pos).column.get()
+```
+
+per-row hot path 不再从 `Block` 里按 name 或 position 查找 column，而是直接使用预先收集好的 `IColumn *`：
+
+```text
+compareRowWithOwnedKey(sort_columns, row, heap.worst().key)
+materializeOwnedKey(sort_columns, row)
+```
+
+heap 中始终只保存 owned key，不保存 `Block`、`IColumn *` 或 historical row reference。当前 row 只有在 `update()` 函数内部比较时可以直接读 `stage1_block` column。这样既避免 retain historical blocks，又避免给明显不可能进入 local TopK 的 rows 做 key copy。
+
+当前 block 内如果某个 entry 被后续更优 row eviction，可以通过 `block_sequence` 和 `row_index_in_stage1_block` 将对应的 `topn_candidate_filter` 置为 false。历史 block 的 entry 被 eviction 时不能回撤，因为 rows 已经输出给上层 TopN。
+
 ## Filter 坐标
 
 MSLM 中存在多套 filter 坐标：
@@ -428,7 +540,7 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 
 - `dt_enable_multi_stage_late_materialization` 为 true。
 - 当前查询已经满足 MSLM 启用条件。
-- 查询包含 TopN，且 TopN 的 child 是单表 `PhysicalTableScan`。
+- 查询包含 TopN，且 TopN 的 child 直接是单表 `PhysicalTableScan`。
 - `PhysicalTableScan` 已经包含 residual `filter_conditions`，并且 TableScan 已经包含 pushed-down filters。
 - TopN 上方仍保留全局 TopN executor。
 - TopN 的 `LIMIT` 和 `OFFSET` 是常量。
@@ -437,6 +549,7 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 - ORDER BY 只包含 direct `ColumnRef` plain columns。
 - ORDER BY columns 能通过 `ColumnRef` index 映射到 `table_scan.getColumns()` 中的 column id，并能进一步映射到 DeltaMerge `ColumnDefine`。
 - ORDER BY columns 不需要 TableScan 后 extra cast。
+- ORDER BY columns 类型属于第一版白名单：integer numeric、decimal、date、datetime、float。
 - ORDER BY columns 可加入 Stage 1 column set。
 - ORDER BY columns 数量不超过阈值。
 - final rest columns 数量满足 MSLM 的宽列启发式要求。
@@ -455,9 +568,11 @@ order_by_column_count <= 4
 - ORDER BY expression 不是 plain column。
 - ORDER BY plain column 无法按 column id 映射到 storage column。
 - ORDER BY column 需要 TableScan 后 extra cast。
+- ORDER BY column 类型不在第一版白名单内，例如 String、FixedString、Timestamp、Time/Duration、JSON、Enum/Set、Bit、Vector、Array、Tuple、Map。
 - ORDER BY 依赖 generated column。
 - TopN 包含非常大的 offset。
 - TopN 之前仍有无法进入 MSLM 的 Selection。
+- TopN 的 child 是 `PhysicalFilter`，即 residual Selection 没有被合并到 `PhysicalTableScan::filter_conditions`。
 - 查询没有 residual Selection，或者没有 TableScan pushed-down filters。
 - 查询要求 keep order 且和 MSLM 现有限制冲突。
 - TopN 位于 join、aggregation、window 等算子之上。
@@ -468,26 +583,50 @@ order_by_column_count <= 4
 
 - 明确 MVP 范围。
 - 增加基于 DAG request 的单测设计。
-- 明确 ties、offset、NULL ordering、collation 的测试覆盖。
+- 明确 ties、offset、NULL ordering、float NaN、以及 unsupported string/collation 的禁用测试覆盖。
 
-### Step 2: Plan 识别和元数据传递
+### Step 2: Pattern 识别和元数据传递
 
-- 在 `PhysicalTopN` 构建阶段识别 child 是 `PhysicalTableScan` 的模式。
+- 在 `PhysicalTopN` 构建阶段识别 child 直接是 `PhysicalTableScan` 的模式。
 - 只在 `PhysicalTableScan` 同时包含 residual `filter_conditions` 和 TableScan pushed-down filters 时启用。
+- 如果 child 是 `PhysicalFilter` 或其他节点，第一版不尝试穿透或重写。
 - 构造 `MSLMTopNDescription`。
-- 按 `ColumnRef` index -> `table_scan.getColumns()[index].id` -> `ColumnDefine` 的顺序对齐 order-by columns。
-- 将 order-by columns 合并到 MSLM Stage 1 column set。
-- 使用 Stage 1 storage column names 构造 storage-side `SortDescription`。
 - 保留上层 TopN executor，不修改最终 DAG 语义。
 
-### Step 3: RunningLocalTopN 状态
+### Step 3: ORDER BY ColumnRef 对齐和 Stage 1 列集合
 
-- 实现 per-stream bounded TopK heap。
-- 支持 plain column sort key 提取。
+- 要求每个 `ByItem.expr` 都是 direct `ColumnRef`。
+- 按 `ColumnRef` index -> `table_scan.getColumns()[index].id` -> `ColumnDefine` 的顺序对齐 order-by columns。
+- 禁用无法按 column id 映射到 storage column 的 order-by columns。
+- 禁用 generated column、virtual `_tidb_tid`、需要 TableScan 后 extra cast 的 order-by columns。
+- 校验 order-by column 类型属于第一版白名单。
+- 将 order-by columns 按 ColumnID 合并到 MSLM Stage 1 column set。
+- final rest columns 按 ColumnID 排除 Stage 1 columns，避免 order-by columns 重复读取。
+- 使用 Stage 1 storage column names 构造 storage-side order-by descriptor。
+
+### Step 4: Typed Sort Key Infrastructure
+
+- 实现 fixed-size typed `SortKeyField` / `OwnedSortKey`。
+- 实现 `SortKeyColumnDesc`，保存 column position、direction、type kind、extractor 和 comparator。
+- 支持 integer numeric、decimal、date、datetime、float。
+- Nullable wrapper 统一由 extractor/comparator 处理，nested type 必须在白名单内。
+- integer numeric、date、datetime 直接比较底层整数值。
+- decimal 直接比较底层 decimal value。
+- float comparator 使用现有 `CompareHelper` 的 NaN 语义。
+- 不支持 String/FixedString/collation，不使用 order-preserving bytes key，不直接使用通用 `Field` 作为 heap key。
+
+### Step 5: RunningLocalTopN 状态和算法
+
+- 实现 per-stream bounded TopK heap，heap top 是当前 worst candidate。
+- heap entry 只保存 owned key、stream/block sequence 和 current-block row index。
+- 每个 block 开始时预先收集 order-by column pointers，per-row compare 不从 `Block` 查找 column。
+- 支持 row-to-owned lazy compare，只在 row 需要进入 heap 时 materialize owned key。
 - 支持 ASC/DESC、NULL ordering 和 TiFlash 现有 comparator 语义。
 - 生成当前 block 的 `topn_candidate_filter`。
+- 支持 current block eviction，将被当前 block 后续更优 row 挤出的 candidate filter 清掉。
+- 历史 block entry eviction 不做回撤。
 
-### Step 4: 接入 MSLM final rest 读取
+### Step 6: 接入 MSLM final rest 读取
 
 - 在 residual filter 之后执行 running local TopN。
 - 将 `residual_filter` 和 `topn_candidate_filter` 合并。
@@ -496,7 +635,7 @@ order_by_column_count <= 4
 - 当 `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。
 - 当 `residual_filter && topn_candidate_filter` 部分通过时，compose 回 Stage 0 原始坐标并读取 final rest columns。
 
-### Step 5: Runtime Stats 和可观测性
+### Step 7: Runtime Stats 和可观测性
 
 - 保留 `stage1_output_rows` internal counter，表示 residual Selection 通过的真实行数。
 - 增加 `topn_candidate_rows` internal counter，表示 `residual_filter && topn_candidate_filter` 之后实际读取 final rest columns 的行数。
@@ -575,7 +714,8 @@ order_by_column_count <= 4
 
 - MVP 只支持 plain columns。
 - 复用现有排序比较逻辑。
-- 对 ASC/DESC、NULL、collation 做单测。
+- 对 ASC/DESC、NULL、float NaN、decimal/date/datetime 做正确性单测。
+- 对 String/FixedString/collation 做禁用路径单测。
 
 ### ORDER BY 列对齐错误
 

From b0f4fce9e5536034ab4199ef84a1a7fc9da625c8 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Mon, 27 Jul 2026 22:28:17 +0800
Subject: [PATCH 05/18] Implement running local TopN for MSLM

---
 dbms/src/Debug/MockStorage.cpp                |  18 +-
 dbms/src/Debug/MockStorage.h                  |   7 +-
 .../Coprocessor/DAGStorageInterpreter.cpp     |  15 +-
 .../Flash/Coprocessor/DAGStorageInterpreter.h |   5 +-
 .../Planner/Plans/PhysicalMockTableScan.cpp   |  15 +-
 .../Planner/Plans/PhysicalMockTableScan.h     |   9 +
 .../Flash/Planner/Plans/PhysicalTableScan.cpp |  20 +-
 .../Flash/Planner/Plans/PhysicalTableScan.h   |   7 +
 dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp |  94 +++++
 .../Flash/tests/gtest_executors_with_dm.cpp   | 136 +++++++
 .../DeltaMerge/Filter/PushDownFilter.h        |   8 +-
 ...ageLateMaterializationBlockInputStream.cpp |  40 ++
 ...StageLateMaterializationBlockInputStream.h |   1 +
 ...ultiStageLateMaterializationRuntimeStats.h |   1 +
 .../MultiStageLateMaterializationTopN.cpp     | 341 ++++++++++++++++++
 .../MultiStageLateMaterializationTopN.h       | 139 +++++++
 dbms/src/Storages/SelectQueryInfo.cpp         |   2 +
 dbms/src/Storages/SelectQueryInfo.h           |   3 +
 dbms/src/Storages/StorageDeltaMerge.cpp       | 154 +++++++-
 19 files changed, 996 insertions(+), 19 deletions(-)
 create mode 100644 dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
 create mode 100644 dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h

diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp
index 9a2c1d04db3..cf346e16e42 100644
--- a/dbms/src/Debug/MockStorage.cpp
+++ b/dbms/src/Debug/MockStorage.cpp
@@ -376,7 +376,8 @@ BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge(
     std::vector<int> runtime_filter_ids,
     int rf_max_wait_time_ms,
     const google::protobuf::RepeatedPtrField<tipb::Expr> * pushed_down_filters,
-    const TiDB::ColumnInfos * table_scan_column_infos)
+    const TiDB::ColumnInfos * table_scan_column_infos,
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn)
 {
     static const google::protobuf::RepeatedPtrField<tipb::Expr> empty_pushed_down_filters{};
     static const auto empty_ann_query_info = tipb::ANNQueryInfo{};
@@ -400,6 +401,7 @@ BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge(
             runtime_filter_ids,
             rf_max_wait_time_ms,
             context.getTimezoneInfo());
+        query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn;
         BlockInputStreams ins = storage->read(
             column_names,
             query_info,
@@ -476,7 +478,8 @@ void MockStorage::buildExecFromDeltaMerge(
     int rf_max_wait_time_ms,
     const google::protobuf::RepeatedPtrField<tipb::Expr> * pushed_down_filters,
     const String & table_scan_executor_id,
-    const TiDB::ColumnInfos * table_scan_column_infos)
+    const TiDB::ColumnInfos * table_scan_column_infos,
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn)
 {
     static const google::protobuf::RepeatedPtrField<tipb::Expr> empty_pushed_down_filters{};
     static const auto empty_ann_query_info = tipb::ANNQueryInfo{};
@@ -511,9 +514,13 @@ void MockStorage::buildExecFromDeltaMerge(
                         &multi_stage_late_materialization_runtime_stats->stage0_output_rows));
                 dag_context->setExecutorRowsOverride(
                     filter_conditions->executor_id,
-                    std::shared_ptr<std::atomic<UInt64>>(
-                        multi_stage_late_materialization_runtime_stats,
-                        &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+                    multi_stage_late_materialization_topn != nullptr
+                        ? std::shared_ptr<std::atomic<UInt64>>(
+                            multi_stage_late_materialization_runtime_stats,
+                            &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
+                        : std::shared_ptr<std::atomic<UInt64>>(
+                            multi_stage_late_materialization_runtime_stats,
+                            &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
             }
         }
         query_info.dag_query = std::make_unique<DAGQueryInfo>(
@@ -526,6 +533,7 @@ void MockStorage::buildExecFromDeltaMerge(
             context.getTimezoneInfo());
         query_info.enable_multi_stage_late_materialization = enable_multi_stage_late_materialization;
         query_info.multi_stage_late_materialization_runtime_stats = multi_stage_late_materialization_runtime_stats;
+        query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn;
         storage->read(
             exec_context_,
             group_builder,
diff --git a/dbms/src/Debug/MockStorage.h b/dbms/src/Debug/MockStorage.h
index 1259b9ad1dc..d149b8565a2 100644
--- a/dbms/src/Debug/MockStorage.h
+++ b/dbms/src/Debug/MockStorage.h
@@ -20,6 +20,7 @@
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
 #include <Operators/Operator.h>
 #include <Storages/DeltaMerge/ColumnDefine_fwd.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 #include <TiDB/Schema/TiDB_fwd.h>
 #include <common/types.h>
 
@@ -107,7 +108,8 @@ class MockStorage
         std::vector<int> runtime_filter_ids = std::vector<int>(),
         int rf_max_wait_time_ms = 0,
         const google::protobuf::RepeatedPtrField<tipb::Expr> * pushed_down_filters = nullptr,
-        const TiDB::ColumnInfos * scan_column_infos = nullptr);
+        const TiDB::ColumnInfos * scan_column_infos = nullptr,
+        const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn = nullptr);
 
     void buildExecFromDeltaMerge(
         PipelineExecutorContext & exec_context_,
@@ -121,7 +123,8 @@ class MockStorage
         int rf_max_wait_time_ms = 0,
         const google::protobuf::RepeatedPtrField<tipb::Expr> * pushed_down_filters = nullptr,
         const String & table_scan_executor_id = "",
-        const TiDB::ColumnInfos * scan_column_infos = nullptr);
+        const TiDB::ColumnInfos * scan_column_infos = nullptr,
+        const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn = nullptr);
 
     bool tableExistsForDeltaMerge(Int64 table_id);
 
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index 0a5307cb0c3..5c6852fe7b8 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -382,11 +382,13 @@ DAGStorageInterpreter::DAGStorageInterpreter(
     Context & context_,
     const TiDBTableScan & table_scan_,
     const FilterConditions & filter_conditions_,
-    size_t max_streams_)
+    size_t max_streams_,
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn_)
     : context(context_)
     , table_scan(table_scan_)
     , filter_conditions(filter_conditions_)
     , max_streams(max_streams_)
+    , multi_stage_late_materialization_topn(multi_stage_late_materialization_topn_)
     , log(Logger::get(context.getDAGContext()->log ? context.getDAGContext()->log->identifier() : ""))
     , logical_table_id(table_scan.getLogicalTableID())
     , tmt(context.getTMTContext())
@@ -1052,9 +1054,13 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
                 &multi_stage_late_materialization_runtime_stats->stage0_output_rows));
         dagContext().setExecutorRowsOverride(
             filter_conditions.executor_id,
-            std::shared_ptr<std::atomic<UInt64>>(
-                multi_stage_late_materialization_runtime_stats,
-                &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+            multi_stage_late_materialization_topn != nullptr
+                ? std::shared_ptr<std::atomic<UInt64>>(
+                    multi_stage_late_materialization_runtime_stats,
+                    &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
+                : std::shared_ptr<std::atomic<UInt64>>(
+                    multi_stage_late_materialization_runtime_stats,
+                    &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
     }
 
     auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo {
@@ -1074,6 +1080,7 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
         query_info.is_fast_scan = table_scan.isFastScan();
         query_info.enable_multi_stage_late_materialization = enable_multi_stage_late_materialization;
         query_info.multi_stage_late_materialization_runtime_stats = multi_stage_late_materialization_runtime_stats;
+        query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn;
         return query_info;
     };
     RUNTIME_CHECK_MSG(mvcc_query_info->scan_context != nullptr, "Unexpected null scan_context");
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h
index 40dcdd72631..dc0bf1e3e0d 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h
@@ -22,6 +22,7 @@
 #include <Flash/Coprocessor/RemoteRequest.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 #include <Storages/DeltaMerge/Remote/DisaggSnapshot_fwd.h>
 #include <Storages/KVStore/Read/LearnerRead.h>
 #include <Storages/KVStore/Read/RegionException.h>
@@ -47,7 +48,8 @@ class DAGStorageInterpreter
         Context & context_,
         const TiDBTableScan & table_scan,
         const FilterConditions & filter_conditions_,
-        size_t max_streams_);
+        size_t max_streams_,
+        const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn_ = nullptr);
 
     ~DAGStorageInterpreter();
 
@@ -146,6 +148,7 @@ class DAGStorageInterpreter
     const TiDBTableScan & table_scan;
     const FilterConditions & filter_conditions;
     const size_t max_streams;
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn;
     LoggerPtr log;
 
     /// derived from other members, doesn't change during DAGStorageInterpreter's lifetime
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp
index 6597e82a52f..72b013838da 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp
@@ -73,7 +73,8 @@ std::pair<NamesAndTypes, BlockInputStreams> mockSchemaAndStreams(
             table_scan.getRuntimeFilterIDs(),
             10000,
             &table_scan.getPushedDownFilters(),
-            use_table_scan_columns ? &table_scan.getColumns() : nullptr));
+            use_table_scan_columns ? &table_scan.getColumns() : nullptr,
+            nullptr));
     }
     else
     {
@@ -184,7 +185,8 @@ void PhysicalMockTableScan::buildPipelineExecGroupImpl(
             rf_max_wait_time_ms,
             &pushed_down_filters,
             execId(),
-            use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr);
+            use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr,
+            multi_stage_late_materialization_topn);
         for (size_t i = 0; i < group_builder.concurrency(); ++i)
         {
             if (auto * source_op = dynamic_cast<UnorderedSourceOp *>(group_builder.getCurBuilder(i).source_op.get()))
@@ -240,7 +242,8 @@ bool PhysicalMockTableScan::setFilterConditions(
             runtime_filter_ids,
             10000,
             &pushed_down_filters,
-            use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr));
+            use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr,
+            multi_stage_late_materialization_topn));
 
         return true;
     }
@@ -261,6 +264,12 @@ const String & PhysicalMockTableScan::getFilterConditionsId() const
     return filter_conditions.executor_id;
 }
 
+void PhysicalMockTableScan::setMultiStageLateMaterializationTopN(
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn)
+{
+    multi_stage_late_materialization_topn = topn;
+}
+
 void PhysicalMockTableScan::buildRuntimeFilterInLocalStream(Context & context)
 {
     for (const auto & local_stream : mock_streams)
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h
index 2b2c2d29548..246297727c1 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h
+++ b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h
@@ -19,6 +19,7 @@
 #include <Flash/Coprocessor/RuntimeFilterMgr.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Flash/Planner/Plans/PhysicalLeaf.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 #include <tipb/executor.pb.h>
 
 namespace DB
@@ -63,6 +64,12 @@ class PhysicalMockTableScan : public PhysicalLeaf
 
     const String & getFilterConditionsId() const;
 
+    const TiDB::ColumnInfos & getTableScanColumns() const { return table_scan_columns; }
+
+    const google::protobuf::RepeatedPtrField<tipb::Expr> & getPushedDownFilters() const { return pushed_down_filters; }
+
+    void setMultiStageLateMaterializationTopN(const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn);
+
 private:
     void buildBlockInputStreamImpl(DAGPipeline & pipeline, Context & /*context*/, size_t /*max_streams*/) override;
 
@@ -92,6 +99,8 @@ class PhysicalMockTableScan : public PhysicalLeaf
 
     TiDB::ColumnInfos table_scan_columns;
 
+    DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn;
+
     bool use_table_scan_columns_for_delta_merge;
 
     const int rf_max_wait_time_ms = 10000;
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp
index 6165bbe448e..2076961ef08 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp
@@ -114,7 +114,12 @@ void PhysicalTableScan::buildBlockInputStreamImpl(DAGPipeline & pipeline, Contex
     }
     else
     {
-        DAGStorageInterpreter storage_interpreter(context, tidb_table_scan, filter_conditions, max_streams);
+        DAGStorageInterpreter storage_interpreter(
+            context,
+            tidb_table_scan,
+            filter_conditions,
+            max_streams,
+            multi_stage_late_materialization_topn);
         storage_interpreter.execute(pipeline);
     }
     buildProjection(pipeline);
@@ -137,7 +142,12 @@ void PhysicalTableScan::buildPipeline(
     }
     else
     {
-        DAGStorageInterpreter storage_interpreter(context, tidb_table_scan, filter_conditions, context.getMaxStreams());
+        DAGStorageInterpreter storage_interpreter(
+            context,
+            tidb_table_scan,
+            filter_conditions,
+            context.getMaxStreams(),
+            multi_stage_late_materialization_topn);
         storage_interpreter.execute(exec_context, pipeline_exec_builder);
     }
     buildProjection(exec_context, pipeline_exec_builder);
@@ -214,4 +224,10 @@ const String & PhysicalTableScan::getFilterConditionsId() const
     RUNTIME_CHECK(hasFilterConditions());
     return filter_conditions.executor_id;
 }
+
+void PhysicalTableScan::setMultiStageLateMaterializationTopN(
+    const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn)
+{
+    multi_stage_late_materialization_topn = topn;
+}
 } // namespace DB
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h
index 284b75d0c73..bd654def8a8 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h
@@ -18,6 +18,7 @@
 #include <Flash/Coprocessor/FilterConditions.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Flash/Planner/Plans/PhysicalLeaf.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 #include <tipb/executor.pb.h>
 
 namespace DB
@@ -48,6 +49,10 @@ class PhysicalTableScan : public PhysicalLeaf
 
     const String & getFilterConditionsId() const;
 
+    const TiDBTableScan & getTiDBTableScan() const { return tidb_table_scan; }
+
+    void setMultiStageLateMaterializationTopN(const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn);
+
     void buildPipeline(PipelineBuilder & builder, Context & context, PipelineExecutorContext & exec_context) override;
 
 private:
@@ -67,6 +72,8 @@ class PhysicalTableScan : public PhysicalLeaf
 
     TiDBTableScan tidb_table_scan;
 
+    DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn;
+
     Block sample_block;
 
     PipelineExecGroupBuilder pipeline_exec_builder;
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
index e6b3e7022f8..c2320dcf2e8 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
@@ -13,18 +13,111 @@
 // limitations under the License.
 
 #include <Common/Logger.h>
+#include <Flash/Coprocessor/DAGCodec.h>
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Coprocessor/DAGExpressionAnalyzer.h>
 #include <Flash/Coprocessor/DAGPipeline.h>
+#include <Flash/Coprocessor/DAGUtils.h>
 #include <Flash/Coprocessor/InterpreterUtils.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
 #include <Flash/Planner/FinalizeHelper.h>
 #include <Flash/Planner/PhysicalPlanHelper.h>
+#include <Flash/Planner/Plans/PhysicalMockTableScan.h>
+#include <Flash/Planner/Plans/PhysicalTableScan.h>
 #include <Flash/Planner/Plans/PhysicalTopN.h>
 #include <Interpreters/Context.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
+#include <fmt/core.h>
 
 namespace DB
 {
+namespace
+{
+DM::MultiStageLateMaterializationTopNDescriptionPtr tryBuildMultiStageLateMaterializationTopN(
+    const tipb::TopN & top_n,
+    const TiDB::ColumnInfos & table_scan_columns,
+    const google::protobuf::RepeatedPtrField<tipb::Expr> & pushed_down_filters,
+    bool has_filter_conditions,
+    const LoggerPtr & log)
+{
+    auto disable = [&](const String & reason) -> DM::MultiStageLateMaterializationTopNDescriptionPtr {
+        LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason={}", reason);
+        return nullptr;
+    };
+
+    if (!has_filter_conditions)
+        return disable("no residual filter conditions");
+    if (pushed_down_filters.empty())
+        return disable("no stage0 pushed down filters");
+    if (top_n.limit() == 0)
+        return disable("topk is zero");
+    if (top_n.limit() > DM::multi_stage_late_materialization_topn_max_topk)
+        return disable(fmt::format("topk is too large: {}", top_n.limit()));
+    if (top_n.order_by_size() > static_cast<int>(DM::multi_stage_late_materialization_topn_max_order_by_columns))
+        return disable(fmt::format("too many order by columns: {}", top_n.order_by_size()));
+
+    auto desc = std::make_shared<DM::MultiStageLateMaterializationTopNDescription>();
+    desc->topk = top_n.limit();
+    desc->order_by_columns.reserve(top_n.order_by_size());
+
+    for (const auto & by_item : top_n.order_by())
+    {
+        if (!isColumnExpr(by_item.expr()))
+            return disable("order by expression is not direct ColumnRef");
+
+        const auto column_index = decodeDAGInt64(by_item.expr().val());
+        if (column_index < 0 || column_index >= static_cast<Int64>(table_scan_columns.size()))
+            return disable(fmt::format(
+                "order by ColumnRef index out of range, index={}, table_scan_column_size={}",
+                column_index,
+                table_scan_columns.size()));
+
+        const auto & column = table_scan_columns[column_index];
+        if (column.hasGeneratedColumnFlag())
+            return disable(fmt::format("order by generated column is unsupported, column_id={}", column.id));
+        if (column.id == ExtraTableIDColumnID || column.id == ExtraCommitTSColumnID)
+            return disable(fmt::format("order by virtual column is unsupported, column_id={}", column.id));
+
+        desc->order_by_columns.push_back(
+            DM::MultiStageLateMaterializationTopNOrderByColumn{column.id, by_item.desc() ? -1 : 1});
+    }
+
+    return desc;
+}
+
+void tryAttachMultiStageLateMaterializationTopN(
+    const tipb::TopN & top_n,
+    const PhysicalPlanNodePtr & child,
+    const LoggerPtr & log)
+{
+    if (auto table_scan = std::dynamic_pointer_cast<PhysicalTableScan>(child))
+    {
+        auto topn = tryBuildMultiStageLateMaterializationTopN(
+            top_n,
+            table_scan->getTiDBTableScan().getColumns(),
+            table_scan->getTiDBTableScan().getPushedDownFilters(),
+            table_scan->hasFilterConditions(),
+            log);
+        table_scan->setMultiStageLateMaterializationTopN(topn);
+        return;
+    }
+
+    if (auto table_scan = std::dynamic_pointer_cast<PhysicalMockTableScan>(child))
+    {
+        auto topn = tryBuildMultiStageLateMaterializationTopN(
+            top_n,
+            table_scan->getTableScanColumns(),
+            table_scan->getPushedDownFilters(),
+            table_scan->hasFilterConditions(),
+            log);
+        table_scan->setMultiStageLateMaterializationTopN(topn);
+        return;
+    }
+
+    LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason=child is not direct table scan");
+}
+} // namespace
+
 PhysicalPlanNodePtr PhysicalTopN::build(
     const Context & context,
     const String & executor_id,
@@ -45,6 +138,7 @@ PhysicalPlanNodePtr PhysicalTopN::build(
 
     auto order_columns = analyzer.buildOrderColumns(before_sort_actions, top_n.order_by());
     SortDescription order_descr = getSortDescription(order_columns, top_n.order_by());
+    tryAttachMultiStageLateMaterializationTopN(top_n, child, log);
 
     auto physical_top_n = std::make_shared<PhysicalTopN>(
         executor_id,
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 8e6019b7774..3f985238bf2 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -149,6 +149,8 @@ tipb::TableScan * findMutableTableScan(tipb::Executor * executor)
 {
     switch (executor->tp())
     {
+    case tipb::ExecType::TypeTopN:
+        return findMutableTableScan(executor->mutable_topn()->mutable_child());
     case tipb::ExecType::TypeTableScan:
         return executor->mutable_tbl_scan();
     case tipb::ExecType::TypeSelection:
@@ -164,6 +166,8 @@ tipb::Selection * findMutableSelection(tipb::Executor * executor)
 {
     switch (executor->tp())
     {
+    case tipb::ExecType::TypeTopN:
+        return findMutableSelection(executor->mutable_topn()->mutable_child());
     case tipb::ExecType::TypeSelection:
         return executor->mutable_selection();
     case tipb::ExecType::TypeProjection:
@@ -364,6 +368,29 @@ std::shared_ptr<tipb::DAGRequest> buildDAGRequestWithPushedDownFilterAndProjecti
     return request;
 }
 
+std::shared_ptr<tipb::DAGRequest> buildDAGRequestWithPushedDownFilterAndTopN(
+    MockDAGRequestContext & context,
+    const String & table_name,
+    const ASTPtr & pushed_down_filter,
+    const ASTPtr & residual_filter,
+    const String & order_by_column,
+    bool is_desc,
+    UInt64 limit)
+{
+    auto pushed_down_filter_request = context.scan("test_db", table_name).filter(pushed_down_filter).build(context);
+    RUNTIME_CHECK(pushed_down_filter_request->root_executor().tp() == tipb::ExecType::TypeSelection);
+    RUNTIME_CHECK(pushed_down_filter_request->root_executor().selection().conditions_size() == 1);
+    const auto & pushed_down_filter_expr = pushed_down_filter_request->root_executor().selection().conditions(0);
+
+    auto request = context.scan("test_db", table_name)
+                       .filter(residual_filter)
+                       .topN(order_by_column, is_desc, limit)
+                       .build(context);
+    *findMutableTableScan(request->mutable_root_executor())->add_pushed_down_filter_conditions()
+        = pushed_down_filter_expr;
+    return request;
+}
+
 ASTPtr buildAllGreaterThanMinusOneFilter(size_t begin_col_id, size_t end_col_id)
 {
     RUNTIME_CHECK(begin_col_id < end_col_id);
@@ -583,6 +610,115 @@ try
 }
 CATCH
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNAsc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "c1",
+        /*is_desc=*/false,
+        /*limit=*/2);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto expected = executeStreams(&disabled_dag_context);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto actual = executeStreams(&enabled_dag_context);
+    ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
+        << "\n  expect_block: \n"
+        << getColumnsContent(expected) << "\n actual_block: \n"
+        << getColumnsContent(actual);
+
+    auto table_scan_rows = enabled_dag_context.getExecutorRowsOverride("table_scan_0");
+    auto selection_rows = enabled_dag_context.getExecutorRowsOverride("selection_1");
+    ASSERT_NE(table_scan_rows, nullptr);
+    ASSERT_NE(selection_rows, nullptr);
+    ASSERT_EQ(table_scan_rows->load(), 16);
+    ASSERT_GE(selection_rows->load(), 2);
+    ASSERT_LT(selection_rows->load(), 16);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForStringOrderBy)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    MockColumnInfoVec columns;
+    columns.push_back({"c0", TiDB::TP::TypeLongLong});
+    columns.push_back({"c1", TiDB::TP::TypeLongLong});
+    columns.push_back({"s", TiDB::TP::TypeString});
+    for (size_t col_id = 2; col_id <= 14; ++col_id)
+        columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong});
+
+    constexpr size_t rows = 32;
+    std::vector<Int64> c0_values;
+    std::vector<Int64> c1_values;
+    std::vector<String> s_values;
+    c0_values.reserve(rows);
+    c1_values.reserve(rows);
+    s_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+    {
+        c0_values.push_back(row);
+        c1_values.push_back(row % 8);
+        s_values.push_back(fmt::format("s{:02}", 31 - row));
+    }
+
+    ColumnsWithTypeAndName data;
+    data.emplace_back(toVec<Int64>("c0", c0_values));
+    data.emplace_back(toVec<Int64>("c1", c1_values));
+    data.emplace_back(toVec<String>("s", s_values));
+    for (size_t col_id = 2; col_id <= 14; ++col_id)
+    {
+        std::vector<Int64> values;
+        values.reserve(rows);
+        for (size_t row = 0; row < rows; ++row)
+            values.push_back(static_cast<Int64>(col_id * 1000 + row));
+        data.emplace_back(toVec<Int64>(fmt::format("c{}", col_id), values));
+    }
+    context.addMockDeltaMerge(
+        {"test_db", "multi_stage_lm_string_topn"},
+        columns,
+        data,
+        /*concurrency_hint=*/4);
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_string_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "s",
+        /*is_desc=*/false,
+        /*limit=*/2);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto expected = executeStreams(&disabled_dag_context);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto actual = executeStreams(&enabled_dag_context);
+    ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
+        << "\n  expect_block: \n"
+        << getColumnsContent(expected) << "\n actual_block: \n"
+        << getColumnsContent(actual);
+
+    assertMultiStageRowsOverride(enabled_dag_context, 16);
+}
+CATCH
+
 TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationEnabledAtRestColumnThreshold)
 try
 {
diff --git a/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h b/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h
index 538531afe0a..a813a47a08f 100644
--- a/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h
+++ b/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h
@@ -22,6 +22,7 @@
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Interpreters/ExpressionActions.h>
 #include <Storages/DeltaMerge/Filter/RSOperator.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 
 namespace DB
 {
@@ -45,7 +46,8 @@ class PushDownFilter
         const ColumnDefinesPtr & filter_columns_,
         const String filter_column_name_,
         const ExpressionActionsPtr & extra_cast_,
-        const ColumnDefinesPtr & columns_after_cast_)
+        const ColumnDefinesPtr & columns_after_cast_,
+        const MultiStageLateMaterializationTopNDescriptionPtr & topn_ = nullptr)
         : rs_operator(rs_operator_)
         , before_where(before_where_)
         , project_after_where(project_after_where_)
@@ -53,10 +55,12 @@ class PushDownFilter
         , filter_columns(filter_columns_)
         , extra_cast(extra_cast_)
         , columns_after_cast(columns_after_cast_)
+        , topn(topn_)
     {}
 
     explicit PushDownFilter(const RSOperatorPtr & rs_operator_)
         : rs_operator(rs_operator_)
+        , topn(nullptr)
     {}
 
     Poco::JSON::Object::Ptr toJSONObject() const;
@@ -94,6 +98,8 @@ class PushDownFilter
     const ExpressionActionsPtr extra_cast;
     // If the extra_cast is not null, the types of the columns may be changed
     const ColumnDefinesPtr columns_after_cast;
+    // Optional storage-side TopN metadata used only by multi-stage late materialization.
+    const MultiStageLateMaterializationTopNDescriptionPtr topn;
 };
 
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index cbb6002f0fd..2b462dcc239 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -57,6 +57,10 @@ MultiStageLateMaterializationBlockInputStream::MultiStageLateMaterializationBloc
           buildResidualFilterHeader(residual_filter_),
           residual_filter_->before_where,
           residual_filter_->filter_column_name)
+    , running_topn(
+          residual_filter_->topn != nullptr
+              ? std::make_unique<RunningLocalTopN>(*residual_filter_->topn, *residual_filter_->filter_columns)
+              : nullptr)
     , log(Logger::get(NAME, req_id_))
 {
     RUNTIME_CHECK(residual_filter != nullptr);
@@ -326,6 +330,9 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         if (runtime_stats)
             runtime_stats->stage1_output_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
 
+        if (running_topn == nullptr && runtime_stats)
+            runtime_stats->topn_candidate_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
+
         if (residual_passed_rows == 0)
         {
             ++late_mode_blocks;
@@ -333,6 +340,39 @@ Block MultiStageLateMaterializationBlockInputStream::read()
             continue;
         }
 
+        if (running_topn != nullptr)
+        {
+            auto topn_result = running_topn->update(stage1_block, residual_filter_ptr, residual_passed_rows);
+            if (runtime_stats)
+                runtime_stats->topn_candidate_rows.fetch_add(topn_result.passed_count, std::memory_order_relaxed);
+
+            if (topn_result.passed_count == 0)
+            {
+                ++late_mode_blocks;
+                skipNextBlockOrRead(final_rest_stream, "final_rest");
+                continue;
+            }
+
+            if (topn_result.passed_count == stage1_block.rows())
+            {
+                ++direct_mode_blocks;
+                return buildDirectModeBlock(
+                    stage1_block,
+                    stage0_filter_ptr,
+                    effective_stage0_filter.passed_count,
+                    nullptr,
+                    topn_result.passed_count);
+            }
+
+            ++late_mode_blocks;
+            return buildLateModeBlock(
+                stage1_block,
+                stage0_filter_ptr,
+                stage0_block.rows(),
+                topn_result.filter,
+                topn_result.passed_count);
+        }
+
         if (residual_passed_rows == stage1_block.rows())
         {
             ++direct_mode_blocks;
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
index 90ac71bc11d..c7bde1736a1 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
@@ -101,6 +101,7 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
     BitmapFilterPtr bitmap_filter;
     MultiStageLateMaterializationRuntimeStatsPtr runtime_stats;
     FilterTransformAction residual_filter_action;
+    std::unique_ptr<RunningLocalTopN> running_topn;
 
     size_t late_mode_blocks = 0;
     size_t direct_mode_blocks = 0;
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index 9bc7d221adc..20d0a9a0ad4 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -26,6 +26,7 @@ struct MultiStageLateMaterializationRuntimeStats
 {
     std::atomic<UInt64> stage0_output_rows{0};
     std::atomic<UInt64> stage1_output_rows{0};
+    std::atomic<UInt64> topn_candidate_rows{0};
 };
 
 using MultiStageLateMaterializationRuntimeStatsPtr = std::shared_ptr<MultiStageLateMaterializationRuntimeStats>;
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
new file mode 100644
index 00000000000..507fdcbe210
--- /dev/null
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -0,0 +1,341 @@
+// Copyright 2026 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnDecimal.h>
+#include <Columns/ColumnNullable.h>
+#include <Columns/ColumnVector.h>
+#include <Common/Exception.h>
+#include <Common/typeid_cast.h>
+#include <DataTypes/DataTypeDate.h>
+#include <DataTypes/DataTypeDateTime.h>
+#include <DataTypes/DataTypeDecimal.h>
+#include <DataTypes/DataTypeNullable.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Functions/FunctionHelpers.h>
+#include <Storages/DeltaMerge/DeltaMergeDefines.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
+
+#include <algorithm>
+
+namespace DB::DM
+{
+namespace
+{
+struct SortKeyTypeInfo
+{
+    SortKeyKind kind;
+    void (*extract_field)(const IColumn & column, size_t row, SortKeyField & field);
+};
+
+const IColumn & unwrapNullableColumn(const IColumn & column, size_t row, bool & is_null)
+{
+    if (const auto * nullable_column = typeid_cast<const ColumnNullable *>(&column))
+    {
+        if (nullable_column->isNullAt(row))
+        {
+            is_null = true;
+            return nullable_column->getNestedColumn();
+        }
+        return nullable_column->getNestedColumn();
+    }
+
+    is_null = false;
+    return column;
+}
+
+template <typename ColumnType>
+const ColumnType & castColumn(const IColumn & column)
+{
+    const auto * typed_column = typeid_cast<const ColumnType *>(&column);
+    RUNTIME_CHECK_MSG(
+        typed_column != nullptr,
+        "Unexpected order by column type in TopN-enhanced multi-stage late materialization, column_type={}",
+        column.getName());
+    return *typed_column;
+}
+
+template <typename ColumnValueType, typename StorageValueType, SortKeyKind kind, typename ColumnType>
+void extractFieldImpl(const IColumn & column, size_t row, SortKeyField & field)
+{
+    field.kind = kind;
+
+    bool is_null = false;
+    const auto & nested_column = unwrapNullableColumn(column, row, is_null);
+    field.is_null = is_null;
+    if (is_null)
+        return;
+
+    const auto & data = castColumn<ColumnType>(nested_column).getData();
+    field.value = static_cast<StorageValueType>(data[row]);
+}
+
+template <typename DecimalType, SortKeyKind kind>
+void extractDecimalFieldImpl(const IColumn & column, size_t row, SortKeyField & field)
+{
+    field.kind = kind;
+
+    bool is_null = false;
+    const auto & nested_column = unwrapNullableColumn(column, row, is_null);
+    field.is_null = is_null;
+    if (is_null)
+        return;
+
+    const auto & data = castColumn<ColumnDecimal<DecimalType>>(nested_column).getData();
+    field.value = data[row];
+}
+
+template <typename T>
+int compareValue(const T & lhs, const T & rhs)
+{
+    return lhs > rhs ? 1 : (lhs < rhs ? -1 : 0);
+}
+
+int compareSortKeyFields(const SortKeyField & lhs, const SortKeyField & rhs)
+{
+    RUNTIME_CHECK(lhs.kind == rhs.kind);
+
+    if (lhs.is_null || rhs.is_null)
+    {
+        if (lhs.is_null && rhs.is_null)
+            return 0;
+        return lhs.is_null ? -1 : 1;
+    }
+
+    switch (lhs.kind)
+    {
+    case SortKeyKind::Int64:
+        return compareValue(std::get<Int64>(lhs.value), std::get<Int64>(rhs.value));
+    case SortKeyKind::UInt64:
+    case SortKeyKind::Date:
+    case SortKeyKind::DateTime:
+        return compareValue(std::get<UInt64>(lhs.value), std::get<UInt64>(rhs.value));
+    case SortKeyKind::Float32:
+        return CompareHelper<Float32>::compare(std::get<Float32>(lhs.value), std::get<Float32>(rhs.value), -1);
+    case SortKeyKind::Float64:
+        return CompareHelper<Float64>::compare(std::get<Float64>(lhs.value), std::get<Float64>(rhs.value), -1);
+    case SortKeyKind::Decimal32:
+        return compareValue(std::get<Decimal32>(lhs.value), std::get<Decimal32>(rhs.value));
+    case SortKeyKind::Decimal64:
+        return compareValue(std::get<Decimal64>(lhs.value), std::get<Decimal64>(rhs.value));
+    case SortKeyKind::Decimal128:
+        return compareValue(std::get<Decimal128>(lhs.value), std::get<Decimal128>(rhs.value));
+    case SortKeyKind::Decimal256:
+        return compareValue(std::get<Decimal256>(lhs.value), std::get<Decimal256>(rhs.value));
+    }
+    RUNTIME_CHECK_MSG(false, "Unexpected sort key kind");
+    return 0;
+}
+
+SortKeyTypeInfo getSortKeyTypeInfo(const DataTypePtr & type)
+{
+    const auto type_not_null = removeNullable(type);
+    if (checkDataType<DataTypeInt8>(type_not_null.get()))
+        return {SortKeyKind::Int64, extractFieldImpl<Int8, Int64, SortKeyKind::Int64, ColumnVector<Int8>>};
+    if (checkDataType<DataTypeInt16>(type_not_null.get()))
+        return {SortKeyKind::Int64, extractFieldImpl<Int16, Int64, SortKeyKind::Int64, ColumnVector<Int16>>};
+    if (checkDataType<DataTypeInt32>(type_not_null.get()))
+        return {SortKeyKind::Int64, extractFieldImpl<Int32, Int64, SortKeyKind::Int64, ColumnVector<Int32>>};
+    if (checkDataType<DataTypeInt64>(type_not_null.get()))
+        return {SortKeyKind::Int64, extractFieldImpl<Int64, Int64, SortKeyKind::Int64, ColumnVector<Int64>>};
+    if (checkDataType<DataTypeUInt8>(type_not_null.get()))
+        return {SortKeyKind::UInt64, extractFieldImpl<UInt8, UInt64, SortKeyKind::UInt64, ColumnVector<UInt8>>};
+    if (checkDataType<DataTypeUInt16>(type_not_null.get()))
+        return {SortKeyKind::UInt64, extractFieldImpl<UInt16, UInt64, SortKeyKind::UInt64, ColumnVector<UInt16>>};
+    if (checkDataType<DataTypeUInt32>(type_not_null.get()))
+        return {SortKeyKind::UInt64, extractFieldImpl<UInt32, UInt64, SortKeyKind::UInt64, ColumnVector<UInt32>>};
+    if (checkDataType<DataTypeUInt64>(type_not_null.get()))
+        return {SortKeyKind::UInt64, extractFieldImpl<UInt64, UInt64, SortKeyKind::UInt64, ColumnVector<UInt64>>};
+    if (checkDataType<DataTypeFloat32>(type_not_null.get()))
+        return {SortKeyKind::Float32, extractFieldImpl<Float32, Float32, SortKeyKind::Float32, ColumnVector<Float32>>};
+    if (checkDataType<DataTypeFloat64>(type_not_null.get()))
+        return {SortKeyKind::Float64, extractFieldImpl<Float64, Float64, SortKeyKind::Float64, ColumnVector<Float64>>};
+    if (checkDataType<DataTypeDecimal32>(type_not_null.get()))
+        return {SortKeyKind::Decimal32, extractDecimalFieldImpl<Decimal32, SortKeyKind::Decimal32>};
+    if (checkDataType<DataTypeDecimal64>(type_not_null.get()))
+        return {SortKeyKind::Decimal64, extractDecimalFieldImpl<Decimal64, SortKeyKind::Decimal64>};
+    if (checkDataType<DataTypeDecimal128>(type_not_null.get()))
+        return {SortKeyKind::Decimal128, extractDecimalFieldImpl<Decimal128, SortKeyKind::Decimal128>};
+    if (checkDataType<DataTypeDecimal256>(type_not_null.get()))
+        return {SortKeyKind::Decimal256, extractDecimalFieldImpl<Decimal256, SortKeyKind::Decimal256>};
+    if (checkDataType<DataTypeDate>(type_not_null.get()))
+        return {SortKeyKind::Date, extractFieldImpl<UInt16, UInt64, SortKeyKind::Date, ColumnVector<UInt16>>};
+    if (checkDataType<DataTypeDateTime>(type_not_null.get()))
+        return {SortKeyKind::DateTime, extractFieldImpl<UInt32, UInt64, SortKeyKind::DateTime, ColumnVector<UInt32>>};
+
+    throw Exception("Unsupported order by type for TopN-enhanced multi-stage late materialization: " + type->getName());
+}
+} // namespace
+
+bool RunningLocalTopN::HeapComparator::operator()(const HeapEntry & lhs, const HeapEntry & rhs) const
+{
+    RUNTIME_CHECK(owner != nullptr);
+    return owner->compareOwnedKeys(lhs.key, rhs.key) < 0;
+}
+
+RunningLocalTopN::RunningLocalTopN(
+    const MultiStageLateMaterializationTopNDescription & description,
+    const ColumnDefines & stage1_columns)
+    : topk(description.topk)
+    , heap(HeapComparator{this})
+{
+    RUNTIME_CHECK(topk > 0);
+    RUNTIME_CHECK(description.order_by_columns.size() <= multi_stage_late_materialization_topn_max_order_by_columns);
+
+    sort_key_columns.reserve(description.order_by_columns.size());
+    for (const auto & order_by_column : description.order_by_columns)
+    {
+        const auto column_it = std::find_if(stage1_columns.begin(), stage1_columns.end(), [&](const auto & column) {
+            return column.id == order_by_column.column_id;
+        });
+        RUNTIME_CHECK_MSG(
+            column_it != stage1_columns.end(),
+            "Order by column is not found in stage1 columns, column_id={}",
+            order_by_column.column_id);
+
+        auto type_info = getSortKeyTypeInfo(column_it->type);
+        sort_key_columns.push_back(SortKeyColumnDesc{
+            .column_id = order_by_column.column_id,
+            .column_pos = static_cast<size_t>(column_it - stage1_columns.begin()),
+            .kind = type_info.kind,
+            .direction = order_by_column.direction,
+            .extract_field = type_info.extract_field,
+        });
+    }
+}
+
+int RunningLocalTopN::compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const
+{
+    RUNTIME_CHECK(lhs.size == rhs.size);
+    RUNTIME_CHECK(lhs.size == sort_key_columns.size());
+    for (size_t i = 0; i < lhs.size; ++i)
+    {
+        if (const auto cmp = compareSortKeyFields(lhs.fields[i], rhs.fields[i]); cmp != 0)
+            return cmp * sort_key_columns[i].direction;
+    }
+    return 0;
+}
+
+int RunningLocalTopN::compareRowWithOwnedKey(
+    const std::vector<const IColumn *> & sort_columns,
+    size_t row,
+    const OwnedSortKey & rhs) const
+{
+    RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size());
+    RUNTIME_CHECK(rhs.size == sort_key_columns.size());
+
+    SortKeyField lhs_field;
+    for (size_t i = 0; i < sort_key_columns.size(); ++i)
+    {
+        sort_key_columns[i].extract_field(*sort_columns[i], row, lhs_field);
+        if (const auto cmp = compareSortKeyFields(lhs_field, rhs.fields[i]); cmp != 0)
+            return cmp * sort_key_columns[i].direction;
+    }
+    return 0;
+}
+
+OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row) const
+{
+    RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size());
+
+    OwnedSortKey key;
+    key.size = sort_key_columns.size();
+    for (size_t i = 0; i < sort_key_columns.size(); ++i)
+        sort_key_columns[i].extract_field(*sort_columns[i], row, key.fields[i]);
+    return key;
+}
+
+RunningLocalTopNUpdateResult RunningLocalTopN::update(
+    const Block & stage1_block,
+    const IColumn::Filter * residual_filter,
+    size_t residual_passed_rows)
+{
+    ++current_block_sequence;
+
+    const auto rows = stage1_block.rows();
+    RunningLocalTopNUpdateResult result;
+    result.filter.resize_fill(rows, 0);
+
+    if (residual_passed_rows == 0)
+        return result;
+
+    std::vector<ColumnPtr> materialized_const_columns;
+    std::vector<const IColumn *> sort_columns;
+    sort_columns.reserve(sort_key_columns.size());
+    for (const auto & desc : sort_key_columns)
+    {
+        auto column = stage1_block.getByPosition(desc.column_pos).column;
+        if (auto full_column = column->convertToFullColumnIfConst())
+        {
+            materialized_const_columns.push_back(std::move(full_column));
+            column = materialized_const_columns.back();
+        }
+        sort_columns.push_back(column.get());
+    }
+
+    auto markCandidate = [&](size_t row) {
+        if (result.filter[row] == 0)
+        {
+            result.filter[row] = 1;
+            ++result.passed_count;
+        }
+    };
+    auto unmarkCurrentBlockCandidate = [&](const HeapEntry & entry) {
+        if (entry.block_sequence == current_block_sequence && result.filter[entry.row_index_in_stage1_block] != 0)
+        {
+            result.filter[entry.row_index_in_stage1_block] = 0;
+            --result.passed_count;
+        }
+    };
+
+    for (size_t row = 0; row < rows; ++row)
+    {
+        if (residual_filter != nullptr && (*residual_filter)[row] == 0)
+            continue;
+
+        if (heap.size() < topk)
+        {
+            markCandidate(row);
+            heap.push(HeapEntry{
+                .key = materializeOwnedKey(sort_columns, row),
+                .block_sequence = current_block_sequence,
+                .row_index_in_stage1_block = static_cast<UInt32>(row),
+            });
+            continue;
+        }
+
+        const auto cmp = compareRowWithOwnedKey(sort_columns, row, heap.top().key);
+        if (cmp < 0)
+        {
+            auto evicted = heap.top();
+            heap.pop();
+            unmarkCurrentBlockCandidate(evicted);
+
+            markCandidate(row);
+            heap.push(HeapEntry{
+                .key = materializeOwnedKey(sort_columns, row),
+                .block_sequence = current_block_sequence,
+                .row_index_in_stage1_block = static_cast<UInt32>(row),
+            });
+        }
+        else if (cmp == 0)
+        {
+            markCandidate(row);
+        }
+    }
+
+    return result;
+}
+
+} // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
new file mode 100644
index 00000000000..c45d6797bc8
--- /dev/null
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
@@ -0,0 +1,139 @@
+// Copyright 2026 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Columns/IColumn.h>
+#include <Common/Decimal.h>
+#include <Core/Block.h>
+#include <Storages/DeltaMerge/ColumnDefine_fwd.h>
+#include <Storages/KVStore/Types.h>
+#include <common/types.h>
+
+#include <array>
+#include <memory>
+#include <queue>
+#include <variant>
+#include <vector>
+
+namespace DB::DM
+{
+
+inline constexpr UInt64 multi_stage_late_materialization_topn_max_topk = 4096;
+inline constexpr size_t multi_stage_late_materialization_topn_max_order_by_columns = 4;
+
+struct MultiStageLateMaterializationTopNOrderByColumn
+{
+    ColumnID column_id;
+    int direction;
+};
+
+struct MultiStageLateMaterializationTopNDescription
+{
+    UInt64 topk = 0;
+    std::vector<MultiStageLateMaterializationTopNOrderByColumn> order_by_columns;
+};
+
+using MultiStageLateMaterializationTopNDescriptionPtr = std::shared_ptr<MultiStageLateMaterializationTopNDescription>;
+
+enum class SortKeyKind
+{
+    Int64,
+    UInt64,
+    Float32,
+    Float64,
+    Decimal32,
+    Decimal64,
+    Decimal128,
+    Decimal256,
+    Date,
+    DateTime,
+};
+
+using SortKeyValue = std::variant<Int64, UInt64, Float32, Float64, Decimal32, Decimal64, Decimal128, Decimal256>;
+
+struct SortKeyField
+{
+    bool is_null = false;
+    SortKeyKind kind = SortKeyKind::Int64;
+    SortKeyValue value = Int64{};
+};
+
+struct OwnedSortKey
+{
+    std::array<SortKeyField, multi_stage_late_materialization_topn_max_order_by_columns> fields;
+    size_t size = 0;
+};
+
+struct RunningLocalTopNUpdateResult
+{
+    IColumn::Filter filter;
+    size_t passed_count = 0;
+};
+
+class RunningLocalTopN
+{
+public:
+    RunningLocalTopN(
+        const MultiStageLateMaterializationTopNDescription & description,
+        const ColumnDefines & stage1_columns);
+
+    RunningLocalTopNUpdateResult update(
+        const Block & stage1_block,
+        const IColumn::Filter * residual_filter,
+        size_t residual_passed_rows);
+
+    size_t heapSize() const { return heap.size(); }
+
+private:
+    struct SortKeyColumnDesc
+    {
+        using ExtractFieldFn = void (*)(const IColumn & column, size_t row, SortKeyField & field);
+
+        ColumnID column_id = EmptyColumnID;
+        size_t column_pos = 0;
+        SortKeyKind kind = SortKeyKind::Int64;
+        int direction = 1;
+        ExtractFieldFn extract_field = nullptr;
+    };
+
+    struct HeapEntry
+    {
+        OwnedSortKey key;
+        UInt64 block_sequence = 0;
+        UInt32 row_index_in_stage1_block = 0;
+    };
+
+    struct HeapComparator
+    {
+        const RunningLocalTopN * owner = nullptr;
+
+        bool operator()(const HeapEntry & lhs, const HeapEntry & rhs) const;
+    };
+
+    using Heap = std::priority_queue<HeapEntry, std::vector<HeapEntry>, HeapComparator>;
+
+    int compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const;
+    int compareRowWithOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row, const OwnedSortKey & rhs)
+        const;
+    OwnedSortKey materializeOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row) const;
+
+private:
+    UInt64 topk;
+    std::vector<SortKeyColumnDesc> sort_key_columns;
+    UInt64 current_block_sequence = 0;
+    Heap heap;
+};
+
+} // namespace DB::DM
diff --git a/dbms/src/Storages/SelectQueryInfo.cpp b/dbms/src/Storages/SelectQueryInfo.cpp
index 2ba3b656e08..84325257d1f 100644
--- a/dbms/src/Storages/SelectQueryInfo.cpp
+++ b/dbms/src/Storages/SelectQueryInfo.cpp
@@ -34,6 +34,7 @@ SelectQueryInfo::SelectQueryInfo(const SelectQueryInfo & rhs)
     , has_multiple_partitions(rhs.has_multiple_partitions)
     , enable_multi_stage_late_materialization(rhs.enable_multi_stage_late_materialization)
     , multi_stage_late_materialization_runtime_stats(rhs.multi_stage_late_materialization_runtime_stats)
+    , multi_stage_late_materialization_topn(rhs.multi_stage_late_materialization_topn)
 {}
 
 SelectQueryInfo::SelectQueryInfo(SelectQueryInfo && rhs) noexcept
@@ -47,6 +48,7 @@ SelectQueryInfo::SelectQueryInfo(SelectQueryInfo && rhs) noexcept
     , has_multiple_partitions(rhs.has_multiple_partitions)
     , enable_multi_stage_late_materialization(rhs.enable_multi_stage_late_materialization)
     , multi_stage_late_materialization_runtime_stats(std::move(rhs.multi_stage_late_materialization_runtime_stats))
+    , multi_stage_late_materialization_topn(std::move(rhs.multi_stage_late_materialization_topn))
 {}
 
 } // namespace DB
diff --git a/dbms/src/Storages/SelectQueryInfo.h b/dbms/src/Storages/SelectQueryInfo.h
index 2d6b6e318da..ec4090bd330 100644
--- a/dbms/src/Storages/SelectQueryInfo.h
+++ b/dbms/src/Storages/SelectQueryInfo.h
@@ -37,6 +37,8 @@ namespace DM
 {
 struct MultiStageLateMaterializationRuntimeStats;
 using MultiStageLateMaterializationRuntimeStatsPtr = std::shared_ptr<MultiStageLateMaterializationRuntimeStats>;
+struct MultiStageLateMaterializationTopNDescription;
+using MultiStageLateMaterializationTopNDescriptionPtr = std::shared_ptr<MultiStageLateMaterializationTopNDescription>;
 } // namespace DM
 
 
@@ -62,6 +64,7 @@ struct SelectQueryInfo
     bool has_multiple_partitions = false;
     bool enable_multi_stage_late_materialization = false;
     DM::MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats;
+    DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn;
 
     SelectQueryInfo();
     ~SelectQueryInfo();
diff --git a/dbms/src/Storages/StorageDeltaMerge.cpp b/dbms/src/Storages/StorageDeltaMerge.cpp
index 046b11d2fe5..d5ea4064ca7 100644
--- a/dbms/src/Storages/StorageDeltaMerge.cpp
+++ b/dbms/src/Storages/StorageDeltaMerge.cpp
@@ -22,6 +22,11 @@
 #include <Core/Defines.h>
 #include <DataStreams/IBlockOutputStream.h>
 #include <DataStreams/OneBlockInputStream.h>
+#include <DataTypes/DataTypeDate.h>
+#include <DataTypes/DataTypeDateTime.h>
+#include <DataTypes/DataTypeDecimal.h>
+#include <DataTypes/DataTypeNullable.h>
+#include <DataTypes/DataTypesNumber.h>
 #include <DataTypes/isSupportedDataTypeCast.h>
 #include <Databases/IDatabase.h>
 #include <Debug/MockTiDB.h>
@@ -29,6 +34,7 @@
 #include <Flash/Coprocessor/DAGQueryInfo.h>
 #include <Flash/Coprocessor/InterpreterUtils.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
+#include <Functions/FunctionHelpers.h>
 #include <Interpreters/Context.h>
 #include <Parsers/ASTCreateQuery.h>
 #include <Parsers/ASTExpressionList.h>
@@ -792,6 +798,134 @@ RuntimeFilteList parseRuntimeFilterList(
     return runtime_filter_list;
 }
 
+bool isSupportedMultiStageLateMaterializationTopNOrderByType(const DataTypePtr & type)
+{
+    const auto type_not_null = removeNullable(type);
+    return checkDataType<DataTypeUInt8>(type_not_null.get()) //
+        || checkDataType<DataTypeUInt16>(type_not_null.get()) //
+        || checkDataType<DataTypeUInt32>(type_not_null.get()) //
+        || checkDataType<DataTypeUInt64>(type_not_null.get()) //
+        || checkDataType<DataTypeInt8>(type_not_null.get()) //
+        || checkDataType<DataTypeInt16>(type_not_null.get()) //
+        || checkDataType<DataTypeInt32>(type_not_null.get()) //
+        || checkDataType<DataTypeInt64>(type_not_null.get()) //
+        || checkDataType<DataTypeFloat32>(type_not_null.get()) //
+        || checkDataType<DataTypeFloat64>(type_not_null.get()) //
+        || checkDataType<DataTypeDecimal32>(type_not_null.get()) //
+        || checkDataType<DataTypeDecimal64>(type_not_null.get()) //
+        || checkDataType<DataTypeDecimal128>(type_not_null.get()) //
+        || checkDataType<DataTypeDecimal256>(type_not_null.get()) //
+        || checkDataType<DataTypeDate>(type_not_null.get()) //
+        || checkDataType<DataTypeDateTime>(type_not_null.get());
+}
+
+const TiDB::ColumnInfo * findSourceColumnByID(const TiDB::ColumnInfos & source_columns, ColumnID column_id)
+{
+    const auto it = std::find_if(source_columns.begin(), source_columns.end(), [&](const auto & column) {
+        return column.id == column_id;
+    });
+    return it == source_columns.end() ? nullptr : &*it;
+}
+
+DM::MultiStageLateMaterializationTopNDescriptionPtr buildStorageMultiStageLateMaterializationTopN(
+    const SelectQueryInfo & query_info,
+    const DM::ColumnDefines & final_columns_to_read,
+    const DM::ColumnDefinesPtr & residual_filter_columns,
+    const LoggerPtr & tracing_logger,
+    DM::ColumnDefinesPtr & merged_stage1_columns)
+{
+    auto disable = [&](const String & reason) -> DM::MultiStageLateMaterializationTopNDescriptionPtr {
+        LOG_DEBUG(tracing_logger, "Disable TopN-enhanced multi-stage late materialization, reason={}", reason);
+        return nullptr;
+    };
+
+    if (!query_info.multi_stage_late_materialization_topn)
+        return nullptr;
+    if (query_info.multi_stage_late_materialization_topn->topk == 0)
+        return disable("topk is zero");
+    if (query_info.multi_stage_late_materialization_topn->topk > DM::multi_stage_late_materialization_topn_max_topk)
+        return disable(fmt::format("topk is too large: {}", query_info.multi_stage_late_materialization_topn->topk));
+    if (query_info.multi_stage_late_materialization_topn->order_by_columns.empty())
+        return disable("empty order by columns");
+    if (query_info.multi_stage_late_materialization_topn->order_by_columns.size()
+        > DM::multi_stage_late_materialization_topn_max_order_by_columns)
+        return disable(fmt::format(
+            "too many order by columns: {}",
+            query_info.multi_stage_late_materialization_topn->order_by_columns.size()));
+
+    std::unordered_map<ColumnID, DM::ColumnDefine> final_column_map;
+    final_column_map.reserve(final_columns_to_read.size());
+    for (const auto & column : final_columns_to_read)
+        final_column_map.emplace(column.id, column);
+
+    merged_stage1_columns = std::make_shared<DM::ColumnDefines>(*residual_filter_columns);
+    std::unordered_set<ColumnID> stage1_column_ids;
+    stage1_column_ids.reserve(merged_stage1_columns->size());
+    for (const auto & column : *merged_stage1_columns)
+        stage1_column_ids.insert(column.id);
+
+    auto storage_topn = std::make_shared<DM::MultiStageLateMaterializationTopNDescription>();
+    storage_topn->topk = query_info.multi_stage_late_materialization_topn->topk;
+    storage_topn->order_by_columns.reserve(query_info.multi_stage_late_materialization_topn->order_by_columns.size());
+
+    for (const auto & order_by_column : query_info.multi_stage_late_materialization_topn->order_by_columns)
+    {
+        const auto column_id = order_by_column.column_id;
+        if (column_id == ExtraTableIDColumnID || column_id == ExtraCommitTSColumnID)
+            return disable(fmt::format("order by virtual column is unsupported, column_id={}", column_id));
+
+        const auto * source_column = findSourceColumnByID(query_info.dag_query->source_columns, column_id);
+        if (source_column == nullptr)
+            return disable(
+                fmt::format("order by column is not found in table scan source columns, column_id={}", column_id));
+        if (source_column->hasGeneratedColumnFlag())
+            return disable(fmt::format("order by generated column is unsupported, column_id={}", column_id));
+        if (source_column->tp == TiDB::TypeTimestamp || source_column->tp == TiDB::TypeTime)
+            return disable(fmt::format(
+                "order by column type needs or may need table scan extra cast, column_id={}, tp={}",
+                column_id,
+                static_cast<int>(source_column->tp)));
+
+        const auto final_column_it = final_column_map.find(column_id);
+        if (final_column_it == final_column_map.end())
+            return disable(
+                fmt::format("order by column is not found in storage columns_to_read, column_id={}", column_id));
+        if (!isSupportedMultiStageLateMaterializationTopNOrderByType(final_column_it->second.type))
+            return disable(fmt::format(
+                "unsupported order by storage type, column_id={}, type={}",
+                column_id,
+                final_column_it->second.type->getName()));
+
+        if (stage1_column_ids.insert(column_id).second)
+            merged_stage1_columns->push_back(final_column_it->second);
+        storage_topn->order_by_columns.push_back(order_by_column);
+    }
+
+    const auto final_rest_col_cnt = final_columns_to_read.size() - merged_stage1_columns->size();
+    if (final_rest_col_cnt == 0)
+        return disable("no final rest columns after adding order by columns");
+    if (final_rest_col_cnt < 10)
+        return disable(fmt::format("too few final rest columns after adding order by columns: {}", final_rest_col_cnt));
+    if (final_rest_col_cnt < 2 * merged_stage1_columns->size())
+    {
+        return disable(fmt::format(
+            "final rest columns are not wide enough after adding order by columns, final_rest_col_cnt={}, "
+            "stage1_col_cnt={}",
+            final_rest_col_cnt,
+            merged_stage1_columns->size()));
+    }
+
+    LOG_DEBUG(
+        tracing_logger,
+        "Enable TopN-enhanced multi-stage late materialization, topk={} order_by_columns={} stage1_columns={} "
+        "final_rest_columns={}",
+        storage_topn->topk,
+        storage_topn->order_by_columns.size(),
+        merged_stage1_columns->size(),
+        final_rest_col_cnt);
+    return storage_topn;
+}
+
 DM::PushDownFilterPtr buildMultiStageLateMaterializationFilter(
     const SelectQueryInfo & query_info,
     const DM::PushDownFilterPtr & stage0_filter,
@@ -832,7 +966,25 @@ DM::PushDownFilterPtr buildMultiStageLateMaterializationFilter(
         "Build multi-stage late materialization residual filter, stage1_filter_columns={} final_columns_to_read={}",
         residual_filter->filter_columns->size(),
         final_columns_to_read.size());
-    return residual_filter;
+    DM::ColumnDefinesPtr merged_stage1_columns;
+    auto storage_topn = buildStorageMultiStageLateMaterializationTopN(
+        query_info,
+        final_columns_to_read,
+        residual_filter->filter_columns,
+        tracing_logger,
+        merged_stage1_columns);
+    if (storage_topn == nullptr)
+        return residual_filter;
+
+    return std::make_shared<DM::PushDownFilter>(
+        residual_filter->rs_operator,
+        residual_filter->before_where,
+        residual_filter->project_after_where,
+        merged_stage1_columns,
+        residual_filter->filter_column_name,
+        residual_filter->extra_cast,
+        residual_filter->columns_after_cast,
+        storage_topn);
 }
 } // namespace
 

From 2c86fad4df73595ce1213a23803f19d1e911147d Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Tue, 28 Jul 2026 14:01:39 +0800
Subject: [PATCH 06/18] Refine multi-stage late materialization TopN

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 dbms/src/Debug/MockStorage.cpp                |   2 +-
 .../Coprocessor/DAGStorageInterpreter.cpp     |   2 +-
 dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp |   5 +-
 .../Flash/tests/gtest_executors_with_dm.cpp   | 147 ++++++++++---
 dbms/src/Interpreters/Settings.h              |   2 +-
 .../MultiStageLateMaterializationTopN.cpp     | 197 +++++++++++++-----
 .../MultiStageLateMaterializationTopN.h       |  21 +-
 dbms/src/Storages/StorageDeltaMerge.cpp       |   6 +-
 ...-07-25-multi-stage-late-materialization.md |   6 +-
 .../2026-07-26-running-local-topn-for-mslm.md |  21 +-
 10 files changed, 307 insertions(+), 102 deletions(-)

diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp
index cf346e16e42..03e72a2ddc3 100644
--- a/dbms/src/Debug/MockStorage.cpp
+++ b/dbms/src/Debug/MockStorage.cpp
@@ -159,7 +159,7 @@ bool shouldEnableMultiStageLateMaterializationForMockDeltaMerge(
     const TiDB::ColumnInfos & scan_column_infos)
 {
     const auto & settings = context.getSettingsRef();
-    if (!settings.dt_enable_multi_stage_late_materialization)
+    if (settings.dt_enable_multi_stage_late_materialization == 0)
         return false;
     if (filter_conditions == nullptr || !filter_conditions->hasValue())
         return false;
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index 5c6852fe7b8..7502210d29a 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -1797,7 +1797,7 @@ std::pair<Names, std::vector<UInt8>> DAGStorageInterpreter::getColumnsForTableSc
 
 bool DAGStorageInterpreter::shouldEnableMultiStageLateMaterialization() const
 {
-    if (!context.getSettingsRef().dt_enable_multi_stage_late_materialization)
+    if (context.getSettingsRef().dt_enable_multi_stage_late_materialization == 0)
         return false;
 
     auto disable = [&](const String & reason) {
diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
index c2320dcf2e8..09d344e5cfa 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
@@ -138,7 +138,10 @@ PhysicalPlanNodePtr PhysicalTopN::build(
 
     auto order_columns = analyzer.buildOrderColumns(before_sort_actions, top_n.order_by());
     SortDescription order_descr = getSortDescription(order_columns, top_n.order_by());
-    tryAttachMultiStageLateMaterializationTopN(top_n, child, log);
+    if (context.getSettingsRef().dt_enable_multi_stage_late_materialization == 2)
+        tryAttachMultiStageLateMaterializationTopN(top_n, child, log);
+    else
+        LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason=setting disabled");
 
     auto physical_top_n = std::make_shared<PhysicalTopN>(
         executor_id,
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 3f985238bf2..626b25d640c 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -145,6 +145,10 @@ class ExecutorsWithDMTestRunner : public DB::tests::ExecutorTest
 
 namespace
 {
+constexpr UInt64 MSLMModeDisabled = 0;
+constexpr UInt64 MSLMModeSelection = 1;
+constexpr UInt64 MSLMModeTopN = 2;
+
 tipb::TableScan * findMutableTableScan(tipb::Executor * executor)
 {
     switch (executor->tp())
@@ -562,13 +566,13 @@ try
         lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast<Int64>(3)))), lit(Field(static_cast<Int64>(0)))));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -592,13 +596,13 @@ try
         lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast<Int64>(8)))), lit(Field(static_cast<Int64>(0)))));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -625,13 +629,98 @@ try
         /*is_desc=*/false,
         /*limit=*/2);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
+    DAGContext selection_only_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto selection_only_actual = executeStreams(&selection_only_dag_context);
+    ASSERT_TRUE(columnsEqual(expected, selection_only_actual, /*_restrict=*/false))
+        << "\n  expect_block: \n"
+        << getColumnsContent(expected) << "\n actual_block: \n"
+        << getColumnsContent(selection_only_actual);
+    assertMultiStageRowsOverride(selection_only_dag_context, 16);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto actual = executeStreams(&enabled_dag_context);
+    ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
+        << "\n  expect_block: \n"
+        << getColumnsContent(expected) << "\n actual_block: \n"
+        << getColumnsContent(actual);
+
+    auto table_scan_rows = enabled_dag_context.getExecutorRowsOverride("table_scan_0");
+    auto selection_rows = enabled_dag_context.getExecutorRowsOverride("selection_1");
+    ASSERT_NE(table_scan_rows, nullptr);
+    ASSERT_NE(selection_rows, nullptr);
+    ASSERT_EQ(table_scan_rows->load(), 16);
+    ASSERT_GE(selection_rows->load(), 2);
+    ASSERT_LT(selection_rows->load(), 16);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateTime)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    MockColumnInfoVec columns;
+    columns.push_back({"c0", TiDB::TP::TypeLongLong});
+    columns.push_back({"ts_col", TiDB::TP::TypeDatetime});
+    columns.push_back({"c1", TiDB::TP::TypeLongLong});
+    for (size_t col_id = 2; col_id <= 14; ++col_id)
+        columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong});
+
+    constexpr size_t rows = 32;
+    std::vector<Int64> c0_values;
+    ColumnWithNullableMyDateTime ts_values;
+    std::vector<Int64> c1_values;
+    c0_values.reserve(rows);
+    ts_values.reserve(rows);
+    c1_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+    {
+        c0_values.push_back(row);
+        ts_values.push_back(MyDateTime(2020, 1, 1, static_cast<UInt32>(row), 0, 0, 0).toPackedUInt());
+        c1_values.push_back(row % 8);
+    }
+
+    ColumnsWithTypeAndName data;
+    data.emplace_back(toVec<Int64>("c0", c0_values));
+    data.emplace_back(toNullableVec<MyDateTime>("ts_col", ts_values));
+    data.emplace_back(toVec<Int64>("c1", c1_values));
+    for (size_t col_id = 2; col_id <= 14; ++col_id)
+    {
+        std::vector<Int64> values;
+        values.reserve(rows);
+        for (size_t row = 0; row < rows; ++row)
+            values.push_back(static_cast<Int64>(col_id * 1000 + row));
+        data.emplace_back(toVec<Int64>(fmt::format("c{}", col_id), values));
+    }
+    context.addMockDeltaMerge(
+        {"test_db", "multi_stage_lm_datetime_topn"},
+        columns,
+        data,
+        /*concurrency_hint=*/4);
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_datetime_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ts_col",
+        /*is_desc=*/false,
+        /*limit=*/2);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    auto expected = executeStreams(&disabled_dag_context);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -703,11 +792,11 @@ try
         /*is_desc=*/false,
         /*limit=*/2);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -731,13 +820,13 @@ try
         lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         buildAllGreaterThanMinusOneFilter(/*begin_col_id=*/1, /*end_col_id=*/6));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -763,13 +852,13 @@ try
         lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         lt(col("c1"), lit(Field(static_cast<Int64>(4)))));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -799,13 +888,13 @@ try
         lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         buildAllGreaterThanMinusOneFilter(/*begin_col_id=*/1, /*end_col_id=*/9));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
@@ -840,11 +929,11 @@ try
     for (const auto & condition : pushed_down_filter_request->root_executor().selection().conditions())
         *table_scan->add_pushed_down_filter_conditions() = condition;
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -914,7 +1003,7 @@ try
     *findMutableTableScan(request->mutable_root_executor())->add_pushed_down_filter_conditions()
         = pushed_down_filter_expr;
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     const auto selected_rows = selectRows(rows, [](size_t) { return true; });
     auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(
@@ -923,7 +1012,7 @@ try
         "project_2_");
     assertColumnsEqual(expected, executeStreams(&disabled_dag_context));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -994,7 +1083,7 @@ try
         }));
     auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_ts_col, "table_scan_0_");
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -1064,7 +1153,7 @@ try
         tipb::ScalarFuncSig::LTDuration,
         context.getCollation());
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     const auto selected_rows = selectRows(rows, [](size_t row) { return row % 8 > 2 && row % 8 < 6; });
     auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(
@@ -1073,7 +1162,7 @@ try
         "project_2_");
     assertColumnsEqual(expected, executeStreams(&disabled_dag_context));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -1146,11 +1235,11 @@ try
         context.getCollation());
     request->set_time_zone_name("Asia/Shanghai");
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto expected = executeStreams(&disabled_dag_context);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -1175,13 +1264,13 @@ try
         gt(col("c0"), lit(Field(static_cast<Int64>(16)))),
         lt(col("c0"), lit(Field(static_cast<Int64>(24)))));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     const auto selected_rows = selectRows(32, [](size_t row) { return row > 16 && row < 24; });
     auto expected = buildExpectedPayloadColumns(selected_rows);
     assertColumnsEqual(expected, executeStreams(&disabled_dag_context));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -1249,11 +1338,11 @@ try
         buildNullableInt64Values(selected_rows, [](Int64 row) { return row + 1; }));
     auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_commit_ts_col);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
     DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     assertColumnsEqual(expected, executeStreams(&disabled_dag_context));
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
     assertColumnsEqual(expected, actual);
@@ -1280,7 +1369,7 @@ try
         gt(col("c1"), lit(Field(static_cast<Int64>(-1)))));
     rewriteTableScanColumnAsExtraTableID(request.get(), /*column_index=*/1);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     auto actual = executeStreams(&enabled_dag_context);
 
@@ -1304,7 +1393,7 @@ try
         eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast<Int64>(3)))), lit(Field(static_cast<Int64>(0)))));
     request->set_collect_execution_summaries(true);
 
-    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true;
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
     DAGContext dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
     executeStreams(&dag_context);
 
diff --git a/dbms/src/Interpreters/Settings.h b/dbms/src/Interpreters/Settings.h
index 703dabdf632..a449aef7c3f 100644
--- a/dbms/src/Interpreters/Settings.h
+++ b/dbms/src/Interpreters/Settings.h
@@ -226,7 +226,7 @@ struct Settings
     M(SettingUInt64, dt_max_sharing_column_bytes_for_all, 2048 * Constant::MB, "Memory limitation for data sharing of all requests, include those sharing blocks in block queue. 0 means disable data sharing")                         \
     M(SettingUInt64, dt_max_sharing_column_count, 5, "Deprecated")                                                                                                                                                                      \
     M(SettingBool, dt_enable_bitmap_filter, true, "Use bitmap filter to read data or not")                                                                                                                                              \
-    M(SettingBool, dt_enable_multi_stage_late_materialization, true, "Enable multi-stage late materialization for residual filters in DeltaTree Engine.")                                                                               \
+    M(SettingUInt64, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, 2: TopN enhanced.")                                                    \
     M(SettingDouble, dt_read_thread_count_scale, 2.0, "Number of read thread = number of logical cpu cores * dt_read_thread_count_scale.  Only has meaning at server startup.")                                                         \
     M(SettingDouble, io_thread_count_scale, 5.0, "Number of thread of IOThreadPool = number of logical cpu cores * io_thread_count_scale.  Only has meaning at server startup.")                                                        \
     M(SettingUInt64, init_thread_count_scale, 100, "Number of thread = number of logical cpu cores * init_thread_count_scale. It just works for thread pool for initStores and loadMetadata. Only has meaning at server startup.")      \
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
index 507fdcbe210..bb37b3e7a0a 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -21,6 +21,8 @@
 #include <DataTypes/DataTypeDate.h>
 #include <DataTypes/DataTypeDateTime.h>
 #include <DataTypes/DataTypeDecimal.h>
+#include <DataTypes/DataTypeMyDate.h>
+#include <DataTypes/DataTypeMyDateTime.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <Functions/FunctionHelpers.h>
@@ -28,6 +30,7 @@
 #include <Storages/DeltaMerge/MultiStageLateMaterializationTopN.h>
 
 #include <algorithm>
+#include <utility>
 
 namespace DB::DM
 {
@@ -36,23 +39,17 @@ namespace
 struct SortKeyTypeInfo
 {
     SortKeyKind kind;
-    void (*extract_field)(const IColumn & column, size_t row, SortKeyField & field);
+    SortKeyColumnView (*build_column_view)(const IColumn & column);
+    void (*extract_field)(const SortKeyColumnView & column, size_t row, SortKeyField & field);
+    int (*compare_column_with_owned)(const SortKeyColumnView & column, size_t row, const SortKeyField & field);
 };
 
-const IColumn & unwrapNullableColumn(const IColumn & column, size_t row, bool & is_null)
+std::pair<const IColumn *, const IColumn::Filter *> unwrapNullableColumn(const IColumn & column)
 {
     if (const auto * nullable_column = typeid_cast<const ColumnNullable *>(&column))
-    {
-        if (nullable_column->isNullAt(row))
-        {
-            is_null = true;
-            return nullable_column->getNestedColumn();
-        }
-        return nullable_column->getNestedColumn();
-    }
+        return {&nullable_column->getNestedColumn(), &nullable_column->getNullMapData()};
 
-    is_null = false;
-    return column;
+    return {&column, nullptr};
 }
 
 template <typename ColumnType>
@@ -66,33 +63,50 @@ const ColumnType & castColumn(const IColumn & column)
     return *typed_column;
 }
 
+template <typename ColumnType>
+SortKeyColumnView buildSortKeyColumnViewImpl(const IColumn & column)
+{
+    const auto [nested_column, null_map] = unwrapNullableColumn(column);
+    return {
+        .data = &castColumn<ColumnType>(*nested_column).getData(),
+        .null_map = null_map,
+    };
+}
+
+template <typename ColumnType>
+const typename ColumnType::Container & getColumnData(const SortKeyColumnView & column)
+{
+    return *static_cast<const typename ColumnType::Container *>(column.data);
+}
+
+bool isNullAt(const IColumn::Filter * null_map, size_t row)
+{
+    return null_map != nullptr && (*null_map)[row] != 0;
+}
+
 template <typename ColumnValueType, typename StorageValueType, SortKeyKind kind, typename ColumnType>
-void extractFieldImpl(const IColumn & column, size_t row, SortKeyField & field)
+void extractFieldImpl(const SortKeyColumnView & column, size_t row, SortKeyField & field)
 {
     field.kind = kind;
 
-    bool is_null = false;
-    const auto & nested_column = unwrapNullableColumn(column, row, is_null);
-    field.is_null = is_null;
-    if (is_null)
+    field.is_null = isNullAt(column.null_map, row);
+    if (field.is_null)
         return;
 
-    const auto & data = castColumn<ColumnType>(nested_column).getData();
+    const auto & data = getColumnData<ColumnType>(column);
     field.value = static_cast<StorageValueType>(data[row]);
 }
 
 template <typename DecimalType, SortKeyKind kind>
-void extractDecimalFieldImpl(const IColumn & column, size_t row, SortKeyField & field)
+void extractDecimalFieldImpl(const SortKeyColumnView & column, size_t row, SortKeyField & field)
 {
     field.kind = kind;
 
-    bool is_null = false;
-    const auto & nested_column = unwrapNullableColumn(column, row, is_null);
-    field.is_null = is_null;
-    if (is_null)
+    field.is_null = isNullAt(column.null_map, row);
+    if (field.is_null)
         return;
 
-    const auto & data = castColumn<ColumnDecimal<DecimalType>>(nested_column).getData();
+    const auto & data = getColumnData<ColumnDecimal<DecimalType>>(column);
     field.value = data[row];
 }
 
@@ -102,6 +116,88 @@ int compareValue(const T & lhs, const T & rhs)
     return lhs > rhs ? 1 : (lhs < rhs ? -1 : 0);
 }
 
+int compareNull(bool lhs_is_null, bool rhs_is_null)
+{
+    if (lhs_is_null && rhs_is_null)
+        return 0;
+    return lhs_is_null ? -1 : 1;
+}
+
+template <typename StorageValueType, typename ColumnType>
+int compareColumnWithOwnedFieldImpl(
+    const SortKeyColumnView & column,
+    size_t row,
+    const SortKeyField & rhs)
+{
+    const bool lhs_is_null = isNullAt(column.null_map, row);
+    if (lhs_is_null || rhs.is_null)
+        return compareNull(lhs_is_null, rhs.is_null);
+
+    const auto & data = getColumnData<ColumnType>(column);
+    return compareValue(static_cast<StorageValueType>(data[row]), std::get<StorageValueType>(rhs.value));
+}
+
+template <typename ColumnValueType, typename ColumnType>
+int compareFloatColumnWithOwnedFieldImpl(
+    const SortKeyColumnView & column,
+    size_t row,
+    const SortKeyField & rhs)
+{
+    const bool lhs_is_null = isNullAt(column.null_map, row);
+    if (lhs_is_null || rhs.is_null)
+        return compareNull(lhs_is_null, rhs.is_null);
+
+    const auto & data = getColumnData<ColumnType>(column);
+    return CompareHelper<ColumnValueType>::compare(data[row], std::get<ColumnValueType>(rhs.value), -1);
+}
+
+template <typename DecimalType>
+int compareDecimalColumnWithOwnedFieldImpl(
+    const SortKeyColumnView & column,
+    size_t row,
+    const SortKeyField & rhs)
+{
+    const bool lhs_is_null = isNullAt(column.null_map, row);
+    if (lhs_is_null || rhs.is_null)
+        return compareNull(lhs_is_null, rhs.is_null);
+
+    const auto & data = getColumnData<ColumnDecimal<DecimalType>>(column);
+    return compareValue(data[row], std::get<DecimalType>(rhs.value));
+}
+
+template <typename ColumnValueType, typename StorageValueType, SortKeyKind kind, typename ColumnType>
+SortKeyTypeInfo makeSortKeyTypeInfo()
+{
+    return {
+        .kind = kind,
+        .build_column_view = buildSortKeyColumnViewImpl<ColumnType>,
+        .extract_field = extractFieldImpl<ColumnValueType, StorageValueType, kind, ColumnType>,
+        .compare_column_with_owned = compareColumnWithOwnedFieldImpl<StorageValueType, ColumnType>,
+    };
+}
+
+template <typename ColumnValueType, SortKeyKind kind, typename ColumnType>
+SortKeyTypeInfo makeFloatSortKeyTypeInfo()
+{
+    return {
+        .kind = kind,
+        .build_column_view = buildSortKeyColumnViewImpl<ColumnType>,
+        .extract_field = extractFieldImpl<ColumnValueType, ColumnValueType, kind, ColumnType>,
+        .compare_column_with_owned = compareFloatColumnWithOwnedFieldImpl<ColumnValueType, ColumnType>,
+    };
+}
+
+template <typename DecimalType, SortKeyKind kind>
+SortKeyTypeInfo makeDecimalSortKeyTypeInfo()
+{
+    return {
+        .kind = kind,
+        .build_column_view = buildSortKeyColumnViewImpl<ColumnDecimal<DecimalType>>,
+        .extract_field = extractDecimalFieldImpl<DecimalType, kind>,
+        .compare_column_with_owned = compareDecimalColumnWithOwnedFieldImpl<DecimalType>,
+    };
+}
+
 int compareSortKeyFields(const SortKeyField & lhs, const SortKeyField & rhs)
 {
     RUNTIME_CHECK(lhs.kind == rhs.kind);
@@ -142,37 +238,41 @@ SortKeyTypeInfo getSortKeyTypeInfo(const DataTypePtr & type)
 {
     const auto type_not_null = removeNullable(type);
     if (checkDataType<DataTypeInt8>(type_not_null.get()))
-        return {SortKeyKind::Int64, extractFieldImpl<Int8, Int64, SortKeyKind::Int64, ColumnVector<Int8>>};
+        return makeSortKeyTypeInfo<Int8, Int64, SortKeyKind::Int64, ColumnVector<Int8>>();
     if (checkDataType<DataTypeInt16>(type_not_null.get()))
-        return {SortKeyKind::Int64, extractFieldImpl<Int16, Int64, SortKeyKind::Int64, ColumnVector<Int16>>};
+        return makeSortKeyTypeInfo<Int16, Int64, SortKeyKind::Int64, ColumnVector<Int16>>();
     if (checkDataType<DataTypeInt32>(type_not_null.get()))
-        return {SortKeyKind::Int64, extractFieldImpl<Int32, Int64, SortKeyKind::Int64, ColumnVector<Int32>>};
+        return makeSortKeyTypeInfo<Int32, Int64, SortKeyKind::Int64, ColumnVector<Int32>>();
     if (checkDataType<DataTypeInt64>(type_not_null.get()))
-        return {SortKeyKind::Int64, extractFieldImpl<Int64, Int64, SortKeyKind::Int64, ColumnVector<Int64>>};
+        return makeSortKeyTypeInfo<Int64, Int64, SortKeyKind::Int64, ColumnVector<Int64>>();
     if (checkDataType<DataTypeUInt8>(type_not_null.get()))
-        return {SortKeyKind::UInt64, extractFieldImpl<UInt8, UInt64, SortKeyKind::UInt64, ColumnVector<UInt8>>};
+        return makeSortKeyTypeInfo<UInt8, UInt64, SortKeyKind::UInt64, ColumnVector<UInt8>>();
     if (checkDataType<DataTypeUInt16>(type_not_null.get()))
-        return {SortKeyKind::UInt64, extractFieldImpl<UInt16, UInt64, SortKeyKind::UInt64, ColumnVector<UInt16>>};
+        return makeSortKeyTypeInfo<UInt16, UInt64, SortKeyKind::UInt64, ColumnVector<UInt16>>();
     if (checkDataType<DataTypeUInt32>(type_not_null.get()))
-        return {SortKeyKind::UInt64, extractFieldImpl<UInt32, UInt64, SortKeyKind::UInt64, ColumnVector<UInt32>>};
+        return makeSortKeyTypeInfo<UInt32, UInt64, SortKeyKind::UInt64, ColumnVector<UInt32>>();
     if (checkDataType<DataTypeUInt64>(type_not_null.get()))
-        return {SortKeyKind::UInt64, extractFieldImpl<UInt64, UInt64, SortKeyKind::UInt64, ColumnVector<UInt64>>};
+        return makeSortKeyTypeInfo<UInt64, UInt64, SortKeyKind::UInt64, ColumnVector<UInt64>>();
     if (checkDataType<DataTypeFloat32>(type_not_null.get()))
-        return {SortKeyKind::Float32, extractFieldImpl<Float32, Float32, SortKeyKind::Float32, ColumnVector<Float32>>};
+        return makeFloatSortKeyTypeInfo<Float32, SortKeyKind::Float32, ColumnVector<Float32>>();
     if (checkDataType<DataTypeFloat64>(type_not_null.get()))
-        return {SortKeyKind::Float64, extractFieldImpl<Float64, Float64, SortKeyKind::Float64, ColumnVector<Float64>>};
+        return makeFloatSortKeyTypeInfo<Float64, SortKeyKind::Float64, ColumnVector<Float64>>();
     if (checkDataType<DataTypeDecimal32>(type_not_null.get()))
-        return {SortKeyKind::Decimal32, extractDecimalFieldImpl<Decimal32, SortKeyKind::Decimal32>};
+        return makeDecimalSortKeyTypeInfo<Decimal32, SortKeyKind::Decimal32>();
     if (checkDataType<DataTypeDecimal64>(type_not_null.get()))
-        return {SortKeyKind::Decimal64, extractDecimalFieldImpl<Decimal64, SortKeyKind::Decimal64>};
+        return makeDecimalSortKeyTypeInfo<Decimal64, SortKeyKind::Decimal64>();
     if (checkDataType<DataTypeDecimal128>(type_not_null.get()))
-        return {SortKeyKind::Decimal128, extractDecimalFieldImpl<Decimal128, SortKeyKind::Decimal128>};
+        return makeDecimalSortKeyTypeInfo<Decimal128, SortKeyKind::Decimal128>();
     if (checkDataType<DataTypeDecimal256>(type_not_null.get()))
-        return {SortKeyKind::Decimal256, extractDecimalFieldImpl<Decimal256, SortKeyKind::Decimal256>};
+        return makeDecimalSortKeyTypeInfo<Decimal256, SortKeyKind::Decimal256>();
     if (checkDataType<DataTypeDate>(type_not_null.get()))
-        return {SortKeyKind::Date, extractFieldImpl<UInt16, UInt64, SortKeyKind::Date, ColumnVector<UInt16>>};
+        return makeSortKeyTypeInfo<UInt16, UInt64, SortKeyKind::Date, ColumnVector<UInt16>>();
     if (checkDataType<DataTypeDateTime>(type_not_null.get()))
-        return {SortKeyKind::DateTime, extractFieldImpl<UInt32, UInt64, SortKeyKind::DateTime, ColumnVector<UInt32>>};
+        return makeSortKeyTypeInfo<UInt32, UInt64, SortKeyKind::DateTime, ColumnVector<UInt32>>();
+    if (checkDataType<DataTypeMyDate>(type_not_null.get()))
+        return makeSortKeyTypeInfo<UInt64, UInt64, SortKeyKind::Date, ColumnVector<UInt64>>();
+    if (checkDataType<DataTypeMyDateTime>(type_not_null.get()))
+        return makeSortKeyTypeInfo<UInt64, UInt64, SortKeyKind::DateTime, ColumnVector<UInt64>>();
 
     throw Exception("Unsupported order by type for TopN-enhanced multi-stage late materialization: " + type->getName());
 }
@@ -210,7 +310,9 @@ RunningLocalTopN::RunningLocalTopN(
             .column_pos = static_cast<size_t>(column_it - stage1_columns.begin()),
             .kind = type_info.kind,
             .direction = order_by_column.direction,
+            .build_column_view = type_info.build_column_view,
             .extract_field = type_info.extract_field,
+            .compare_column_with_owned = type_info.compare_column_with_owned,
         });
     }
 }
@@ -228,31 +330,30 @@ int RunningLocalTopN::compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSort
 }
 
 int RunningLocalTopN::compareRowWithOwnedKey(
-    const std::vector<const IColumn *> & sort_columns,
+    const std::vector<SortKeyColumnView> & sort_columns,
     size_t row,
     const OwnedSortKey & rhs) const
 {
     RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size());
     RUNTIME_CHECK(rhs.size == sort_key_columns.size());
 
-    SortKeyField lhs_field;
     for (size_t i = 0; i < sort_key_columns.size(); ++i)
     {
-        sort_key_columns[i].extract_field(*sort_columns[i], row, lhs_field);
-        if (const auto cmp = compareSortKeyFields(lhs_field, rhs.fields[i]); cmp != 0)
+        const auto cmp = sort_key_columns[i].compare_column_with_owned(sort_columns[i], row, rhs.fields[i]);
+        if (cmp != 0)
             return cmp * sort_key_columns[i].direction;
     }
     return 0;
 }
 
-OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row) const
+OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row) const
 {
     RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size());
 
     OwnedSortKey key;
     key.size = sort_key_columns.size();
     for (size_t i = 0; i < sort_key_columns.size(); ++i)
-        sort_key_columns[i].extract_field(*sort_columns[i], row, key.fields[i]);
+        sort_key_columns[i].extract_field(sort_columns[i], row, key.fields[i]);
     return key;
 }
 
@@ -271,7 +372,7 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         return result;
 
     std::vector<ColumnPtr> materialized_const_columns;
-    std::vector<const IColumn *> sort_columns;
+    std::vector<SortKeyColumnView> sort_columns;
     sort_columns.reserve(sort_key_columns.size());
     for (const auto & desc : sort_key_columns)
     {
@@ -281,7 +382,7 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
             materialized_const_columns.push_back(std::move(full_column));
             column = materialized_const_columns.back();
         }
-        sort_columns.push_back(column.get());
+        sort_columns.push_back(desc.build_column_view(*column));
     }
 
     auto markCandidate = [&](size_t row) {
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
index c45d6797bc8..8dbb183bdd9 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
@@ -30,7 +30,7 @@
 namespace DB::DM
 {
 
-inline constexpr UInt64 multi_stage_late_materialization_topn_max_topk = 4096;
+inline constexpr UInt64 multi_stage_late_materialization_topn_max_topk = 2048;
 inline constexpr size_t multi_stage_late_materialization_topn_max_order_by_columns = 4;
 
 struct MultiStageLateMaterializationTopNOrderByColumn
@@ -82,6 +82,12 @@ struct RunningLocalTopNUpdateResult
     size_t passed_count = 0;
 };
 
+struct SortKeyColumnView
+{
+    const void * data = nullptr;
+    const IColumn::Filter * null_map = nullptr;
+};
+
 class RunningLocalTopN
 {
 public:
@@ -99,13 +105,20 @@ class RunningLocalTopN
 private:
     struct SortKeyColumnDesc
     {
-        using ExtractFieldFn = void (*)(const IColumn & column, size_t row, SortKeyField & field);
+        using BuildColumnViewFn = SortKeyColumnView (*)(const IColumn & column);
+        using ExtractFieldFn = void (*)(const SortKeyColumnView & column, size_t row, SortKeyField & field);
+        using CompareColumnWithOwnedFn = int (*)(
+            const SortKeyColumnView & column,
+            size_t row,
+            const SortKeyField & field);
 
         ColumnID column_id = EmptyColumnID;
         size_t column_pos = 0;
         SortKeyKind kind = SortKeyKind::Int64;
         int direction = 1;
+        BuildColumnViewFn build_column_view = nullptr;
         ExtractFieldFn extract_field = nullptr;
+        CompareColumnWithOwnedFn compare_column_with_owned = nullptr;
     };
 
     struct HeapEntry
@@ -125,9 +138,9 @@ class RunningLocalTopN
     using Heap = std::priority_queue<HeapEntry, std::vector<HeapEntry>, HeapComparator>;
 
     int compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const;
-    int compareRowWithOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row, const OwnedSortKey & rhs)
+    int compareRowWithOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row, const OwnedSortKey & rhs)
         const;
-    OwnedSortKey materializeOwnedKey(const std::vector<const IColumn *> & sort_columns, size_t row) const;
+    OwnedSortKey materializeOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row) const;
 
 private:
     UInt64 topk;
diff --git a/dbms/src/Storages/StorageDeltaMerge.cpp b/dbms/src/Storages/StorageDeltaMerge.cpp
index d5ea4064ca7..a1b93a6e728 100644
--- a/dbms/src/Storages/StorageDeltaMerge.cpp
+++ b/dbms/src/Storages/StorageDeltaMerge.cpp
@@ -25,6 +25,8 @@
 #include <DataTypes/DataTypeDate.h>
 #include <DataTypes/DataTypeDateTime.h>
 #include <DataTypes/DataTypeDecimal.h>
+#include <DataTypes/DataTypeMyDate.h>
+#include <DataTypes/DataTypeMyDateTime.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <DataTypes/isSupportedDataTypeCast.h>
@@ -816,7 +818,9 @@ bool isSupportedMultiStageLateMaterializationTopNOrderByType(const DataTypePtr &
         || checkDataType<DataTypeDecimal128>(type_not_null.get()) //
         || checkDataType<DataTypeDecimal256>(type_not_null.get()) //
         || checkDataType<DataTypeDate>(type_not_null.get()) //
-        || checkDataType<DataTypeDateTime>(type_not_null.get());
+        || checkDataType<DataTypeDateTime>(type_not_null.get()) //
+        || checkDataType<DataTypeMyDate>(type_not_null.get()) //
+        || checkDataType<DataTypeMyDateTime>(type_not_null.get());
 }
 
 const TiDB::ColumnInfo * findSourceColumnByID(const TiDB::ColumnInfos & source_columns, ColumnID column_id)
diff --git a/docs/design/2026-07-25-multi-stage-late-materialization.md b/docs/design/2026-07-25-multi-stage-late-materialization.md
index e55e90e4174..4422b4b3631 100644
--- a/docs/design/2026-07-25-multi-stage-late-materialization.md
+++ b/docs/design/2026-07-25-multi-stage-late-materialization.md
@@ -188,7 +188,7 @@ DirectMode 不需要新增 `full_stream`。它本质上还是两次读取：
 
 启用条件：
 
-- TiFlash setting `dt_enable_multi_stage_late_materialization` 为 true。
+- TiFlash setting `dt_enable_multi_stage_late_materialization >= 1`。
 - 当前 query 使用 pipeline 执行模式。
 - TableScan 已存在 pushed down filters，即可以形成 Stage 0。
 - TableScan 上方存在 residual Selection，即 `filter_conditions.hasValue()`。
@@ -458,13 +458,13 @@ if (filter_conditions.hasValue() && multi_stage_lm_consumed_filter_conditions)
 改动点：
 
 - 在 TiFlash pipeline TableScan 构建路径中识别 candidate：
-  - `dt_enable_multi_stage_late_materialization` 为 true。
+  - `dt_enable_multi_stage_late_materialization >= 1`。
   - `filter_conditions.hasValue()`。
   - TableScan 已经有 pushed down filters。
   - `generated_column_infos` 为空。
   - `read_opts.keep_order == false`。
   - 列数满足启用规则。
-- 增加 TiFlash setting `dt_enable_multi_stage_late_materialization`，并增加一个内部 flag，例如 `multi_stage_lm_enabled`。
+- 增加 TiFlash setting `dt_enable_multi_stage_late_materialization`，其中 `0` 表示禁用，`1` 表示启用 selection MSLM，`2` 表示启用 TopN-enhanced MSLM；并增加一个内部 flag，例如 `multi_stage_lm_enabled`。
 - 将 residual filter conditions 传入 DeltaMerge read 层。
 - 当 `multi_stage_lm_enabled` 为 true 时，跳过 `executePushedDownFilter(...)`。
 - 同时把当前 `group_builder.getCurProfileInfos()` 注册给 `filter_conditions.executor_id`。
diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index f8c7027ceb2..abdd89d4f34 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -74,7 +74,7 @@ LIMIT k;
 
 - 只支持已经满足现有 MSLM 启用条件的查询。也就是说，TableScan 必须已有 pushed-down filters 作为 Stage 0，且 TableScan 上方必须有 residual Selection 作为 Stage 1。
 - 只支持 `ORDER BY` plain columns，即 TiPB `ByItem.expr` 必须直接是 `ColumnRef`。
-- 只支持 constant `LIMIT` 和 constant `OFFSET`。
+- 只支持 TiDB 已经能下发为 constant TopN limit 的查询。SQL 层的 `OFFSET` 会在 TiDB pushdown TopN 时合并进 TiPB `TopN.limit`。
 - 只支持 TiDB 下发的单表 `TopN -> Selection -> TableScan`。在 TiFlash physical plan 中，这一层 `Selection` 必须已经被合并到 `PhysicalTableScan::filter_conditions`，所以第一版实际只识别 `PhysicalTopN -> PhysicalTableScan(with filter_conditions and pushed_down_filters)`。
 - 第一版不穿透 `PhysicalFilter`。如果 TiFlash physical plan 中仍然存在独立的 `PhysicalFilter -> PhysicalTableScan`，说明 residual Selection 没有进入当前 MSLM Stage 1 识别范围，TopN-enhanced MSLM 直接禁用。
 - 不移除上层 TopN。
@@ -86,17 +86,13 @@ LIMIT k;
 - 阶段二支持有 TableScan pushed filter 但没有 residual Selection 的 `TopN -> TableScan`。这要求把 residual filter 变成 optional，Stage 1 只读取 order-by columns。
 - 阶段三支持完全没有 filter 的 `TopN -> TableScan`。这需要把 MSLM 从当前的 filter-driven 模型泛化为 order-by-column-driven 模型，Stage 1 order-by stream 需要成为 block driver，工程风险更高。
 
-其中本地 TopN 使用的 K 为：
+其中本地 TopN 使用的 K 为 TiPB `TopN.limit`：
 
 ```text
-K = LIMIT + OFFSET
+K = tipb.TopN.limit
 ```
 
-如果没有 OFFSET，则：
-
-```text
-K = LIMIT
-```
+对 SQL 层的 `LIMIT offset, count`，TiDB 下发 pushed-down TopN 前会把 `PhysicalTopN.Count` 改成 `offset + count`，所以 TiFlash 侧不再额外处理 offset。
 
 ## 核心思路
 
@@ -202,7 +198,7 @@ TableScan
 struct MSLMTopNDescription
 {
     SortDescription storage_sort_description;
-    UInt64 topk; // limit + offset
+    UInt64 topk; // tipb.TopN.limit
     ColumnDefines order_by_columns;
 };
 ```
@@ -538,13 +534,12 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 
 必须满足：
 
-- `dt_enable_multi_stage_late_materialization` 为 true。
+- `dt_enable_multi_stage_late_materialization == 2`。
 - 当前查询已经满足 MSLM 启用条件。
 - 查询包含 TopN，且 TopN 的 child 直接是单表 `PhysicalTableScan`。
 - `PhysicalTableScan` 已经包含 residual `filter_conditions`，并且 TableScan 已经包含 pushed-down filters。
 - TopN 上方仍保留全局 TopN executor。
-- TopN 的 `LIMIT` 和 `OFFSET` 是常量。
-- `topk = limit + offset` 未超过阈值。
+- TiPB `TopN.limit` 已经包含 SQL 层 offset，并且未超过阈值。
 - TopN 之前的 Selection predicates 全部在 MSLM Stage 0 或 Stage 1 内执行。
 - ORDER BY 只包含 direct `ColumnRef` plain columns。
 - ORDER BY columns 能通过 `ColumnRef` index 映射到 `table_scan.getColumns()` 中的 column id，并能进一步映射到 DeltaMerge `ColumnDefine`。
@@ -557,7 +552,7 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 建议第一版阈值：
 
 ```text
-topk <= 4096
+topk <= 2048
 order_by_column_count <= 4
 ```
 

From 4bf257de7911a7e19f6222fcc693c0f2d8a8bb82 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Tue, 28 Jul 2026 14:52:54 +0800
Subject: [PATCH 07/18] Refine multi-stage late materialization logging

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 dbms/src/Debug/MockStorage.cpp                |  4 +-
 .../Coprocessor/DAGStorageInterpreter.cpp     |  5 +-
 ...ageLateMaterializationBlockInputStream.cpp | 10 +--
 ...ultiStageLateMaterializationRuntimeStats.h | 61 +++++++++++++++++++
 dbms/src/Storages/DeltaMerge/Segment.cpp      |  2 +-
 5 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp
index 03e72a2ddc3..97cb10cb25b 100644
--- a/dbms/src/Debug/MockStorage.cpp
+++ b/dbms/src/Debug/MockStorage.cpp
@@ -504,7 +504,9 @@ void MockStorage::buildExecFromDeltaMerge(
         if (enable_multi_stage_late_materialization)
         {
             multi_stage_late_materialization_runtime_stats
-                = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>();
+                = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
+                    fmt::format("mock table_scan_executor_id={}", table_scan_executor_id),
+                    multi_stage_late_materialization_topn != nullptr);
             if (auto * dag_context = context.getDAGContext(); dag_context != nullptr && !table_scan_executor_id.empty())
             {
                 dag_context->setExecutorRowsOverride(
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index 7502210d29a..ba7d5c437ea 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -1045,8 +1045,9 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
     DM::MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats;
     if (enable_multi_stage_late_materialization)
     {
-        multi_stage_late_materialization_runtime_stats
-            = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>();
+        multi_stage_late_materialization_runtime_stats = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
+            fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID()),
+            multi_stage_late_materialization_topn != nullptr);
         dagContext().setExecutorRowsOverride(
             table_scan.getTableScanExecutorID(),
             std::shared_ptr<std::atomic<UInt64>>(
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index 2b462dcc239..6c60df24b3e 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -401,11 +401,11 @@ void MultiStageLateMaterializationBlockInputStream::logSummary()
         return;
     summary_logged = true;
 
-    LOG_INFO(
-        log,
-        "Multi-stage late materialization finished, late_mode_blocks={} direct_mode_blocks={}",
-        late_mode_blocks,
-        direct_mode_blocks);
+    if (runtime_stats)
+        runtime_stats->finishStream(
+            late_mode_blocks,
+            direct_mode_blocks,
+            running_topn != nullptr ? running_topn->heapSize() : 0);
 }
 
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index 20d0a9a0ad4..c0e0651aeab 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <common/logger_useful.h>
 #include <common/types.h>
 
 #include <atomic>
@@ -24,9 +25,69 @@ namespace DB::DM
 
 struct MultiStageLateMaterializationRuntimeStats
 {
+    MultiStageLateMaterializationRuntimeStats() = default;
+
+    MultiStageLateMaterializationRuntimeStats(const String & log_id_, bool topn_enabled_)
+        : topn_enabled(topn_enabled_)
+        , log(Logger::get("MultiStageLateMaterialization", log_id_))
+    {}
+
+    ~MultiStageLateMaterializationRuntimeStats()
+    {
+        try
+        {
+            logSummary();
+        }
+        catch (...)
+        {
+        }
+    }
+
+    void finishStream(UInt64 stream_late_mode_blocks, UInt64 stream_direct_mode_blocks, UInt64 stream_topn_heap_size)
+    {
+        finished_streams.fetch_add(1, std::memory_order_relaxed);
+        late_mode_blocks.fetch_add(stream_late_mode_blocks, std::memory_order_relaxed);
+        direct_mode_blocks.fetch_add(stream_direct_mode_blocks, std::memory_order_relaxed);
+        topn_heap_size_sum.fetch_add(stream_topn_heap_size, std::memory_order_relaxed);
+    }
+
+    void logSummary() const
+    {
+        if (log == nullptr)
+            return;
+
+        const auto stage1_rows = stage1_output_rows.load(std::memory_order_relaxed);
+        const auto topn_candidate_rows_for_log
+            = topn_enabled ? topn_candidate_rows.load(std::memory_order_relaxed) : 0;
+        const auto topn_filtered_rows
+            = topn_enabled && stage1_rows >= topn_candidate_rows_for_log ? stage1_rows - topn_candidate_rows_for_log
+                                                                          : 0;
+
+        LOG_INFO(
+            log,
+            "Multi-stage late materialization finished, streams={} late_mode_blocks={} direct_mode_blocks={} "
+            "stage0_output_rows={} stage1_output_rows={} topn_enabled={} topn_candidate_rows={} "
+            "topn_filtered_rows={} topn_heap_size_sum={}",
+            finished_streams.load(std::memory_order_relaxed),
+            late_mode_blocks.load(std::memory_order_relaxed),
+            direct_mode_blocks.load(std::memory_order_relaxed),
+            stage0_output_rows.load(std::memory_order_relaxed),
+            stage1_rows,
+            topn_enabled,
+            topn_candidate_rows_for_log,
+            topn_filtered_rows,
+            topn_heap_size_sum.load(std::memory_order_relaxed));
+    }
+
     std::atomic<UInt64> stage0_output_rows{0};
     std::atomic<UInt64> stage1_output_rows{0};
     std::atomic<UInt64> topn_candidate_rows{0};
+    std::atomic<UInt64> finished_streams{0};
+    std::atomic<UInt64> late_mode_blocks{0};
+    std::atomic<UInt64> direct_mode_blocks{0};
+    std::atomic<UInt64> topn_heap_size_sum{0};
+    bool topn_enabled = false;
+    LoggerPtr log;
 };
 
 using MultiStageLateMaterializationRuntimeStatsPtr = std::shared_ptr<MultiStageLateMaterializationRuntimeStats>;
diff --git a/dbms/src/Storages/DeltaMerge/Segment.cpp b/dbms/src/Storages/DeltaMerge/Segment.cpp
index 31abaeee776..95f751c5cb2 100644
--- a/dbms/src/Storages/DeltaMerge/Segment.cpp
+++ b/dbms/src/Storages/DeltaMerge/Segment.cpp
@@ -3732,7 +3732,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
             expected_block_size,
             ReadTag::Query);
 
-        LOG_INFO(
+        LOG_DEBUG(
             segment_snap->log,
             "Use multi-stage late materialization, stage0_filter_columns={} stage1_filter_columns={} "
             "final_rest_columns={}",

From e39d2084c6caa0d0632e5d1b640e47ad5e22a3b8 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Tue, 28 Jul 2026 23:24:05 +0800
Subject: [PATCH 08/18] Add adaptive disable for MSLM local TopN

---
 ...ageLateMaterializationBlockInputStream.cpp |  57 +++++-
 ...StageLateMaterializationBlockInputStream.h |   9 +
 ...ultiStageLateMaterializationRuntimeStats.h |  26 ++-
 .../MultiStageLateMaterializationTopN.h       |  16 +-
 .../gtest_skippable_block_input_stream.cpp    | 172 ++++++++++++++++++
 .../2026-07-26-running-local-topn-for-mslm.md |  28 ++-
 6 files changed, 287 insertions(+), 21 deletions(-)

diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index 6c60df24b3e..b997f410a80 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -17,6 +17,8 @@
 #include <Storages/DeltaMerge/DeltaMergeHelpers.h>
 #include <Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h>
 
+#include <algorithm>
+
 namespace DB::DM
 {
 namespace
@@ -282,6 +284,53 @@ Block MultiStageLateMaterializationBlockInputStream::buildDirectModeBlock(
     return full_block;
 }
 
+void MultiStageLateMaterializationBlockInputStream::updateTopNAdaptiveState(
+    UInt64 residual_passed_rows,
+    UInt64 topn_candidate_rows)
+{
+    if (running_topn == nullptr || topn_adaptive_disabled || residual_passed_rows == 0)
+        return;
+
+    const auto warmup_rows
+        = std::max(running_topn->topK() * 4, multi_stage_late_materialization_topn_adaptive_rows_threshold);
+    if (!topn_adaptive_warmed_up)
+    {
+        topn_adaptive_warmup_observed_rows += residual_passed_rows;
+        if (topn_adaptive_warmup_observed_rows < warmup_rows)
+            return;
+
+        topn_adaptive_warmed_up = true;
+        return;
+    }
+
+    topn_adaptive_input_rows += residual_passed_rows;
+    topn_adaptive_candidate_rows += topn_candidate_rows;
+
+    const auto min_input_rows_to_check
+        = std::max(running_topn->topK() * 4, multi_stage_late_materialization_topn_adaptive_rows_threshold);
+    if (topn_adaptive_input_rows < min_input_rows_to_check)
+        return;
+
+    if (topn_adaptive_candidate_rows * multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator
+        < topn_adaptive_input_rows * multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator)
+    {
+        return;
+    }
+
+    topn_adaptive_disabled = true;
+    LOG_INFO(
+        log,
+        "Disable running local TopN for multi-stage late materialization adaptively, warmup_rows={} input_rows={} "
+        "candidate_rows={} min_input_rows_to_check={} heap_size={} disable_ratio={}/{}",
+        topn_adaptive_warmup_observed_rows,
+        topn_adaptive_input_rows,
+        topn_adaptive_candidate_rows,
+        min_input_rows_to_check,
+        running_topn->heapSize(),
+        multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator,
+        multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator);
+}
+
 Block MultiStageLateMaterializationBlockInputStream::read()
 {
     while (true)
@@ -330,7 +379,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         if (runtime_stats)
             runtime_stats->stage1_output_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
 
-        if (running_topn == nullptr && runtime_stats)
+        if (!shouldUseRunningTopN() && runtime_stats)
             runtime_stats->topn_candidate_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
 
         if (residual_passed_rows == 0)
@@ -340,11 +389,12 @@ Block MultiStageLateMaterializationBlockInputStream::read()
             continue;
         }
 
-        if (running_topn != nullptr)
+        if (shouldUseRunningTopN())
         {
             auto topn_result = running_topn->update(stage1_block, residual_filter_ptr, residual_passed_rows);
             if (runtime_stats)
                 runtime_stats->topn_candidate_rows.fetch_add(topn_result.passed_count, std::memory_order_relaxed);
+            updateTopNAdaptiveState(residual_passed_rows, topn_result.passed_count);
 
             if (topn_result.passed_count == 0)
             {
@@ -405,7 +455,8 @@ void MultiStageLateMaterializationBlockInputStream::logSummary()
         runtime_stats->finishStream(
             late_mode_blocks,
             direct_mode_blocks,
-            running_topn != nullptr ? running_topn->heapSize() : 0);
+            running_topn != nullptr ? running_topn->heapSize() : 0,
+            topn_adaptive_disabled);
 }
 
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
index c7bde1736a1..00a3abf3c08 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
@@ -90,6 +90,10 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
         const IColumn::Filter * residual_filter,
         size_t residual_passed_rows);
 
+    bool shouldUseRunningTopN() const { return running_topn != nullptr && !topn_adaptive_disabled; }
+
+    void updateTopNAdaptiveState(UInt64 residual_passed_rows, UInt64 topn_candidate_rows);
+
     void logSummary();
 
 private:
@@ -105,6 +109,11 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
 
     size_t late_mode_blocks = 0;
     size_t direct_mode_blocks = 0;
+    UInt64 topn_adaptive_warmup_observed_rows = 0;
+    UInt64 topn_adaptive_input_rows = 0;
+    UInt64 topn_adaptive_candidate_rows = 0;
+    bool topn_adaptive_warmed_up = false;
+    bool topn_adaptive_disabled = false;
     bool summary_logged = false;
 
     const LoggerPtr log;
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index c0e0651aeab..7b7a19d2c03 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -39,16 +39,21 @@ struct MultiStageLateMaterializationRuntimeStats
             logSummary();
         }
         catch (...)
-        {
-        }
+        {}
     }
 
-    void finishStream(UInt64 stream_late_mode_blocks, UInt64 stream_direct_mode_blocks, UInt64 stream_topn_heap_size)
+    void finishStream(
+        UInt64 stream_late_mode_blocks,
+        UInt64 stream_direct_mode_blocks,
+        UInt64 stream_topn_heap_size,
+        bool stream_topn_adaptive_disabled)
     {
         finished_streams.fetch_add(1, std::memory_order_relaxed);
         late_mode_blocks.fetch_add(stream_late_mode_blocks, std::memory_order_relaxed);
         direct_mode_blocks.fetch_add(stream_direct_mode_blocks, std::memory_order_relaxed);
         topn_heap_size_sum.fetch_add(stream_topn_heap_size, std::memory_order_relaxed);
+        if (stream_topn_adaptive_disabled)
+            topn_adaptive_disabled_streams.fetch_add(1, std::memory_order_relaxed);
     }
 
     void logSummary() const
@@ -57,17 +62,16 @@ struct MultiStageLateMaterializationRuntimeStats
             return;
 
         const auto stage1_rows = stage1_output_rows.load(std::memory_order_relaxed);
-        const auto topn_candidate_rows_for_log
-            = topn_enabled ? topn_candidate_rows.load(std::memory_order_relaxed) : 0;
-        const auto topn_filtered_rows
-            = topn_enabled && stage1_rows >= topn_candidate_rows_for_log ? stage1_rows - topn_candidate_rows_for_log
-                                                                          : 0;
+        const auto topn_candidate_rows_for_log = topn_enabled ? topn_candidate_rows.load(std::memory_order_relaxed) : 0;
+        const auto topn_filtered_rows = topn_enabled && stage1_rows >= topn_candidate_rows_for_log
+            ? stage1_rows - topn_candidate_rows_for_log
+            : 0;
 
         LOG_INFO(
             log,
             "Multi-stage late materialization finished, streams={} late_mode_blocks={} direct_mode_blocks={} "
             "stage0_output_rows={} stage1_output_rows={} topn_enabled={} topn_candidate_rows={} "
-            "topn_filtered_rows={} topn_heap_size_sum={}",
+            "topn_filtered_rows={} topn_heap_size_sum={} topn_adaptive_disabled_streams={}",
             finished_streams.load(std::memory_order_relaxed),
             late_mode_blocks.load(std::memory_order_relaxed),
             direct_mode_blocks.load(std::memory_order_relaxed),
@@ -76,7 +80,8 @@ struct MultiStageLateMaterializationRuntimeStats
             topn_enabled,
             topn_candidate_rows_for_log,
             topn_filtered_rows,
-            topn_heap_size_sum.load(std::memory_order_relaxed));
+            topn_heap_size_sum.load(std::memory_order_relaxed),
+            topn_adaptive_disabled_streams.load(std::memory_order_relaxed));
     }
 
     std::atomic<UInt64> stage0_output_rows{0};
@@ -86,6 +91,7 @@ struct MultiStageLateMaterializationRuntimeStats
     std::atomic<UInt64> late_mode_blocks{0};
     std::atomic<UInt64> direct_mode_blocks{0};
     std::atomic<UInt64> topn_heap_size_sum{0};
+    std::atomic<UInt64> topn_adaptive_disabled_streams{0};
     bool topn_enabled = false;
     LoggerPtr log;
 };
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
index 8dbb183bdd9..7d4a2d97b4f 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h
@@ -32,6 +32,9 @@ namespace DB::DM
 
 inline constexpr UInt64 multi_stage_late_materialization_topn_max_topk = 2048;
 inline constexpr size_t multi_stage_late_materialization_topn_max_order_by_columns = 4;
+inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_rows_threshold = 8192;
+inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator = 7;
+inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator = 10;
 
 struct MultiStageLateMaterializationTopNOrderByColumn
 {
@@ -101,16 +104,15 @@ class RunningLocalTopN
         size_t residual_passed_rows);
 
     size_t heapSize() const { return heap.size(); }
+    UInt64 topK() const { return topk; }
 
 private:
     struct SortKeyColumnDesc
     {
         using BuildColumnViewFn = SortKeyColumnView (*)(const IColumn & column);
         using ExtractFieldFn = void (*)(const SortKeyColumnView & column, size_t row, SortKeyField & field);
-        using CompareColumnWithOwnedFn = int (*)(
-            const SortKeyColumnView & column,
-            size_t row,
-            const SortKeyField & field);
+        using CompareColumnWithOwnedFn
+            = int (*)(const SortKeyColumnView & column, size_t row, const SortKeyField & field);
 
         ColumnID column_id = EmptyColumnID;
         size_t column_pos = 0;
@@ -138,8 +140,10 @@ class RunningLocalTopN
     using Heap = std::priority_queue<HeapEntry, std::vector<HeapEntry>, HeapComparator>;
 
     int compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const;
-    int compareRowWithOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row, const OwnedSortKey & rhs)
-        const;
+    int compareRowWithOwnedKey(
+        const std::vector<SortKeyColumnView> & sort_columns,
+        size_t row,
+        const OwnedSortKey & rhs) const;
     OwnedSortKey materializeOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row) const;
 
 private:
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
index cc51ac457b2..ec69668f3e0 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
@@ -96,6 +96,14 @@ std::vector<UInt8> makeResidualValues(size_t rows, size_t passed_rows)
     return values;
 }
 
+std::vector<UInt8> makeTopNAdaptiveValues(size_t rows, size_t candidate_rows)
+{
+    RUNTIME_CHECK(candidate_rows <= rows);
+    std::vector<UInt8> values(rows, 3);
+    std::fill(values.begin(), values.begin() + candidate_rows, 1);
+    return values;
+}
+
 size_t drainMultiStageStream(const BlockInputStreamPtr & stream)
 {
     size_t rows = 0;
@@ -276,6 +284,30 @@ PushDownFilterPtr makeResidualFilterForMultiStageTest()
         nullptr);
 }
 
+PushDownFilterPtr makeResidualFilterWithTopNForMultiStageTest(UInt64 topk)
+{
+    auto filter_columns = std::make_shared<ColumnDefines>(makeMultiStageFilterColumns());
+    auto actions = std::make_shared<ExpressionActions>(toEmptyBlock(*filter_columns).getNamesAndTypes());
+    actions->add(ExpressionAction::copyColumn(MULTI_STAGE_FILTER_COL_NAME, MULTI_STAGE_FILTER_TMP_COL_NAME));
+
+    auto topn = std::make_shared<MultiStageLateMaterializationTopNDescription>();
+    topn->topk = topk;
+    topn->order_by_columns.push_back(MultiStageLateMaterializationTopNOrderByColumn{
+        .column_id = MULTI_STAGE_FILTER_COL_ID,
+        .direction = 1,
+    });
+
+    return std::make_shared<PushDownFilter>(
+        EMPTY_RS_OPERATOR,
+        actions,
+        nullptr,
+        filter_columns,
+        MULTI_STAGE_FILTER_TMP_COL_NAME,
+        nullptr,
+        nullptr,
+        topn);
+}
+
 void ensureMultiStageTestFunctionsRegistered()
 {
     static std::once_flag once;
@@ -1007,6 +1039,146 @@ try
 }
 CATCH
 
+TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDoesNotDisableDuringWarmup)
+try
+{
+    constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
+
+    std::vector<IColumn::Filter> stage0_filters;
+    std::vector<std::vector<UInt8>> filter_values;
+    std::vector<std::vector<Int64>> rest_values;
+    stage0_filters.emplace_back(warmup_rows, 1);
+    filter_values.emplace_back(warmup_rows, 1);
+    rest_values.emplace_back(makeRestValues(0, warmup_rows));
+
+    auto stage0_stream = std::make_shared<DeterministicStage0FilterInputStream>(std::move(stage0_filters));
+    auto stage1_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageFilterColumns(),
+        filter_values,
+        rest_values);
+    auto final_rest_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageRestColumns(),
+        filter_values,
+        rest_values);
+
+    auto bitmap_filter = std::make_shared<BitmapFilter>(warmup_rows, 1);
+    auto runtime_stats = std::make_shared<MultiStageLateMaterializationRuntimeStats>();
+    auto stream = std::make_shared<MultiStageLateMaterializationBlockInputStream>(
+        makeMultiStageColumnsToRead(),
+        stage0_stream,
+        stage1_stream,
+        final_rest_stream,
+        makeResidualFilterWithTopNForMultiStageTest(/*topk=*/2),
+        bitmap_filter,
+        "test",
+        runtime_stats);
+
+    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows);
+    ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows);
+    ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows);
+    ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0);
+}
+CATCH
+
+TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDoesNotDisableBelowThreshold)
+try
+{
+    constexpr UInt64 topk = 2;
+    constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
+    constexpr size_t second_block_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
+    constexpr size_t below_threshold_candidate_rows = second_block_rows
+            * multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator
+            / multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator
+        - 1;
+
+    std::vector<IColumn::Filter> stage0_filters;
+    std::vector<std::vector<UInt8>> filter_values;
+    std::vector<std::vector<Int64>> rest_values;
+    stage0_filters.emplace_back(warmup_rows, 1);
+    filter_values.emplace_back(warmup_rows, 2);
+    rest_values.emplace_back(makeRestValues(0, warmup_rows));
+    stage0_filters.emplace_back(second_block_rows, 1);
+    filter_values.emplace_back(makeTopNAdaptiveValues(second_block_rows, below_threshold_candidate_rows));
+    rest_values.emplace_back(makeRestValues(100000, second_block_rows));
+
+    auto stage0_stream = std::make_shared<DeterministicStage0FilterInputStream>(std::move(stage0_filters));
+    auto stage1_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageFilterColumns(),
+        filter_values,
+        rest_values);
+    auto final_rest_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageRestColumns(),
+        filter_values,
+        rest_values);
+
+    auto bitmap_filter = std::make_shared<BitmapFilter>(warmup_rows + second_block_rows, 1);
+    auto runtime_stats = std::make_shared<MultiStageLateMaterializationRuntimeStats>();
+    auto stream = std::make_shared<MultiStageLateMaterializationBlockInputStream>(
+        makeMultiStageColumnsToRead(),
+        stage0_stream,
+        stage1_stream,
+        final_rest_stream,
+        makeResidualFilterWithTopNForMultiStageTest(topk),
+        bitmap_filter,
+        "test",
+        runtime_stats);
+
+    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows + below_threshold_candidate_rows);
+    ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows + below_threshold_candidate_rows);
+    ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0);
+}
+CATCH
+
+TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDisable)
+try
+{
+    constexpr UInt64 topk = 2;
+    constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
+    constexpr size_t second_block_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
+
+    std::vector<IColumn::Filter> stage0_filters;
+    std::vector<std::vector<UInt8>> filter_values;
+    std::vector<std::vector<Int64>> rest_values;
+    stage0_filters.emplace_back(warmup_rows, 1);
+    filter_values.emplace_back(warmup_rows, 1);
+    rest_values.emplace_back(makeRestValues(0, warmup_rows));
+    stage0_filters.emplace_back(second_block_rows, 1);
+    filter_values.emplace_back(second_block_rows, 1);
+    rest_values.emplace_back(makeRestValues(100000, second_block_rows));
+
+    auto stage0_stream = std::make_shared<DeterministicStage0FilterInputStream>(std::move(stage0_filters));
+    auto stage1_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageFilterColumns(),
+        filter_values,
+        rest_values);
+    auto final_rest_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageRestColumns(),
+        filter_values,
+        rest_values);
+
+    auto bitmap_filter = std::make_shared<BitmapFilter>(warmup_rows + second_block_rows, 1);
+    auto runtime_stats = std::make_shared<MultiStageLateMaterializationRuntimeStats>();
+    auto stream = std::make_shared<MultiStageLateMaterializationBlockInputStream>(
+        makeMultiStageColumnsToRead(),
+        stage0_stream,
+        stage1_stream,
+        final_rest_stream,
+        makeResidualFilterWithTopNForMultiStageTest(topk),
+        bitmap_filter,
+        "test",
+        runtime_stats);
+
+    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows + second_block_rows);
+    ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 1);
+}
+CATCH
+
 TEST_F(SkippableBlockInputStreamTest, InMemory2)
 try
 {
diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index abdd89d4f34..44c3504a69f 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -134,12 +134,30 @@ Final:
 
 `topn_candidate_filter` 只用于减少 final rest columns 的读取。上层 TopN executor 仍然保留，因此 storage 内输出可以是全局 TopN 的 superset。
 
-当前 MSLM 已经不再使用 adaptive mode。TopN-enhanced MSLM 因此不需要根据 residual 过滤率在 direct/late 两种策略之间做启发式选择，而是只保留下面三个自然分支：
+当前 MSLM 已经不再使用旧的 direct/late adaptive mode。TopN-enhanced MSLM 因此不需要根据 residual 过滤率在 direct/late 两种策略之间做启发式选择，而是只保留下面三个自然分支：
 
 - `residual_filter && topn_candidate_filter` 全过滤时，skip final rest stream。
 - `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。
 - 其他部分通过场景，把 Stage 1 坐标的 filter compose 回 Stage 0 原始坐标，再用 combined filter 读取 final rest columns。
 
+但是 running local TopN 本身有额外代价：Stage 1 需要读取 order-by columns，执行层需要做 comparator、heap 维护和 candidate filter 合并。如果 TopN 几乎不能减少 final rest columns 的读取，就可能带来性能回退。
+
+因此第一版增加一个运行时 one-way adaptive disable：
+
+```text
+开始时启用 running local TopN
+先用 warmup rows 建立初始 TopN 状态
+warmup 后，持续统计后续 residual input rows 和 TopN candidate rows
+如果后续累计 residual_input_rows 达到最小检查门槛，且 topn_candidate_rows / residual_input_rows 过高
+  后续 blocks 禁用 TopN filter
+  后续 final rest 只使用 residual filter
+禁用后不再重新启用
+```
+
+这个 adaptive 只做运行时止损，不改变正确性：禁用后输出的是 residual Selection 后的更多 rows，仍然是上层全局 TopN 的 superset。上层 TopN executor 保留，因此最终结果不变。
+
+需要注意，这个 runtime disable 不能避免已经加入 Stage 1 的 order-by columns IO，因为 Stage 1 schema 在 stream 构建时已经固定。它主要减少后续 comparator、heap 维护和 TopN candidate filter 处理成本。避免明显不划算的 order-by column 读取仍然依赖静态启发式 gating。
+
 ## 正确性
 
 ### Local TopK 覆盖 Global TopK
@@ -554,6 +572,9 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 ```text
 topk <= 2048
 order_by_column_count <= 4
+runtime adaptive warmup rows >= max(topk * 4, 8192)
+after warmup, check only when cumulative residual_input_rows >= max(topk * 4, 8192)
+after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >= 0.7
 ```
 
 实际阈值可以根据 benchmark 调整。
@@ -634,10 +655,11 @@ order_by_column_count <= 4
 
 - 保留 `stage1_output_rows` internal counter，表示 residual Selection 通过的真实行数。
 - 增加 `topn_candidate_rows` internal counter，表示 `residual_filter && topn_candidate_filter` 之后实际读取 final rest columns 的行数。
+- 增加 runtime one-way adaptive disable 统计，表示有多少 streams 因 TopN candidate ratio 过高而禁用后续 TopN filter。
 - 在 TopN-enhanced MSLM 启用时，可以用 `topn_candidate_rows` overwrite residual Selection executor 的 `actRows`，让 `EXPLAIN ANALYZE` 直接展示 final rest columns materialized rows。
 - 该 overwrite 会改变 Selection `actRows` 的展示语义：它不再表示纯 residual Selection 后的逻辑行数，而是表示 storage 内输出给上层 TopN 的 candidate rows。
 - TableScan executor 的 `actRows` 仍然使用 `stage0_output_rows`。
-- 可选增加 debug log，输出 residual passed rows、TopN candidate rows、heap size、filtered rows。
+- 可选增加 debug log，输出 residual passed rows、TopN candidate rows、heap size、filtered rows、adaptive disable 状态。
 
 ## 测试方案
 
@@ -740,6 +762,8 @@ order_by_column_count <= 4
 
 - 使用保守启发式启用。
 - `topk` 和 order-by column count 设置上限。
+- 使用 runtime one-way adaptive disable，在 candidate ratio 过高时禁用后续 TopN filter。
+- adaptive disable 不能避免 Stage 1 已经额外读取的 order-by columns，因此静态 gating 仍然是主要防线。
 - 对 ties 多的场景 benchmark。
 
 ### Heap 持有历史 Block 导致内存累积

From c7fa510ed5fca4ed8a3bd769745dd673c353e7df Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 10:26:28 +0800
Subject: [PATCH 09/18] Fix flaky MSLM TopN tests

---
 dbms/src/Flash/tests/gtest_executors_with_dm.cpp | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 626b25d640c..9a0e2dfd409 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -630,13 +630,13 @@ try
         /*limit=*/2);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
-    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto expected = executeStreams(&disabled_dag_context);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
     ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection;
-    DAGContext selection_only_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext selection_only_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto selection_only_actual = executeStreams(&selection_only_dag_context);
     ASSERT_TRUE(columnsEqual(expected, selection_only_actual, /*_restrict=*/false))
         << "\n  expect_block: \n"
@@ -645,7 +645,7 @@ try
     assertMultiStageRowsOverride(selection_only_dag_context, 16);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
-    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
         << "\n  expect_block: \n"
@@ -717,11 +717,11 @@ try
         /*limit=*/2);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
-    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto expected = executeStreams(&disabled_dag_context);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
-    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
         << "\n  expect_block: \n"
@@ -793,11 +793,11 @@ try
         /*limit=*/2);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
-    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto expected = executeStreams(&disabled_dag_context);
 
     context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
-    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4);
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
     auto actual = executeStreams(&enabled_dag_context);
     ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false))
         << "\n  expect_block: \n"

From a209e82ee10b6a3b9adc3913d9ccfe010a778831 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 10:48:32 +0800
Subject: [PATCH 10/18] Format MSLM TopN changes

---
 .../Coprocessor/DAGStorageInterpreter.cpp     | 23 ++++++++++---------
 .../MultiStageLateMaterializationTopN.cpp     | 18 ++++-----------
 2 files changed, 17 insertions(+), 24 deletions(-)

diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index ba7d5c437ea..c4a1a043eb8 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -371,10 +371,10 @@ String genErrMsgForLocalRead(const KeyspaceID keyspace_id, const TableID & table
     return table_id == logical_table_id
         ? fmt::format("(while creating read sources from storage, keyspace_id={} table_id={})", keyspace_id, table_id)
         : fmt::format(
-            "(while creating read sources from storage, keyspace_id={} table_id={} logical_table_id={})",
-            keyspace_id,
-            table_id,
-            logical_table_id);
+              "(while creating read sources from storage, keyspace_id={} table_id={} logical_table_id={})",
+              keyspace_id,
+              table_id,
+              logical_table_id);
 }
 } // namespace
 
@@ -1045,9 +1045,10 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
     DM::MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats;
     if (enable_multi_stage_late_materialization)
     {
-        multi_stage_late_materialization_runtime_stats = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
-            fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID()),
-            multi_stage_late_materialization_topn != nullptr);
+        multi_stage_late_materialization_runtime_stats
+            = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
+                fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID()),
+                multi_stage_late_materialization_topn != nullptr);
         dagContext().setExecutorRowsOverride(
             table_scan.getTableScanExecutorID(),
             std::shared_ptr<std::atomic<UInt64>>(
@@ -1057,11 +1058,11 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
             filter_conditions.executor_id,
             multi_stage_late_materialization_topn != nullptr
                 ? std::shared_ptr<std::atomic<UInt64>>(
-                    multi_stage_late_materialization_runtime_stats,
-                    &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
+                      multi_stage_late_materialization_runtime_stats,
+                      &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
                 : std::shared_ptr<std::atomic<UInt64>>(
-                    multi_stage_late_materialization_runtime_stats,
-                    &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+                      multi_stage_late_materialization_runtime_stats,
+                      &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
     }
 
     auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo {
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
index bb37b3e7a0a..6cb85474c46 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -124,10 +124,7 @@ int compareNull(bool lhs_is_null, bool rhs_is_null)
 }
 
 template <typename StorageValueType, typename ColumnType>
-int compareColumnWithOwnedFieldImpl(
-    const SortKeyColumnView & column,
-    size_t row,
-    const SortKeyField & rhs)
+int compareColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs)
 {
     const bool lhs_is_null = isNullAt(column.null_map, row);
     if (lhs_is_null || rhs.is_null)
@@ -138,10 +135,7 @@ int compareColumnWithOwnedFieldImpl(
 }
 
 template <typename ColumnValueType, typename ColumnType>
-int compareFloatColumnWithOwnedFieldImpl(
-    const SortKeyColumnView & column,
-    size_t row,
-    const SortKeyField & rhs)
+int compareFloatColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs)
 {
     const bool lhs_is_null = isNullAt(column.null_map, row);
     if (lhs_is_null || rhs.is_null)
@@ -152,10 +146,7 @@ int compareFloatColumnWithOwnedFieldImpl(
 }
 
 template <typename DecimalType>
-int compareDecimalColumnWithOwnedFieldImpl(
-    const SortKeyColumnView & column,
-    size_t row,
-    const SortKeyField & rhs)
+int compareDecimalColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs)
 {
     const bool lhs_is_null = isNullAt(column.null_map, row);
     if (lhs_is_null || rhs.is_null)
@@ -346,7 +337,8 @@ int RunningLocalTopN::compareRowWithOwnedKey(
     return 0;
 }
 
-OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row) const
+OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector<SortKeyColumnView> & sort_columns, size_t row)
+    const
 {
     RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size());
 

From c9eb74ff022845a52b5ea0b288a857127a251405 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 13:15:40 +0800
Subject: [PATCH 11/18] Format DAGStorageInterpreter with llvm 17

---
 .../Flash/Coprocessor/DAGStorageInterpreter.cpp  | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index c4a1a043eb8..eae4d13a2ba 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -371,10 +371,10 @@ String genErrMsgForLocalRead(const KeyspaceID keyspace_id, const TableID & table
     return table_id == logical_table_id
         ? fmt::format("(while creating read sources from storage, keyspace_id={} table_id={})", keyspace_id, table_id)
         : fmt::format(
-              "(while creating read sources from storage, keyspace_id={} table_id={} logical_table_id={})",
-              keyspace_id,
-              table_id,
-              logical_table_id);
+            "(while creating read sources from storage, keyspace_id={} table_id={} logical_table_id={})",
+            keyspace_id,
+            table_id,
+            logical_table_id);
 }
 } // namespace
 
@@ -1058,11 +1058,11 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
             filter_conditions.executor_id,
             multi_stage_late_materialization_topn != nullptr
                 ? std::shared_ptr<std::atomic<UInt64>>(
-                      multi_stage_late_materialization_runtime_stats,
-                      &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
+                    multi_stage_late_materialization_runtime_stats,
+                    &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
                 : std::shared_ptr<std::atomic<UInt64>>(
-                      multi_stage_late_materialization_runtime_stats,
-                      &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+                    multi_stage_late_materialization_runtime_stats,
+                    &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
     }
 
     auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo {

From 2f799be166a6a9345c60c1814daec07d4ac2381d Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 15:52:57 +0800
Subject: [PATCH 12/18] Clamp multi-stage late materialization mode

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 .../Flash/tests/gtest_executors_with_dm.cpp   | 32 ++++++++
 dbms/src/Interpreters/Settings.h              |  2 +-
 dbms/src/Interpreters/SettingsCommon.h        | 76 +++++++++++++++++++
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 9a0e2dfd409..4707f10ca05 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -149,6 +149,38 @@ constexpr UInt64 MSLMModeDisabled = 0;
 constexpr UInt64 MSLMModeSelection = 1;
 constexpr UInt64 MSLMModeTopN = 2;
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationSettingClampsInvalidMode)
+try
+{
+    Settings settings;
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+
+    settings.dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled);
+
+    settings.dt_enable_multi_stage_late_materialization = MSLMModeSelection;
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeSelection);
+
+    settings.dt_enable_multi_stage_late_materialization = MSLMModeTopN;
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+
+    settings.dt_enable_multi_stage_late_materialization = 999;
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+
+    settings.set("dt_enable_multi_stage_late_materialization", Field(static_cast<UInt64>(3)));
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+
+    settings.set("dt_enable_multi_stage_late_materialization", String("999"));
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+
+    settings.set("dt_enable_multi_stage_late_materialization", Field(static_cast<Int64>(-1)));
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled);
+
+    settings.set("dt_enable_multi_stage_late_materialization", String("-1"));
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled);
+}
+CATCH
+
 tipb::TableScan * findMutableTableScan(tipb::Executor * executor)
 {
     switch (executor->tp())
diff --git a/dbms/src/Interpreters/Settings.h b/dbms/src/Interpreters/Settings.h
index a449aef7c3f..0b78ab82703 100644
--- a/dbms/src/Interpreters/Settings.h
+++ b/dbms/src/Interpreters/Settings.h
@@ -226,7 +226,7 @@ struct Settings
     M(SettingUInt64, dt_max_sharing_column_bytes_for_all, 2048 * Constant::MB, "Memory limitation for data sharing of all requests, include those sharing blocks in block queue. 0 means disable data sharing")                         \
     M(SettingUInt64, dt_max_sharing_column_count, 5, "Deprecated")                                                                                                                                                                      \
     M(SettingBool, dt_enable_bitmap_filter, true, "Use bitmap filter to read data or not")                                                                                                                                              \
-    M(SettingUInt64, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, 2: TopN enhanced.")                                                    \
+    M(SettingMultiStageLateMaterializationMode, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, 2: TopN enhanced. Values outside [0, 2] are clamped to the nearest boundary.") \
     M(SettingDouble, dt_read_thread_count_scale, 2.0, "Number of read thread = number of logical cpu cores * dt_read_thread_count_scale.  Only has meaning at server startup.")                                                         \
     M(SettingDouble, io_thread_count_scale, 5.0, "Number of thread of IOThreadPool = number of logical cpu cores * io_thread_count_scale.  Only has meaning at server startup.")                                                        \
     M(SettingUInt64, init_thread_count_scale, 100, "Number of thread = number of logical cpu cores * init_thread_count_scale. It just works for thread pool for initStores and loadMetadata. Only has meaning at server startup.")      \
diff --git a/dbms/src/Interpreters/SettingsCommon.h b/dbms/src/Interpreters/SettingsCommon.h
index 87ebe447b91..e83b9518466 100644
--- a/dbms/src/Interpreters/SettingsCommon.h
+++ b/dbms/src/Interpreters/SettingsCommon.h
@@ -102,6 +102,82 @@ using SettingUInt64 = SettingInt<UInt64>;
 using SettingInt64 = SettingInt<Int64>;
 using SettingBool = SettingInt<bool>;
 
+struct SettingMultiStageLateMaterializationMode
+{
+public:
+    bool changed = false;
+
+    SettingMultiStageLateMaterializationMode(UInt64 x = 0) // NOLINT(google-explicit-constructor)
+        : value(clampMode(x))
+    {}
+
+    SettingMultiStageLateMaterializationMode(const SettingMultiStageLateMaterializationMode & setting)
+        : value(setting.value.load())
+    {}
+
+    operator UInt64() const { return value.load(); } // NOLINT(google-explicit-constructor)
+
+    SettingMultiStageLateMaterializationMode & operator=(UInt64 x)
+    {
+        set(x);
+        return *this;
+    }
+
+    SettingMultiStageLateMaterializationMode & operator=(const SettingMultiStageLateMaterializationMode & setting)
+    {
+        set(setting.value.load());
+        return *this;
+    }
+
+    String toString() const { return DB::toString(value.load()); }
+
+    void set(UInt64 x)
+    {
+        value.store(clampMode(x));
+        changed = true;
+    }
+
+    void set(const Field & x)
+    {
+        if (x.getType() == Field::Types::Int64)
+        {
+            const auto mode = safeGet<Int64>(x);
+            set(mode < 0 ? 0 : static_cast<UInt64>(mode));
+            return;
+        }
+        if (x.getType() == Field::Types::Float64)
+        {
+            const auto mode = safeGet<Float64>(x);
+            set(mode < 0 ? 0 : static_cast<UInt64>(mode));
+            return;
+        }
+        set(applyVisitor(FieldVisitorConvertToNumber<UInt64>(), x));
+    }
+
+    void set(const String & x) { set(!x.empty() && x[0] == '-' ? 0 : parse<UInt64>(x)); }
+
+    void set(ReadBuffer & buf)
+    {
+        UInt64 x = 0;
+        readVarT(x, buf);
+        set(x);
+    }
+
+    UInt64 get() const { return value.load(); }
+
+    void write(WriteBuffer & buf) const { writeVarT(value.load(), buf); }
+
+private:
+    static UInt64 clampMode(UInt64 x) { return x > 2 ? 2 : x; }
+
+    std::atomic<UInt64> value;
+};
+
+ALWAYS_INLINE inline auto format_as(SettingMultiStageLateMaterializationMode s)
+{
+    return s.get();
+}
+
 
 /** Unlike SettingUInt64, supports the value of 'auto' - the number of processor cores without taking into account SMT.
   * A value of 0 is also treated as auto.

From f8c724d82d950ee099af9882b98f9a7a9d8879b0 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 18:03:26 +0800
Subject: [PATCH 13/18] Refine multi-stage late materialization TopN

---
 dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp |   3 +
 .../Flash/tests/gtest_executors_with_dm.cpp   | 456 +++++++++++++++++-
 .../MultiStageLateMaterializationTopN.cpp     |   6 +-
 .../gtest_skippable_block_input_stream.cpp    |  19 +-
 dbms/src/Storages/StorageDeltaMerge.cpp       |   4 +-
 5 files changed, 471 insertions(+), 17 deletions(-)

diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
index 09d344e5cfa..aadf3993fc0 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
@@ -65,6 +65,9 @@ DM::MultiStageLateMaterializationTopNDescriptionPtr tryBuildMultiStageLateMateri
         if (!isColumnExpr(by_item.expr()))
             return disable("order by expression is not direct ColumnRef");
 
+        // The ColumnRef index is relative to TopN's input schema. TopN-enhanced MSLM is only attached
+        // when the child is a direct TableScan, whose output schema is generated from table_scan_columns
+        // with the same order, so the index can be used to locate the original table scan column here.
         const auto column_index = decodeDAGInt64(by_item.expr().val());
         if (column_index < 0 || column_index >= static_cast<Int64>(table_scan_columns.size()))
             return disable(fmt::format(
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 4707f10ca05..29bb9e0051d 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -334,6 +334,42 @@ void assertColumnsEqual(const ColumnsWithTypeAndName & expected, const ColumnsWi
         << getColumnsContent(actual);
 }
 
+void assertTopNEnhancedMultiStageRowsOverride(
+    DAGContext & dag_context,
+    UInt64 expected_table_scan_rows,
+    UInt64 min_selection_rows)
+{
+    auto table_scan_rows = dag_context.getExecutorRowsOverride("table_scan_0");
+    auto selection_rows = dag_context.getExecutorRowsOverride("selection_1");
+    ASSERT_NE(table_scan_rows, nullptr);
+    ASSERT_NE(selection_rows, nullptr);
+    ASSERT_EQ(table_scan_rows->load(), expected_table_scan_rows);
+    ASSERT_GE(selection_rows->load(), min_selection_rows);
+    ASSERT_LT(selection_rows->load(), expected_table_scan_rows);
+}
+
+void assertTopNEnhancedMultiStageMatchesDisabled(
+    ExecutorsWithDMTestRunner & runner,
+    MockDAGRequestContext & context,
+    const std::shared_ptr<tipb::DAGRequest> & request,
+    const String & log_identifier,
+    UInt64 expected_table_scan_rows,
+    UInt64 min_selection_rows)
+{
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
+    DAGContext disabled_dag_context(*request, log_identifier, /*concurrency=*/1);
+    auto expected = runner.executeStreams(&disabled_dag_context);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
+    DAGContext enabled_dag_context(*request, log_identifier, /*concurrency=*/1);
+    auto actual = runner.executeStreams(&enabled_dag_context);
+
+    assertColumnsEqual(expected, actual);
+    assertTopNEnhancedMultiStageRowsOverride(enabled_dag_context, expected_table_scan_rows, min_selection_rows);
+}
+
 void rewriteTableScanColumnAsCommitTS(tipb::DAGRequest * request, Int32 column_index)
 {
     auto * table_scan = findMutableTableScan(request->mutable_root_executor());
@@ -409,8 +445,7 @@ std::shared_ptr<tipb::DAGRequest> buildDAGRequestWithPushedDownFilterAndTopN(
     const String & table_name,
     const ASTPtr & pushed_down_filter,
     const ASTPtr & residual_filter,
-    const String & order_by_column,
-    bool is_desc,
+    MockOrderByItemVec order_by_items,
     UInt64 limit)
 {
     auto pushed_down_filter_request = context.scan("test_db", table_name).filter(pushed_down_filter).build(context);
@@ -420,13 +455,31 @@ std::shared_ptr<tipb::DAGRequest> buildDAGRequestWithPushedDownFilterAndTopN(
 
     auto request = context.scan("test_db", table_name)
                        .filter(residual_filter)
-                       .topN(order_by_column, is_desc, limit)
+                       .topN(std::move(order_by_items), static_cast<int>(limit))
                        .build(context);
     *findMutableTableScan(request->mutable_root_executor())->add_pushed_down_filter_conditions()
         = pushed_down_filter_expr;
     return request;
 }
 
+std::shared_ptr<tipb::DAGRequest> buildDAGRequestWithPushedDownFilterAndTopN(
+    MockDAGRequestContext & context,
+    const String & table_name,
+    const ASTPtr & pushed_down_filter,
+    const ASTPtr & residual_filter,
+    const String & order_by_column,
+    bool is_desc,
+    UInt64 limit)
+{
+    return buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        table_name,
+        pushed_down_filter,
+        residual_filter,
+        MockOrderByItemVec{{order_by_column, is_desc}},
+        limit);
+}
+
 ASTPtr buildAllGreaterThanMinusOneFilter(size_t begin_col_id, size_t end_col_id)
 {
     RUNTIME_CHECK(begin_col_id < end_col_id);
@@ -467,6 +520,49 @@ void addMockMultiStageLateMaterializationTable(
     context.addMockDeltaMerge({"test_db", table_name}, columns, data, /*concurrency_hint=*/4);
 }
 
+void addMockMultiStageLateMaterializationTopNOrderByTable(
+    MockDAGRequestContext & context,
+    const String & table_name,
+    TiDB::TP order_by_type,
+    ColumnWithTypeAndName order_by_column)
+{
+    constexpr size_t rows = 32;
+    MockColumnInfoVec columns;
+    ColumnsWithTypeAndName data;
+
+    columns.push_back({"c0", TiDB::TP::TypeLongLong});
+    columns.push_back({"ord", order_by_type});
+    columns.push_back({"c1", TiDB::TP::TypeLongLong});
+
+    std::vector<Int64> c0_values;
+    std::vector<Int64> c1_values;
+    c0_values.reserve(rows);
+    c1_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+    {
+        c0_values.push_back(static_cast<Int64>(row));
+        c1_values.push_back(static_cast<Int64>(row % 8));
+    }
+
+    data.emplace_back(toVec<Int64>("c0", c0_values));
+    data.emplace_back(std::move(order_by_column));
+    data.emplace_back(toVec<Int64>("c1", c1_values));
+
+    for (size_t col_id = 2; col_id <= 14; ++col_id)
+    {
+        const auto name = fmt::format("c{}", col_id);
+        columns.push_back({name, TiDB::TP::TypeLongLong});
+
+        std::vector<Int64> values;
+        values.reserve(rows);
+        for (size_t row = 0; row < rows; ++row)
+            values.push_back(static_cast<Int64>(col_id * 1000 + row));
+        data.emplace_back(toVec<Int64>(name, values));
+    }
+
+    context.addMockDeltaMerge({"test_db", table_name}, columns, data, /*concurrency_hint=*/4);
+}
+
 void rewriteComparisonRightLiteral(
     tipb::Expr * condition,
     tipb::ScalarFuncSig sig,
@@ -517,6 +613,27 @@ void assertMultiStageRowsOverride(DAGContext & dag_context, UInt64 expected_sele
     ASSERT_EQ(selection_rows->load(), expected_selection_rows);
     ASSERT_GE(table_scan_rows->load(), selection_rows->load());
 }
+
+void assertTopNEnhancedMultiStageFallsBackToSelection(
+    ExecutorsWithDMTestRunner & runner,
+    MockDAGRequestContext & context,
+    const std::shared_ptr<tipb::DAGRequest> & request,
+    const String & log_identifier,
+    UInt64 expected_selection_rows)
+{
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
+    DAGContext disabled_dag_context(*request, log_identifier, /*concurrency=*/1);
+    auto expected = runner.executeStreams(&disabled_dag_context);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
+    DAGContext enabled_dag_context(*request, log_identifier, /*concurrency=*/1);
+    auto actual = runner.executeStreams(&enabled_dag_context);
+
+    assertColumnsEqual(expected, actual);
+    assertMultiStageRowsOverride(enabled_dag_context, expected_selection_rows);
+}
 } // namespace
 
 TEST_F(ExecutorsWithDMTestRunner, Basic)
@@ -694,6 +811,55 @@ try
 }
 CATCH
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "c1",
+        /*is_desc=*/true,
+        /*limit=*/2);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/2);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNMultiColumnOrderBy)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        MockOrderByItemVec{{"c1", false}, {"c2", true}},
+        /*limit=*/3);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/3);
+}
+CATCH
+
 TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateTime)
 try
 {
@@ -770,6 +936,290 @@ try
 }
 CATCH
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNNullableOrderByNullOrdering)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_nullable_topn",
+        TiDB::TP::TypeLongLong,
+        toNullableVec<Int64>("ord", {{}, 100, {}, 90, 80, 70, 60, 50, 40, 30, 20, 10, 9,  8,  7,  6,
+                                     5,  4,   3,  2,  1,  0,  -1, -2, -3, -4, -5, -6, -7, -8, -9, -10}));
+
+    for (const auto is_desc : {false, true})
+    {
+        auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+            context,
+            "multi_stage_lm_nullable_topn",
+            lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+            gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+            "ord",
+            is_desc,
+            /*limit=*/2);
+
+        assertTopNEnhancedMultiStageMatchesDisabled(
+            *this,
+            context,
+            request,
+            dag_context_ptr->log->identifier(),
+            /*expected_table_scan_rows=*/16,
+            /*min_selection_rows=*/2);
+    }
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNFloatDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    std::vector<Float64> order_by_values;
+    order_by_values.reserve(32);
+    for (size_t row = 0; row < 32; ++row)
+        order_by_values.push_back(static_cast<Float64>(row) + 0.5);
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_float_topn",
+        TiDB::TP::TypeDouble,
+        toVec<Float64>("ord", order_by_values));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_float_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/true,
+        /*limit=*/3);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/3);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDecimalDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    std::vector<DecimalField<Decimal32>> order_by_values;
+    order_by_values.reserve(32);
+    for (size_t row = 0; row < 32; ++row)
+        order_by_values.emplace_back(static_cast<Int32>(row * 10), 0);
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_decimal_topn",
+        TiDB::TP::TypeNewDecimal,
+        toVec<Decimal32>("ord", order_by_values));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_decimal_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/true,
+        /*limit=*/3);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/3);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    std::vector<typename TypeTraits<MyDate>::FieldType> order_by_values;
+    order_by_values.reserve(32);
+    for (size_t row = 0; row < 32; ++row)
+        order_by_values.push_back(MyDate(2020, 1, static_cast<UInt32>(row % 28 + 1)).toPackedUInt());
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_date_topn",
+        TiDB::TP::TypeDate,
+        toVec<MyDate>("ord", order_by_values));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_date_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/true,
+        /*limit=*/3);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/3);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateTimeDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    std::vector<std::optional<typename TypeTraits<MyDateTime>::FieldType>> order_by_values;
+    order_by_values.reserve(32);
+    for (size_t row = 0; row < 32; ++row)
+        order_by_values.push_back(MyDateTime(2020, 1, 1, static_cast<UInt32>(row % 24), 0, 0, 0).toPackedUInt());
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_datetime_desc_topn",
+        TiDB::TP::TypeDatetime,
+        toNullableVec<MyDateTime>("ord", order_by_values));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_datetime_desc_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/true,
+        /*limit=*/3);
+
+    assertTopNEnhancedMultiStageMatchesDisabled(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_table_scan_rows=*/16,
+        /*min_selection_rows=*/3);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNTimestampDesc)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    std::vector<std::optional<typename TypeTraits<MyDateTime>::FieldType>> order_by_values;
+    order_by_values.reserve(32);
+    for (size_t row = 0; row < 32; ++row)
+        order_by_values.push_back(MyDateTime(2020, 1, 1, static_cast<UInt32>(row % 24), 0, 0, 0).toPackedUInt());
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_timestamp_topn",
+        TiDB::TP::TypeTimestamp,
+        toNullableVec<MyDateTime>("ord", order_by_values));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_timestamp_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/true,
+        /*limit=*/3);
+    request->set_time_zone_name("Asia/Shanghai");
+
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNVersionColumn)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    MockColumnInfoVec columns;
+    columns.push_back({"pk", TiDB::TP::TypeLongLong});
+    columns.push_back({"commit_ts", TiDB::TP::TypeLongLong});
+    for (size_t col_id = 1; col_id <= 14; ++col_id)
+        columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong});
+
+    constexpr size_t rows = 32;
+    std::vector<Int64> pk_values;
+    std::vector<Int64> commit_ts_dummy_values;
+    pk_values.reserve(rows);
+    commit_ts_dummy_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+    {
+        pk_values.push_back(row);
+        commit_ts_dummy_values.push_back(-1);
+    }
+
+    ColumnsWithTypeAndName data;
+    data.emplace_back(toVec<Int64>("pk", pk_values));
+    data.emplace_back(toVec<Int64>("commit_ts", commit_ts_dummy_values));
+    for (size_t col_id = 1; col_id <= 14; ++col_id)
+    {
+        std::vector<Int64> values;
+        values.reserve(rows);
+        for (size_t row = 0; row < rows; ++row)
+        {
+            if (col_id == 1)
+                values.push_back(row % 8);
+            else
+                values.push_back(static_cast<Int64>(col_id * 1000 + row));
+        }
+        data.emplace_back(toVec<Int64>(fmt::format("c{}", col_id), values));
+    }
+    context.addMockDeltaMerge(
+        {"test_db", "multi_stage_lm_version_topn"},
+        columns,
+        data,
+        /*concurrency_hint=*/4);
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_version_topn",
+        lt(col("pk"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "commit_ts",
+        /*is_desc=*/true,
+        /*limit=*/3);
+    rewriteTableScanColumnAsCommitTS(request.get(), /*column_index=*/1);
+
+    const auto selected_rows = std::vector<Int64>{15, 14, 13};
+    auto expected_commit_ts_col = toNullableVec<Int64>(
+        MutableSupport::version_column_name,
+        buildNullableInt64Values(selected_rows, [](Int64 row) { return row + 1; }));
+    auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_commit_ts_col);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled;
+    DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
+    assertColumnsEqual(expected, executeStreams(&disabled_dag_context));
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr);
+    ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr);
+
+    context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN;
+    DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1);
+    assertColumnsEqual(expected, executeStreams(&enabled_dag_context));
+    assertMultiStageRowsOverride(enabled_dag_context, 16);
+}
+CATCH
+
 TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForStringOrderBy)
 try
 {
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
index 6cb85474c46..40899bbd463 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -409,6 +409,8 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         }
 
         const auto cmp = compareRowWithOwnedKey(sort_columns, row, heap.top().key);
+        // Equal sort keys are not kept as extra candidates. This matches the existing local TopN behavior,
+        // where ties are broken locally by the entries already selected by the partial TopN.
         if (cmp < 0)
         {
             auto evicted = heap.top();
@@ -422,10 +424,6 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
                 .row_index_in_stage1_block = static_cast<UInt32>(row),
             });
         }
-        else if (cmp == 0)
-        {
-            markCandidate(row);
-        }
     }
 
     return result;
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
index ec69668f3e0..c05eb82cfbd 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
@@ -1042,6 +1042,7 @@ CATCH
 TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDoesNotDisableDuringWarmup)
 try
 {
+    constexpr UInt64 topk = 2;
     constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold;
 
     std::vector<IColumn::Filter> stage0_filters;
@@ -1068,15 +1069,15 @@ try
         stage0_stream,
         stage1_stream,
         final_rest_stream,
-        makeResidualFilterWithTopNForMultiStageTest(/*topk=*/2),
+        makeResidualFilterWithTopNForMultiStageTest(topk),
         bitmap_filter,
         "test",
         runtime_stats);
 
-    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows);
+    ASSERT_EQ(drainMultiStageStream(stream), topk);
     ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows);
     ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows);
-    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk);
     ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0);
 }
 CATCH
@@ -1124,15 +1125,15 @@ try
         "test",
         runtime_stats);
 
-    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows + below_threshold_candidate_rows);
+    ASSERT_EQ(drainMultiStageStream(stream), topk * 2);
     ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows);
     ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows);
-    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows + below_threshold_candidate_rows);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk * 2);
     ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0);
 }
 CATCH
 
-TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDisable)
+TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNTieRowsDoNotExpandCandidates)
 try
 {
     constexpr UInt64 topk = 2;
@@ -1171,11 +1172,11 @@ try
         "test",
         runtime_stats);
 
-    ASSERT_EQ(drainMultiStageStream(stream), warmup_rows + second_block_rows);
+    ASSERT_EQ(drainMultiStageStream(stream), topk);
     ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows);
     ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows);
-    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), warmup_rows + second_block_rows);
-    ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 1);
+    ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk);
+    ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0);
 }
 CATCH
 
diff --git a/dbms/src/Storages/StorageDeltaMerge.cpp b/dbms/src/Storages/StorageDeltaMerge.cpp
index a1b93a6e728..63e2835da17 100644
--- a/dbms/src/Storages/StorageDeltaMerge.cpp
+++ b/dbms/src/Storages/StorageDeltaMerge.cpp
@@ -980,8 +980,10 @@ DM::PushDownFilterPtr buildMultiStageLateMaterializationFilter(
     if (storage_topn == nullptr)
         return residual_filter;
 
+    // This filter object only carries residual filter actions and the stage1 columns. The actual rough-set
+    // skipping of MSLM streams must use the original query-level rs_operator from the stage0 filter.
     return std::make_shared<DM::PushDownFilter>(
-        residual_filter->rs_operator,
+        DM::EMPTY_RS_OPERATOR,
         residual_filter->before_where,
         residual_filter->project_after_where,
         merged_stage1_columns,

From ee4cf8ef6cbf67544f2c83f41e496f03ec794317 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Wed, 29 Jul 2026 22:32:48 +0800
Subject: [PATCH 14/18] Add MSLM scan context reporting

---
 contrib/tipb                                  |   2 +-
 .../Coprocessor/DAGStorageInterpreter.cpp     |   8 +-
 .../DeltaMerge/File/DMFilePackFilter.cpp      |   6 +-
 .../Storages/DeltaMerge/File/DMFileReader.cpp |  38 +-
 ...ageLateMaterializationBlockInputStream.cpp |  20 +-
 ...ultiStageLateMaterializationRuntimeStats.h | 139 ++++++-
 .../MultiStageLateMaterializationTopN.cpp     |  10 +-
 dbms/src/Storages/DeltaMerge/ReadMode.h       |   4 +-
 dbms/src/Storages/DeltaMerge/ScanContext.cpp  | 343 +++++++++++++++++-
 dbms/src/Storages/DeltaMerge/ScanContext.h    |  43 ++-
 dbms/src/Storages/DeltaMerge/Segment.cpp      |  10 +-
 .../tests/gtest_dm_delta_merge_store.cpp      |   9 +-
 .../2026-07-26-running-local-topn-for-mslm.md | 242 +++++++++++-
 13 files changed, 778 insertions(+), 96 deletions(-)

diff --git a/contrib/tipb b/contrib/tipb
index fe7badb76b6..474dc14add1 160000
--- a/contrib/tipb
+++ b/contrib/tipb
@@ -1 +1 @@
-Subproject commit fe7badb76b660abaedd8706accd809809c475280
+Subproject commit 474dc14add1ef2d1c77bc6f0a42f93546583323f
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index eae4d13a2ba..6d96fb851b2 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -1059,10 +1059,16 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
             multi_stage_late_materialization_topn != nullptr
                 ? std::shared_ptr<std::atomic<UInt64>>(
                     multi_stage_late_materialization_runtime_stats,
-                    &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
+                    &multi_stage_late_materialization_runtime_stats->final_rest_input_rows)
                 : std::shared_ptr<std::atomic<UInt64>>(
                     multi_stage_late_materialization_runtime_stats,
                     &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+        if (auto scan_context_it = dagContext().scan_context_map.find(table_scan.getTableScanExecutorID());
+            scan_context_it != dagContext().scan_context_map.end() && scan_context_it->second != nullptr)
+        {
+            scan_context_it->second->setMultiStageLateMaterializationRuntimeStats(
+                multi_stage_late_materialization_runtime_stats);
+        }
     }
 
     auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo {
diff --git a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
index 1ffa12e0cb6..8d2fbc9857c 100644
--- a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
+++ b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
@@ -185,8 +185,10 @@ DMFilePackFilterResultPtr DMFilePackFilter::load(ReadTag read_tag)
     // In table scanning, DMFilePackFilter of a DMFile may be created several times:
     // 1. When building MVCC bitmap (ReadTag::MVCC).
     // 2. When building LM filter stream (ReadTag::LMFilter).
-    // 3. When building MSLM stage1 residual filter stream (ReadTag::MSLMStage1Filter).
-    // 4. When building stream of other columns (ReadTag::Query).
+    // 3. When building MSLM pushed filter stream (ReadTag::MSLMPushedFilter).
+    // 4. When building MSLM candidate stream (ReadTag::MSLMCandidate).
+    // 5. When building MSLM final rest stream (ReadTag::MSLMFinalRest).
+    // 6. When building stream of other columns (ReadTag::Query).
     // Only need to count the filter result once.
     // TODO: We can create DMFilePackFilter at the beginning and pass it to the stages described above.
     if (read_tag == ReadTag::Query)
diff --git a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
index 3a3126a346b..e9307564663 100644
--- a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
+++ b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
@@ -138,7 +138,8 @@ size_t DMFileReader::skipNextBlock()
         return 0;
 
     addSkippedRows(read_rows);
-    scan_context->late_materialization_skip_rows += read_rows;
+    if (read_tag == ReadTag::Query)
+        scan_context->late_materialization_skip_rows += read_rows;
     return read_rows;
 }
 
@@ -285,7 +286,7 @@ bool DMFileReader::isCacheableColumn(const ColumnDefine & cd)
 Block DMFileReader::read()
 {
     Stopwatch watch;
-    SCOPE_EXIT(scan_context->total_dmfile_read_time_ns += watch.elapsed(););
+    SCOPE_EXIT(scan_context->addDMFileReadTime(watch.elapsed(), read_tag););
 
     /// 1. Skip filtered out packs.
     if (size_t skip_rows; !getSkippedRows(skip_rows))
@@ -694,41 +695,12 @@ bool DMFileReader::getCachedPacks(
 
 void DMFileReader::addScannedRows(UInt64 rows)
 {
-    switch (read_tag)
-    {
-    case ReadTag::Query:
-        scan_context->dmfile_data_scanned_rows += rows;
-        break;
-    case ReadTag::MVCC:
-        scan_context->dmfile_mvcc_scanned_rows += rows;
-        break;
-    case ReadTag::LMFilter:
-        scan_context->dmfile_lm_filter_scanned_rows += rows;
-        break;
-    case ReadTag::MSLMStage1Filter:
-        scan_context->dmfile_mslm_stage1_filter_scanned_rows += rows;
-        break;
-    default:
-        break;
-    }
+    scan_context->addDMFileScannedRows(rows, read_tag);
 }
 
 void DMFileReader::addSkippedRows(UInt64 rows)
 {
-    switch (read_tag)
-    {
-    case ReadTag::Query:
-        scan_context->dmfile_data_skipped_rows += rows;
-        break;
-    case ReadTag::MVCC:
-        scan_context->dmfile_mvcc_skipped_rows += rows;
-        break;
-    case ReadTag::LMFilter:
-        scan_context->dmfile_lm_filter_skipped_rows += rows;
-        break;
-    default:
-        break;
-    }
+    scan_context->addDMFileSkippedRows(rows, read_tag);
 }
 
 void DMFileReader::initAllMatchBlockInfo()
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index b997f410a80..48a37664cf1 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -346,9 +346,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
 
         auto effective_stage0_filter = buildStage0EffectiveFilter(stage0_block, stage0_filter);
         if (runtime_stats)
-            runtime_stats->stage0_output_rows.fetch_add(
-                effective_stage0_filter.passed_count,
-                std::memory_order_relaxed);
+            runtime_stats->recordPushedFilter(stage0_block.rows(), effective_stage0_filter.passed_count);
         if (effective_stage0_filter.passed_count == 0)
         {
             skipNextBlockOrRead(stage1_filter_stream, "stage1_filter");
@@ -377,10 +375,15 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         FilterPtr residual_filter_ptr = nullptr;
         const auto residual_passed_rows = executeResidualFilter(stage1_block, filter_eval_block, residual_filter_ptr);
         if (runtime_stats)
-            runtime_stats->stage1_output_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
+            runtime_stats->recordResidualFilter(stage1_block.rows(), residual_passed_rows);
 
         if (!shouldUseRunningTopN() && runtime_stats)
-            runtime_stats->topn_candidate_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed);
+        {
+            if (running_topn != nullptr)
+                runtime_stats->recordRunningTopNBypass(residual_passed_rows);
+            else
+                runtime_stats->recordFinalRestInputRows(residual_passed_rows);
+        }
 
         if (residual_passed_rows == 0)
         {
@@ -393,7 +396,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         {
             auto topn_result = running_topn->update(stage1_block, residual_filter_ptr, residual_passed_rows);
             if (runtime_stats)
-                runtime_stats->topn_candidate_rows.fetch_add(topn_result.passed_count, std::memory_order_relaxed);
+                runtime_stats->recordRunningTopN(residual_passed_rows, topn_result.passed_count);
             updateTopNAdaptiveState(residual_passed_rows, topn_result.passed_count);
 
             if (topn_result.passed_count == 0)
@@ -456,7 +459,10 @@ void MultiStageLateMaterializationBlockInputStream::logSummary()
             late_mode_blocks,
             direct_mode_blocks,
             running_topn != nullptr ? running_topn->heapSize() : 0,
-            topn_adaptive_disabled);
+            topn_adaptive_disabled,
+            topn_adaptive_warmup_observed_rows,
+            topn_adaptive_input_rows,
+            topn_adaptive_candidate_rows);
 }
 
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index 7b7a19d2c03..fc65f103de0 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -46,44 +46,150 @@ struct MultiStageLateMaterializationRuntimeStats
         UInt64 stream_late_mode_blocks,
         UInt64 stream_direct_mode_blocks,
         UInt64 stream_topn_heap_size,
-        bool stream_topn_adaptive_disabled)
+        bool stream_topn_adaptive_disabled,
+        UInt64 stream_topn_adaptive_warmup_rows,
+        UInt64 stream_topn_adaptive_input_rows,
+        UInt64 stream_topn_adaptive_candidate_rows)
     {
         finished_streams.fetch_add(1, std::memory_order_relaxed);
         late_mode_blocks.fetch_add(stream_late_mode_blocks, std::memory_order_relaxed);
         direct_mode_blocks.fetch_add(stream_direct_mode_blocks, std::memory_order_relaxed);
         topn_heap_size_sum.fetch_add(stream_topn_heap_size, std::memory_order_relaxed);
+        topn_adaptive_warmup_rows.fetch_add(stream_topn_adaptive_warmup_rows, std::memory_order_relaxed);
+        topn_adaptive_post_warmup_input_rows.fetch_add(stream_topn_adaptive_input_rows, std::memory_order_relaxed);
+        topn_adaptive_post_warmup_candidate_rows.fetch_add(
+            stream_topn_adaptive_candidate_rows,
+            std::memory_order_relaxed);
         if (stream_topn_adaptive_disabled)
             topn_adaptive_disabled_streams.fetch_add(1, std::memory_order_relaxed);
     }
 
+    void recordPushedFilter(UInt64 input_rows, UInt64 selected_rows)
+    {
+        pushed_filter_input_rows.fetch_add(input_rows, std::memory_order_relaxed);
+        pushed_filter_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed);
+        pushed_filter_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed);
+
+        // Keep the old counter name for existing actRows/tests.
+        stage0_output_rows.fetch_add(selected_rows, std::memory_order_relaxed);
+    }
+
+    void recordResidualFilter(UInt64 input_rows, UInt64 selected_rows)
+    {
+        residual_filter_input_rows.fetch_add(input_rows, std::memory_order_relaxed);
+        residual_filter_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed);
+        residual_filter_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed);
+
+        // Keep the old counter name for existing actRows/tests.
+        stage1_output_rows.fetch_add(selected_rows, std::memory_order_relaxed);
+    }
+
+    void recordFinalRestInputRows(UInt64 rows)
+    {
+        final_rest_input_rows.fetch_add(rows, std::memory_order_relaxed);
+
+        // Keep the old counter name for existing actRows/tests.
+        topn_candidate_rows.fetch_add(rows, std::memory_order_relaxed);
+    }
+
+    void recordRunningTopN(UInt64 input_rows, UInt64 selected_rows)
+    {
+        topn_enabled.store(true, std::memory_order_relaxed);
+        running_topn_input_rows.fetch_add(input_rows, std::memory_order_relaxed);
+        running_topn_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed);
+        running_topn_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed);
+        recordFinalRestInputRows(selected_rows);
+    }
+
+    void recordRunningTopNBypass(UInt64 rows)
+    {
+        topn_enabled.store(true, std::memory_order_relaxed);
+        running_topn_input_rows.fetch_add(rows, std::memory_order_relaxed);
+        running_topn_bypass_rows.fetch_add(rows, std::memory_order_relaxed);
+        recordFinalRestInputRows(rows);
+    }
+
+    void merge(const MultiStageLateMaterializationRuntimeStats & other)
+    {
+        pushed_filter_input_rows.fetch_add(other.pushed_filter_input_rows.load(std::memory_order_relaxed));
+        pushed_filter_selected_rows.fetch_add(other.pushed_filter_selected_rows.load(std::memory_order_relaxed));
+        pushed_filter_filtered_rows.fetch_add(other.pushed_filter_filtered_rows.load(std::memory_order_relaxed));
+        residual_filter_input_rows.fetch_add(other.residual_filter_input_rows.load(std::memory_order_relaxed));
+        residual_filter_selected_rows.fetch_add(other.residual_filter_selected_rows.load(std::memory_order_relaxed));
+        residual_filter_filtered_rows.fetch_add(other.residual_filter_filtered_rows.load(std::memory_order_relaxed));
+        final_rest_input_rows.fetch_add(other.final_rest_input_rows.load(std::memory_order_relaxed));
+        running_topn_input_rows.fetch_add(other.running_topn_input_rows.load(std::memory_order_relaxed));
+        running_topn_selected_rows.fetch_add(other.running_topn_selected_rows.load(std::memory_order_relaxed));
+        running_topn_bypass_rows.fetch_add(other.running_topn_bypass_rows.load(std::memory_order_relaxed));
+        running_topn_filtered_rows.fetch_add(other.running_topn_filtered_rows.load(std::memory_order_relaxed));
+
+        stage0_output_rows.fetch_add(other.stage0_output_rows.load(std::memory_order_relaxed));
+        stage1_output_rows.fetch_add(other.stage1_output_rows.load(std::memory_order_relaxed));
+        topn_candidate_rows.fetch_add(other.topn_candidate_rows.load(std::memory_order_relaxed));
+        finished_streams.fetch_add(other.finished_streams.load(std::memory_order_relaxed));
+        late_mode_blocks.fetch_add(other.late_mode_blocks.load(std::memory_order_relaxed));
+        direct_mode_blocks.fetch_add(other.direct_mode_blocks.load(std::memory_order_relaxed));
+        topn_heap_size_sum.fetch_add(other.topn_heap_size_sum.load(std::memory_order_relaxed));
+        topn_adaptive_warmup_rows.fetch_add(other.topn_adaptive_warmup_rows.load(std::memory_order_relaxed));
+        topn_adaptive_post_warmup_input_rows.fetch_add(
+            other.topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed));
+        topn_adaptive_post_warmup_candidate_rows.fetch_add(
+            other.topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed));
+        topn_adaptive_disabled_streams.fetch_add(other.topn_adaptive_disabled_streams.load(std::memory_order_relaxed));
+        topn_enabled.store(
+            topn_enabled.load(std::memory_order_relaxed) || other.topn_enabled.load(std::memory_order_relaxed),
+            std::memory_order_relaxed);
+    }
+
     void logSummary() const
     {
         if (log == nullptr)
             return;
 
-        const auto stage1_rows = stage1_output_rows.load(std::memory_order_relaxed);
-        const auto topn_candidate_rows_for_log = topn_enabled ? topn_candidate_rows.load(std::memory_order_relaxed) : 0;
-        const auto topn_filtered_rows = topn_enabled && stage1_rows >= topn_candidate_rows_for_log
-            ? stage1_rows - topn_candidate_rows_for_log
-            : 0;
-
         LOG_INFO(
             log,
             "Multi-stage late materialization finished, streams={} late_mode_blocks={} direct_mode_blocks={} "
-            "stage0_output_rows={} stage1_output_rows={} topn_enabled={} topn_candidate_rows={} "
-            "topn_filtered_rows={} topn_heap_size_sum={} topn_adaptive_disabled_streams={}",
+            "pushed_filter_input_rows={} pushed_filter_selected_rows={} pushed_filter_filtered_rows={} "
+            "residual_filter_input_rows={} residual_filter_selected_rows={} residual_filter_filtered_rows={} "
+            "final_rest_input_rows={} topn_enabled={} running_topn_input_rows={} running_topn_selected_rows={} "
+            "running_topn_bypass_rows={} running_topn_filtered_rows={} running_topn_heap_size_sum={} "
+            "topn_adaptive_warmup_rows={} topn_adaptive_post_warmup_input_rows={} "
+            "topn_adaptive_post_warmup_candidate_rows={} topn_adaptive_disabled_streams={}",
             finished_streams.load(std::memory_order_relaxed),
             late_mode_blocks.load(std::memory_order_relaxed),
             direct_mode_blocks.load(std::memory_order_relaxed),
-            stage0_output_rows.load(std::memory_order_relaxed),
-            stage1_rows,
-            topn_enabled,
-            topn_candidate_rows_for_log,
-            topn_filtered_rows,
+            pushed_filter_input_rows.load(std::memory_order_relaxed),
+            pushed_filter_selected_rows.load(std::memory_order_relaxed),
+            pushed_filter_filtered_rows.load(std::memory_order_relaxed),
+            residual_filter_input_rows.load(std::memory_order_relaxed),
+            residual_filter_selected_rows.load(std::memory_order_relaxed),
+            residual_filter_filtered_rows.load(std::memory_order_relaxed),
+            final_rest_input_rows.load(std::memory_order_relaxed),
+            topn_enabled.load(std::memory_order_relaxed),
+            running_topn_input_rows.load(std::memory_order_relaxed),
+            running_topn_selected_rows.load(std::memory_order_relaxed),
+            running_topn_bypass_rows.load(std::memory_order_relaxed),
+            running_topn_filtered_rows.load(std::memory_order_relaxed),
             topn_heap_size_sum.load(std::memory_order_relaxed),
+            topn_adaptive_warmup_rows.load(std::memory_order_relaxed),
+            topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed),
+            topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed),
             topn_adaptive_disabled_streams.load(std::memory_order_relaxed));
     }
 
+    std::atomic<UInt64> pushed_filter_input_rows{0};
+    std::atomic<UInt64> pushed_filter_selected_rows{0};
+    std::atomic<UInt64> pushed_filter_filtered_rows{0};
+    std::atomic<UInt64> residual_filter_input_rows{0};
+    std::atomic<UInt64> residual_filter_selected_rows{0};
+    std::atomic<UInt64> residual_filter_filtered_rows{0};
+    std::atomic<UInt64> final_rest_input_rows{0};
+    std::atomic<UInt64> running_topn_input_rows{0};
+    std::atomic<UInt64> running_topn_selected_rows{0};
+    std::atomic<UInt64> running_topn_bypass_rows{0};
+    std::atomic<UInt64> running_topn_filtered_rows{0};
+
+    // Legacy names kept because actRows overrides and existing unit tests still refer to them.
     std::atomic<UInt64> stage0_output_rows{0};
     std::atomic<UInt64> stage1_output_rows{0};
     std::atomic<UInt64> topn_candidate_rows{0};
@@ -91,8 +197,11 @@ struct MultiStageLateMaterializationRuntimeStats
     std::atomic<UInt64> late_mode_blocks{0};
     std::atomic<UInt64> direct_mode_blocks{0};
     std::atomic<UInt64> topn_heap_size_sum{0};
+    std::atomic<UInt64> topn_adaptive_warmup_rows{0};
+    std::atomic<UInt64> topn_adaptive_post_warmup_input_rows{0};
+    std::atomic<UInt64> topn_adaptive_post_warmup_candidate_rows{0};
     std::atomic<UInt64> topn_adaptive_disabled_streams{0};
-    bool topn_enabled = false;
+    std::atomic<bool> topn_enabled{false};
     LoggerPtr log;
 };
 
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
index 40899bbd463..2e579aa670d 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -377,14 +377,14 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         sort_columns.push_back(desc.build_column_view(*column));
     }
 
-    auto markCandidate = [&](size_t row) {
+    auto mark_candidate = [&](size_t row) {
         if (result.filter[row] == 0)
         {
             result.filter[row] = 1;
             ++result.passed_count;
         }
     };
-    auto unmarkCurrentBlockCandidate = [&](const HeapEntry & entry) {
+    auto unmark_current_block_candidate = [&](const HeapEntry & entry) {
         if (entry.block_sequence == current_block_sequence && result.filter[entry.row_index_in_stage1_block] != 0)
         {
             result.filter[entry.row_index_in_stage1_block] = 0;
@@ -399,7 +399,7 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
 
         if (heap.size() < topk)
         {
-            markCandidate(row);
+            mark_candidate(row);
             heap.push(HeapEntry{
                 .key = materializeOwnedKey(sort_columns, row),
                 .block_sequence = current_block_sequence,
@@ -415,9 +415,9 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         {
             auto evicted = heap.top();
             heap.pop();
-            unmarkCurrentBlockCandidate(evicted);
+            unmark_current_block_candidate(evicted);
 
-            markCandidate(row);
+            mark_candidate(row);
             heap.push(HeapEntry{
                 .key = materializeOwnedKey(sort_columns, row),
                 .block_sequence = current_block_sequence,
diff --git a/dbms/src/Storages/DeltaMerge/ReadMode.h b/dbms/src/Storages/DeltaMerge/ReadMode.h
index 607e0788a62..791f3c45767 100644
--- a/dbms/src/Storages/DeltaMerge/ReadMode.h
+++ b/dbms/src/Storages/DeltaMerge/ReadMode.h
@@ -45,7 +45,9 @@ enum class ReadTag
     Query, // Read columns required by queries.
     MVCC, // Read columns to build MVCC bitmap.
     LMFilter, // Read columns required by late-materialization filter.
-    MSLMStage1Filter, // Read columns required by multi-stage late-materialization residual filter.
+    MSLMPushedFilter, // Read columns required by multi-stage late-materialization pushed filter.
+    MSLMCandidate, // Read columns required by multi-stage late-materialization residual/order-by candidate read.
+    MSLMFinalRest, // Read final rest columns for multi-stage late-materialization.
 };
 
 enum class ReadRUType
diff --git a/dbms/src/Storages/DeltaMerge/ScanContext.cpp b/dbms/src/Storages/DeltaMerge/ScanContext.cpp
index 5e7e4366fa3..c3f7700462f 100644
--- a/dbms/src/Storages/DeltaMerge/ScanContext.cpp
+++ b/dbms/src/Storages/DeltaMerge/ScanContext.cpp
@@ -25,6 +25,73 @@
 
 namespace DB::DM
 {
+namespace
+{
+static constexpr uint64_t NS_TO_MS = 1'000'000;
+
+bool isQueryReadTag(ReadTag read_tag)
+{
+    switch (read_tag)
+    {
+    case ReadTag::Query:
+    case ReadTag::LMFilter:
+    case ReadTag::MSLMPushedFilter:
+    case ReadTag::MSLMCandidate:
+    case ReadTag::MSLMFinalRest:
+        return true;
+    default:
+        return false;
+    }
+}
+
+void serializeReadStage(const MSLMReadStageScanContext & read_stage, tipb::TiFlashMSLMReadStageContext * proto)
+{
+    proto->set_dmfile_scanned_rows(read_stage.dmfile_scanned_rows);
+    proto->set_dmfile_skipped_rows(read_stage.dmfile_skipped_rows);
+    proto->set_read_bytes(read_stage.read_bytes);
+    proto->set_read_time_ms(read_stage.read_time_ns / NS_TO_MS);
+}
+
+void mergeReadStage(MSLMReadStageScanContext & read_stage, const tipb::TiFlashMSLMReadStageContext & proto)
+{
+    read_stage.dmfile_scanned_rows += proto.dmfile_scanned_rows();
+    read_stage.dmfile_skipped_rows += proto.dmfile_skipped_rows();
+    read_stage.read_bytes += proto.read_bytes();
+    read_stage.read_time_ns += proto.read_time_ms() * NS_TO_MS;
+}
+
+void serializeFilter(
+    uint64_t input_rows,
+    uint64_t selected_rows,
+    uint64_t filtered_rows,
+    tipb::TiFlashMSLMFilterContext * proto)
+{
+    proto->set_input_rows(input_rows);
+    proto->set_selected_rows(selected_rows);
+    proto->set_filtered_rows(filtered_rows);
+}
+
+Poco::JSON::Object::Ptr readStageToJson(const MSLMReadStageScanContext & read_stage)
+{
+    static constexpr double NS_TO_MS_SCALE = 1'000'000.0;
+    Poco::JSON::Object::Ptr json = new Poco::JSON::Object();
+    json->set("dmfile_scanned_rows", read_stage.dmfile_scanned_rows.load());
+    json->set("dmfile_skipped_rows", read_stage.dmfile_skipped_rows.load());
+    json->set("read_bytes", read_stage.read_bytes.load());
+    json->set("read_time", fmt::format("{:.3f}ms", read_stage.read_time_ns.load() / NS_TO_MS_SCALE));
+    return json;
+}
+
+Poco::JSON::Object::Ptr filterToJson(uint64_t input_rows, uint64_t selected_rows, uint64_t filtered_rows)
+{
+    Poco::JSON::Object::Ptr json = new Poco::JSON::Object();
+    json->set("input_rows", input_rows);
+    json->set("selected_rows", selected_rows);
+    json->set("filtered_rows", filtered_rows);
+    return json;
+}
+} // namespace
+
 void ScanContext::setRegionNumOfCurrentInstance(uint64_t region_num)
 {
     region_num_of_instance[current_instance_id] = region_num;
@@ -103,7 +170,6 @@ String ScanContext::toJson() const
     json->set("dmfile_mvcc_skipped_rows", dmfile_mvcc_skipped_rows.load());
     json->set("dmfile_lm_filter_scanned_rows", dmfile_lm_filter_scanned_rows.load());
     json->set("dmfile_lm_filter_skipped_rows", dmfile_lm_filter_skipped_rows.load());
-    json->set("dmfile_mslm_stage1_filter_scanned_rows", dmfile_mslm_stage1_filter_scanned_rows.load());
     json->set("dmfile_read_time", fmt::format("{:.3f}ms", total_dmfile_read_time_ns.load() / NS_TO_MS_SCALE));
 
     json->set(
@@ -191,6 +257,48 @@ String ScanContext::toJson() const
         json->set("pushdown", pushdown_executor->toJSONObject());
     }
 
+    if (hasMultiStageLateMaterializationContext())
+    {
+        const auto & stats = getMultiStageLateMaterializationRuntimeStats();
+        Poco::JSON::Object::Ptr mslm = new Poco::JSON::Object();
+        mslm->set("streams", stats.finished_streams.load());
+        mslm->set("late_mode_blocks", stats.late_mode_blocks.load());
+        mslm->set("direct_mode_blocks", stats.direct_mode_blocks.load());
+        mslm->set("pushed_filter_read", readStageToJson(mslm_pushed_filter_read));
+        mslm->set("candidate_read", readStageToJson(mslm_candidate_read));
+        mslm->set("final_rest_read", readStageToJson(mslm_final_rest_read));
+        mslm->set(
+            "pushed_filter",
+            filterToJson(
+                stats.pushed_filter_input_rows.load(),
+                stats.pushed_filter_selected_rows.load(),
+                stats.pushed_filter_filtered_rows.load()));
+        mslm->set(
+            "residual_filter",
+            filterToJson(
+                stats.residual_filter_input_rows.load(),
+                stats.residual_filter_selected_rows.load(),
+                stats.residual_filter_filtered_rows.load()));
+        if (stats.topn_enabled.load())
+        {
+            Poco::JSON::Object::Ptr running_topn = new Poco::JSON::Object();
+            running_topn->set("input_rows", stats.running_topn_input_rows.load());
+            running_topn->set("selected_rows", stats.running_topn_selected_rows.load());
+            running_topn->set("bypass_rows", stats.running_topn_bypass_rows.load());
+            running_topn->set("filtered_rows", stats.running_topn_filtered_rows.load());
+            running_topn->set("heap_size_sum", stats.topn_heap_size_sum.load());
+            running_topn->set("adaptive_warmup_rows", stats.topn_adaptive_warmup_rows.load());
+            running_topn->set("adaptive_post_warmup_input_rows", stats.topn_adaptive_post_warmup_input_rows.load());
+            running_topn->set(
+                "adaptive_post_warmup_candidate_rows",
+                stats.topn_adaptive_post_warmup_candidate_rows.load());
+            running_topn->set("adaptive_disabled_streams", stats.topn_adaptive_disabled_streams.load());
+            mslm->set("running_topn", running_topn);
+        }
+        mslm->set("final_rest_input_rows", stats.final_rest_input_rows.load());
+        json->set("multi_stage_late_materialization", mslm);
+    }
+
     std::stringstream buf;
     json->stringify(buf);
     return buf.str();
@@ -275,7 +383,7 @@ std::optional<LACBytesCollector> ScanContext::newLACBytesCollector(ReadTag read_
 {
     if (resource_group_name.empty())
         return std::nullopt;
-    if (read_tag != ReadTag::Query && read_tag != ReadTag::LMFilter && read_tag != ReadTag::MSLMStage1Filter)
+    if (!isQueryReadTag(read_tag))
         return std::nullopt;
     return LACBytesCollector(resource_group_name);
 }
@@ -285,8 +393,7 @@ void ScanContext::addUserReadBytes(
     ReadTag read_tag,
     std::optional<LACBytesCollector> & lac_bytes_collector)
 {
-    if (read_tag != ReadTag::Query && read_tag != ReadTag::LMFilter && read_tag != ReadTag::MSLMStage1Filter
-        && read_tag != ReadTag::MVCC)
+    if (!isQueryReadTag(read_tag) && read_tag != ReadTag::MVCC)
         return;
     if (read_tag == ReadTag::MVCC)
     {
@@ -297,6 +404,8 @@ void ScanContext::addUserReadBytes(
     else
     {
         query_read_bytes += bytes;
+        if (auto * read_stage = getMutableMSLMReadStage(read_tag); read_stage != nullptr)
+            read_stage->read_bytes += bytes;
         if (query_read_bytes_counter)
             query_read_bytes_counter->Increment(bytes);
         if (lac_bytes_collector)
@@ -304,4 +413,230 @@ void ScanContext::addUserReadBytes(
     }
 }
 
+bool ScanContext::hasMultiStageLateMaterializationContext() const
+{
+    return multi_stage_late_materialization_enabled.load(std::memory_order_relaxed)
+        || multi_stage_late_materialization_runtime_stats != nullptr;
+}
+
+const MultiStageLateMaterializationRuntimeStats & ScanContext::getMultiStageLateMaterializationRuntimeStats() const
+{
+    if (multi_stage_late_materialization_runtime_stats)
+        return *multi_stage_late_materialization_runtime_stats;
+    return merged_multi_stage_late_materialization_runtime_stats;
+}
+
+MultiStageLateMaterializationRuntimeStats & ScanContext::getMergedMultiStageLateMaterializationRuntimeStats()
+{
+    multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+    return merged_multi_stage_late_materialization_runtime_stats;
+}
+
+MSLMReadStageScanContext * ScanContext::getMutableMSLMReadStage(ReadTag read_tag)
+{
+    switch (read_tag)
+    {
+    case ReadTag::MSLMPushedFilter:
+        multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+        return &mslm_pushed_filter_read;
+    case ReadTag::MSLMCandidate:
+        multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+        return &mslm_candidate_read;
+    case ReadTag::MSLMFinalRest:
+        multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+        return &mslm_final_rest_read;
+    default:
+        return nullptr;
+    }
+}
+
+const MSLMReadStageScanContext * ScanContext::getMSLMReadStage(ReadTag read_tag) const
+{
+    switch (read_tag)
+    {
+    case ReadTag::MSLMPushedFilter:
+        return &mslm_pushed_filter_read;
+    case ReadTag::MSLMCandidate:
+        return &mslm_candidate_read;
+    case ReadTag::MSLMFinalRest:
+        return &mslm_final_rest_read;
+    default:
+        return nullptr;
+    }
+}
+
+void ScanContext::addDMFileReadTime(uint64_t ns, ReadTag read_tag)
+{
+    total_dmfile_read_time_ns += ns;
+    if (auto * read_stage = getMutableMSLMReadStage(read_tag); read_stage != nullptr)
+        read_stage->read_time_ns += ns;
+}
+
+void ScanContext::addDMFileScannedRows(uint64_t rows, ReadTag read_tag)
+{
+    switch (read_tag)
+    {
+    case ReadTag::Query:
+        dmfile_data_scanned_rows += rows;
+        break;
+    case ReadTag::MVCC:
+        dmfile_mvcc_scanned_rows += rows;
+        break;
+    case ReadTag::LMFilter:
+        dmfile_lm_filter_scanned_rows += rows;
+        break;
+    case ReadTag::MSLMPushedFilter:
+    case ReadTag::MSLMCandidate:
+    case ReadTag::MSLMFinalRest:
+        getMutableMSLMReadStage(read_tag)->dmfile_scanned_rows += rows;
+        break;
+    default:
+        break;
+    }
+}
+
+void ScanContext::addDMFileSkippedRows(uint64_t rows, ReadTag read_tag)
+{
+    switch (read_tag)
+    {
+    case ReadTag::Query:
+        dmfile_data_skipped_rows += rows;
+        break;
+    case ReadTag::MVCC:
+        dmfile_mvcc_skipped_rows += rows;
+        break;
+    case ReadTag::LMFilter:
+        dmfile_lm_filter_skipped_rows += rows;
+        break;
+    case ReadTag::MSLMPushedFilter:
+    case ReadTag::MSLMCandidate:
+    case ReadTag::MSLMFinalRest:
+        getMutableMSLMReadStage(read_tag)->dmfile_skipped_rows += rows;
+        break;
+    default:
+        break;
+    }
+}
+
+void ScanContext::setMultiStageLateMaterializationRuntimeStats(
+    const MultiStageLateMaterializationRuntimeStatsPtr & stats)
+{
+    multi_stage_late_materialization_runtime_stats = stats;
+    if (stats != nullptr)
+        multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+}
+
+void ScanContext::serializeMultiStageLateMaterialization(tipb::TiFlashScanContext & proto) const
+{
+    if (!hasMultiStageLateMaterializationContext())
+        return;
+
+    const auto & stats = getMultiStageLateMaterializationRuntimeStats();
+    auto * mslm = proto.mutable_multi_stage_late_materialization();
+    mslm->set_streams(stats.finished_streams.load(std::memory_order_relaxed));
+    mslm->set_late_mode_blocks(stats.late_mode_blocks.load(std::memory_order_relaxed));
+    mslm->set_direct_mode_blocks(stats.direct_mode_blocks.load(std::memory_order_relaxed));
+    serializeReadStage(mslm_pushed_filter_read, mslm->mutable_pushed_filter_read());
+    serializeReadStage(mslm_candidate_read, mslm->mutable_candidate_read());
+    serializeReadStage(mslm_final_rest_read, mslm->mutable_final_rest_read());
+    serializeFilter(
+        stats.pushed_filter_input_rows.load(std::memory_order_relaxed),
+        stats.pushed_filter_selected_rows.load(std::memory_order_relaxed),
+        stats.pushed_filter_filtered_rows.load(std::memory_order_relaxed),
+        mslm->mutable_pushed_filter());
+    serializeFilter(
+        stats.residual_filter_input_rows.load(std::memory_order_relaxed),
+        stats.residual_filter_selected_rows.load(std::memory_order_relaxed),
+        stats.residual_filter_filtered_rows.load(std::memory_order_relaxed),
+        mslm->mutable_residual_filter());
+    if (stats.topn_enabled.load(std::memory_order_relaxed))
+    {
+        auto * topn = mslm->mutable_running_topn();
+        topn->set_input_rows(stats.running_topn_input_rows.load(std::memory_order_relaxed));
+        topn->set_selected_rows(stats.running_topn_selected_rows.load(std::memory_order_relaxed));
+        topn->set_bypass_rows(stats.running_topn_bypass_rows.load(std::memory_order_relaxed));
+        topn->set_filtered_rows(stats.running_topn_filtered_rows.load(std::memory_order_relaxed));
+        topn->set_heap_size_sum(stats.topn_heap_size_sum.load(std::memory_order_relaxed));
+        topn->set_adaptive_warmup_rows(stats.topn_adaptive_warmup_rows.load(std::memory_order_relaxed));
+        topn->set_adaptive_post_warmup_input_rows(
+            stats.topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed));
+        topn->set_adaptive_post_warmup_candidate_rows(
+            stats.topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed));
+        topn->set_adaptive_disabled_streams(stats.topn_adaptive_disabled_streams.load(std::memory_order_relaxed));
+    }
+    mslm->set_final_rest_input_rows(stats.final_rest_input_rows.load(std::memory_order_relaxed));
+}
+
+void ScanContext::deserializeMultiStageLateMaterialization(const tipb::TiFlashScanContext & proto)
+{
+    if (!proto.has_multi_stage_late_materialization())
+        return;
+    mergeMultiStageLateMaterialization(proto);
+}
+
+void ScanContext::mergeMultiStageLateMaterialization(const ScanContext & other)
+{
+    if (!other.hasMultiStageLateMaterializationContext())
+        return;
+
+    multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+    mslm_pushed_filter_read.merge(other.mslm_pushed_filter_read);
+    mslm_candidate_read.merge(other.mslm_candidate_read);
+    mslm_final_rest_read.merge(other.mslm_final_rest_read);
+    merged_multi_stage_late_materialization_runtime_stats.merge(other.getMultiStageLateMaterializationRuntimeStats());
+}
+
+void ScanContext::mergeMultiStageLateMaterialization(const tipb::TiFlashScanContext & other)
+{
+    if (!other.has_multi_stage_late_materialization())
+        return;
+
+    const auto & mslm = other.multi_stage_late_materialization();
+    multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed);
+
+    if (mslm.has_pushed_filter_read())
+        mergeReadStage(mslm_pushed_filter_read, mslm.pushed_filter_read());
+    if (mslm.has_candidate_read())
+        mergeReadStage(mslm_candidate_read, mslm.candidate_read());
+    if (mslm.has_final_rest_read())
+        mergeReadStage(mslm_final_rest_read, mslm.final_rest_read());
+
+    auto & stats = getMergedMultiStageLateMaterializationRuntimeStats();
+    stats.finished_streams += mslm.streams();
+    stats.late_mode_blocks += mslm.late_mode_blocks();
+    stats.direct_mode_blocks += mslm.direct_mode_blocks();
+    if (mslm.has_pushed_filter())
+    {
+        const auto & pushed_filter = mslm.pushed_filter();
+        stats.pushed_filter_input_rows += pushed_filter.input_rows();
+        stats.pushed_filter_selected_rows += pushed_filter.selected_rows();
+        stats.pushed_filter_filtered_rows += pushed_filter.filtered_rows();
+        stats.stage0_output_rows += pushed_filter.selected_rows();
+    }
+    if (mslm.has_residual_filter())
+    {
+        const auto & residual_filter = mslm.residual_filter();
+        stats.residual_filter_input_rows += residual_filter.input_rows();
+        stats.residual_filter_selected_rows += residual_filter.selected_rows();
+        stats.residual_filter_filtered_rows += residual_filter.filtered_rows();
+        stats.stage1_output_rows += residual_filter.selected_rows();
+    }
+    stats.final_rest_input_rows += mslm.final_rest_input_rows();
+    stats.topn_candidate_rows += mslm.final_rest_input_rows();
+    if (mslm.has_running_topn())
+    {
+        const auto & running_topn = mslm.running_topn();
+        stats.topn_enabled.store(true, std::memory_order_relaxed);
+        stats.running_topn_input_rows += running_topn.input_rows();
+        stats.running_topn_selected_rows += running_topn.selected_rows();
+        stats.running_topn_bypass_rows += running_topn.bypass_rows();
+        stats.running_topn_filtered_rows += running_topn.filtered_rows();
+        stats.topn_heap_size_sum += running_topn.heap_size_sum();
+        stats.topn_adaptive_warmup_rows += running_topn.adaptive_warmup_rows();
+        stats.topn_adaptive_post_warmup_input_rows += running_topn.adaptive_post_warmup_input_rows();
+        stats.topn_adaptive_post_warmup_candidate_rows += running_topn.adaptive_post_warmup_candidate_rows();
+        stats.topn_adaptive_disabled_streams += running_topn.adaptive_disabled_streams();
+    }
+}
+
 } // namespace DB::DM
diff --git a/dbms/src/Storages/DeltaMerge/ScanContext.h b/dbms/src/Storages/DeltaMerge/ScanContext.h
index 14e1b5e4e5e..0eaa038e5d7 100644
--- a/dbms/src/Storages/DeltaMerge/ScanContext.h
+++ b/dbms/src/Storages/DeltaMerge/ScanContext.h
@@ -18,6 +18,7 @@
 #include <Common/TiFlashMetrics.h>
 #include <Flash/ResourceControl/LocalAdmissionController.h>
 #include <Poco/Util/AbstractConfiguration.h>
+#include <Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h>
 #include <Storages/DeltaMerge/ReadMode.h>
 #include <Storages/DeltaMerge/ScanContext_fwd.h>
 #include <common/types.h>
@@ -34,6 +35,23 @@ namespace DB::DM
 {
 class PushDownFilter;
 using PushDownFilterPtr = std::shared_ptr<PushDownFilter>;
+
+struct MSLMReadStageScanContext
+{
+    std::atomic<uint64_t> dmfile_scanned_rows{0};
+    std::atomic<uint64_t> dmfile_skipped_rows{0};
+    std::atomic<uint64_t> read_bytes{0};
+    std::atomic<uint64_t> read_time_ns{0};
+
+    void merge(const MSLMReadStageScanContext & other)
+    {
+        dmfile_scanned_rows += other.dmfile_scanned_rows;
+        dmfile_skipped_rows += other.dmfile_skipped_rows;
+        read_bytes += other.read_bytes;
+        read_time_ns += other.read_time_ns;
+    }
+};
+
 /// ScanContext is used to record statistical information in table scan for current query.
 /// For each table scan(one executor id), there is only one ScanContext.
 /// ScanContext helps to collect the statistical information of the table scan to show in `EXPLAIN ANALYZE`.
@@ -46,7 +64,12 @@ class ScanContext
     std::atomic<uint64_t> dmfile_mvcc_skipped_rows{0};
     std::atomic<uint64_t> dmfile_lm_filter_scanned_rows{0};
     std::atomic<uint64_t> dmfile_lm_filter_skipped_rows{0};
-    std::atomic<uint64_t> dmfile_mslm_stage1_filter_scanned_rows{0};
+    std::atomic<bool> multi_stage_late_materialization_enabled{false};
+    MSLMReadStageScanContext mslm_pushed_filter_read;
+    MSLMReadStageScanContext mslm_candidate_read;
+    MSLMReadStageScanContext mslm_final_rest_read;
+    MultiStageLateMaterializationRuntimeStats merged_multi_stage_late_materialization_runtime_stats;
+    MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats;
     std::atomic<uint64_t> total_dmfile_read_time_ns{0};
 
     std::atomic<uint64_t> total_rs_pack_filter_check_time_ns{0};
@@ -130,6 +153,7 @@ class ScanContext
         dmfile_mvcc_skipped_rows = tiflash_scan_context_pb.dmfile_mvcc_skipped_rows();
         dmfile_lm_filter_scanned_rows = tiflash_scan_context_pb.dmfile_lm_filter_scanned_rows();
         dmfile_lm_filter_skipped_rows = tiflash_scan_context_pb.dmfile_lm_filter_skipped_rows();
+        deserializeMultiStageLateMaterialization(tiflash_scan_context_pb);
         total_rs_pack_filter_check_time_ns = tiflash_scan_context_pb.total_dmfile_rs_check_ms() * 1000000;
         // TODO: rs_pack_filter_none, rs_pack_filter_some, rs_pack_filter_all,rs_pack_filter_all_null
         // rs_dmfile_read_with_all
@@ -186,6 +210,7 @@ class ScanContext
         tiflash_scan_context_pb.set_dmfile_mvcc_skipped_rows(dmfile_mvcc_skipped_rows);
         tiflash_scan_context_pb.set_dmfile_lm_filter_scanned_rows(dmfile_lm_filter_scanned_rows);
         tiflash_scan_context_pb.set_dmfile_lm_filter_skipped_rows(dmfile_lm_filter_skipped_rows);
+        serializeMultiStageLateMaterialization(tiflash_scan_context_pb);
         tiflash_scan_context_pb.set_total_dmfile_rs_check_ms(total_rs_pack_filter_check_time_ns / 1000000);
         // TODO: pack_filter_none, pack_filter_some, pack_filter_all
         tiflash_scan_context_pb.set_total_dmfile_read_ms(total_dmfile_read_time_ns / 1000000);
@@ -241,7 +266,7 @@ class ScanContext
         dmfile_mvcc_skipped_rows += other.dmfile_mvcc_skipped_rows;
         dmfile_lm_filter_scanned_rows += other.dmfile_lm_filter_scanned_rows;
         dmfile_lm_filter_skipped_rows += other.dmfile_lm_filter_skipped_rows;
-        dmfile_mslm_stage1_filter_scanned_rows += other.dmfile_mslm_stage1_filter_scanned_rows;
+        mergeMultiStageLateMaterialization(other);
         total_rs_pack_filter_check_time_ns += other.total_rs_pack_filter_check_time_ns;
         rs_pack_filter_none += other.rs_pack_filter_none;
         rs_pack_filter_some += other.rs_pack_filter_some;
@@ -303,6 +328,7 @@ class ScanContext
         dmfile_mvcc_skipped_rows += other.dmfile_mvcc_skipped_rows();
         dmfile_lm_filter_scanned_rows += other.dmfile_lm_filter_scanned_rows();
         dmfile_lm_filter_skipped_rows += other.dmfile_lm_filter_skipped_rows();
+        mergeMultiStageLateMaterialization(other);
         total_rs_pack_filter_check_time_ns += other.total_dmfile_rs_check_ms() * 1000000;
         // TODO: rs_pack_filter_none, rs_pack_filter_some, rs_pack_filter_all, rs_pack_filter_all_null
         // rs_dmfile_read_with_all
@@ -360,9 +386,22 @@ class ScanContext
     // LACBytesCollector is not thread-safe, to avoid locking, we create a new one for each stream.
     std::optional<LACBytesCollector> newLACBytesCollector(ReadTag read_tag);
     void addUserReadBytes(size_t bytes, ReadTag read_tag, std::optional<LACBytesCollector> & lac_bytes_collector);
+    void addDMFileReadTime(uint64_t ns, ReadTag read_tag);
+    void addDMFileScannedRows(uint64_t rows, ReadTag read_tag);
+    void addDMFileSkippedRows(uint64_t rows, ReadTag read_tag);
+    void setMultiStageLateMaterializationRuntimeStats(const MultiStageLateMaterializationRuntimeStatsPtr & stats);
     uint64_t userReadBytes() const { return query_read_bytes + mvcc_read_bytes; }
 
 private:
+    bool hasMultiStageLateMaterializationContext() const;
+    const MultiStageLateMaterializationRuntimeStats & getMultiStageLateMaterializationRuntimeStats() const;
+    MultiStageLateMaterializationRuntimeStats & getMergedMultiStageLateMaterializationRuntimeStats();
+    MSLMReadStageScanContext * getMutableMSLMReadStage(ReadTag read_tag);
+    const MSLMReadStageScanContext * getMSLMReadStage(ReadTag read_tag) const;
+    void serializeMultiStageLateMaterialization(tipb::TiFlashScanContext & proto) const;
+    void deserializeMultiStageLateMaterialization(const tipb::TiFlashScanContext & proto);
+    void mergeMultiStageLateMaterialization(const ScanContext & other);
+    void mergeMultiStageLateMaterialization(const tipb::TiFlashScanContext & other);
     void serializeRegionNumOfInstance(tipb::TiFlashScanContext & proto) const;
     void deserializeRegionNumberOfInstance(const tipb::TiFlashScanContext & proto);
     void mergeRegionNumberOfInstance(const ScanContext & other);
diff --git a/dbms/src/Storages/DeltaMerge/Segment.cpp b/dbms/src/Storages/DeltaMerge/Segment.cpp
index 95f751c5cb2..9fbe02fb691 100644
--- a/dbms/src/Storages/DeltaMerge/Segment.cpp
+++ b/dbms/src/Storages/DeltaMerge/Segment.cpp
@@ -3627,6 +3627,8 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
     size_t expected_block_size)
 {
     const auto & filter_columns = filter->filter_columns;
+    const bool enable_multi_stage_late_materialization
+        = multi_stage_late_materialization_filter && multi_stage_late_materialization_filter->before_where;
     BlockInputStreamPtr filter_column_stream = getConcatSkippableBlockInputStream(
         bitmap_filter,
         segment_snap,
@@ -3636,7 +3638,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
         filter->rs_operator,
         start_ts,
         expected_block_size,
-        ReadTag::LMFilter);
+        enable_multi_stage_late_materialization ? ReadTag::MSLMPushedFilter : ReadTag::LMFilter);
 
     if (unlikely(filter_columns->size() == columns_to_read.size()))
     {
@@ -3682,7 +3684,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
         dm_context.tracing_id);
     filter_column_stream->setExtraInfo("push down filter");
 
-    if (multi_stage_late_materialization_filter && multi_stage_late_materialization_filter->before_where)
+    if (enable_multi_stage_late_materialization)
     {
         const auto & stage1_filter_columns = multi_stage_late_materialization_filter->filter_columns;
         RUNTIME_CHECK(stage1_filter_columns != nullptr);
@@ -3705,7 +3707,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
             filter->rs_operator,
             start_ts,
             expected_block_size,
-            ReadTag::MSLMStage1Filter);
+            ReadTag::MSLMCandidate);
 
         auto final_rest_columns_to_read = std::make_shared<ColumnDefines>(columns_to_read);
         for (const auto & col : *stage1_filter_columns)
@@ -3730,7 +3732,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream(
             filter->rs_operator,
             start_ts,
             expected_block_size,
-            ReadTag::Query);
+            ReadTag::MSLMFinalRest);
 
         LOG_DEBUG(
             segment_snap->log,
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp
index 71e6d953c92..cbccfbd477d 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp
@@ -268,9 +268,12 @@ try
     in->readSuffix();
 
     ASSERT_EQ(output_rows, rows);
-    ASSERT_EQ(scan_context->dmfile_lm_filter_scanned_rows, rows);
-    ASSERT_EQ(scan_context->dmfile_mslm_stage1_filter_scanned_rows, rows);
-    ASSERT_NE(scan_context->toJson().find("dmfile_mslm_stage1_filter_scanned_rows"), String::npos);
+    ASSERT_EQ(scan_context->dmfile_lm_filter_scanned_rows, 0);
+    ASSERT_EQ(scan_context->mslm_pushed_filter_read.dmfile_scanned_rows, rows);
+    ASSERT_EQ(scan_context->mslm_candidate_read.dmfile_scanned_rows, rows);
+    ASSERT_EQ(scan_context->mslm_final_rest_read.dmfile_scanned_rows, rows);
+    ASSERT_NE(scan_context->toJson().find("multi_stage_late_materialization"), String::npos);
+    ASSERT_NE(scan_context->toJson().find("candidate_read"), String::npos);
 }
 CATCH
 
diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index 44c3504a69f..b73a78228fa 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -17,6 +17,7 @@
 - [Running Local TopN 状态](#running-local-topn-状态)
 - [Filter 坐标](#filter-坐标)
 - [启用规则](#启用规则)
+- [可观测性](#可观测性)
 - [实现计划](#实现计划)
 - [测试方案](#测试方案)
 - [性能评估](#性能评估)
@@ -147,8 +148,8 @@ Final:
 ```text
 开始时启用 running local TopN
 先用 warmup rows 建立初始 TopN 状态
-warmup 后，持续统计后续 residual input rows 和 TopN candidate rows
-如果后续累计 residual_input_rows 达到最小检查门槛，且 topn_candidate_rows / residual_input_rows 过高
+warmup 后，持续统计后续 `adaptive_post_warmup_input_rows` 和 `adaptive_post_warmup_candidate_rows`
+如果后续累计 input rows 达到最小检查门槛，且 candidate rows / input rows 过高
   后续 blocks 禁用 TopN filter
   后续 final rest 只使用 residual filter
 禁用后不再重新启用
@@ -156,7 +157,7 @@ warmup 后，持续统计后续 residual input rows 和 TopN candidate rows
 
 这个 adaptive 只做运行时止损，不改变正确性：禁用后输出的是 residual Selection 后的更多 rows，仍然是上层全局 TopN 的 superset。上层 TopN executor 保留，因此最终结果不变。
 
-需要注意，这个 runtime disable 不能避免已经加入 Stage 1 的 order-by columns IO，因为 Stage 1 schema 在 stream 构建时已经固定。它主要减少后续 comparator、heap 维护和 TopN candidate filter 处理成本。避免明显不划算的 order-by column 读取仍然依赖静态启发式 gating。
+需要注意，这个 runtime disable 不能避免已经加入 Stage 1 的 order-by columns IO，因为 Stage 1 schema 在 stream 构建时已经固定。它主要减少后续 comparator、heap 维护和 TopN filter 处理成本。避免明显不划算的 order-by column 读取仍然依赖静态启发式 gating。
 
 ## 正确性
 
@@ -269,7 +270,7 @@ final_rest_columns = final_columns_to_read - stage1_columns
 
 如果 order-by column 已经属于 residual filter columns，则 Stage 1 只读一次。如果 order-by column 也是最终输出列，则 final rest columns 中必须排除该列，避免重复读取。
 
-第一版不支持没有 residual Selection 的 `TopN -> TableScan`。这种场景没有 Selection executor 可以承载 candidate rows 的 `EXPLAIN ANALYZE` 展示，也需要让 residual filter 变成 optional，后续单独设计。
+第一版不支持没有 residual Selection 的 `TopN -> TableScan`。这种场景需要将 candidate read 从 residual-filter-driven 扩展为 order-by-column-driven；scan context 的 `residual_filter` 已经按 optional 设计，可以兼容后续扩展。
 
 plain column 还需要保证 storage 内比较语义和上层 TopN 一致。第一版建议保守禁用下面场景：
 
@@ -573,8 +574,8 @@ combined_filter = compose(stage0_filter, stage1_final_filter)
 topk <= 2048
 order_by_column_count <= 4
 runtime adaptive warmup rows >= max(topk * 4, 8192)
-after warmup, check only when cumulative residual_input_rows >= max(topk * 4, 8192)
-after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >= 0.7
+after warmup, check only when adaptive_post_warmup_input_rows >= max(topk * 4, 8192)
+after warmup, disable if adaptive_post_warmup_candidate_rows / adaptive_post_warmup_input_rows >= 0.7
 ```
 
 实际阈值可以根据 benchmark 调整。
@@ -593,6 +594,208 @@ after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >=
 - 查询要求 keep order 且和 MSLM 现有限制冲突。
 - TopN 位于 join、aggregation、window 等算子之上。
 
+## 可观测性
+
+`EXPLAIN ANALYZE` 中各 executor 的 `actRows` 语义保持现状：
+
+- TableScan executor 的 `actRows` 展示 Stage 0 pushed filter 之后的 rows。
+- 没有启用 running local TopN 时，Selection executor 的 `actRows` 展示 residual Selection 之后的 rows。
+- 启用 running local TopN 时，Selection executor 的 `actRows` 展示 MSLM 最终需要读取 final rest columns 的 rows。
+
+MSLM 内部细节通过 TiFlash scan context 展示。原 single-stage late materialization 继续使用现有字段，不改变语义：
+
+- `lm_skip_rows`
+- `dmfile_lm_filter_scanned_rows`
+- `dmfile_lm_filter_skipped_rows`
+
+MSLM 不复用这些旧 LM 字段，而是在 `TiFlashScanContext` 中新增独立的 `multi_stage_late_materialization` context。这样可以避免 single-stage LM 和 MSLM 的统计混在一起解释。
+
+建议 TiPB 结构如下：
+
+```proto
+message TiFlashScanContext {
+    // Existing fields unchanged.
+    optional uint64 dmfile_lm_filter_scanned_rows = 33;
+    optional uint64 dmfile_lm_filter_skipped_rows = 34;
+
+    // Only set when multi-stage late materialization is enabled.
+    optional TiFlashMultiStageLateMaterializationScanContext multi_stage_late_materialization = 35;
+
+    // Existing vector index fields unchanged.
+    optional uint64 total_vector_idx_load_from_s3 = 100;
+    // ...
+}
+
+message TiFlashMultiStageLateMaterializationScanContext {
+    optional uint64 streams = 1;
+    optional uint64 late_mode_blocks = 2;
+    optional uint64 direct_mode_blocks = 3;
+
+    // Physical read stages. These fields describe DMFile/storage read cost only.
+    optional TiFlashMSLMReadStageContext pushed_filter_read = 10;
+    optional TiFlashMSLMReadStageContext candidate_read = 11;
+    optional TiFlashMSLMReadStageContext final_rest_read = 12;
+
+    // Logical row reducers. Each reducer is independently optional.
+    optional TiFlashMSLMFilterContext pushed_filter = 20;
+    optional TiFlashMSLMFilterContext residual_filter = 21;
+    optional TiFlashMSLMRunningTopNContext running_topn = 22;
+
+    // Logical rows that actually need final-rest materialization.
+    optional uint64 final_rest_input_rows = 30;
+}
+
+message TiFlashMSLMReadStageContext {
+    optional uint64 dmfile_scanned_rows = 1;
+    optional uint64 dmfile_skipped_rows = 2;
+    optional uint64 read_bytes = 3;
+    optional uint64 read_time_ms = 4;
+}
+
+message TiFlashMSLMFilterContext {
+    optional uint64 input_rows = 1;
+    optional uint64 selected_rows = 2;
+    optional uint64 filtered_rows = 3;
+}
+
+message TiFlashMSLMRunningTopNContext {
+    optional uint64 input_rows = 1;
+    optional uint64 selected_rows = 2;
+    optional uint64 bypass_rows = 3;
+    optional uint64 filtered_rows = 4;
+    optional uint64 heap_size_sum = 5;
+
+    optional uint64 adaptive_warmup_rows = 10;
+    optional uint64 adaptive_post_warmup_input_rows = 11;
+    optional uint64 adaptive_post_warmup_candidate_rows = 12;
+    optional uint64 adaptive_disabled_streams = 13;
+}
+```
+
+### Logical Reducer 语义
+
+`pushed_filter` 表示 pushed filter 这个逻辑 reducer。只要 MSLM 中存在 pushed filter，就填该字段：
+
+- `input_rows`: 进入 pushed filter 的逻辑 rows。
+- `selected_rows`: pushed filter 之后留下、会进入后续 MSLM stage 的逻辑 rows。
+- `filtered_rows`: 被 pushed filter 过滤掉的逻辑 rows。
+
+`residual_filter` 表示 residual Selection 这个逻辑 reducer。没有 residual filter 时不填该字段，不用 `0` 表示不存在：
+
+- `input_rows`: 进入 residual filter 的逻辑 rows。当前实现中通常等于 `pushed_filter.selected_rows`。
+- `selected_rows`: residual filter 之后留下、会进入后续 TopN 或 final rest 读取的逻辑 rows。
+- `filtered_rows`: 被 residual filter 过滤掉的逻辑 rows。
+
+filter reducer 的 rows 满足：
+
+```text
+filter.input_rows = filter.selected_rows + filter.filtered_rows
+```
+
+`running_topn` 表示 running local TopN 这个逻辑 reducer。没有 running local TopN 时不填该字段：
+
+- `input_rows`: 进入 running local TopN 的逻辑 rows。
+- `selected_rows`: 被 running local TopN 选中、需要继续读取 final rest columns 的 rows。
+- `bypass_rows`: runtime adaptive disable 之后，不再经过 TopN filter、直接进入 final rest 读取的 rows。
+- `filtered_rows`: 被 running local TopN 过滤掉、因此不需要读取 final rest columns 的 rows。
+- `heap_size_sum`: 所有 streams 的 heap size 总和。它不是 `selected_rows`。
+
+`running_topn` 的 rows 满足：
+
+```text
+running_topn.input_rows
+  = running_topn.selected_rows
+  + running_topn.bypass_rows
+  + running_topn.filtered_rows
+```
+
+`final_rest_input_rows` 表示最终需要读取 final rest columns 的逻辑 rows。它和物理读行数不同：
+
+```text
+with running_topn:
+  final_rest_input_rows = running_topn.selected_rows + running_topn.bypass_rows
+
+without running_topn and with residual_filter:
+  final_rest_input_rows = residual_filter.selected_rows
+
+without running_topn and without residual_filter:
+  final_rest_input_rows = pushed_filter.selected_rows
+```
+
+如果未来支持没有 pushed filter、直接由 TopN 驱动的 MSLM，`pushed_filter` 和 `residual_filter` 可以都不存在，`running_topn.input_rows` 表示进入 TopN 的候选 rows，`final_rest_input_rows` 仍然等于 `running_topn.selected_rows + running_topn.bypass_rows`。
+
+### Physical Read Stage 语义
+
+`pushed_filter_read`、`candidate_read`、`final_rest_read` 只表示物理读成本，不承担 logical input/output rows 语义。因此 read stage 中不设置 `input_rows/output_rows` 字段。
+
+- `pushed_filter_read`: 为 pushed filter columns 做的物理读取。
+- `candidate_read`: 为产生 final-rest candidate row set 做的物理读取。当前实现中，它读取 residual filter columns 和 order-by columns；未来 topn-only 路径中，它可以只读取 order-by columns。
+- `final_rest_read`: 为最终输出所需 rest columns 做的物理读取。
+
+这些 read stage 的字段含义：
+
+- `dmfile_scanned_rows`: DMFile 层实际 scan 的 rows。
+- `dmfile_skipped_rows`: DMFile 层跳过的 rows。
+- `read_bytes`: 该 read stage 实际读取的 bytes。
+- `read_time_ms`: 该 read stage 的读取耗时。
+
+`final_rest_input_rows` 和 `final_rest_read.dmfile_scanned_rows` 不能混用：
+
+```text
+final_rest_input_rows:
+  logical rows that need final-rest materialization.
+
+final_rest_read.dmfile_scanned_rows:
+  physical rows scanned by DMFile for final-rest columns. It may be larger than
+  final_rest_input_rows because of pack/block-level read amplification.
+```
+
+### ReadTag 拆分
+
+为了让 MSLM 不污染原 single-stage LM 的统计，TiFlash 内部建议拆分 read tag：
+
+```cpp
+ReadTag::LMFilter              // existing single-stage LM filter read
+ReadTag::MSLMPushedFilter      // MSLM pushed filter read
+ReadTag::MSLMCandidate         // MSLM residual/order-by candidate read
+ReadTag::MSLMFinalRest         // MSLM final rest read
+```
+
+其中 `ReadTag::LMFilter` 继续填旧字段；`MSLMPushedFilter`、`MSLMCandidate`、`MSLMFinalRest` 只填新的 `multi_stage_late_materialization` context。RU 统计上这些 MSLM read tags 仍然属于 query read。
+
+### Optional 字段规则
+
+不存在或无意义的字段不填，不用 `0` 表示不存在。`0` 只表示该阶段或 reducer 真实存在，但计数结果为 0。
+
+例如没有 residual filter 时：
+
+```json
+{
+  "multi_stage_late_materialization": {
+    "pushed_filter": {
+      "input_rows": 1000000,
+      "selected_rows": 120000,
+      "filtered_rows": 880000
+    },
+    "final_rest_input_rows": 120000
+  }
+}
+```
+
+而不是：
+
+```json
+{
+  "multi_stage_late_materialization": {
+    "residual_filter": {
+      "input_rows": 0,
+      "selected_rows": 0,
+      "filtered_rows": 0
+    }
+  }
+}
+```
+
 ## 实现计划
 
 ### Step 1: 文档和测试用例设计
@@ -653,13 +856,15 @@ after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >=
 
 ### Step 7: Runtime Stats 和可观测性
 
-- 保留 `stage1_output_rows` internal counter，表示 residual Selection 通过的真实行数。
-- 增加 `topn_candidate_rows` internal counter，表示 `residual_filter && topn_candidate_filter` 之后实际读取 final rest columns 的行数。
+- 保留 logical filter counters：`pushed_filter.input_rows/selected_rows/filtered_rows`、`residual_filter.input_rows/selected_rows/filtered_rows`。
+- 增加 `final_rest_input_rows` internal counter，表示实际需要读取 final rest columns 的逻辑行数。
+- 增加 running TopN counters：`input_rows`、`selected_rows`、`bypass_rows`、`filtered_rows`、`heap_size_sum`。
 - 增加 runtime one-way adaptive disable 统计，表示有多少 streams 因 TopN candidate ratio 过高而禁用后续 TopN filter。
-- 在 TopN-enhanced MSLM 启用时，可以用 `topn_candidate_rows` overwrite residual Selection executor 的 `actRows`，让 `EXPLAIN ANALYZE` 直接展示 final rest columns materialized rows。
-- 该 overwrite 会改变 Selection `actRows` 的展示语义：它不再表示纯 residual Selection 后的逻辑行数，而是表示 storage 内输出给上层 TopN 的 candidate rows。
-- TableScan executor 的 `actRows` 仍然使用 `stage0_output_rows`。
-- 可选增加 debug log，输出 residual passed rows、TopN candidate rows、heap size、filtered rows、adaptive disable 状态。
+- 在 running local TopN 启用时，用 `final_rest_input_rows` overwrite residual Selection executor 的 `actRows`，让 `EXPLAIN ANALYZE` 直接展示 final rest columns materialized rows。
+- 未启用 running local TopN 时，Selection executor 的 `actRows` 仍展示 residual Selection 后的逻辑行数。
+- TableScan executor 的 `actRows` 仍然使用 pushed filter output rows。
+- 增加 `TiFlashScanContext.multi_stage_late_materialization`，不要复用原 single-stage LM 的 `lm_*` / `dmfile_lm_filter_*` 字段。
+- 增加 debug log，输出 residual passed rows、TopN selected/bypass/filtered rows、heap size、adaptive disable 状态。
 
 ## 测试方案
 
@@ -683,8 +888,8 @@ after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >=
 - ties 场景比较，验证结果满足 SQL order semantics。
 - Stage 0 全过滤。
 - residual filter 全过滤。
-- TopN candidate filter 全过滤。
-- TopN candidate filter 全通过。
+- TopN filter 全过滤。
+- TopN filter 全通过。
 - ORDER BY plain column 按 column id 对齐到 storage column，避免依赖 child schema name。
 - ORDER BY column 需要 TableScan 后 extra cast 时禁用优化。
 - 没有 residual Selection 时禁用优化。
@@ -694,7 +899,7 @@ after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >=
 
 - MSLM 原有测试必须全部通过。
 - TopN 上方的 projection 仍然正确。
-- TopN-enhanced MSLM 启用时，EXPLAIN ANALYZE 中 residual Selection `actRows` 展示 TopN candidate rows，即 final rest columns materialized rows。
+- TopN-enhanced MSLM 启用时，EXPLAIN ANALYZE 中 residual Selection `actRows` 展示 `final_rest_input_rows`，即 final rest columns materialized rows。
 - 禁用条件下必须走原路径。
 
 ## 性能评估
@@ -719,7 +924,8 @@ after warmup, disable if cumulative topn_candidate_rows / residual_input_rows >=
 - DMFile read bytes。
 - CPU time。
 - memory usage。
-- candidate rows output to upper TopN。
+- `final_rest_input_rows`。
+- `running_topn.selected_rows/bypass_rows/filtered_rows`。
 
 ## 风险
 
@@ -785,7 +991,7 @@ Running local TopN 不能撤回之前已经输出的 rows。极端数据分布
 缓解：
 
 - 上层 TopN 保留，保证正确性。
-- 通过 stats 观察 candidate rows 数量。
+- 通过 stats 观察 `final_rest_input_rows` 和 `running_topn.selected_rows/bypass_rows/filtered_rows`。
 - 后续考虑更强的 candidate row id materialization。
 
 ## 备选方案
@@ -830,5 +1036,5 @@ Running local TopN 是 block-local TopN 的增强版，复杂度增加有限，
 - ties 场景是否需要为了结果稳定性保留所有 equal-worst rows。
 - 第一版是否禁用 variable-length order-by columns，还是实现 per-stream memory guard。
 - 是否需要支持 ORDER BY expression，以及 expression action 如何复用。
-- 是否需要在 EXPLAIN ANALYZE 中暴露 TopN-enhanced MSLM 的 candidate rows。
+- 是否需要在 `EXPLAIN ANALYZE` 中展示 scan context 之外的更细粒度 MSLM TopN 信息。
 - 是否需要在 cost model 中引入列宽和 TopN selectivity。

From dedc049efef7ed8afbd0edbaf2a7dd73a4dd96de Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Thu, 30 Jul 2026 10:10:29 +0800
Subject: [PATCH 15/18] Fix MSLM TopN fallback stats and type checks

---
 dbms/src/Debug/MockStorage.cpp                |  19 +-
 dbms/src/Debug/MockStorage.h                  |  18 ++
 .../Coprocessor/DAGStorageInterpreter.cpp     |  13 +-
 .../Flash/tests/gtest_executors_with_dm.cpp   | 189 +++++++++++++++++-
 ...ageLateMaterializationBlockInputStream.cpp |   2 +
 ...ultiStageLateMaterializationRuntimeStats.h |  11 +-
 dbms/src/Storages/DeltaMerge/ScanContext.cpp  |   2 +-
 dbms/src/Storages/StorageDeltaMerge.cpp       |  26 ++-
 8 files changed, 252 insertions(+), 28 deletions(-)

diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp
index 97cb10cb25b..57a7599bd40 100644
--- a/dbms/src/Debug/MockStorage.cpp
+++ b/dbms/src/Debug/MockStorage.cpp
@@ -505,10 +505,10 @@ void MockStorage::buildExecFromDeltaMerge(
         {
             multi_stage_late_materialization_runtime_stats
                 = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
-                    fmt::format("mock table_scan_executor_id={}", table_scan_executor_id),
-                    multi_stage_late_materialization_topn != nullptr);
+                    fmt::format("mock table_scan_executor_id={}", table_scan_executor_id));
             if (auto * dag_context = context.getDAGContext(); dag_context != nullptr && !table_scan_executor_id.empty())
             {
+                dag_context->scan_context_map[table_scan_executor_id] = query_info.mvcc_query_info->scan_context;
                 dag_context->setExecutorRowsOverride(
                     table_scan_executor_id,
                     std::shared_ptr<std::atomic<UInt64>>(
@@ -516,14 +516,12 @@ void MockStorage::buildExecFromDeltaMerge(
                         &multi_stage_late_materialization_runtime_stats->stage0_output_rows));
                 dag_context->setExecutorRowsOverride(
                     filter_conditions->executor_id,
-                    multi_stage_late_materialization_topn != nullptr
-                        ? std::shared_ptr<std::atomic<UInt64>>(
-                            multi_stage_late_materialization_runtime_stats,
-                            &multi_stage_late_materialization_runtime_stats->topn_candidate_rows)
-                        : std::shared_ptr<std::atomic<UInt64>>(
-                            multi_stage_late_materialization_runtime_stats,
-                            &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+                    std::shared_ptr<std::atomic<UInt64>>(
+                        multi_stage_late_materialization_runtime_stats,
+                        &multi_stage_late_materialization_runtime_stats->final_rest_input_rows));
             }
+            query_info.mvcc_query_info->scan_context->setMultiStageLateMaterializationRuntimeStats(
+                multi_stage_late_materialization_runtime_stats);
         }
         query_info.dag_query = std::make_unique<DAGQueryInfo>(
             filter_conditions->conditions,
@@ -617,6 +615,8 @@ void MockStorage::addTableInfoForDeltaMerge(const String & name, const MockColum
         TiDB::ColumnInfo ret;
         ret.name = column.name;
         ret.tp = column.type;
+        ret.collate = column.collate;
+        ret.elems = column.elems;
 
         if (!column.nullable)
             ret.setNotNullFlag();
@@ -900,6 +900,7 @@ TiDB::ColumnInfos mockColumnInfosToTiDBColumnInfos(const MockColumnInfoVec & moc
         column_info.name = mock_column_info.name;
         column_info.tp = mock_column_info.type;
         column_info.collate = mock_column_info.collate;
+        column_info.elems = mock_column_info.elems;
         column_info.id = col_id++;
         // TODO: find a way to assign decimal field's flen.
         if (column_info.tp == TiDB::TP::TypeNewDecimal)
diff --git a/dbms/src/Debug/MockStorage.h b/dbms/src/Debug/MockStorage.h
index d149b8565a2..84ea3a6919d 100644
--- a/dbms/src/Debug/MockStorage.h
+++ b/dbms/src/Debug/MockStorage.h
@@ -27,6 +27,8 @@
 #include <atomic>
 #include <memory>
 #include <unordered_map>
+#include <utility>
+#include <vector>
 
 namespace DB
 {
@@ -37,10 +39,26 @@ struct SelectQueryInfo;
 
 struct MockColumnInfo
 {
+    MockColumnInfo() = default;
+
+    MockColumnInfo(
+        String name_,
+        TiDB::TP type_,
+        bool nullable_ = true,
+        Poco::Dynamic::Var collate_ = {},
+        std::vector<std::pair<std::string, Int16>> elems_ = {})
+        : name(std::move(name_))
+        , type(type_)
+        , nullable(nullable_)
+        , collate(std::move(collate_))
+        , elems(std::move(elems_))
+    {}
+
     String name;
     TiDB::TP type;
     bool nullable = true;
     Poco::Dynamic::Var collate{}; // default empty means no collation.
+    std::vector<std::pair<std::string, Int16>> elems;
 };
 using MockColumnInfoVec = std::vector<MockColumnInfo>;
 using TableInfo = TiDB::TableInfo;
diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
index 6d96fb851b2..154450989a2 100644
--- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp
@@ -1047,8 +1047,7 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
     {
         multi_stage_late_materialization_runtime_stats
             = std::make_shared<DM::MultiStageLateMaterializationRuntimeStats>(
-                fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID()),
-                multi_stage_late_materialization_topn != nullptr);
+                fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID()));
         dagContext().setExecutorRowsOverride(
             table_scan.getTableScanExecutorID(),
             std::shared_ptr<std::atomic<UInt64>>(
@@ -1056,13 +1055,9 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
                 &multi_stage_late_materialization_runtime_stats->stage0_output_rows));
         dagContext().setExecutorRowsOverride(
             filter_conditions.executor_id,
-            multi_stage_late_materialization_topn != nullptr
-                ? std::shared_ptr<std::atomic<UInt64>>(
-                    multi_stage_late_materialization_runtime_stats,
-                    &multi_stage_late_materialization_runtime_stats->final_rest_input_rows)
-                : std::shared_ptr<std::atomic<UInt64>>(
-                    multi_stage_late_materialization_runtime_stats,
-                    &multi_stage_late_materialization_runtime_stats->stage1_output_rows));
+            std::shared_ptr<std::atomic<UInt64>>(
+                multi_stage_late_materialization_runtime_stats,
+                &multi_stage_late_materialization_runtime_stats->final_rest_input_rows));
         if (auto scan_context_it = dagContext().scan_context_map.find(table_scan.getTableScanExecutorID());
             scan_context_it != dagContext().scan_context_map.end() && scan_context_it->second != nullptr)
         {
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 29bb9e0051d..9e7b11ccf05 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -14,11 +14,13 @@
 
 #include <Common/MyDuration.h>
 #include <Common/MyTime.h>
+#include <DataTypes/DataTypeEnum.h>
 #include <Debug/MockExecutor/AstToPB.h>
 #include <Debug/MockStorage.h>
 #include <Flash/Statistics/ExecutorStatisticsCollector.h>
 #include <Interpreters/Context.h>
 #include <Operators/OperatorProfileInfo.h>
+#include <Storages/DeltaMerge/ScanContext.h>
 #include <Storages/KVStore/Types.h>
 #include <Storages/MutableSupport.h>
 #include <TestUtils/ExecutorTestUtils.h>
@@ -334,6 +336,21 @@ void assertColumnsEqual(const ColumnsWithTypeAndName & expected, const ColumnsWi
         << getColumnsContent(actual);
 }
 
+DM::MultiStageLateMaterializationRuntimeStatsPtr getMultiStageRuntimeStats(DAGContext & dag_context)
+{
+    DM::MultiStageLateMaterializationRuntimeStatsPtr stats;
+    for (const auto & scan_context_entry : dag_context.scan_context_map)
+    {
+        const auto & scan_context = scan_context_entry.second;
+        if (scan_context == nullptr || scan_context->multi_stage_late_materialization_runtime_stats == nullptr)
+            continue;
+        RUNTIME_CHECK_MSG(stats == nullptr, "Found multiple table scan contexts with multi-stage runtime stats");
+        stats = scan_context->multi_stage_late_materialization_runtime_stats;
+    }
+    RUNTIME_CHECK_MSG(stats != nullptr, "Cannot find table scan context with multi-stage runtime stats");
+    return stats;
+}
+
 void assertTopNEnhancedMultiStageRowsOverride(
     DAGContext & dag_context,
     UInt64 expected_table_scan_rows,
@@ -346,6 +363,12 @@ void assertTopNEnhancedMultiStageRowsOverride(
     ASSERT_EQ(table_scan_rows->load(), expected_table_scan_rows);
     ASSERT_GE(selection_rows->load(), min_selection_rows);
     ASSERT_LT(selection_rows->load(), expected_table_scan_rows);
+
+    auto stats = getMultiStageRuntimeStats(dag_context);
+    ASSERT_NE(stats, nullptr);
+    ASSERT_TRUE(stats->topn_enabled.load());
+    ASSERT_GT(stats->running_topn_input_rows.load(), 0);
+    ASSERT_EQ(selection_rows->load(), stats->final_rest_input_rows.load());
 }
 
 void assertTopNEnhancedMultiStageMatchesDisabled(
@@ -523,7 +546,7 @@ void addMockMultiStageLateMaterializationTable(
 void addMockMultiStageLateMaterializationTopNOrderByTable(
     MockDAGRequestContext & context,
     const String & table_name,
-    TiDB::TP order_by_type,
+    const MockColumnInfo & order_by_info,
     ColumnWithTypeAndName order_by_column)
 {
     constexpr size_t rows = 32;
@@ -531,7 +554,7 @@ void addMockMultiStageLateMaterializationTopNOrderByTable(
     ColumnsWithTypeAndName data;
 
     columns.push_back({"c0", TiDB::TP::TypeLongLong});
-    columns.push_back({"ord", order_by_type});
+    columns.push_back(order_by_info);
     columns.push_back({"c1", TiDB::TP::TypeLongLong});
 
     std::vector<Int64> c0_values;
@@ -563,6 +586,31 @@ void addMockMultiStageLateMaterializationTopNOrderByTable(
     context.addMockDeltaMerge({"test_db", table_name}, columns, data, /*concurrency_hint=*/4);
 }
 
+void addMockMultiStageLateMaterializationTopNOrderByTable(
+    MockDAGRequestContext & context,
+    const String & table_name,
+    TiDB::TP order_by_type,
+    ColumnWithTypeAndName order_by_column)
+{
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        table_name,
+        MockColumnInfo{"ord", order_by_type},
+        std::move(order_by_column));
+}
+
+ColumnWithTypeAndName toEnum16Vec(
+    const String & name,
+    const std::vector<Int16> & values,
+    const DataTypeEnum16::Values & enum_values)
+{
+    auto type = std::make_shared<DataTypeEnum16>(enum_values);
+    auto column = type->createColumn();
+    for (const auto value : values)
+        column->insert(Field(static_cast<Int64>(value)));
+    return {std::move(column), type, name};
+}
+
 void rewriteComparisonRightLiteral(
     tipb::Expr * condition,
     tipb::ScalarFuncSig sig,
@@ -612,6 +660,23 @@ void assertMultiStageRowsOverride(DAGContext & dag_context, UInt64 expected_sele
     ASSERT_NE(selection_rows, nullptr);
     ASSERT_EQ(selection_rows->load(), expected_selection_rows);
     ASSERT_GE(table_scan_rows->load(), selection_rows->load());
+
+    auto stats = getMultiStageRuntimeStats(dag_context);
+    ASSERT_NE(stats, nullptr);
+    ASSERT_EQ(selection_rows->load(), stats->final_rest_input_rows.load());
+}
+
+void assertTopNEnhancedMultiStageFallbackStats(DAGContext & dag_context, UInt64 expected_selection_rows)
+{
+    auto stats = getMultiStageRuntimeStats(dag_context);
+    ASSERT_NE(stats, nullptr);
+    ASSERT_FALSE(stats->topn_enabled.load());
+    ASSERT_EQ(stats->running_topn_input_rows.load(), 0);
+    ASSERT_EQ(stats->running_topn_selected_rows.load(), 0);
+    ASSERT_EQ(stats->running_topn_bypass_rows.load(), 0);
+    ASSERT_EQ(stats->running_topn_filtered_rows.load(), 0);
+    ASSERT_EQ(stats->stage1_output_rows.load(), expected_selection_rows);
+    ASSERT_EQ(stats->final_rest_input_rows.load(), expected_selection_rows);
 }
 
 void assertTopNEnhancedMultiStageFallsBackToSelection(
@@ -633,6 +698,7 @@ void assertTopNEnhancedMultiStageFallsBackToSelection(
 
     assertColumnsEqual(expected, actual);
     assertMultiStageRowsOverride(enabled_dag_context, expected_selection_rows);
+    assertTopNEnhancedMultiStageFallbackStats(enabled_dag_context, expected_selection_rows);
 }
 } // namespace
 
@@ -1287,6 +1353,96 @@ try
         << getColumnsContent(actual);
 
     assertMultiStageRowsOverride(enabled_dag_context, 16);
+    assertTopNEnhancedMultiStageFallbackStats(enabled_dag_context, 16);
+}
+CATCH
+
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForBitSetEnumOrderBy)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    constexpr size_t rows = 32;
+    std::vector<UInt64> uint64_order_by_values;
+    uint64_order_by_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+        uint64_order_by_values.push_back(static_cast<UInt64>(rows - row));
+
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_bit_topn",
+        TiDB::TP::TypeBit,
+        toVec<UInt64>("ord", uint64_order_by_values));
+    auto bit_request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_bit_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/false,
+        /*limit=*/2);
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        bit_request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
+
+    MockColumnInfo set_info{"ord", TiDB::TP::TypeSet};
+    for (Int16 value = 1; value <= 6; ++value)
+        set_info.elems.emplace_back(fmt::format("set_{}", value), value);
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_set_topn",
+        set_info,
+        toVec<UInt64>("ord", uint64_order_by_values));
+    auto set_request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_set_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/false,
+        /*limit=*/2);
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        set_request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
+
+    MockColumnInfo enum_info{"ord", TiDB::TP::TypeEnum};
+    DataTypeEnum16::Values enum_values;
+    enum_values.reserve(rows);
+    std::vector<Int16> enum_order_by_values;
+    enum_order_by_values.reserve(rows);
+    for (size_t row = 0; row < rows; ++row)
+    {
+        const auto value = static_cast<Int16>(row + 1);
+        enum_info.elems.emplace_back(fmt::format("enum_{}", value), value);
+        enum_values.emplace_back(fmt::format("enum_{}", value), value);
+        enum_order_by_values.push_back(static_cast<Int16>(rows - row));
+    }
+    addMockMultiStageLateMaterializationTopNOrderByTable(
+        context,
+        "multi_stage_lm_enum_topn",
+        enum_info,
+        toEnum16Vec("ord", enum_order_by_values, enum_values));
+    auto enum_request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_enum_topn",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "ord",
+        /*is_desc=*/false,
+        /*limit=*/2);
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        enum_request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
 }
 CATCH
 
@@ -1353,6 +1509,35 @@ try
 }
 CATCH
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNFallsBackWhenRestColumnsBelowThreshold)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    addMockMultiStageLateMaterializationTable(
+        context,
+        "multi_stage_lm_topn_rest_below_threshold",
+        /*column_count=*/11);
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm_topn_rest_below_threshold",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "c2",
+        /*is_desc=*/false,
+        /*limit=*/2);
+
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
+}
+CATCH
+
 TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationDisabledWhenRestColumnsRatioBelowThreshold)
 try
 {
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index 48a37664cf1..0aa1dee7569 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -68,6 +68,8 @@ MultiStageLateMaterializationBlockInputStream::MultiStageLateMaterializationBloc
     RUNTIME_CHECK(residual_filter != nullptr);
     RUNTIME_CHECK(residual_filter->before_where != nullptr);
     RUNTIME_CHECK(residual_filter->filter_columns != nullptr);
+    if (running_topn != nullptr && runtime_stats)
+        runtime_stats->recordRunningTopNEnabled();
 }
 
 Block MultiStageLateMaterializationBlockInputStream::buildResidualFilterHeader(
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index fc65f103de0..b72725717d7 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -27,9 +27,8 @@ struct MultiStageLateMaterializationRuntimeStats
 {
     MultiStageLateMaterializationRuntimeStats() = default;
 
-    MultiStageLateMaterializationRuntimeStats(const String & log_id_, bool topn_enabled_)
-        : topn_enabled(topn_enabled_)
-        , log(Logger::get("MultiStageLateMaterialization", log_id_))
+    explicit MultiStageLateMaterializationRuntimeStats(const String & log_id_)
+        : log(Logger::get("MultiStageLateMaterialization", log_id_))
     {}
 
     ~MultiStageLateMaterializationRuntimeStats()
@@ -92,9 +91,11 @@ struct MultiStageLateMaterializationRuntimeStats
         topn_candidate_rows.fetch_add(rows, std::memory_order_relaxed);
     }
 
+    void recordRunningTopNEnabled() { topn_enabled.store(true, std::memory_order_relaxed); }
+
     void recordRunningTopN(UInt64 input_rows, UInt64 selected_rows)
     {
-        topn_enabled.store(true, std::memory_order_relaxed);
+        recordRunningTopNEnabled();
         running_topn_input_rows.fetch_add(input_rows, std::memory_order_relaxed);
         running_topn_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed);
         running_topn_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed);
@@ -103,7 +104,7 @@ struct MultiStageLateMaterializationRuntimeStats
 
     void recordRunningTopNBypass(UInt64 rows)
     {
-        topn_enabled.store(true, std::memory_order_relaxed);
+        recordRunningTopNEnabled();
         running_topn_input_rows.fetch_add(rows, std::memory_order_relaxed);
         running_topn_bypass_rows.fetch_add(rows, std::memory_order_relaxed);
         recordFinalRestInputRows(rows);
diff --git a/dbms/src/Storages/DeltaMerge/ScanContext.cpp b/dbms/src/Storages/DeltaMerge/ScanContext.cpp
index c3f7700462f..97bd28eb752 100644
--- a/dbms/src/Storages/DeltaMerge/ScanContext.cpp
+++ b/dbms/src/Storages/DeltaMerge/ScanContext.cpp
@@ -27,7 +27,7 @@ namespace DB::DM
 {
 namespace
 {
-static constexpr uint64_t NS_TO_MS = 1'000'000;
+constexpr uint64_t NS_TO_MS = 1'000'000;
 
 bool isQueryReadTag(ReadTag read_tag)
 {
diff --git a/dbms/src/Storages/StorageDeltaMerge.cpp b/dbms/src/Storages/StorageDeltaMerge.cpp
index 63e2835da17..b2b6b408fb6 100644
--- a/dbms/src/Storages/StorageDeltaMerge.cpp
+++ b/dbms/src/Storages/StorageDeltaMerge.cpp
@@ -823,6 +823,28 @@ bool isSupportedMultiStageLateMaterializationTopNOrderByType(const DataTypePtr &
         || checkDataType<DataTypeMyDateTime>(type_not_null.get());
 }
 
+bool isSupportedMultiStageLateMaterializationTopNSourceType(TiDB::TP type)
+{
+    switch (type)
+    {
+    case TiDB::TypeTiny:
+    case TiDB::TypeShort:
+    case TiDB::TypeLong:
+    case TiDB::TypeLongLong:
+    case TiDB::TypeInt24:
+    case TiDB::TypeFloat:
+    case TiDB::TypeDouble:
+    case TiDB::TypeDecimal:
+    case TiDB::TypeNewDecimal:
+    case TiDB::TypeDate:
+    case TiDB::TypeNewDate:
+    case TiDB::TypeDatetime:
+        return true;
+    default:
+        return false;
+    }
+}
+
 const TiDB::ColumnInfo * findSourceColumnByID(const TiDB::ColumnInfos & source_columns, ColumnID column_id)
 {
     const auto it = std::find_if(source_columns.begin(), source_columns.end(), [&](const auto & column) {
@@ -884,9 +906,9 @@ DM::MultiStageLateMaterializationTopNDescriptionPtr buildStorageMultiStageLateMa
                 fmt::format("order by column is not found in table scan source columns, column_id={}", column_id));
         if (source_column->hasGeneratedColumnFlag())
             return disable(fmt::format("order by generated column is unsupported, column_id={}", column_id));
-        if (source_column->tp == TiDB::TypeTimestamp || source_column->tp == TiDB::TypeTime)
+        if (!isSupportedMultiStageLateMaterializationTopNSourceType(source_column->tp))
             return disable(fmt::format(
-                "order by column type needs or may need table scan extra cast, column_id={}, tp={}",
+                "unsupported order by source type, column_id={}, tp={}",
                 column_id,
                 static_cast<int>(source_column->tp)));
 

From df223b52fcf75b96daac3a7bf5a23893e0ace1ef Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Thu, 30 Jul 2026 15:54:10 +0800
Subject: [PATCH 16/18] Address MSLM TopN review comments

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp |  2 +-
 .../Flash/tests/gtest_executors_with_dm.cpp   | 14 +---
 dbms/src/Interpreters/Settings.h              |  2 +-
 dbms/src/Interpreters/SettingsCommon.h        | 76 -------------------
 .../DeltaMerge/File/DMFilePackFilter.cpp      |  7 +-
 .../Storages/DeltaMerge/File/DMFileReader.cpp |  2 +-
 dbms/src/Storages/DeltaMerge/ReadMode.h       |  7 ++
 .../DeltaMerge/tests/gtest_dm_file.cpp        | 49 ++++++++++++
 .../2026-07-26-running-local-topn-for-mslm.md | 12 +--
 9 files changed, 74 insertions(+), 97 deletions(-)

diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
index aadf3993fc0..cc53a12f6d0 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
@@ -141,7 +141,7 @@ PhysicalPlanNodePtr PhysicalTopN::build(
 
     auto order_columns = analyzer.buildOrderColumns(before_sort_actions, top_n.order_by());
     SortDescription order_descr = getSortDescription(order_columns, top_n.order_by());
-    if (context.getSettingsRef().dt_enable_multi_stage_late_materialization == 2)
+    if (context.getSettingsRef().dt_enable_multi_stage_late_materialization >= 2)
         tryAttachMultiStageLateMaterializationTopN(top_n, child, log);
     else
         LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason=setting disabled");
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 9e7b11ccf05..7e5dbd6c45c 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -151,7 +151,7 @@ constexpr UInt64 MSLMModeDisabled = 0;
 constexpr UInt64 MSLMModeSelection = 1;
 constexpr UInt64 MSLMModeTopN = 2;
 
-TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationSettingClampsInvalidMode)
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationSettingKeepsUInt64Mode)
 try
 {
     Settings settings;
@@ -167,19 +167,13 @@ try
     ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
 
     settings.dt_enable_multi_stage_late_materialization = 999;
-    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 999);
 
     settings.set("dt_enable_multi_stage_late_materialization", Field(static_cast<UInt64>(3)));
-    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 3);
 
     settings.set("dt_enable_multi_stage_late_materialization", String("999"));
-    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN);
-
-    settings.set("dt_enable_multi_stage_late_materialization", Field(static_cast<Int64>(-1)));
-    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled);
-
-    settings.set("dt_enable_multi_stage_late_materialization", String("-1"));
-    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled);
+    ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 999);
 }
 CATCH
 
diff --git a/dbms/src/Interpreters/Settings.h b/dbms/src/Interpreters/Settings.h
index 0b78ab82703..583ef02d162 100644
--- a/dbms/src/Interpreters/Settings.h
+++ b/dbms/src/Interpreters/Settings.h
@@ -226,7 +226,7 @@ struct Settings
     M(SettingUInt64, dt_max_sharing_column_bytes_for_all, 2048 * Constant::MB, "Memory limitation for data sharing of all requests, include those sharing blocks in block queue. 0 means disable data sharing")                         \
     M(SettingUInt64, dt_max_sharing_column_count, 5, "Deprecated")                                                                                                                                                                      \
     M(SettingBool, dt_enable_bitmap_filter, true, "Use bitmap filter to read data or not")                                                                                                                                              \
-    M(SettingMultiStageLateMaterializationMode, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, 2: TopN enhanced. Values outside [0, 2] are clamped to the nearest boundary.") \
+    M(SettingUInt64, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, >=2: TopN enhanced.") \
     M(SettingDouble, dt_read_thread_count_scale, 2.0, "Number of read thread = number of logical cpu cores * dt_read_thread_count_scale.  Only has meaning at server startup.")                                                         \
     M(SettingDouble, io_thread_count_scale, 5.0, "Number of thread of IOThreadPool = number of logical cpu cores * io_thread_count_scale.  Only has meaning at server startup.")                                                        \
     M(SettingUInt64, init_thread_count_scale, 100, "Number of thread = number of logical cpu cores * init_thread_count_scale. It just works for thread pool for initStores and loadMetadata. Only has meaning at server startup.")      \
diff --git a/dbms/src/Interpreters/SettingsCommon.h b/dbms/src/Interpreters/SettingsCommon.h
index e83b9518466..87ebe447b91 100644
--- a/dbms/src/Interpreters/SettingsCommon.h
+++ b/dbms/src/Interpreters/SettingsCommon.h
@@ -102,82 +102,6 @@ using SettingUInt64 = SettingInt<UInt64>;
 using SettingInt64 = SettingInt<Int64>;
 using SettingBool = SettingInt<bool>;
 
-struct SettingMultiStageLateMaterializationMode
-{
-public:
-    bool changed = false;
-
-    SettingMultiStageLateMaterializationMode(UInt64 x = 0) // NOLINT(google-explicit-constructor)
-        : value(clampMode(x))
-    {}
-
-    SettingMultiStageLateMaterializationMode(const SettingMultiStageLateMaterializationMode & setting)
-        : value(setting.value.load())
-    {}
-
-    operator UInt64() const { return value.load(); } // NOLINT(google-explicit-constructor)
-
-    SettingMultiStageLateMaterializationMode & operator=(UInt64 x)
-    {
-        set(x);
-        return *this;
-    }
-
-    SettingMultiStageLateMaterializationMode & operator=(const SettingMultiStageLateMaterializationMode & setting)
-    {
-        set(setting.value.load());
-        return *this;
-    }
-
-    String toString() const { return DB::toString(value.load()); }
-
-    void set(UInt64 x)
-    {
-        value.store(clampMode(x));
-        changed = true;
-    }
-
-    void set(const Field & x)
-    {
-        if (x.getType() == Field::Types::Int64)
-        {
-            const auto mode = safeGet<Int64>(x);
-            set(mode < 0 ? 0 : static_cast<UInt64>(mode));
-            return;
-        }
-        if (x.getType() == Field::Types::Float64)
-        {
-            const auto mode = safeGet<Float64>(x);
-            set(mode < 0 ? 0 : static_cast<UInt64>(mode));
-            return;
-        }
-        set(applyVisitor(FieldVisitorConvertToNumber<UInt64>(), x));
-    }
-
-    void set(const String & x) { set(!x.empty() && x[0] == '-' ? 0 : parse<UInt64>(x)); }
-
-    void set(ReadBuffer & buf)
-    {
-        UInt64 x = 0;
-        readVarT(x, buf);
-        set(x);
-    }
-
-    UInt64 get() const { return value.load(); }
-
-    void write(WriteBuffer & buf) const { writeVarT(value.load(), buf); }
-
-private:
-    static UInt64 clampMode(UInt64 x) { return x > 2 ? 2 : x; }
-
-    std::atomic<UInt64> value;
-};
-
-ALWAYS_INLINE inline auto format_as(SettingMultiStageLateMaterializationMode s)
-{
-    return s.get();
-}
-
 
 /** Unlike SettingUInt64, supports the value of 'auto' - the number of processor cores without taking into account SMT.
   * A value of 0 is also treated as auto.
diff --git a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
index 8d2fbc9857c..45e41cebb1b 100644
--- a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
+++ b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp
@@ -87,7 +87,7 @@ DMFilePackFilterResultPtr DMFilePackFilter::load(ReadTag read_tag)
     SCOPE_EXIT({ scan_context->total_rs_pack_filter_check_time_ns += watch.elapsed(); });
     const size_t pack_count = dmfile->getPacks();
     DMFilePackFilterResult result(index_cache, read_limiter, pack_count);
-    result.param.record_trim_metrics = read_tag == ReadTag::Query;
+    result.param.record_trim_metrics = isQueryRoughSetStatsOwner(read_tag);
     auto read_all_packs = (rowkey_ranges.size() == 1 && rowkey_ranges[0].all()) || rowkey_ranges.empty();
     if (!read_all_packs)
     {
@@ -189,9 +189,10 @@ DMFilePackFilterResultPtr DMFilePackFilter::load(ReadTag read_tag)
     // 4. When building MSLM candidate stream (ReadTag::MSLMCandidate).
     // 5. When building MSLM final rest stream (ReadTag::MSLMFinalRest).
     // 6. When building stream of other columns (ReadTag::Query).
-    // Only need to count the filter result once.
+    // Query owns these statistics for normal/LM reads. MSLMPushedFilter owns them for MSLM because it always
+    // traverses the logical MSLM input.
     // TODO: We can create DMFilePackFilter at the beginning and pass it to the stages described above.
-    if (read_tag == ReadTag::Query)
+    if (isQueryRoughSetStatsOwner(read_tag))
     {
         scan_context->rs_pack_filter_none += none_count;
         scan_context->rs_pack_filter_some += some_count;
diff --git a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
index e9307564663..f25515164ad 100644
--- a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
+++ b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp
@@ -163,7 +163,7 @@ std::pair<size_t, RSResult> DMFileReader::getReadRows()
     }
 
     next_row_offset += read_rows;
-    if (read_tag == ReadTag::Query && last_pack_res.allMatch())
+    if (isQueryRoughSetStatsOwner(read_tag) && last_pack_res.allMatch())
         scan_context->rs_dmfile_read_with_all += next_pack_id - start_pack_id;
     return {read_rows, last_pack_res};
 }
diff --git a/dbms/src/Storages/DeltaMerge/ReadMode.h b/dbms/src/Storages/DeltaMerge/ReadMode.h
index 791f3c45767..aa3ce574012 100644
--- a/dbms/src/Storages/DeltaMerge/ReadMode.h
+++ b/dbms/src/Storages/DeltaMerge/ReadMode.h
@@ -50,6 +50,13 @@ enum class ReadTag
     MSLMFinalRest, // Read final rest columns for multi-stage late-materialization.
 };
 
+/// A DMFile may be opened once for every physical MSLM read stage, but query-level
+/// rough-set statistics must be recorded by exactly one stage.
+inline constexpr bool isQueryRoughSetStatsOwner(ReadTag read_tag)
+{
+    return read_tag == ReadTag::Query || read_tag == ReadTag::MSLMPushedFilter;
+}
+
 enum class ReadRUType
 {
     MVCC_ESTIMATE,
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp
index 3c4324ab3c4..d69fd348a2f 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp
@@ -1705,6 +1705,55 @@ RSOperatorPtr toRSFilter(const ColumnDefine & cd, const HandleRange & range)
 }
 } // namespace
 
+TEST_P(DMFileTest, MSLMRoughSetStatsRecordedOnce)
+try
+{
+    auto cols = DMTestEnv::getDefaultColumns();
+    ColumnDefine i64_cd(2, "i64", typeFromString("Int64"));
+    cols->push_back(i64_cd);
+    reload(cols);
+
+    constexpr Int64 num_rows = 128;
+    {
+        auto block = DMTestEnv::prepareSimpleWriteBlock(0, num_rows, false);
+        block.insert(DB::tests::createColumn<Int64>(createNumbers<Int64>(0, num_rows), i64_cd.name, i64_cd.id));
+        auto stream = std::make_shared<DMFileBlockOutputStream>(dbContext(), dm_file, *cols);
+        stream->writePrefix();
+        stream->write(block, {});
+        stream->writeSuffix();
+    }
+
+    const auto filter = toRSFilter(i64_cd, HandleRange{0, num_rows});
+    const auto read_all = [&](ReadTag read_tag, const ScanContextPtr & scan_context) {
+        DMFileBlockInputStreamBuilder builder(dbContext());
+        auto stream = builder.setColumnCache(column_cache)
+                          .setRSOperator(filter)
+                          .setReadTag(read_tag)
+                          .build(dm_file, *cols, RowKeyRanges{RowKeyRange::newAll(false, 1)}, scan_context);
+        while (stream->read()) {}
+    };
+    const auto pack_filter_count = [](const ScanContextPtr & scan_context) {
+        return scan_context->rs_pack_filter_none.load() + scan_context->rs_pack_filter_some.load()
+            + scan_context->rs_pack_filter_all.load() + scan_context->rs_pack_filter_all_null.load();
+    };
+
+    auto query_context = std::make_shared<ScanContext>();
+    read_all(ReadTag::Query, query_context);
+    EXPECT_EQ(pack_filter_count(query_context), 1);
+    EXPECT_EQ(query_context->rs_dmfile_read_with_all, 1);
+
+    auto mslm_context = std::make_shared<ScanContext>();
+    read_all(ReadTag::MSLMPushedFilter, mslm_context);
+    EXPECT_EQ(pack_filter_count(mslm_context), 1);
+    EXPECT_EQ(mslm_context->rs_dmfile_read_with_all, 1);
+
+    read_all(ReadTag::MSLMCandidate, mslm_context);
+    read_all(ReadTag::MSLMFinalRest, mslm_context);
+    EXPECT_EQ(pack_filter_count(mslm_context), 1);
+    EXPECT_EQ(mslm_context->rs_dmfile_read_with_all, 1);
+}
+CATCH
+
 TEST_P(DMFileTest, ReadFilteredByRoughSetFilter)
 try
 {
diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md
index b73a78228fa..f53fb413bb1 100644
--- a/docs/design/2026-07-26-running-local-topn-for-mslm.md
+++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md
@@ -163,19 +163,21 @@ warmup 后，持续统计后续 `adaptive_post_warmup_input_rows` 和 `adaptive_
 
 ### Local TopK 覆盖 Global TopK
 
-对于任意一个 stream，如果某一行 `r` 是全局 TopK 结果中的一行，那么 `r` 必然属于它所在 stream 的 local TopK。
+如果没有 ties，对于任意一个 stream，如果某一行 `r` 是全局 TopK 结果中的一行，那么 `r` 必然属于它所在 stream 的 local TopK。
 
 证明：
 
 如果 `r` 不属于所在 stream 的 local TopK，说明这个 stream 中至少有 K 行按照相同 comparator 排在 `r` 前面。这 K 行在全局也排在 `r` 前面，因此 `r` 不可能是全局 TopK，矛盾。
 
-所以：
+所以在没有 ties 时：
 
 ```text
 global TopK subset union(each stream local TopK)
 ```
 
-第一版保留上层全局 TopN executor，因此只要 storage 内输出包含全局 TopK 的 superset，最终结果就是正确的。
+有 ties 时，local TopK 不一定保留和全局 TopK 完全相同的一组物理 rows，因为相同 sort key 的 rows 可以在不同执行路径中选择不同 tie winners。此时需要的正确性条件不是固定物理 rows 的严格 subset，而是每个 stream 输出足够多的 rows，且这些 rows 的 sort key 不劣于该 stream 内第 K 个 sort key。这样所有 stream 输出的 union 仍然包含足够多的 rows，其 sort key 不劣于全局第 K 个 sort key，上层 TopN 可以在 SQL ties 语义下得到合法结果。
+
+第一版保留上层全局 TopN executor，因此只要 storage 内输出包含上层 TopN 生成合法结果所需的 candidate rows，最终结果就是正确的。
 
 ### Running Local TopN 的 Streaming Superset
 
@@ -195,7 +197,7 @@ final local TopK 的 superset
 
 当 local TopN heap 已满时，如果当前 row 按 comparator 不优于 heap worst row，则该 row 可以被跳过。原因是后续扫描只会增加竞争者，不会让这个 row 的排序位置变得更靠前。
 
-对于和 heap worst row 相等的 ties，第一版采用保守策略：相等 rows 作为 candidate 输出，但不一定放入 bounded heap。这会降低 ties 较多场景下的剪枝效果，但可以避免因为 ties 选择不同物理 rows 引起不必要的行为差异。
+对于和 heap worst row 相等的 ties，第一版不会继续输出这些相等 rows 作为 candidate，而是让 heap 中已有 entries 成为本 stream 的 local tie winners。这样可以避免 ties 大量扩展 candidate rows 和 final-rest IO。代价是 ties 下可能选择不同的物理 rows，但 SQL `ORDER BY` 未提供完整 tie-break key 时本来就不保证相同 sort key 内的物理 row 选择，这也属于本设计的非目标。
 
 ## 执行流程
 
@@ -403,7 +405,7 @@ final rest columns 按现有 MSLM 逻辑读取。
 上层 TopN 保留，保证最终结果正确。
 ```
 
-Ties 也需要避免破坏内存上界。第一版 heap 仍最多保存 `topk` 个 entries。与 heap worst 相等的 rows 可以作为当前 block candidate 输出，但不插入 heap。这样 ties 较多时剪枝效果会下降，但 heap 内存保持 bounded。
+Ties 也需要避免破坏内存和 IO 上界。第一版 heap 仍最多保存 `topk` 个 entries；与 heap worst 相等的 rows 不插入 heap，也不作为当前 block candidate 输出。这样 ties 较多时 candidate rows 不会被 ties 放大，heap 内存保持 bounded，但不同执行路径可能选择不同的物理 tie winners。
 
 ## Running Local TopN 状态
 

From 4bec38ebcf54a1fa84a2ee0cbbeafae30ad90455 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Thu, 30 Jul 2026 17:49:42 +0800
Subject: [PATCH 17/18] Optimize running TopN heap update

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 .../MultiStageLateMaterializationTopN.cpp     | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
index 2e579aa670d..9b90cebf250 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp
@@ -358,7 +358,6 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
 
     const auto rows = stage1_block.rows();
     RunningLocalTopNUpdateResult result;
-    result.filter.resize_fill(rows, 0);
 
     if (residual_passed_rows == 0)
         return result;
@@ -377,15 +376,26 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         sort_columns.push_back(desc.build_column_view(*column));
     }
 
+    auto ensure_filter = [&]() -> IColumn::Filter & {
+        if (result.filter.empty())
+            result.filter.resize_fill(rows, 0);
+        return result.filter;
+    };
     auto mark_candidate = [&](size_t row) {
-        if (result.filter[row] == 0)
+        auto & filter = ensure_filter();
+        if (filter[row] == 0)
         {
-            result.filter[row] = 1;
+            filter[row] = 1;
             ++result.passed_count;
         }
     };
     auto unmark_current_block_candidate = [&](const HeapEntry & entry) {
-        if (entry.block_sequence == current_block_sequence && result.filter[entry.row_index_in_stage1_block] != 0)
+        if (entry.block_sequence != current_block_sequence)
+            return;
+
+        RUNTIME_CHECK(!result.filter.empty());
+        RUNTIME_CHECK(entry.row_index_in_stage1_block < result.filter.size());
+        if (result.filter[entry.row_index_in_stage1_block] != 0)
         {
             result.filter[entry.row_index_in_stage1_block] = 0;
             --result.passed_count;
@@ -413,9 +423,8 @@ RunningLocalTopNUpdateResult RunningLocalTopN::update(
         // where ties are broken locally by the entries already selected by the partial TopN.
         if (cmp < 0)
         {
-            auto evicted = heap.top();
+            unmark_current_block_candidate(heap.top());
             heap.pop();
-            unmark_current_block_candidate(evicted);
 
             mark_candidate(row);
             heap.push(HeapEntry{

From 890944abe4a2c6f6c6c9f6caf2517f008d2c3410 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Thu, 30 Jul 2026 23:12:57 +0800
Subject: [PATCH 18/18] Refine MSLM TopN stats flush and partition guard

---
 dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp |  3 +
 .../Flash/tests/gtest_executors_with_dm.cpp   | 28 ++++++
 ...ageLateMaterializationBlockInputStream.cpp | 70 +++++++++++---
 ...StageLateMaterializationBlockInputStream.h | 14 ++-
 ...ultiStageLateMaterializationRuntimeStats.h | 93 +++++++++++++++++++
 .../gtest_skippable_block_input_stream.cpp    | 49 ++++++++++
 6 files changed, 242 insertions(+), 15 deletions(-)

diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
index cc53a12f6d0..3249b3d5e6d 100644
--- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
+++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp
@@ -55,6 +55,9 @@ DM::MultiStageLateMaterializationTopNDescriptionPtr tryBuildMultiStageLateMateri
         return disable(fmt::format("topk is too large: {}", top_n.limit()));
     if (top_n.order_by_size() > static_cast<int>(DM::multi_stage_late_materialization_topn_max_order_by_columns))
         return disable(fmt::format("too many order by columns: {}", top_n.order_by_size()));
+    if (top_n.partition_by_size() != 0)
+        return disable(
+            fmt::format("partition TopN is unsupported, partition_by_columns={}", top_n.partition_by_size()));
 
     auto desc = std::make_shared<DM::MultiStageLateMaterializationTopNDescription>();
     desc->topk = top_n.limit();
diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
index 7e5dbd6c45c..c5ce0a57e04 100644
--- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
+++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp
@@ -1206,6 +1206,34 @@ try
 }
 CATCH
 
+TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForPartitionBy)
+try
+{
+    enablePipeline(true);
+    context.context->setSetting("max_block_size", Field(static_cast<UInt64>(64)));
+
+    auto request = buildDAGRequestWithPushedDownFilterAndTopN(
+        context,
+        "multi_stage_lm",
+        lt(col("c0"), lit(Field(static_cast<Int64>(16)))),
+        gt(col("c1"), lit(Field(static_cast<Int64>(-1)))),
+        "c1",
+        /*is_desc=*/false,
+        /*limit=*/2);
+
+    auto * topn = request->mutable_root_executor()->mutable_topn();
+    ASSERT_EQ(topn->order_by_size(), 1);
+    *topn->add_partition_by() = topn->order_by(0);
+
+    assertTopNEnhancedMultiStageFallsBackToSelection(
+        *this,
+        context,
+        request,
+        dag_context_ptr->log->identifier(),
+        /*expected_selection_rows=*/16);
+}
+CATCH
+
 TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNVersionColumn)
 try
 {
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
index 0aa1dee7569..16874b53cf1 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp
@@ -23,6 +23,9 @@ namespace DB::DM
 {
 namespace
 {
+constexpr UInt64 runtime_stats_flush_logical_rows = static_cast<UInt64>(DEFAULT_MERGE_BLOCK_SIZE) * 64;
+constexpr UInt64 runtime_stats_flush_blocks = 64;
+
 void filterBlock(Block & block, const IColumn::Filter & filter, size_t passed_count)
 {
     if (!block)
@@ -72,6 +75,11 @@ MultiStageLateMaterializationBlockInputStream::MultiStageLateMaterializationBloc
         runtime_stats->recordRunningTopNEnabled();
 }
 
+MultiStageLateMaterializationBlockInputStream::~MultiStageLateMaterializationBlockInputStream()
+{
+    flushRuntimeStatsNoThrow();
+}
+
 Block MultiStageLateMaterializationBlockInputStream::buildResidualFilterHeader(
     const PushDownFilterPtr & residual_filter)
 {
@@ -342,17 +350,18 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         stage0_block = stage0_filter_stream->read(stage0_filter, true);
         if (!stage0_block)
         {
-            logSummary();
+            finishRuntimeStats();
             return {};
         }
 
+        ++runtime_stats_pending_blocks;
         auto effective_stage0_filter = buildStage0EffectiveFilter(stage0_block, stage0_filter);
-        if (runtime_stats)
-            runtime_stats->recordPushedFilter(stage0_block.rows(), effective_stage0_filter.passed_count);
+        local_runtime_stats.recordPushedFilter(stage0_block.rows(), effective_stage0_filter.passed_count);
         if (effective_stage0_filter.passed_count == 0)
         {
             skipNextBlockOrRead(stage1_filter_stream, "stage1_filter");
             skipNextBlockOrRead(final_rest_stream, "final_rest");
+            flushRuntimeStatsIfNeeded();
             continue;
         }
 
@@ -376,41 +385,42 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         Block filter_eval_block;
         FilterPtr residual_filter_ptr = nullptr;
         const auto residual_passed_rows = executeResidualFilter(stage1_block, filter_eval_block, residual_filter_ptr);
-        if (runtime_stats)
-            runtime_stats->recordResidualFilter(stage1_block.rows(), residual_passed_rows);
+        local_runtime_stats.recordResidualFilter(stage1_block.rows(), residual_passed_rows);
 
-        if (!shouldUseRunningTopN() && runtime_stats)
+        if (!shouldUseRunningTopN())
         {
             if (running_topn != nullptr)
-                runtime_stats->recordRunningTopNBypass(residual_passed_rows);
+                local_runtime_stats.recordRunningTopNBypass(residual_passed_rows);
             else
-                runtime_stats->recordFinalRestInputRows(residual_passed_rows);
+                local_runtime_stats.recordFinalRestInputRows(residual_passed_rows);
         }
 
         if (residual_passed_rows == 0)
         {
             ++late_mode_blocks;
             skipNextBlockOrRead(final_rest_stream, "final_rest");
+            flushRuntimeStatsIfNeeded();
             continue;
         }
 
         if (shouldUseRunningTopN())
         {
             auto topn_result = running_topn->update(stage1_block, residual_filter_ptr, residual_passed_rows);
-            if (runtime_stats)
-                runtime_stats->recordRunningTopN(residual_passed_rows, topn_result.passed_count);
+            local_runtime_stats.recordRunningTopN(residual_passed_rows, topn_result.passed_count);
             updateTopNAdaptiveState(residual_passed_rows, topn_result.passed_count);
 
             if (topn_result.passed_count == 0)
             {
                 ++late_mode_blocks;
                 skipNextBlockOrRead(final_rest_stream, "final_rest");
+                flushRuntimeStatsIfNeeded();
                 continue;
             }
 
             if (topn_result.passed_count == stage1_block.rows())
             {
                 ++direct_mode_blocks;
+                flushRuntimeStatsIfNeeded();
                 return buildDirectModeBlock(
                     stage1_block,
                     stage0_filter_ptr,
@@ -420,6 +430,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
             }
 
             ++late_mode_blocks;
+            flushRuntimeStatsIfNeeded();
             return buildLateModeBlock(
                 stage1_block,
                 stage0_filter_ptr,
@@ -431,6 +442,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
         if (residual_passed_rows == stage1_block.rows())
         {
             ++direct_mode_blocks;
+            flushRuntimeStatsIfNeeded();
             return buildDirectModeBlock(
                 stage1_block,
                 stage0_filter_ptr,
@@ -441,6 +453,7 @@ Block MultiStageLateMaterializationBlockInputStream::read()
 
         RUNTIME_CHECK(residual_filter_ptr != nullptr);
         ++late_mode_blocks;
+        flushRuntimeStatsIfNeeded();
         return buildLateModeBlock(
             stage1_block,
             stage0_filter_ptr,
@@ -450,12 +463,43 @@ Block MultiStageLateMaterializationBlockInputStream::read()
     }
 }
 
-void MultiStageLateMaterializationBlockInputStream::logSummary()
+void MultiStageLateMaterializationBlockInputStream::flushRuntimeStats()
+{
+    if (local_runtime_stats.empty())
+    {
+        runtime_stats_pending_blocks = 0;
+        return;
+    }
+    if (runtime_stats)
+        runtime_stats->merge(local_runtime_stats);
+    local_runtime_stats.reset();
+    runtime_stats_pending_blocks = 0;
+}
+
+void MultiStageLateMaterializationBlockInputStream::flushRuntimeStatsNoThrow()
+{
+    try
+    {
+        finishRuntimeStats();
+    }
+    catch (...)
+    {}
+}
+
+void MultiStageLateMaterializationBlockInputStream::flushRuntimeStatsIfNeeded()
+{
+    if (local_runtime_stats.pushed_filter_input_rows >= runtime_stats_flush_logical_rows
+        || runtime_stats_pending_blocks >= runtime_stats_flush_blocks)
+        flushRuntimeStats();
+}
+
+void MultiStageLateMaterializationBlockInputStream::finishRuntimeStats()
 {
-    if (summary_logged)
+    if (runtime_stats_finished)
         return;
-    summary_logged = true;
+    runtime_stats_finished = true;
 
+    flushRuntimeStats();
     if (runtime_stats)
         runtime_stats->finishStream(
             late_mode_blocks,
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
index 00a3abf3c08..68f0a132c0f 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h
@@ -40,6 +40,8 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
         const String & req_id_,
         const MultiStageLateMaterializationRuntimeStatsPtr & runtime_stats_ = nullptr);
 
+    ~MultiStageLateMaterializationBlockInputStream() override;
+
     String getName() const override { return NAME; }
 
     Block getHeader() const override { return header; }
@@ -94,7 +96,13 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
 
     void updateTopNAdaptiveState(UInt64 residual_passed_rows, UInt64 topn_candidate_rows);
 
-    void logSummary();
+    void flushRuntimeStats();
+
+    void flushRuntimeStatsNoThrow();
+
+    void flushRuntimeStatsIfNeeded();
+
+    void finishRuntimeStats();
 
 private:
     Block header;
@@ -104,6 +112,8 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
     PushDownFilterPtr residual_filter;
     BitmapFilterPtr bitmap_filter;
     MultiStageLateMaterializationRuntimeStatsPtr runtime_stats;
+    MultiStageLateMaterializationRuntimeStatsDelta local_runtime_stats;
+    UInt64 runtime_stats_pending_blocks = 0;
     FilterTransformAction residual_filter_action;
     std::unique_ptr<RunningLocalTopN> running_topn;
 
@@ -114,7 +124,7 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream
     UInt64 topn_adaptive_candidate_rows = 0;
     bool topn_adaptive_warmed_up = false;
     bool topn_adaptive_disabled = false;
-    bool summary_logged = false;
+    bool runtime_stats_finished = false;
 
     const LoggerPtr log;
 };
diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
index b72725717d7..4403e03c2f1 100644
--- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
+++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h
@@ -23,6 +23,77 @@
 namespace DB::DM
 {
 
+struct MultiStageLateMaterializationRuntimeStatsDelta
+{
+    void reset() { *this = {}; }
+
+    bool empty() const
+    {
+        return pushed_filter_input_rows == 0 && residual_filter_input_rows == 0 && final_rest_input_rows == 0
+            && running_topn_input_rows == 0;
+    }
+
+    void recordPushedFilter(UInt64 input_rows, UInt64 selected_rows)
+    {
+        pushed_filter_input_rows += input_rows;
+        pushed_filter_selected_rows += selected_rows;
+        pushed_filter_filtered_rows += input_rows - selected_rows;
+
+        // Keep the old counter name for existing actRows/tests.
+        stage0_output_rows += selected_rows;
+    }
+
+    void recordResidualFilter(UInt64 input_rows, UInt64 selected_rows)
+    {
+        residual_filter_input_rows += input_rows;
+        residual_filter_selected_rows += selected_rows;
+        residual_filter_filtered_rows += input_rows - selected_rows;
+
+        // Keep the old counter name for existing actRows/tests.
+        stage1_output_rows += selected_rows;
+    }
+
+    void recordFinalRestInputRows(UInt64 rows)
+    {
+        final_rest_input_rows += rows;
+
+        // Keep the old counter name for existing actRows/tests.
+        topn_candidate_rows += rows;
+    }
+
+    void recordRunningTopN(UInt64 input_rows, UInt64 selected_rows)
+    {
+        running_topn_input_rows += input_rows;
+        running_topn_selected_rows += selected_rows;
+        running_topn_filtered_rows += input_rows - selected_rows;
+        recordFinalRestInputRows(selected_rows);
+    }
+
+    void recordRunningTopNBypass(UInt64 rows)
+    {
+        running_topn_input_rows += rows;
+        running_topn_bypass_rows += rows;
+        recordFinalRestInputRows(rows);
+    }
+
+    UInt64 pushed_filter_input_rows = 0;
+    UInt64 pushed_filter_selected_rows = 0;
+    UInt64 pushed_filter_filtered_rows = 0;
+    UInt64 residual_filter_input_rows = 0;
+    UInt64 residual_filter_selected_rows = 0;
+    UInt64 residual_filter_filtered_rows = 0;
+    UInt64 final_rest_input_rows = 0;
+    UInt64 running_topn_input_rows = 0;
+    UInt64 running_topn_selected_rows = 0;
+    UInt64 running_topn_bypass_rows = 0;
+    UInt64 running_topn_filtered_rows = 0;
+
+    // Legacy names kept because actRows overrides and existing unit tests still refer to them.
+    UInt64 stage0_output_rows = 0;
+    UInt64 stage1_output_rows = 0;
+    UInt64 topn_candidate_rows = 0;
+};
+
 struct MultiStageLateMaterializationRuntimeStats
 {
     MultiStageLateMaterializationRuntimeStats() = default;
@@ -110,6 +181,28 @@ struct MultiStageLateMaterializationRuntimeStats
         recordFinalRestInputRows(rows);
     }
 
+    void merge(const MultiStageLateMaterializationRuntimeStatsDelta & delta)
+    {
+        if (delta.empty())
+            return;
+
+        pushed_filter_input_rows.fetch_add(delta.pushed_filter_input_rows, std::memory_order_relaxed);
+        pushed_filter_selected_rows.fetch_add(delta.pushed_filter_selected_rows, std::memory_order_relaxed);
+        pushed_filter_filtered_rows.fetch_add(delta.pushed_filter_filtered_rows, std::memory_order_relaxed);
+        residual_filter_input_rows.fetch_add(delta.residual_filter_input_rows, std::memory_order_relaxed);
+        residual_filter_selected_rows.fetch_add(delta.residual_filter_selected_rows, std::memory_order_relaxed);
+        residual_filter_filtered_rows.fetch_add(delta.residual_filter_filtered_rows, std::memory_order_relaxed);
+        final_rest_input_rows.fetch_add(delta.final_rest_input_rows, std::memory_order_relaxed);
+        running_topn_input_rows.fetch_add(delta.running_topn_input_rows, std::memory_order_relaxed);
+        running_topn_selected_rows.fetch_add(delta.running_topn_selected_rows, std::memory_order_relaxed);
+        running_topn_bypass_rows.fetch_add(delta.running_topn_bypass_rows, std::memory_order_relaxed);
+        running_topn_filtered_rows.fetch_add(delta.running_topn_filtered_rows, std::memory_order_relaxed);
+
+        stage0_output_rows.fetch_add(delta.stage0_output_rows, std::memory_order_relaxed);
+        stage1_output_rows.fetch_add(delta.stage1_output_rows, std::memory_order_relaxed);
+        topn_candidate_rows.fetch_add(delta.topn_candidate_rows, std::memory_order_relaxed);
+    }
+
     void merge(const MultiStageLateMaterializationRuntimeStats & other)
     {
         pushed_filter_input_rows.fetch_add(other.pushed_filter_input_rows.load(std::memory_order_relaxed));
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
index c05eb82cfbd..54e72b9afe3 100644
--- a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
+++ b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp
@@ -721,6 +721,55 @@ try
 }
 CATCH
 
+TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationRuntimeStatsFlushesByBlockCount)
+try
+{
+    constexpr size_t blocks_before_flush = 64;
+
+    std::vector<IColumn::Filter> stage0_filters;
+    std::vector<std::vector<UInt8>> filter_values;
+    std::vector<std::vector<Int64>> rest_values;
+    for (size_t i = 0; i < blocks_before_flush; ++i)
+    {
+        stage0_filters.emplace_back(makeFilter({1}));
+        filter_values.emplace_back(std::vector<UInt8>{1});
+        rest_values.emplace_back(std::vector<Int64>{static_cast<Int64>(i)});
+    }
+
+    auto stage0_stream = std::make_shared<DeterministicStage0FilterInputStream>(std::move(stage0_filters));
+    auto stage1_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageFilterColumns(),
+        filter_values,
+        rest_values);
+    auto final_rest_stream = std::make_shared<DeterministicSkippableBlockInputStream>(
+        makeMultiStageRestColumns(),
+        filter_values,
+        rest_values);
+
+    auto bitmap_filter = std::make_shared<BitmapFilter>(blocks_before_flush, 1);
+    auto runtime_stats = std::make_shared<MultiStageLateMaterializationRuntimeStats>();
+    auto stream = std::make_shared<MultiStageLateMaterializationBlockInputStream>(
+        makeMultiStageColumnsToRead(),
+        stage0_stream,
+        stage1_stream,
+        final_rest_stream,
+        makeResidualFilterForMultiStageTest(),
+        bitmap_filter,
+        "test",
+        runtime_stats);
+
+    for (size_t i = 0; i < blocks_before_flush - 1; ++i)
+        assertMultiStageRows(stream->read(), {1}, {static_cast<Int64>(i)});
+    ASSERT_EQ(runtime_stats->stage0_output_rows.load(), 0);
+    ASSERT_EQ(runtime_stats->stage1_output_rows.load(), 0);
+
+    assertMultiStageRows(stream->read(), {1}, {static_cast<Int64>(blocks_before_flush - 1)});
+    ASSERT_EQ(runtime_stats->stage0_output_rows.load(), blocks_before_flush);
+    ASSERT_EQ(runtime_stats->stage1_output_rows.load(), blocks_before_flush);
+    ASSERT_EQ(runtime_stats->final_rest_input_rows.load(), blocks_before_flush);
+}
+CATCH
+
 TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationLateMode)
 try
 {