TsDB2/DB/core/query_agg.go

package core

import (
	"bytes"
	"encoding/binary"
	"math/big"
	"slices"

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

// orderKeySpec 预解析的排序键: 排序列不存在于表中时在构建阶段就被丢掉,
// 比较阶段不再做任何字段查找。
type orderKeySpec struct {
	field string
	ft    types.FieldType
	desc  bool
}

// buildOrderSpecs 把 OrderBy 解析成排序键规格 (字段名 + 列类型 + 方向)。
func (q *Query) buildOrderSpecs() []orderKeySpec {
	specs := make([]orderKeySpec, 0, len(q.orderBys))
	for _, ob := range q.orderBys {
		idx := q.table.findFieldIndex(ob.Field)
		if idx < 0 {
			continue
		}
		specs = append(specs, orderKeySpec{
			field: ob.Field,
			ft:    q.table.Field[idx].FieldType,
			desc:  ob.Descending,
		})
	}
	return specs
}

// compareRowsForOrderBy 按 OrderBy 规则比较两行, 返回 -1/0/1。
// 0 表示"所有排序列都无法区分"(字段缺失), 由调用方用稳定排序保持原顺序。
//
// 保留给外部调用方; 查询内部排序走 sortResults (decoration 每行只提取一次键)。
func (q *Query) compareRowsForOrderBy(a, b []types.FieldValue) int {
	return compareRowsBySpecs(a, b, q.buildOrderSpecs())
}

// compareRowsBySpecs 按预解析规格比较两行 (每行先经 lookupFieldValue 取值)。
func compareRowsBySpecs(a, b []types.FieldValue, specs []orderKeySpec) int {
	for _, s := range specs {
		vi, _ := lookupFieldValue(a, []byte(s.field))
		vj, _ := lookupFieldValue(b, []byte(s.field))
		if vi == nil || vj == nil {
			continue
		}
		cmp := compareFieldValues(vi, vj, s.ft)
		if cmp == 0 {
			continue
		}
		if s.desc {
			return -cmp
		}
		return cmp
	}
	return 0
}

// ======================== 排序 ========================

// sortResults 内存排序: decoration 阶段每行按名提取一次排序键, 键随行进数组
// (排序会重排元素, 平行键数组会错位), 比较阶段零查找。
// 行集合可能因 NULL 跳列而宽窄不一, 键按名进行、缺失用 nil 占位,
// 语义与旧比较器完全一致。
func (q *Query) sortResults(results [][]types.FieldValue) {
	specs := q.buildOrderSpecs()
	if len(specs) == 0 {
		return
	}
	dec := make([]struct {
		row  []types.FieldValue
		keys []interface{}
	}, len(results))
	for i := range results {
		k := make([]interface{}, len(specs))
		for j, s := range specs {
			k[j], _ = lookupFieldValue(results[i], []byte(s.field))
		}
		dec[i] = struct {
			row  []types.FieldValue
			keys []interface{}
		}{results[i], k}
	}
	slices.SortStableFunc(dec, func(a, b struct {
		row  []types.FieldValue
		keys []interface{}
	}) int {
		return compareKeys(a.keys, b.keys, specs)
	})
	for i := range dec {
		results[i] = dec[i].row
	}
}

// compareKeys 比较两行预提取的排序键。
func compareKeys(ka, kb []interface{}, specs []orderKeySpec) int {
	for j, s := range specs {
		vi, vj := ka[j], kb[j]
		if vi == nil || vj == nil {
			continue
		}
		cmp := compareFieldValues(vi, vj, s.ft)
		if cmp == 0 {
			continue
		}
		if s.desc {
			return -cmp
		}
		return cmp
	}
	return 0
}

// ======================== 流式 top-k (ORDER BY + LIMIT) ========================

// topkItem 堆内元素: 行 + 预提取排序键 + 输入序号。
// 序号做相等键的决胜 —— 弹出顺序与 slices.SortStableFunc 完全一致 (稳定)。
type topkItem struct {
	pkID uint64
	row  []types.FieldValue
	keys []interface{}
	idx  int
}

// topkHeap 容量上限 limit+offset 的小顶堆: 超过即弹出当前最小, 堆里始终是
// 已见行中的前 k 名。heap.Interface。
type topkHeap struct {
	specs []orderKeySpec
	items []topkItem
}

func (h *topkHeap) Len() int { return len(h.items) }

// Less 定义"输出序的逆序" (大顶堆): 堆顶是【当前集合里最后输出】的元素 ——
// 容量超限时的淘汰对象。相等键时输入序号大者更靠后、先被淘汰 (稳定)。
func (h *topkHeap) Less(i, j int) bool {
	c := compareKeys(h.items[i].keys, h.items[j].keys, h.specs)
	if c != 0 {
		return c > 0
	}
	return h.items[i].idx > h.items[j].idx
}

func (h *topkHeap) Swap(i, j int) { h.items[i], h.items[j] = h.items[j], h.items[i] }

func (h *topkHeap) Push(x interface{}) { h.items = append(h.items, x.(topkItem)) }

func (h *topkHeap) Pop() interface{} {
	old := h.items
	n := len(old)
	it := old[n-1]
	old[n-1] = topkItem{} // 防内存泄漏
	h.items = old[:n-1]
	return it
}

// cmpInt64 / cmpUint64 把三态比较收敛到一处, 避免每个分支重复写一遍 if 链。
func cmpInt64(a, b int64) int {
	switch {
	case a < b:
		return -1
	case a > b:
		return 1
	}
	return 0
}

func cmpUint64(a, b uint64) int {
	switch {
	case a < b:
		return -1
	case a > b:
		return 1
	}
	return 0
}

// compareFieldValues 按【字段类型】比较同一列的两个值, 返回 -1/0/1。
//
// 类型必须由调用方传入, 不能只看 Go 值类型: 同一列都可能读成 int16, 但它既可能来自
// 有符号的 FieldType_Int16 (int16 语义), 也可能来自无符号的 FieldType_Uint16 ——
// 后者的值是 uint16 的位模式被重解释成 int16, 两者序不同 (位模式 0xFFFF: int16 是 -1,
// uint16 是 65535)。旧实现只按 Go 值类型分派, 既分不清这两种, 又整个漏掉了 int8/int64。
//
// 有符号类型比的是【解码后的值】而不是盘上字节: 盘上是保序编码 (符号位已翻转),
// 直接比字节也对, 但这里拿到的是已解码的值, 比数值更直白。
func compareFieldValues(a, b interface{}, ft types.FieldType) int {
	switch ft {
	case types.FieldType_Int8:
		return cmpInt64(int64(a.(int8)), int64(b.(int8)))
	case types.FieldType_Int16:
		return cmpInt64(int64(a.(int16)), int64(b.(int16)))
	case types.FieldType_Int32:
		return cmpInt64(int64(a.(int32)), int64(b.(int32)))
	case types.FieldType_Int64:
		return cmpInt64(a.(int64), b.(int64))
	case types.FieldType_Uint8:
		return cmpUint64(uint64(a.(byte)), uint64(b.(byte)))
	case types.FieldType_Uint16:
		// 无符号: 先把 int16 位模式还原回 uint16 再比, 否则 > 32767 的值会被当成负数
		return cmpUint64(uint64(uint16(a.(int16))), uint64(uint16(b.(int16))))
	case types.FieldType_Uint32:
		return cmpUint64(uint64(uint32(a.(int32))), uint64(uint32(b.(int32))))
	case types.FieldType_Uint64:
		return cmpUint64(a.(uint64), b.(uint64))
	case types.FieldType_DateTime:
		return cmpUint64(a.(uint64), b.(uint64))
	case types.FieldType_Float64:
		// 同一列读出的一定是同一类型, 故 b 也是 float64; 容错地走 toFloat64
		fa, fb := toFloat64(a), toFloat64(b)
		switch {
		case fa < fb:
			return -1
		case fa > fb:
			return 1
		}
		return 0
	case types.FieldType_Float32:
		fa, fb := a.(float32), b.(float32)
		switch {
		case fa < fb:
			return -1
		case fa > fb:
			return 1
		}
		return 0
	case types.FieldType_Bool:
		// false < true (与盘上 0x00 < 0x01 的字节序一致)
		ba, bb := a.(bool), b.(bool)
		switch {
		case ba == bb:
			return 0
		case !ba:
			return -1
		}
		return 1
	case types.FieldType_Year:
		// 同 uint16: 载体是 int16, 但语义是无符号年份
		return cmpUint64(uint64(uint16(a.(int16))), uint64(uint16(b.(int16))))
	case types.FieldType_Decimal:
		// 走数值比较 (自动对齐 scale): 同一列的 scale 必然相同, 但比较函数不该依赖这一点
		return a.(types.Decimal).Cmp(b.(types.Decimal))
	case types.FieldType_Int128:
		na, oka := a.(*big.Int)
		nb, okb := b.(*big.Int)
		if !oka || !okb {
			return 0
		}
		return na.Cmp(nb)
	case types.FieldType_UUID:
		ua, oka := a.(types.UUID)
		ub, okb := b.(types.UUID)
		if !oka || !okb {
			return 0
		}
		return bytes.Compare(ua[:], ub[:])
	}
	// 动态字段与所有定长字符串: 按字节序
	ba, ok := a.([]byte)
	if !ok {
		return 0
	}
	bb, ok := b.([]byte)
	if !ok {
		return 0
	}
	return bytes.Compare(ba, bb)
}

// fieldTypeOf 取某列的存储类型, 供按类型比较/聚合使用; 列不存在时返回 0。
func (q *Query) fieldTypeOf(name string) types.FieldType {
	if i := q.table.findFieldIndex(name); i >= 0 {
		return q.table.Field[i].FieldType
	}
	return 0
}

// ======================== 流式聚合 ========================

// FindOne 流式查询，返回单行 (命中即停，不遍历全表)
func (q *Query) FindOne() ([]types.FieldValue, error) {
	var result []types.FieldValue
	found := false
	err := q.scan(func(row []types.FieldValue) bool {
		result = row
		found = true
		return false // 找到第一行即停止
	})
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, page.ErrNotFound
	}
	return result, nil
}

// Count 流式计数。
//
// 三级加速 (每级都严格保持与全量扫描相同的计数语义):
//  1. 无条件无分页排序 → 逐页数活行 (DataPage::CountLiveRows), 不物化、不解码;
//  2. WHERE 被索引区间完全消费 (remainingCond == nil) → 直接数索引条目,
//     不回表、不反序列化 —— 旧实现这里每个命中行都做一次 ResolveByPK +
//     deserializeRow, "用索引 COUNT" 反而比全表扫描更慢 (20k 行 632ms);
//  3. 其余 → 裸字节扫描计数 (deserializeField 都不需要)。
//     带 OFFSET/LIMIT/ORDER BY 时保持旧行为 (旧实现的 Count 语义上会先被分页截断)。
func (q *Query) Count() (int64, error) {
	paged := q.limit > 0 || q.offset > 0 || len(q.orderBys) > 0
	if !paged && q.whereCond == nil {
		return q.countAllLive(), nil
	}
	if !paged {
		if plan := q.buildScanPlan(); plan != nil && plan.remainingCond == nil {
			var count int64
			for _, kr := range plan.ranges {
				err := 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 {
							count++
						}
						return true
					})
				if err != nil {
					return 0, err
				}
			}
			return count, nil
		}
		var count int64
		err := q.scanRaw(func(_ uint64, _ []byte) bool {
			count++
			return true
		})
		return count, err
	}
	var count int64
	err := q.scan(func(row []types.FieldValue) bool {
		count++
		return true
	})
	return count, err
}

// countAllLive 逐页累加活行数 (无条件 COUNT(*) 的 O(页数) 路径)。
func (q *Query) countAllLive() int64 {
	if q.table.indexMeta == nil {
		return 0
	}
	primMeta, primHandler, ok := q.table.indexMeta.rowStoreMeta()
	if !ok || primMeta == nil || primHandler == nil {
		return 0
	}
	// 处理器实现了免物化计数就用它; 否则退回"读全部行再取 len"
	type liveCounter interface {
		ReadDataPageLiveCount(handle page.HandleDb, dpID uint32) int
	}
	counter, canCount := primHandler.(liveCounter)
	var total int64
	for _, dpID := range primHandler.GetDataPageIDs(q.table.handle, primMeta.EntryPage) {
		if canCount {
			total += int64(counter.ReadDataPageLiveCount(q.table.handle, dpID))
			continue
		}
		total += int64(len(primHandler.ReadDataPageRows(q.table.handle, dpID)))
	}
	return total
}

// Sum 流式求和。
// 每行只解码目标一列 (Table.deserializeField), 不再物化整行; 带 ORDER BY 时
// 排序管线需要整行, 保持旧路径。条件过滤 (whereCond/remainingCond) 在裸字节上
// 进行, 与旧路径同一套 Match。
func (q *Query) Sum(col string) (float64, error) {
	if len(q.orderBys) > 0 {
		var sum float64
		err := q.scan(func(row []types.FieldValue) bool {
			for _, fv := range row {
				if string(fv.FieldName) == col {
					sum += toFloat64(fv.Value)
					break
				}
			}
			return true
		})
		return sum, err
	}
	var sum float64
	err := q.scanRaw(func(_ uint64, rowData []byte) bool {
		if v, ok := q.table.deserializeField(rowData, col); ok {
			sum += toFloat64(v)
		}
		return true
	})
	return sum, err
}

// Max 流式求最大值。
//
// 加速分两档:
//  1. 目标列是某二级索引首列且无 WHERE/分页/排序 → 直接取索引端点
//     (反向 RangeScan 第一条即最大键), O(log N) —— 旧实现对 100 万行
//     线性扫 1.4s, MySQL/SQLite 同场景是 2.6ms/45µs;
//  2. 其余 → 裸字节扫描, 每行只解码目标一列。
func (q *Query) Max(col string) (interface{}, error) {
	if v, hit, err := q.minMaxViaIndex(col, true); hit || err != nil {
		return v, err
	}
	ft := q.fieldTypeOf(col)
	if len(q.orderBys) > 0 {
		var maxVal interface{}
		err := q.scan(func(row []types.FieldValue) bool {
			for _, fv := range row {
				if string(fv.FieldName) == col {
					if maxVal == nil || compareFieldValues(fv.Value, maxVal, ft) > 0 {
						maxVal = fv.Value
					}
					break
				}
			}
			return true
		})
		return maxVal, err
	}
	var maxVal interface{}
	err := q.scanRaw(func(_ uint64, rowData []byte) bool {
		if v, ok := q.table.deserializeField(rowData, col); ok {
			if maxVal == nil || compareFieldValues(v, maxVal, ft) > 0 {
				maxVal = v
			}
		}
		return true
	})
	return maxVal, err
}

// Min 流式求最小值 (加速路径同 Max, 端点取正向第一条)。
func (q *Query) Min(col string) (interface{}, error) {
	if v, hit, err := q.minMaxViaIndex(col, false); hit || err != nil {
		return v, err
	}
	ft := q.fieldTypeOf(col)
	if len(q.orderBys) > 0 {
		var minVal interface{}
		err := q.scan(func(row []types.FieldValue) bool {
			for _, fv := range row {
				if string(fv.FieldName) == col {
					if minVal == nil || compareFieldValues(fv.Value, minVal, ft) < 0 {
						minVal = fv.Value
					}
					break
				}
			}
			return true
		})
		return minVal, err
	}
	var minVal interface{}
	err := q.scanRaw(func(_ uint64, rowData []byte) bool {
		if v, ok := q.table.deserializeField(rowData, col); ok {
			if minVal == nil || compareFieldValues(v, minVal, ft) < 0 {
				minVal = v
			}
		}
		return true
	})
	return minVal, err
}

// minMaxViaIndex 取索引端点求 Max/Min。
//
// 前提: 目标列是某个二级索引的首列, 且查询不带 WHERE/分页/排序 —— 此时
// "全表最大值"就是"该索引里最大的键" (索引键编码与行内字节同源, 字节序即值序,
// 见 encodeIndexKey 的说明), 反向 (Max) / 正向 (Min) RangeScan 的第一条即答案。
// 回调返回 false 提前终止遍历, 成本只有一次树下降。
//
// 返回 (值, 是否命中快速路径, err); 未命中时调用方走全扫路径。
func (q *Query) minMaxViaIndex(col string, wantMax bool) (interface{}, bool, error) {
	if q.whereCond != nil || q.limit > 0 || q.offset > 0 || len(q.orderBys) > 0 {
		return nil, false, nil
	}
	if q.table.indexMeta == nil {
		return nil, false, nil
	}
	for i := range q.table.indexMeta.metas {
		m := &q.table.indexMeta.metas[i]
		if m.Flag&0x80 != 0 || m.IndexType == page.IndexTypePrimary {
			continue
		}
		if q.fieldNameFromMeta(m) != col {
			continue
		}
		if _, ok := q.buildKeyLayout(m); !ok {
			continue // 索引当前不可用 (列被删 / 动态列), 退回全扫
		}
		handler, ok := index.Compose.GetHandler(m.IndexType)
		if !ok {
			continue
		}
		var (
			got   interface{}
			found bool
		)
		err := handler.RangeScan(q.table.handle, m.EntryPage,
			nil, nil, true, true, wantMax,
			func(row []byte) bool {
				if len(row) < 8 {
					return true
				}
				storeRow, ok := q.table.indexMeta.ResolveByPK(binary.BigEndian.Uint64(row[:8]))
				if !ok {
					return true // 端点条目解不出行: 索引端点不可信, 由外层退回全扫
				}
				if v, ok := q.table.deserializeField(storeRow[8:], col); ok {
					got, found = v, true
				}
				return false // 只取端点一条
			})
		if err != nil || !found {
			return nil, false, err
		}
		return got, true, nil
	}
	return nil, false, nil
}

func toFloat64(v interface{}) float64 {
	switch val := v.(type) {
	case float64:
		return val
	case float32:
		return float64(val)
	case uint64:
		return float64(val)
	case int64:
		return float64(val)
	case int:
		return float64(val)
	case int8:
		return float64(val)
	case int16:
		return float64(val)
	case int32:
		return float64(val)
	case byte:
		return float64(val)
	case bool:
		if val {
			return 1
		}
		return 0
	case types.Decimal:
		// 近似: decimal → float64 会失真, 但 Sum/Avg 的返回值类型就是 float64。
		// 需要精确求和的金额场景请用 Decimal.Add 自行累加。
		return val.Float64()
	case *big.Int:
		if val == nil {
			return 0
		}
		f, _ := new(big.Float).SetInt(val).Float64()
		return f
	}
	return 0
}