TsDB2/DB/core/query_exec.go

package core

import (
	"container/heap"
	"encoding/binary"
	"slices"

	"git.ycl.pub/yl32940/TsDB2/DB/core/types"
)

// executeScan 执行索引扫描或全表扫描, 将反序列化后的行传递给回调。
// budget > 0 时最多交付 budget 行 (LIMIT+OFFSET 场景): 二级索引回表据此
// 提前收工, 而不是把命中的上千行全部解出再丢掉 —— Where(hash=重).First()
// 这类点查只允许解 1 行。
func (q *Query) executeScan(callback func(pkID uint64, row []types.FieldValue) bool, plan *scanPlan, budget int) error {
	if plan != nil {
		return q.indexScan(plan, callback, budget)
	}
	return q.fullTableScan(callback)
}

// indexScan 使用索引扫描并按"先物化、后回表"两阶段执行。
//
// 为什么不能像旧实现那样在 RangeScan 回调里直接回表 + 触发用户 callback:
// RangeScan 是沿叶子双向链 (PageHdr.NextPage) 推进的裸游标, 而用户的 Scan 回调
// 常常会对【同一棵索引】做 Update/Delete —— 删旧条目 + 插新条目会触发叶分裂/合并,
// 被合并掉的叶子页会被登记回收, 下一个 Commit/drain 时该页被清零并改写成空闲链节点
// (空闲链同样用 PageHdr.NextPage 串接)。此时游标再读该页的 NextPage, 拿到的就是
// 空闲链指针, 于是跳到一张 type=0 的空页上 (实测 panic: bad page type=0)。
//
// 物化把"遍历"与"修改"彻底分开: 遍历阶段树保持只读, 游标始终有效; 修改发生在
// 遍历结束之后, 不再影响游标。二级索引只物化 pkID(8B/行), 内存开销与结果集成正比
// 且远小于整行物化; rowStore 索引物化的是行的快照 (必须拷贝, 因为它是页内缓冲视图)。
func (q *Query) indexScan(plan *scanPlan, callback func(pkID uint64, row []types.FieldValue) bool, budget int) error {
	// 判断 OrderBy 是否需要反向 (只对单区间有定义; 多区间已强制走排序管线)
	reverse := false
	if len(plan.ranges) == 1 && len(q.orderBys) == 1 {
		idxField := q.fieldNameFromMeta(plan.meta)
		if q.orderBys[0].Field == idxField && q.orderBys[0].Descending {
			reverse = true
		}
	}

	if plan.rowStore {
		// 承载行数据的索引: RangeScan 返回 storeRow = [pkID(8B) + rowData]。
		// storeRow 指向页缓冲, 遍历结束后页可能被回收改写, 必须拷贝一份快照。
		// 多区间 (IN 展开) 时不同值可能编码出重叠区间 → 按 pkID 去重。
		var rows [][]byte
		var seen map[uint64]struct{}
		if len(plan.ranges) > 1 {
			seen = make(map[uint64]struct{}, len(plan.ranges)*4)
		}
		for _, kr := range plan.ranges {
			if err := plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
				kr.startKey, kr.endKey, kr.startInclusive, kr.endInclusive, reverse,
				func(row []byte) bool {
					if len(row) < 8 {
						return true
					}
					if seen != nil {
						pkID := binary.BigEndian.Uint64(row[:8])
						if _, dup := seen[pkID]; dup {
							return true
						}
						seen[pkID] = struct{}{}
					}
					snapshot := make([]byte, len(row))
					copy(snapshot, row)
					rows = append(rows, snapshot)
					return true
				}); err != nil {
				return err
			}
		}
		for _, row := range rows {
			rowData := row[8:]
			if plan.remainingCond != nil {
				if !plan.remainingCond.Match(rowData, q.table) {
					continue
				}
			}
			if !callback(binary.BigEndian.Uint64(row[:8]), q.deserialize(rowData)) {
				return nil
			}
		}
		return nil
	}

	// 二级索引: RangeScan 的输出序即索引键序 (B+ 叶序 / Hash 内部已按键排序),
	// 免排序 ORDER BY 路径依赖它 —— 单区间时原样保留; 只有 IN 多区间才排序去重
	// (此时顺序本来就交给排序管线, 升序输入也是批量回表游标的要求)。
	var pkIDs []uint64
	for _, kr := range plan.ranges {
		if err := plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
			kr.startKey, kr.endKey, kr.startInclusive, kr.endInclusive, reverse,
			func(row []byte) bool {
				if len(row) < 8 {
					return true
				}
				pkIDs = append(pkIDs, binary.BigEndian.Uint64(row[:8]))
				return true
			}); err != nil {
			return err
		}
	}
	if len(plan.ranges) > 1 && len(pkIDs) > 1 {
		slices.Sort(pkIDs)
		// 多区间才可能重复 (同 pkID 命中多个重叠区间); 单区间去重徒增一次扫描
		w := 1
		for i := 1; i < len(pkIDs); i++ {
			if pkIDs[i] != pkIDs[w-1] {
				pkIDs[w] = pkIDs[i]
				w++
			}
		}
		pkIDs = pkIDs[:w]
	}
	q.table.indexMeta.BatchResolveByPK(pkIDs, budget, func(pkID uint64, storeRow []byte) bool {
		rowData := storeRow[8:]
		if plan.remainingCond != nil {
			if !plan.remainingCond.Match(rowData, q.table) {
				return true
			}
		}
		return callback(pkID, q.deserialize(rowData))
	})
	return nil
}

// fullTableScan 全表扫描 (流式回调模式)。与 fullTableScanRaw 同构,
// 回调前多一步 deserialize 整行。
func (q *Query) fullTableScan(callback func(pkID uint64, row []types.FieldValue) bool) error {
	return q.fullTableScanRaw(func(pkID uint64, rowData []byte) bool {
		return callback(pkID, q.deserialize(rowData))
	})
}

// ======================== 裸字节扫描 (聚合专用) ========================

// scanRaw 裸字节版流式扫描: 回调拿 (pkID, 原始行字节), 全程不 deserializeRow。
//
// 背景: Count/Sum/Max/Min 每行只需要"过条件 + 最多一列", 而 scan 管线对每行都
// 分配 []FieldValue 并解码全部活跃列 —— 100 万行的 SUM 里这部分占了绝对大头
// (实测带条件全表扫描每行 ~0.66µs, 整行物化路径 ~1.4µs/行)。
//
// 约束: 只服务内部只读聚合。有 ORDER BY 时排序要比较解码后的列值, 这条路径
// 无能为力 —— 调用方必须先退回整行路径。OFFSET/LIMIT 语义与 scanWithPKID
// 的无排序分支逐行一致。
func (q *Query) scanRaw(callback func(pkID uint64, rowData []byte) bool) error {
	plan := q.buildScanPlan()
	seen := 0
	return q.executeScanRaw(func(pkID uint64, rowData []byte) bool {
		if seen < q.offset {
			seen++
			return true // 跳过 OFFSET
		}
		seen++
		if !callback(pkID, rowData) {
			return false // 用户要求停止
		}
		if q.limit > 0 && seen-q.offset >= q.limit {
			return false // 已满足 LIMIT, 提前终止扫描
		}
		return true
	}, plan)
}

func (q *Query) executeScanRaw(callback func(pkID uint64, rowData []byte) bool, plan *scanPlan) error {
	if plan != nil {
		return q.indexScanRaw(plan, callback)
	}
	return q.fullTableScanRaw(callback)
}

// indexScanRaw 与 indexScan 走同一执行计划, 但回调拿裸行字节, 且遍历中同步回表。
//
// 为什么可以在 RangeScan 回调里直接回表 (indexScan 却必须"先物化、后回表"):
// 物化两阶段是为"用户回调会在遍历中 Update/Delete 同一棵索引"准备的 —— 叶分裂/
// 合并会让裸游标读到被改写的页。本路径的消费者全是只读聚合, 遍历中不存在任何
// 写操作, 游标必然有效; 于是物化 (含 rowStore 的逐行快照拷贝) 一并省掉。
func (q *Query) indexScanRaw(plan *scanPlan, callback func(pkID uint64, rowData []byte) bool) error {
	scanRange := func(kr keyRange) error {
		if plan.rowStore {
			// 承载行数据的索引: 条目即 storeRow = [pkID(8B) + rowData]
			return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
				kr.startKey, kr.endKey, kr.startInclusive, kr.endInclusive, false,
				func(row []byte) bool {
					if len(row) < 8 {
						return true
					}
					rowData := row[8:]
					if plan.remainingCond != nil && !plan.remainingCond.Match(rowData, q.table) {
						return true
					}
					return callback(binary.BigEndian.Uint64(row[:8]), rowData)
				})
		}
		// 二级索引: 条目 = pkID(8B BE), 遍历中同步回表 (只读, 不触碰被遍历的索引结构)
		return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage,
			kr.startKey, kr.endKey, kr.startInclusive, kr.endInclusive, false,
			func(row []byte) bool {
				if len(row) < 8 {
					return true
				}
				storeRow, found := q.table.indexMeta.ResolveByPK(binary.BigEndian.Uint64(row[:8]))
				if !found {
					return true // 行已被删除, 跳过
				}
				rowData := storeRow[8:]
				if plan.remainingCond != nil && !plan.remainingCond.Match(rowData, q.table) {
					return true
				}
				return callback(binary.BigEndian.Uint64(row[:8]), rowData)
			})
	}
	for _, kr := range plan.ranges {
		if err := scanRange(kr); err != nil {
			return err
		}
	}
	return nil
}

// fullTableScanRaw 与 fullTableScan 同构, 回调拿裸行字节。
// ReadDataPageRows 返回的行本就是拷贝快照, rowData 跨回调安全。
func (q *Query) fullTableScanRaw(callback func(pkID uint64, rowData []byte) bool) error {
	// 找承载行数据的索引 (实现 RowStoreHandler 者, 本层不认识索引类型)
	primMeta, primHandler, ok := q.table.indexMeta.rowStoreMeta()
	if !ok || primMeta == nil || primHandler == nil {
		return nil
	}
	dpIDs := primHandler.GetDataPageIDs(q.table.handle, primMeta.EntryPage)
	for _, dpID := range dpIDs {
		for _, storeRow := range primHandler.ReadDataPageRows(q.table.handle, dpID) {
			if len(storeRow) < 8 {
				continue
			}
			rowData := storeRow[8:]
			if q.whereCond != nil {
				if !q.whereCond.Match(rowData, q.table) {
					continue
				}
			}
			if !callback(binary.BigEndian.Uint64(storeRow[:8]), rowData) {
				return nil // 回调要求提前终止
			}
		}
	}
	return nil
}

// topKScan ORDER BY + LIMIT 的流式执行: 扫描全程只保留 limit+offset 行的小顶堆,
// 内存 O(k)、免全量排序 (旧实现先物化全部结果再整体排序, 百万行翻页也要全排)。
// 稳定性: 相等键按输入序号决胜, 输出与 slices.SortStableFunc 后切片逐字节一致。
func (q *Query) topKScan(callback func(pkID uint64, row []types.FieldValue) bool,
	plan *scanPlan, specs []orderKeySpec) error {
	cap := q.limit + q.offset
	if cap < q.limit {
		cap = q.limit // 防 offset 溢出 (理论值, int 实际不会)
	}
	h := &topkHeap{specs: specs}
	idx := 0
	err := q.executeScan(func(pkID uint64, row []types.FieldValue) bool {
		keys := make([]interface{}, len(specs))
		for j, s := range specs {
			keys[j], _ = lookupFieldValue(row, []byte(s.field))
		}
		heap.Push(h, topkItem{pkID: pkID, row: row, keys: keys, idx: idx})
		idx++
		if h.Len() > cap {
			heap.Pop(h)
		}
		return true
	}, plan, 0)
	if err != nil {
		return err
	}
	// 全部弹出: 大顶堆弹出的是"输出序最后"的元素 → items 是输出序的逆序。
	items := make([]topkItem, 0, h.Len())
	for h.Len() > 0 {
		items = append(items, heap.Pop(h).(topkItem))
	}
	if q.offset > 0 {
		if q.offset >= len(items) {
			return nil
		}
		items = items[:len(items)-q.offset] // 弹出序的尾部 = 输出序的头部 → 丢弃
	}
	if q.limit > 0 && len(items) > q.limit {
		items = items[len(items)-q.limit:] // 保留弹出序尾部 = 输出序的前 limit 名
	}
	for i := len(items) - 1; i >= 0; i-- {
		if !callback(items[i].pkID, items[i].row) {
			return nil
		}
	}
	return nil
}