package core import ( "TsDB2/DB/core/index" "TsDB2/DB/core/page" "TsDB2/DB/core/types" "bytes" "encoding/binary" "sort" ) // Query 查询操作对象 type Query struct { table *Table selectCols []string // SELECT 字段列表 whereCond Condition // WHERE 条件 orderBys []SortField // ORDER BY limit int // LIMIT offset int // OFFSET groupBys []string // GROUP BY havingCond Condition // HAVING } type SortField struct { Field string Descending bool } // scanPlan 索引扫描计划 type scanPlan struct { meta *IndexMeta handler index.IndexHandler startKey []byte endKey []byte startInclusive bool endInclusive bool isPrimary bool // 主键/时序索引=true remainingCond Condition // 需要行级过滤的剩余条件 } // Query 返回查询操作对象 func (table *Table) Query() *Query { return &Query{table: table} } func (q *Query) Select(cols ...string) *Query { q.selectCols = cols; 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) GroupBy(cols ...string) *Query { q.groupBys = cols; return q } func (q *Query) Having(cond Condition) *Query { q.havingCond = cond; 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}) return q } // ======================== 索引键编码 ======================== // encodeIndexKey 将条件值编码为索引键字节 (BigEndian, 与 writeFieldValue 一致) func encodeIndexKey(value interface{}, ft types.FieldType) []byte { size := types.GetSize(ft) buf := make([]byte, size) switch ft { case types.FieldType_1byte: buf[0] = value.(byte) case types.FieldType_2byte: binary.BigEndian.PutUint16(buf, uint16(value.(int16))) case types.FieldType_4byte: binary.BigEndian.PutUint32(buf, uint32(value.(int32))) case types.FieldType_8byte: binary.BigEndian.PutUint64(buf, value.(uint64)) default: // 字符串类型 b := value.([]byte) copy(buf, b) } return buf } // ======================== 索引匹配 ======================== // 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 "" } // fieldTypeFromMeta 获取索引第一个字段的类型 func (q *Query) fieldTypeFromMeta(meta *IndexMeta) types.FieldType { fieldIdx := int(meta.FieldIDs[0]) if fieldIdx < len(q.table.Field) { return q.table.Field[fieldIdx].FieldType } return 0 } // condToRange 尝试从条件中提取扫描范围, 返回 (startKey, endKey, startInc, endInc, ok) func (q *Query) condToRange(cond Condition, fieldName string, ft types.FieldType) ([]byte, []byte, bool, bool, Condition) { switch c := cond.(type) { case EqCondition: if c.Field == fieldName { key := encodeIndexKey(c.Value, ft) return key, key, true, true, nil } case GeCondition: if c.Field == fieldName { return encodeIndexKey(c.Value, ft), nil, true, false, nil } case GtCondition: if c.Field == fieldName { return encodeIndexKey(c.Value, ft), nil, false, false, nil } case LeCondition: if c.Field == fieldName { return nil, encodeIndexKey(c.Value, ft), false, true, nil } case LtCondition: if c.Field == fieldName { return nil, encodeIndexKey(c.Value, ft), false, false, nil } case AndCondition: // 合并所有子条件的范围 var sk, ek []byte var si, ei bool = true, true // 默认全包含 var hasStart, hasEnd bool var remaining []Condition for _, sub := range c.Conditions { s, e, sinc, einc, rem2 := q.condToRange(sub, fieldName, ft) if rem2 == nil && (s != nil || e != nil) { // 此子条件可完全用于索引 if s != nil { if !hasStart || (sinc && bytes.Compare(s, sk) > 0) || (!sinc && bytes.Compare(s, sk) >= 0) { sk = s si = sinc hasStart = true } } if e != nil { if !hasEnd || (einc && bytes.Compare(e, ek) < 0) || (!einc && bytes.Compare(e, ek) <= 0) { ek = e ei = einc hasEnd = true } } } else if s != nil || e != nil { // 此子条件部分可用于索引 if s != nil { if !hasStart || (sinc && bytes.Compare(s, sk) > 0) || (!sinc && bytes.Compare(s, sk) >= 0) { sk = s si = sinc hasStart = true } } if e != nil { if !hasEnd || (einc && bytes.Compare(e, ek) < 0) || (!einc && bytes.Compare(e, ek) <= 0) { ek = e ei = einc hasEnd = true } } if rem2 != nil { remaining = append(remaining, rem2) } } else { // 此子条件不能用于索引 remaining = append(remaining, sub) } } if hasStart || hasEnd { var rem Condition if len(remaining) == 1 { rem = remaining[0] } else if len(remaining) > 1 { rem = And(remaining...) } return sk, ek, si, ei, rem } } return nil, nil, false, false, cond } // buildScanPlan 分析 WHERE 条件, 选择最优索引 func (q *Query) buildScanPlan() *scanPlan { if q.whereCond == nil { return nil } if q.table.indexMeta == nil { return nil } var bestPlan *scanPlan for i := range q.table.indexMeta.metas { m := &q.table.indexMeta.metas[i] if m.Flag&0x80 != 0 { continue } isPrimary := m.IndexType == page.IndexTypePrimary || m.IndexType == page.IndexTypeTime fieldName := q.fieldNameFromMeta(m) ft := q.fieldTypeFromMeta(m) if fieldName == "" { continue } sk, ek, si, ei, rem := q.condToRange(q.whereCond, fieldName, ft) if sk == nil && ek == nil { continue // 此索引不匹配 } handler, ok := index.Compose.GetHandler(m.IndexType) if !ok { continue } plan := &scanPlan{ meta: m, handler: handler, startKey: sk, endKey: ek, startInclusive: si, endInclusive: ei, isPrimary: isPrimary, remainingCond: rem, } // 优先选主键/时序索引 (避免回查) if isPrimary || bestPlan == nil { bestPlan = plan } if isPrimary { break // 主键索引最优, 不再找 } } return bestPlan } // ======================== 核心查询 ======================== // 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 { plan := q.buildScanPlan() // 判断是否需要内存排序 (OrderBy 字段与索引不一致时) needSort := false if len(q.orderBys) > 0 { if plan == nil { needSort = true } else { idxField := q.fieldNameFromMeta(plan.meta) if !(len(q.orderBys) == 1 && q.orderBys[0].Field == idxField) { needSort = true } } } // 需要排序时, 必须先收集全部再排序, 无法流式 LIMIT if needSort { var all [][]types.FieldValue err := q.executeScan(func(row []types.FieldValue) bool { all = append(all, row) return true }, plan) if err != nil { return err } q.sortResults(all) // 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 _, row := range all { if !callback(row) { return nil } } return nil } // 不需要排序: 直接流式传递 callback, 支持 LIMIT 提前终止 seen := 0 return q.executeScan(func(row []types.FieldValue) bool { if seen < q.offset { seen++ return true // 跳过 OFFSET } seen++ if !callback(row) { return false // 用户要求停止 } // LIMIT 检查 (0 表示无限制) if q.limit > 0 && seen-q.offset >= q.limit { return false // 已满足, 提前终止扫描 } return true }, plan) } // executeScan 执行索引扫描或全表扫描, 将反序列化后的行传递给回调 func (q *Query) executeScan(callback func(row []types.FieldValue) bool, plan *scanPlan) error { if plan != nil { return q.indexScan(plan, callback) } return q.fullTableScan(callback) } // indexScan 使用索引扫描 (流式回调模式) func (q *Query) indexScan(plan *scanPlan, callback func(row []types.FieldValue) bool) error { // 判断 OrderBy 是否需要反向 reverse := false if len(q.orderBys) == 1 { idxField := q.fieldNameFromMeta(plan.meta) if q.orderBys[0].Field == idxField && q.orderBys[0].Descending { reverse = true } } if plan.isPrimary { // 主键/时序索引: RangeScan 直接返回 storeRow = [pkID(8B) + rowData] return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage, plan.startKey, plan.endKey, plan.startInclusive, plan.endInclusive, reverse, func(row []byte) bool { if len(row) < 8 { return true } rowData := row[8:] if plan.remainingCond != nil { if !plan.remainingCond.Match(rowData, q.table) { return true } } return callback(q.table.deserializeRow(rowData)) }) } // 二级索引: RangeScan 返回 pkID(8B BE), 需回查主键 return plan.handler.RangeScan(q.table.handle, plan.meta.EntryPage, plan.startKey, plan.endKey, plan.startInclusive, plan.endInclusive, reverse, func(row []byte) bool { if len(row) < 8 { return true } pkID := binary.BigEndian.Uint64(row[:8]) storeRow, found := q.table.indexMeta.ResolveByPK(pkID, q.table.Field) if !found { return true } rowData := storeRow[8:] if plan.remainingCond != nil { if !plan.remainingCond.Match(rowData, q.table) { return true } } return callback(q.table.deserializeRow(rowData)) }) } // fullTableScan 全表扫描 (流式回调模式) func (q *Query) fullTableScan(callback func(row []types.FieldValue) bool) error { // 找主键/时序索引 var primMeta *IndexMeta var primHandler index.IndexHandler for i := range q.table.indexMeta.metas { m := &q.table.indexMeta.metas[i] if m.Flag&0x80 != 0 { continue } if m.IndexType == page.IndexTypePrimary || m.IndexType == page.IndexTypeTime { primMeta = m primHandler, _ = index.Compose.GetHandler(m.IndexType) break } } if primMeta == nil || primHandler == nil { return nil } dpIDs := primHandler.GetDataPageIDs(q.table.handle, primMeta.EntryPage) for _, dpID := range dpIDs { rows := primHandler.ReadDataPageRows(q.table.handle, dpID) for _, storeRow := range rows { if len(storeRow) < 8 { continue } rowData := storeRow[8:] if q.whereCond != nil { if !q.whereCond.Match(rowData, q.table) { continue } } if !callback(q.table.deserializeRow(rowData)) { return nil // 回调要求提前终止 } } } return nil } // ======================== 排序 ======================== func (q *Query) sortResults(results [][]types.FieldValue) { sort.SliceStable(results, func(i, j int) bool { for _, ob := range q.orderBys { fi := q.table.findFieldIndex(ob.Field) if fi < 0 { continue } vi := findFieldValue(results[i], ob.Field) vj := findFieldValue(results[j], ob.Field) if vi == nil || vj == nil { continue } cmp := compareFieldValues(vi, vj) if cmp == 0 { continue } if ob.Descending { return cmp > 0 } return cmp < 0 } return false }) } func findFieldValue(row []types.FieldValue, field string) interface{} { for _, fv := range row { if string(fv.FieldName) == field { return fv.Value } } return nil } func compareFieldValues(a, b interface{}) int { switch va := a.(type) { case uint64: vb := b.(uint64) if va < vb { return -1 } if va > vb { return 1 } return 0 case int16: vb := a.(int16) vbb := b.(int16) if vb < vbb { return -1 } if vb > vbb { return 1 } return 0 case int32: vb := b.(int32) if va < vb { return -1 } if va > vb { return 1 } return 0 case byte: vb := b.(byte) if va < vb { return -1 } if va > vb { return 1 } return 0 case []byte: vb := b.([]byte) return bytes.Compare(va, vb) } return 0 } // ======================== 投影 ======================== 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 } // ======================== 流式聚合 ======================== // 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 流式计数 (不收集行,扫描过程中直接累加) func (q *Query) Count() (int64, error) { var count int64 err := q.scan(func(row []types.FieldValue) bool { count++ return true }) return count, err } // Sum 流式求和 (不收集行,扫描过程中直接累加) func (q *Query) Sum(col string) (float64, error) { 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 } // Max 流式求最大值 (不收集行,扫描过程中维护当前最大值) func (q *Query) Max(col string) (interface{}, error) { 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) > 0 { maxVal = fv.Value } break } } return true }) return maxVal, err } // Min 流式求最小值 (不收集行,扫描过程中维护当前最小值) func (q *Query) Min(col string) (interface{}, error) { 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) < 0 { minVal = fv.Value } break } } return true }) return minVal, err } func toFloat64(v interface{}) float64 { switch val := v.(type) { case uint64: return float64(val) case int16: return float64(val) case int32: return float64(val) case byte: return float64(val) } return 0 }