TsDB2/DB/core/query.go

package core

import (
	"slices"

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

// Query 查询操作对象
type Query struct {
	table      *Table
	selectCols []string    // SELECT 字段列表
	whereCond  Condition   // WHERE 条件
	orderBys   []SortField // ORDER BY
	limit      int         // LIMIT
	offset     int         // OFFSET

	// wantMask 投影掩码 (master 下标 → 是否解码), selectCols 与 orderBys 的并集;
	// 排序键从解码后的行里取, ORDER BY 列必须保留在掩码中。nil = 不过滤 (全列)。
	wantMask  []bool
	wantBuilt bool
}

type SortField struct {
	Field      string
	Descending bool
}

// keyRange 索引扫描的单个键区间。
type keyRange struct {
	startKey       []byte
	endKey         []byte
	startInclusive bool
	endInclusive   bool
}

// scanPlan 索引扫描计划。
// ranges 通常为单区间; WHERE ... IN (...) 落在可用分量上时展开成多区间
// (每个值一个精确区间), 空切片表示"结果必为空" (如 IN 空集合)。
type scanPlan struct {
	meta          *IndexMeta
	handler       index.IndexHandler
	ranges        []keyRange
	rowStore      bool      // 该索引是否承载行数据 (主键/时序=true)
	remainingCond Condition // 需要行级过滤的剩余条件
}

// Query 返回查询操作对象
func (table *Table) Query() *Query {
	return &Query{table: table}
}

func (q *Query) Select(cols ...string) *Query {
	q.selectCols = cols
	q.wantBuilt = false // 投影掩码随 selectCols 失效
	return q
}
func (q *Query) Where(cond Condition) *Query { q.whereCond = cond; return q }
func (q *Query) Limit(n int) *Query          { q.limit = n; return q }
func (q *Query) Offset(n int) *Query         { q.offset = n; return q }

func (q *Query) OrderBy(field string, descending ...bool) *Query {
	desc := false
	if len(descending) > 0 {
		desc = descending[0]
	}
	q.orderBys = append(q.orderBys, SortField{Field: field, Descending: desc})
	q.wantBuilt = false // 投影掩码随 orderBys 失效
	return q
}

// rowWithPK 把结果行与它的 pkID 绑在一起 —— 排序/翻页会重排结果集, pkID 必须
// 跟着行一起移动, 否则回调会把 A 行的 pkID 配到 B 行的数据上。
type rowWithPK struct {
	pkID uint64
	row  []types.FieldValue
}

// ======================== 核心查询 ========================

// Find 执行查询，返回多行
func (q *Query) Find() ([][]types.FieldValue, error) {
	var results [][]types.FieldValue
	err := q.scan(func(row []types.FieldValue) bool {
		results = append(results, row)
		return true
	})
	if err != nil {
		return nil, err
	}
	// 投影
	if len(q.selectCols) > 0 {
		results = q.project(results)
	}
	return results, nil
}

// Scan 流式扫描查询, 回调函数逐行处理数据, 返回 false 提前终止
// 适用于不需要全量收集结果的场景, 内存友好
func (q *Query) Scan(callback func(row []types.FieldValue) bool) error {
	return q.scan(callback)
}

// scan 内部流式扫描核心方法
// 自动处理: 索引选择 → 排序 → OFFSET → LIMIT → 回调
func (q *Query) scan(callback func(row []types.FieldValue) bool) error {
	return q.scanWithPKID(func(_ uint64, row []types.FieldValue) bool {
		return callback(row)
	})
}

// ScanWithPKID 流式扫描, 回调同时拿到行主键号 pkID 与列值。
//
// 迁移/导出工具必须用这个而不是 Scan: 行数据里不含 pkID (它是行的 8 字节前缀,
// 不是字段), 只拿字段值就没法还原"这一行的主键号是多少" —— 导出到 SQL 时要么
// 丢掉真实主键, 要么被迫用某个列值冒充 (列值可能是 0、可能被改过)。
func (q *Query) ScanWithPKID(callback func(pkID uint64, row []types.FieldValue) bool) error {
	return q.scanWithPKID(callback)
}

// scanWithPKID 与 scan 完全一致, 但额外把每行的真实 pkID 交给回调。
//
// 为什么需要: 行数据里不含 pkID (它是行的 8 字节前缀, 不是字段), 普通 Scan 回调
// 只拿得到字段值。删除/更新必须按真实 pkID 定位行, 旧实现只能退而求其次地
// "取第一个 uint64 类型的字段值当 pkID" —— 主键列不是行内第一个 8 字节列时,
// 它删掉的是另一个主键指向的行, 匹配上的行反而一条都没动, 而且不报错。
func (q *Query) scanWithPKID(callback func(pkID uint64, row []types.FieldValue) bool) error {
	defer q.table.db.lockRead()() // 读锁: 纯查询可并发; 被 Deleter.Exec 调用时 (外层已持写锁) 同 goroutine 跳过
	plan := q.buildScanPlan()

	// 判断是否需要内存排序 (OrderBy 字段与索引不一致时)
	needSort := false
	if len(q.orderBys) > 0 {
		if plan == nil {
			needSort = true
		} else {
			idxField := q.fieldNameFromMeta(plan.meta)
			// 免排序只适用于单区间: 多区间 (IN 展开) 的输出是"各区间有序、
			// 全局无序"的拼接, 必须交给下面的排序管线。
			if !(len(q.orderBys) == 1 && q.orderBys[0].Field == idxField && len(plan.ranges) == 1) {
				needSort = true
			}
		}
	}

	// 需要排序时, 必须先收集全部再排序, 无法流式 LIMIT
	if needSort {
		specs := q.buildOrderSpecs()
		// ORDER BY + LIMIT: 流式 top-k, 只保留 limit+offset 行的小顶堆,
		// 内存 O(k) 且免全量排序 (旧实现物化全部结果再整体排序)。
		if q.limit > 0 && len(specs) > 0 {
			return q.topKScan(callback, plan, specs)
		}
		var all []rowWithPK
		err := q.executeScan(func(pkID uint64, row []types.FieldValue) bool {
			all = append(all, rowWithPK{pkID, row})
			return true
		}, plan, 0)
		if err != nil {
			return err
		}
		// 排序必须连 pkID 一起搬: 只排 []FieldValue 会让回调拿到错位的 pkID。
		// decoration: 每行只按名提取一次排序键, 键必须随行进数组 —— 排序会重排
		// 切片元素, 平行键数组会因此错位 (索引指向交换后的别的行)。
		// 各行列集合可能因 NULL 跳过而不同, 键按名取、缺失用 nil 占位。
		if len(specs) > 0 {
			dec := make([]struct {
				item rowWithPK
				keys []interface{}
			}, len(all))
			for i := range all {
				k := make([]interface{}, len(specs))
				for j, s := range specs {
					k[j], _ = lookupFieldValue(all[i].row, []byte(s.field))
				}
				dec[i] = struct {
					item rowWithPK
					keys []interface{}
				}{all[i], k}
			}
			slices.SortStableFunc(dec, func(a, b struct {
				item rowWithPK
				keys []interface{}
			}) int {
				return compareKeys(a.keys, b.keys, specs)
			})
			all = all[:0]
			for i := range dec {
				all = append(all, dec[i].item)
			}
		}
		// OFFSET
		if q.offset > 0 {
			if q.offset >= len(all) {
				return nil
			}
			all = all[q.offset:]
		}
		// LIMIT
		if q.limit > 0 && q.limit < len(all) {
			all = all[:q.limit]
		}
		for _, item := range all {
			if !callback(item.pkID, item.row) {
				return nil
			}
		}
		return nil
	}

	// 不需要排序: 直接流式传递 callback, 支持 LIMIT 提前终止。
	// 回表预算 = LIMIT+OFFSET: 达到即收工 (OFFSET 行也要交付给下面的跳行逻辑)。
	budget := 0
	if q.limit > 0 {
		budget = q.limit + q.offset
	}
	seen := 0
	return q.executeScan(func(pkID uint64, row []types.FieldValue) bool {
		if seen < q.offset {
			seen++
			return true // 跳过 OFFSET
		}
		seen++
		if !callback(pkID, row) {
			return false // 用户要求停止
		}
		// LIMIT 检查 (0 表示无限制)
		if q.limit > 0 && seen-q.offset >= q.limit {
			return false // 已满足, 提前终止扫描
		}
		return true
	}, plan, budget)
}

// deserialize 行解码入口: 无投影走全列解码; 有投影只解码选中列 (含 ORDER BY 列,
// 排序键要从解码后的行里取)。掩码按 master 下标, 随查询构建器惰性构建一次。
func (q *Query) deserialize(rowData []byte) []types.FieldValue {
	if len(q.selectCols) == 0 {
		return q.table.deserializeRow(rowData)
	}
	if !q.wantBuilt {
		mask := make([]bool, len(q.table.Field))
		mark := func(name string) {
			if i := q.table.findFieldIndex(name); i >= 0 {
				mask[i] = true
			}
		}
		for _, c := range q.selectCols {
			mark(c)
		}
		for _, ob := range q.orderBys {
			mark(ob.Field)
		}
		q.wantMask = mask
		q.wantBuilt = true
	}
	return q.table.deserializeRowSelect(rowData, q.wantMask)
}

// ======================== 投影 ========================

func (q *Query) project(results [][]types.FieldValue) [][]types.FieldValue {
	colSet := make(map[string]bool)
	for _, c := range q.selectCols {
		colSet[c] = true
	}
	for i, row := range results {
		var filtered []types.FieldValue
		for _, fv := range row {
			if colSet[string(fv.FieldName)] {
				filtered = append(filtered, fv)
			}
		}
		results[i] = filtered
	}
	return results
}