TsDB2/DB/core/query_plan.go

package core

import (
	"bytes"
	"fmt"

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

// ======================== 索引键编码 ========================

// encodeIndexKey 将条件值编码为索引键字节。
//
// 编码与行内存储【同源】(都走 types.EncodeValue): 索引的 B+ 树直接按字节序比较键,
// 因此"索引键字节 == 行内字节"是范围扫描正确的前提。这里不再按类型手写分支 ——
// 以前每加一个类型都要在这里补一份, 漏补的后果是索引选出来的行与过滤结果不一致。
//
// param 是列级类型参数: decimal 是小数位数, 字符串家族是行内宽度码。少了它就没法把
// 条件值编成与行内完全一致的字节 —— decimal 会因 scale 不同而长度/系数不一致,
// 字符串则会按错误的宽度截断或越界, 两种都让"索引选出 5 行、过滤后剩 0 行"。
//
// 键长必须等于 FieldSize(ft, param): B+ 树按【定长拼键】比较, 键长与行内宽度不一致
// 就等于拿两种编码去比大小。
func encodeIndexKey(value interface{}, ft types.FieldType, param byte) []byte {
	if types.IsDynamic(ft) {
		// 动态列不能作为索引分量 (CreateIndex 已拒绝), 这里仅作兜底
		if b, ok := value.([]byte); ok {
			return b
		}
		return nil
	}
	size := types.FieldSize(ft, param)
	if size == 0 {
		// 未注册类型 / 非法列参数: 交给下面 EncodeValue 报出具体原因 (buf 长度为 0)
		panic(fmt.Errorf("%w: 类型 %s 的宽度不可计算 (列级参数 %d)",
			types.ErrValueType, types.TypeNameOf(ft), param))
	}
	buf := make([]byte, size)
	if err := types.EncodeValue(ft, param, value, buf); err != nil {
		// 沿用既有约定: panic 由 encodeIndexKeySafe* 转成 error。
		panic(err)
	}
	return buf
}

// encodeIndexKeySafe 在 encodeIndexKey 外面包一层 panic 恢复。
// 值类型与列不符时 encodeIndexKey 会 panic, 而公开 API 传进来的值不可信,
// 必须转成 error, 不能把用户输入直接变成崩溃。
func encodeIndexKeySafe(value interface{}, ft types.FieldType) (enc []byte, err error) {
	return encodeIndexKeySafeWithParam(value, ft, 0)
}

// encodeIndexKeySafeWithParam 同 encodeIndexKeySafe, 但带上列级类型参数。
// 查询侧拿得到字段元信息时必须走这个版本 (decimal 需要 param 才能编出正确的字节)。
func encodeIndexKeySafeWithParam(value interface{}, ft types.FieldType, param byte) (enc []byte, err error) {
	defer func() {
		if r := recover(); r != nil {
			enc, err = nil, fmt.Errorf("值 %v 与存储类型 %s 不匹配: %v", value, types.TypeNameOf(ft), r)
		}
	}()
	return encodeIndexKey(value, ft, param), nil
}

// ======================== 索引匹配 ========================

// fieldNameFromMeta 获取索引第一个字段的名称
func (q *Query) fieldNameFromMeta(meta *IndexMeta) string {
	fieldIdx := int(meta.FieldIDs[0])
	if fieldIdx < len(q.table.Field) {
		return string(q.table.Field[fieldIdx].FieldName)
	}
	return ""
}

// ======================== 复合索引键规划 ========================
//
// 索引键 = 各分量字段定长拼接 (编码与 extractIndexKey 完全一致), 因此:
//   - 已知前 k 个分量等值 → 键上一段连续区间;
//   - 第 k+1 个分量上的范围条件 → 再收窄这段区间;
//   - 更靠后的分量索引帮不上忙, 原样留作行级后置过滤。
// 下面的代码把 WHERE 按这个模型折叠成 (startKey, endKey, 残余条件)。
//
// 单列索引是 n == 1 的特例, 行为与旧实现一致。

// indexKeyLayout 索引键各分量的偏移与长度 (定长拼接)
type indexKeyLayout struct {
	offsets []int
	sizes   []int
	total   int
}

// cmpKind 叶子条件能为键分量提供的边界方向。
// 取值与 CmpOp 一一对应 (cmpEq == cmpKind(OpEq), ...), classifyLeaf 因此可以
// 直接 cmpKind(t.Op) 转换, 不必再写映射表。
type cmpKind int

const (
	cmpEq = cmpKind(OpEq)
	cmpGe = cmpKind(OpGe)
	cmpGt = cmpKind(OpGt)
	cmpLe = cmpKind(OpLe)
	cmpLt = cmpKind(OpLt)
)

// keyConstraint 某个键分量上的合并边界 (多个 AND 子条件取交集)
type keyConstraint struct {
	lo, hi       []byte // 编码后的分量值; nil 表示该侧无约束
	loInc, hiInc bool
	count        int // 落在该分量上的条件个数, 0 = 无约束
}

// add 把一个边界并入本分量; 返回 false 表示条件自相矛盾。
func (kc *keyConstraint) add(kind cmpKind, enc []byte) bool {
	switch kind {
	case cmpEq:
		if !kc.addLower(enc, true) || !kc.addUpper(enc, true) {
			return false
		}
	case cmpGe:
		if !kc.addLower(enc, true) {
			return false
		}
	case cmpGt:
		if !kc.addLower(enc, false) {
			return false
		}
	case cmpLe:
		if !kc.addUpper(enc, true) {
			return false
		}
	case cmpLt:
		if !kc.addUpper(enc, false) {
			return false
		}
	}
	kc.count++
	return true
}

func (kc *keyConstraint) addLower(enc []byte, inclusive bool) bool {
	if kc.lo == nil {
		kc.lo, kc.loInc = enc, inclusive
		return true
	}
	switch c := bytes.Compare(enc, kc.lo); {
	case c > 0: // 取更高的下界
		kc.lo, kc.loInc = enc, inclusive
	case c == 0: // 同值: 两端都闭才算"含"
		kc.loInc = kc.loInc && inclusive
	}
	return kc.notContradictory()
}

func (kc *keyConstraint) addUpper(enc []byte, inclusive bool) bool {
	if kc.hi == nil {
		kc.hi, kc.hiInc = enc, inclusive
		return true
	}
	switch c := bytes.Compare(enc, kc.hi); {
	case c < 0: // 取更低的上界
		kc.hi, kc.hiInc = enc, inclusive
	case c == 0:
		kc.hiInc = kc.hiInc && inclusive
	}
	return kc.notContradictory()
}

// notContradictory 判断 [lo, hi] 是否还可能有值。
//
// 注意: 这里按编码字节比较, 而定长整数用的是补码 —— 负数的字节序与数值序相反。
// 所以它只在"字节序下已确定矛盾"时才返回 false, 此时真值也必然为空, 反向不成立。
// 调用方因此只把它当作"放弃使用索引"的依据, 绝不据此断言区间非空。
func (kc *keyConstraint) notContradictory() bool {
	if kc.lo == nil || kc.hi == nil {
		return true
	}
	switch c := bytes.Compare(kc.lo, kc.hi); {
	case c > 0:
		return false
	case c == 0:
		return kc.loInc && kc.hiInc
	}
	return true
}

// isEquality 该分量已被钉死为单值 (lo == hi 且两端都闭)
func (kc *keyConstraint) isEquality() bool {
	return kc.count > 0 && kc.lo != nil && kc.hi != nil &&
		kc.loInc && kc.hiInc && bytes.Equal(kc.lo, kc.hi)
}

// buildKeyLayout 依据索引元信息与表字段类型算出键布局。
// 返回 false 表示该索引当前不可用于查询:
//   - 分量列不存在或已被 Drop (新行提不出该分量, 索引里根本没有这批键);
//   - 分量是动态列 (行内只存 6B 堆引用, 值 != 引用, 无法按键定位)。
func (q *Query) buildKeyLayout(meta *IndexMeta) (*indexKeyLayout, bool) {
	n := int(meta.FieldCnt)
	if n <= 0 || n > len(meta.FieldIDs) || n > len(q.table.Field) {
		return nil, false
	}
	lay := &indexKeyLayout{offsets: make([]int, n), sizes: make([]int, n)}
	off := 0
	for k := 0; k < n; k++ {
		fi := int(meta.FieldIDs[k])
		if fi >= len(q.table.Field) {
			return nil, false
		}
		if q.table.Field[fi].IsDropped() {
			return nil, false
		}
		ft := q.table.Field[fi].FieldType
		if types.IsDynamic(ft) {
			return nil, false
		}
		// 键分量宽度必须与行内宽度同源 (含字符串列的宽度码), 否则索引键的偏移切不出该列
		size := int(q.table.Field[fi].Width())
		if size <= 0 {
			return nil, false
		}
		lay.offsets[k] = off
		lay.sizes[k] = size
		off += size
	}
	lay.total = off
	return lay, true
}

// locateComponent 返回字段 master 下标在索引键中的分量位置; -1 表示不属于该索引
func (q *Query) locateComponent(fieldName string, meta *IndexMeta) int {
	fi := q.table.findFieldIndex(fieldName)
	if fi < 0 {
		return -1
	}
	n := int(meta.FieldCnt)
	if n > len(meta.FieldIDs) {
		n = len(meta.FieldIDs)
	}
	for k := 0; k < n; k++ {
		if int(meta.FieldIDs[k]) == fi {
			return k
		}
	}
	return -1
}

// classifyLeaf 判断一个叶子条件能否成为某个键分量的边界。
// comp < 0 表示与该索引无关 (字段不在索引里, 或该条件类型给不出可比边界)。
func (q *Query) classifyLeaf(c Condition, meta *IndexMeta) (comp int, value interface{}, kind cmpKind) {
	switch t := c.(type) {
	case CmpCondition:
		// cmpKind 与 CmpOp 取值一一对应 (见 cmpKind 常量块), 直接转换。
		return q.locateComponent(t.Field, meta), t.Value, cmpKind(t.Op)
	default:
		return -1, nil, cmpEq
	}
}

// flattenAnd 把 AND 树展平成叶子条件。
// OR / NOT / LIKE / IN 不是合取项, 整体当叶子返回 —— 它们给不出索引边界,
// 但必须原样保留下来做行级过滤。
func flattenAnd(cond Condition, out *[]Condition) {
	if and, ok := cond.(AndCondition); ok {
		for _, sub := range and.Conditions {
			flattenAnd(sub, out)
		}
		return
	}
	*out = append(*out, cond)
}

// incBE 把 buf 当作大端无符号整数加一 (原地)。返回 false 表示溢出 ——
// 此时 buf 之后不再存在任何更大的键。
func incBE(buf []byte) bool {
	for i := len(buf) - 1; i >= 0; i-- {
		buf[i]++
		if buf[i] != 0 {
			return true
		}
	}
	return false
}

// fillFF 把 buf[from:] 全置 0xFF (该分量及其后分量的最大值)。
// from 可以是 0 —— 首分量就是范围分量时前缀为空, 此时整个键都无上界。
func fillFF(buf []byte, from int) {
	for i := from; i < len(buf); i++ {
		buf[i] = 0xFF
	}
}

// buildScanPlan 分析 WHERE 条件, 选择最优索引
func (q *Query) buildScanPlan() *scanPlan {
	if q.whereCond == nil {
		return nil
	}
	if q.table.indexMeta == nil {
		return nil
	}

	// 先展平成合取项: 索引区间是按"分量"折叠出来的, 只有 AND 语义才可加。
	var leaves []Condition
	flattenAnd(q.whereCond, &leaves)

	var bestPlan *scanPlan
	for i := range q.table.indexMeta.metas {
		m := &q.table.indexMeta.metas[i]
		if m.Flag&0x80 != 0 {
			continue
		}
		// 主键索引直接承载全部行, 但它的 B+ 树叶键是 pkID, 而不是"建在哪一列上"。
		// planIndex 会用 meta.FieldIDs 编码出的【字段值】当扫描区间 (见 buildKeyLayout),
		// 拿它去比对 pkID 叶键, 语义对不上 —— 只有当该列的值恒等于 pkID 时才偶然正确,
		// 否则会静默地只命中/漏掉若干行 (例如 Delete/Update 的条件扫描)。
		// 因此主键索引不参与列条件规划, 这类查询退回全表扫描 + 行级过滤 (fullTableScan
		// 会施加 whereCond), 结果正确; 按 pkID 的点查/回表仍走 SearchByIndex/ResolveByPK,
		// 不经过这里, 不受影响。
		if m.IndexType == page.IndexTypePrimary {
			continue
		}
		handler, ok := index.Compose.GetHandler(m.IndexType)
		if !ok {
			continue
		}
		plan := q.planIndex(m, handler, leaves)
		if plan == nil {
			continue // 此索引不匹配
		}
		// 该索引是否承载行数据: 实现 RowStoreHandler 即承载, 本层不认识索引类型
		_, plan.rowStore = handler.(index.RowStoreHandler)
		if plan.rowStore {
			return plan // 行存储索引最优 (免回表), 不再找
		}
		if bestPlan == nil {
			bestPlan = plan
		}
	}
	return bestPlan
}

// planIndex 为一个索引推导扫描区间; 该索引用不上时返回 nil。
func (q *Query) planIndex(m *IndexMeta, handler index.IndexHandler, leaves []Condition) *scanPlan {
	lay, ok := q.buildKeyLayout(m)
	if !ok {
		return nil
	}
	n := len(lay.sizes)

	cons := make([]keyConstraint, n)   // 各分量合并后的边界
	compCond := make([][]Condition, n) // 各分量上的原始条件 (用于决定哪些能丢)
	var rest []Condition               // 与索引无关、必须保留的条件

	for _, c := range leaves {
		k, value, kind := q.classifyLeaf(c, m)
		if k < 0 {
			rest = append(rest, c)
			continue
		}
		ft := q.table.Field[m.FieldIDs[k]].FieldType
		// 用 Safe 版本: 条件的值类型与列存储类型不符时 encodeIndexKey 会 panic,
		// 规划阶段不该把用户输入直接变成崩溃 —— 退回行级过滤由 Match 去处理。
		// 带上列级 param (decimal 的 scale / 字符串的宽度码): 索引键必须与行内字节完全一致。
		enc, err := encodeIndexKeySafeWithParam(value, ft, q.table.fieldParam(int(m.FieldIDs[k])))
		if err != nil || len(enc) != lay.sizes[k] {
			// 编码失败, 或长度与分量长度不符 (拼进定长键会错位) → 交给行级过滤
			rest = append(rest, c)
			continue
		}
		if !cons[k].add(kind, enc) {
			// 该分量上的条件互相矛盾 → 结果必为空。
			// 这里保守地"放弃索引"而不是返回空区间: 定长补码编码下负数字节序与
			// 数值序相反, 用字节比较断定"空"对负数不成立, 退回全表扫描更稳。
			return nil
		}
		compCond[k] = append(compCond[k], c)
	}

	// 从左往右消费"等值分量"; 之后最多再消费一个范围分量, 更后面的分量帮不上忙。
	eqLen := 0
	for eqLen < n && cons[eqLen].isEquality() {
		eqLen++
	}
	rangeIdx := -1
	if eqLen < n && cons[eqLen].count > 0 {
		rangeIdx = eqLen
	}

	// IN 展开: 只看"等值前缀的下一个分量", 且其上必须还没有任何边界条件
	// (有范围/等值条件时让位给常规区间规划, IN 由行级过滤兜底)。
	// 每个 IN 值变成该分量取值的精确区间 [前缀||v||00.., 前缀||v||FF..]。
	// 多个 IN 只消费第一个适用的, 其余留在行级过滤 (正确性不受影响)。
	// 注意本段在 cons[0].count 判定【之前】: 只有 IN 条件时 cons 全空,
	// 但它照样能给出一组精确区间。
	for _, c := range leaves {
		in, ok := c.(InCondition)
		if !ok {
			continue
		}
		if eqLen >= n || cons[eqLen].count != 0 {
			continue
		}
		if q.locateComponent(in.Field, m) != eqLen {
			continue
		}
		if plan := q.planInRanges(m, handler, lay, cons, eqLen, rest, compCond, in); plan != nil {
			return plan
		}
		// 返回了 nil = 有值编不出键: 本索引放弃 IN, 让行级过滤兜住正确性,
		// 继续走下面的常规区间规划 (等值前缀等条件仍可用)。
		break
	}

	if cons[0].count == 0 {
		return nil // 首分量没有边界 → 索引给不出连续区间
	}

	startKey := make([]byte, lay.total)
	endKey := make([]byte, lay.total)
	prefix := 0
	for j := 0; j < eqLen; j++ {
		copy(startKey[prefix:], cons[j].lo)
		copy(endKey[prefix:], cons[j].lo)
		prefix += len(cons[j].lo)
	}

	plan := &scanPlan{
		meta: m, handler: handler,
		ranges: []keyRange{{startKey: startKey, endKey: endKey,
			startInclusive: true, endInclusive: true}},
		remainingCond: q.remainCond(rest, compCond, eqLen, rangeIdx),
	}

	// 全部等值 → 精确键 (此时 prefix == total, fillFF 无事可做);
	// 只消费了等值前缀 → 该前缀下的全部键: [前缀||00.., 前缀||FF..]。
	if eqLen == n || rangeIdx < 0 {
		fillFF(endKey, prefix)
		return plan
	}

	rc := &cons[rangeIdx]
	size := lay.sizes[rangeIdx]

	// 下界: [前缀 || 分量值 || 00..]
	if rc.lo != nil {
		copy(startKey[prefix:], rc.lo)
		if !rc.loInc {
			// 严格下界: 要跳过分量 == lo 的全部键, 即把 (前缀||lo) 整体加一
			if !incBE(startKey[:prefix+size]) {
				return nil // 已经是最大键, 该条件不可能有匹配行
			}
		}
	}

	// 上界: [前缀 || 分量值 || FF..]
	kr := &plan.ranges[0]
	if rc.hi == nil {
		fillFF(endKey, prefix)
	} else if rc.hiInc {
		copy(endKey[prefix:], rc.hi)
		fillFF(endKey, prefix+size)
	} else {
		// 严格上界: 分量 == hi 的键全部排除
		copy(endKey[prefix:], rc.hi)
		kr.endInclusive = false
	}
	return plan
}

// planInRanges 把 IN 条件展开成多区间扫描计划。
//
// 返回 nil 表示 IN 无法索引化 (有值编不出键) —— 调用方退回常规规划/全表扫描,
// IN 本身仍在 whereCond 里, 行级过滤保证正确性。
// IN 空集合 → 返回 ranges 为空的计划: 任何区间都不扫, 结果必空 (IN 恒假),
// 与 InCondition.Match 的空集合语义一致。
func (q *Query) planInRanges(m *IndexMeta, handler index.IndexHandler, lay *indexKeyLayout,
	cons []keyConstraint, eqLen int, rest []Condition, compCond [][]Condition, in InCondition) *scanPlan {
	k := eqLen
	size := lay.sizes[k]
	fi := int(m.FieldIDs[k])
	ft := q.table.Field[fi].FieldType
	param := q.table.fieldParam(fi)

	// 等值前缀字节 (与常规规划同一编码: cons[j].lo, isEquality 保证两端闭且相等)
	prefix := make([]byte, lay.offsets[k])
	for j := 0; j < eqLen; j++ {
		copy(prefix[lay.offsets[j]:], cons[j].lo)
	}

	if len(in.Values) == 0 {
		// 恒空: 不扫任何区间。remainingCond 保留全部条件 (无行可过滤, 仅为语义完整)。
		return &scanPlan{meta: m, handler: handler,
			ranges:        nil,
			remainingCond: q.remainCond(rest, compCond, 0, -1)}
	}

	ranges := make([]keyRange, 0, len(in.Values))
	seen := make(map[string]struct{}, len(in.Values)) // 编码后同键的值只留一个区间
	for _, v := range in.Values {
		enc, err := encodeIndexKeySafeWithParam(v, ft, param)
		if err != nil || len(enc) != size {
			return nil // 任一值编不出正确宽度的键 → 整个 IN 退回行级过滤
		}
		// 该分量取值 v 的全部键: [前缀||v||00.., 前缀||v||FF..] (单列索引即 [v, v])
		start := make([]byte, lay.total)
		end := make([]byte, lay.total)
		copy(start, prefix)
		copy(end, prefix)
		copy(start[lay.offsets[k]:], enc)
		copy(end[lay.offsets[k]:], enc)
		fillFF(end, lay.offsets[k]+size)
		// 重复值/编码碰撞值去重: 不去重的话, raw 聚合路径 (Count/Sum 的
		// scanRaw) 会对同一行重复计数 —— 行交付路径有去重, 聚合没有。
		if _, dup := seen[string(start)]; dup {
			continue
		}
		seen[string(start)] = struct{}{}
		ranges = append(ranges, keyRange{startKey: start, endKey: end,
			startInclusive: true, endInclusive: true})
	}
	return &scanPlan{meta: m, handler: handler, ranges: ranges,
		remainingCond: q.remainCond(rest, compCond, eqLen, -1)}
}

// remainCond 汇总索引区间没能表达、仍需逐行过滤的条件。
// 已被区间精确覆盖的分量 (等值前缀 + 那个范围分量) 上的条件可以安全丢掉。
func (q *Query) remainCond(rest []Condition, compCond [][]Condition, eqLen, rangeIdx int) Condition {
	remaining := rest
	for j := range compCond {
		if j < eqLen || j == rangeIdx {
			continue
		}
		remaining = append(remaining, compCond[j]...)
	}
	switch len(remaining) {
	case 0:
		return nil
	case 1:
		return remaining[0]
	default:
		return And(remaining...)
	}
}