TsDB2/DB/core/query_optimizer_test.go

package core

import (
	"sort"
	"testing"

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

// qopNewTable 建测试表: id (uint64 主键) + age (int64, B+ 二级索引 idx_age)
// + score (int64, Hash 二级索引 idx_score) + name (string 32B, 无索引)。
func qopNewTable(t *testing.T, db *DB, name string) *Table {
	t.Helper()
	tbl, err := db.CreateTable(name, []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("id")},
		{FieldType: types.FieldType_Int64, FieldName: []byte("age")},
		{FieldType: types.FieldType_Int64, FieldName: []byte("score")},
		{FieldType: types.FieldType_String, FieldName: []byte("name"), Flag: types.WidthCodeFor(32)},
	})
	if err != nil {
		t.Fatalf("建表 %s: %v", name, err)
	}
	if _, err := tbl.CreateIndex("pk", page.IndexTypePrimary, []byte{0}); err != nil {
		t.Fatalf("建主键索引: %v", err)
	}
	if _, err := tbl.CreateIndex("idx_age", page.IndexTypeBTree, []byte{1}); err != nil {
		t.Fatalf("建 B+ 索引 idx_age: %v", err)
	}
	if _, err := tbl.CreateIndex("idx_score", page.IndexTypeHash, []byte{2}); err != nil {
		t.Fatalf("建 Hash 索引 idx_score: %v", err)
	}
	return tbl
}

// qopInsert 插入一行 (id 显式, age/score/name 按 i 生成)。
func qopInsert(t *testing.T, tbl *Table, id uint64, age, score int64, name string) {
	t.Helper()
	if _, err := tbl.InsertRowWithID(id, []types.FieldValue{
		{FieldName: []byte("id"), Value: id},
		{FieldName: []byte("age"), Value: age},
		{FieldName: []byte("score"), Value: score},
		{FieldName: []byte("name"), Value: name},
	}); err != nil {
		t.Fatalf("插入 id=%d: %v", id, err)
	}
}

// qopPKKey 主键索引键 = pkID 的 8B 大端。
func qopPKKey(id uint64) []byte {
	k := make([]byte, 8)
	types.EncodeValue(types.FieldType_Uint64, 0, id, k)
	return k
}

// qopSeed 灌 n 行: age = i%50, score = i%100, 便于构造多值命中。
func qopSeed(t *testing.T, tbl *Table, n int) {
	t.Helper()
	for i := 1; i <= n; i++ {
		qopInsert(t, tbl, uint64(i), int64(i%50), int64(i%100), "")
	}
}

// collectAges 执行查询并返回 (pkID → age) 映射, 用于与全量预期比对。
func collectAges(t *testing.T, q *Query) map[uint64]int64 {
	t.Helper()
	out := map[uint64]int64{}
	err := q.ScanWithPKID(func(pkID uint64, row []types.FieldValue) bool {
		for _, fv := range row {
			if string(fv.FieldName) == "age" {
				out[pkID] = fv.Value.(int64)
			}
		}
		return true
	})
	if err != nil {
		t.Fatalf("scan: %v", err)
	}
	return out
}

// TestInUsesIndex IN 必须走索引: 规划器为每个值展开一个区间, 而不是退回全表扫描。
func TestInUsesIndex(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 200)

	// 白盒: 规划结果应为多区间计划
	vals := []interface{}{int64(7), int64(23), int64(49)}
	plan := tbl.Query().Where(In("age", vals, 0)).buildScanPlan()
	if plan == nil {
		t.Fatalf("IN(age, 3值) 应生成索引计划, 得到 nil")
	}
	if len(plan.ranges) != 3 {
		t.Fatalf("IN 3 值应展开 3 个区间, 得到 %d", len(plan.ranges))
	}
	if plan.meta == nil || string(plan.meta.IndexName[:7]) != "idx_age" {
		t.Fatalf("计划应选 idx_age, 得到 %+v", plan.meta)
	}

	// 黑盒: 结果集与逐行过滤一致
	got := collectAges(t, tbl.Query().Where(In("age", vals, 0)))
	want := map[uint64]int64{}
	for id := uint64(1); id <= 200; id++ {
		age := int64(id % 50)
		for _, v := range vals {
			if age == v.(int64) {
				want[id] = age
			}
		}
	}
	if len(got) != len(want) {
		t.Fatalf("IN 命中 %d 行, 想要 %d", len(got), len(want))
	}
	for id, age := range want {
		if got[id] != age {
			t.Fatalf("pkID=%d: 得到 age=%d, 想要 %d", id, got[id], age)
		}
	}
}

// TestInOnHashIndex Hash 索引上的等值集合: 每个值一次桶定位。
func TestInOnHashIndex(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 200)

	plan := tbl.Query().Where(In("score", []interface{}{int64(5), int64(66)}, 0)).buildScanPlan()
	if plan == nil {
		t.Fatalf("IN(score, ...) 应生成索引计划")
	}
	if len(plan.ranges) != 2 {
		t.Fatalf("应展开 2 个区间, 得到 %d", len(plan.ranges))
	}

	got := collectAges(t, tbl.Query().Where(In("score", []interface{}{int64(5), int64(66)}, 0)))
	wantCount := 0
	for id := uint64(1); id <= 200; id++ {
		s := int64(id % 100)
		if s == 5 || s == 66 {
			wantCount++
		}
	}
	if len(got) != wantCount {
		t.Fatalf("Hash IN 命中 %d 行, 想要 %d", len(got), wantCount)
	}
}

// TestInEdgeCases 空集合恒假 / 不存在的值 / 重复值不重复返回 / 乱序值。
func TestInEdgeCases(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 100)

	// 空集合: 0 行 (不能退化成查全表)
	if n, err := tbl.Query().Where(In("age", nil, 0)).Count(); err != nil || n != 0 {
		t.Fatalf("空 IN 应得 0 行, 得到 %d (%v)", n, err)
	}
	// 全部不存在的值
	if n, err := tbl.Query().Where(In("age", []interface{}{int64(999), int64(1000)}, 0)).Count(); err != nil || n != 0 {
		t.Fatalf("不存在值的 IN 应得 0 行, 得到 %d", n)
	}
	// 重复值: 每行只返回一次
	dup := []interface{}{int64(7), int64(7), int64(7)}
	got := collectAges(t, tbl.Query().Where(In("age", dup, 0)))
	// age=7 → id%50==7 → id ∈ {7,57} (1..100 内)
	if len(got) != 2 {
		t.Fatalf("重复值的 IN 应命中 2 行, 得到 %d", len(got))
	}
	// 乱序值: 结果集不变
	a := collectAges(t, tbl.Query().Where(In("age", []interface{}{int64(49), int64(7), int64(23)}, 0)))
	b := collectAges(t, tbl.Query().Where(In("age", []interface{}{int64(7), int64(23), int64(49)}, 0)))
	if len(a) != len(b) {
		t.Fatalf("乱序 IN 结果数不一致: %d vs %d", len(a), len(b))
	}
	for id := range a {
		if _, ok := b[id]; !ok {
			t.Fatalf("乱序 IN 结果集不一致: pkID %d 只在一边", id)
		}
	}
}

// TestInWithOtherConditions IN 与其它条件组合: 同列范围条件让位给行级过滤,
// 其它列条件进 remainingCond, 结果都必须正确。
func TestInWithOtherConditions(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 100)

	// IN + 同列范围: 范围条件走索引区间, IN 由行级过滤兜底
	got := collectAges(t, tbl.Query().Where(
		And(In("age", []interface{}{int64(20), int64(25), int64(30)}, 0),
			Gt("age", int64(22), 0))))
	// 期望: age ∈ {25, 30}
	if len(got) != 4 { // id%50==25 → 25,75; id%50==30 → 30,80
		t.Fatalf("IN+范围 应命中 4 行, 得到 %d", len(got))
	}
	for _, age := range got {
		if age != 25 && age != 30 {
			t.Fatalf("出现不该出现的 age=%d", age)
		}
	}

	// IN + 其它列等值 (score 走 hash 或行级过滤)
	got2 := collectAges(t, tbl.Query().Where(
		And(In("age", []interface{}{int64(10), int64(11)}, 0),
			Eq("score", int64(10), 0))))
	// age=10 & score=10 → id%50==10 且 id%100==10 → id=10, 110>100 → 仅 id=10
	if len(got2) != 1 {
		t.Fatalf("IN+其它列 应命中 1 行, 得到 %d", len(got2))
	}
	if got2[10] != 10 {
		t.Fatalf("应为 pkID=10 age=10, 得到 %+v", got2)
	}
}

// TestBatchResolveMatchesPerPK 批量回表与逐个回表结果一致 (含已删除的行号):
// 删除若干行后, 二级索引扫描回表必须自动跳过不存在的 pkID。
func TestBatchResolveMatchesPerPK(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 300)
	// 删除 1/3 的行 (交错)
	for id := uint64(2); id <= 300; id += 3 {
		if err := tbl.DeleteRow("pk", qopPKKey(id)); err != nil {
			t.Fatalf("删除 %d: %v", id, err)
		}
	}
	if err := db.Commit(); err != nil {
		t.Fatalf("commit: %v", err)
	}

	// 大范围扫 + 回表
	got := collectAges(t, tbl.Query().Where(Ge("age", int64(0), 0)))
	want := 0
	for id := uint64(1); id <= 300; id++ {
		if id%3 == 2 {
			continue // 已删除
		}
		want++
	}
	if len(got) != want {
		t.Fatalf("范围扫回表 %d 行, 想要 %d", len(got), want)
	}
	for id := range got {
		if id%3 == 2 {
			t.Fatalf("已删除的行 %d 被回表出来", id)
		}
	}

	// IN 批量回表 (排序去重路径)
	in := collectAges(t, tbl.Query().Where(
		In("age", []interface{}{int64(1), int64(2), int64(3)}, 0)))
	for id := range in {
		if id%3 == 2 {
			t.Fatalf("IN 回表带出已删除行 %d", id)
		}
	}
}

// TestTopKMatchesFullSort ORDER BY + LIMIT 的流式 top-k 与"全量排序再切片"逐行一致,
// 覆盖正/反序、OFFSET、大量重复键 (稳定性) 与 IN 多区间走排序管线的组合。
func TestTopKMatchesFullSort(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 500) // age = i%50: 每个 age 10 行, 天然大量重复键

	full := []uint64{}
	err := tbl.Query().OrderBy("age").ScanWithPKID(func(pkID uint64, _ []types.FieldValue) bool {
		full = append(full, pkID)
		return true
	})
	if err != nil {
		t.Fatalf("full sort: %v", err)
	}

	// 稳定性参照: 相等 age 按扫描顺序 (pkID 升序) 排列
	for _, tc := range []struct {
		name        string
		build       func() *Query
		wantSlice   func() []uint64
	}{
		{
			name:  "asc limit 10",
			build: func() *Query { return tbl.Query().OrderBy("age").Limit(10) },
			wantSlice: func() []uint64 {
				return full[:10]
			},
		},
		{
			name:  "desc limit 20",
			build: func() *Query { return tbl.Query().OrderBy("age", true).Limit(20) },
			wantSlice: func() []uint64 {
				// 稳定降序: 键降序, 相等键保持输入序 (pkID 升序) —
				// 不能对升序结果简单反转 (那会把相等键也倒过来)。
				type pa struct {
					pk  uint64
					age int64
				}
				in := make([]pa, 500)
				for i := 0; i < 500; i++ {
					in[i] = pa{uint64(i + 1), int64((i + 1) % 50)}
				}
				sort.SliceStable(in, func(i, j int) bool { return in[i].age > in[j].age })
				out := make([]uint64, 20)
				for i := range out {
					out[i] = in[i].pk
				}
				return out
			},
		},
		{
			name:  "asc offset 100 limit 50",
			build: func() *Query { return tbl.Query().OrderBy("age").Offset(100).Limit(50) },
			wantSlice: func() []uint64 {
				return full[100:150]
			},
		},
		{
			name: "asc limit 1 (First 场景)",
			build: func() *Query {
				return tbl.Query().OrderBy("age").Limit(1)
			},
			wantSlice: func() []uint64 {
				return full[:1]
			},
		},
	} {
		var got []uint64
		err := tc.build().ScanWithPKID(func(pkID uint64, _ []types.FieldValue) bool {
			got = append(got, pkID)
			return true
		})
		if err != nil {
			t.Fatalf("%s: %v", tc.name, err)
		}
		want := tc.wantSlice()
		if len(got) != len(want) {
			t.Fatalf("%s: %d 行, 想要 %d", tc.name, len(got), len(want))
		}
		for i := range want {
			if got[i] != want[i] {
				t.Fatalf("%s: 第 %d 位 pkID=%d, 想要 %d (稳定性破坏)", tc.name, i, got[i], want[i])
			}
		}
	}

	// IN 多区间 + ORDER BY (强制排序管线) + LIMIT
	var got []uint64
	err = tbl.Query().Where(In("age", []interface{}{int64(3), int64(1), int64(2)}, 0)).
		OrderBy("age").Limit(5).
		ScanWithPKID(func(pkID uint64, _ []types.FieldValue) bool {
			got = append(got, pkID)
			return true
		})
	if err != nil {
		t.Fatalf("IN+ORDER BY+LIMIT: %v", err)
	}
	if len(got) != 5 {
		t.Fatalf("IN+ORDER BY+LIMIT 应得 5 行, 得到 %d", len(got))
	}
	// 期望: age=1 的 10 行 (pkID 1,51,101,151,201,251,301>500→1,51,101,151,201,251... 500 内: 1,51,101,151,201,251,301,351,401,451) 取前 5
	want5 := []uint64{1, 51, 101, 151, 201}
	for i := range want5 {
		if got[i] != want5[i] {
			t.Fatalf("IN+ORDER BY 第 %d 位 = %d, 想要 %d", i, got[i], want5[i])
		}
	}
}

// TestSelectProjectionPushdown 投影下推: 只解码选中列 (+ORDER BY 列),
// 未选中列不出现; 行内 NULL 与字符串列在投影下语义不变。
func TestSelectProjectionPushdown(t *testing.T) {
	db := fixOpenDB(t)
	tbl, err := db.CreateTable("t", []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("id")},
		{FieldType: types.FieldType_Int64, FieldName: []byte("age")},
		{FieldType: types.FieldType_String, FieldName: []byte("name"), Flag: types.WidthCodeFor(32)},
	})
	if err != nil {
		t.Fatalf("建表: %v", err)
	}
	if _, err := tbl.CreateIndex("pk", page.IndexTypePrimary, []byte{0}); err != nil {
		t.Fatalf("建主键索引: %v", err)
	}
	// 交错插入: 部分行 name 为空串, 部分有值
	for i := 1; i <= 20; i++ {
		nm := ""
		if i%2 == 0 {
			nm = "emp"
		}
		if _, err := tbl.InsertRowWithID(uint64(i), []types.FieldValue{
			{FieldName: []byte("id"), Value: uint64(i)},
			{FieldName: []byte("age"), Value: int64(i * 3)},
			{FieldName: []byte("name"), Value: nm},
		}); err != nil {
			t.Fatalf("插入 %d: %v", i, err)
		}
	}
	if err := db.Commit(); err != nil {
		t.Fatalf("commit: %v", err)
	}

	// 只选 name: 行里只有 name 列
	err = tbl.Query().Select("name").Scan(func(row []types.FieldValue) bool {
		if len(row) != 1 || string(row[0].FieldName) != "name" {
			t.Fatalf("投影行应只含 name, 得到 %v", row)
		}
		return true
	})
	if err != nil {
		t.Fatalf("Select(name) scan: %v", err)
	}

	// Select + OrderBy(未选中列): 排序键列必须随投影保留, 排序结果正确
	type pair struct {
		id, age int64
	}
	var got []pair
	err = tbl.Query().Select("id").OrderBy("age", true).Limit(3).Scan(func(row []types.FieldValue) bool {
		var p pair
		for _, fv := range row {
			switch string(fv.FieldName) {
			case "id":
				p.id = int64(fv.Value.(uint64))
			case "age":
				p.age = fv.Value.(int64)
			}
		}
		got = append(got, p)
		return true
	})
	if err != nil {
		t.Fatalf("Select+OrderBy: %v", err)
	}
	// age 倒序前三: id=20,19,18 → age=60,57,54
	want := []pair{{20, 60}, {19, 57}, {18, 54}}
	for i, w := range want {
		if got[i] != w {
			t.Fatalf("第 %d 位 = %+v, 想要 %+v", i, got[i], w)
		}
	}

	// Find + Select: project 后只剩选中列
	rows, err := tbl.Query().Select("age").OrderBy("id").Limit(2).Find()
	if err != nil {
		t.Fatalf("Find+Select: %v", err)
	}
	for _, row := range rows {
		if len(row) != 1 || string(row[0].FieldName) != "age" {
			t.Fatalf("Find+Select 行应只含 age, 得到 %v", row)
		}
	}
}

// TestInCountUsesIndex IN 的 Count 也走索引区间 (不物化行)。
func TestInCountUsesIndex(t *testing.T) {
	db := fixOpenDB(t)
	tbl := qopNewTable(t, db, "t")
	qopSeed(t, tbl, 1000)

	n, err := tbl.Query().Where(In("age", []interface{}{int64(7), int64(8)}, 0)).Count()
	if err != nil {
		t.Fatalf("count: %v", err)
	}
	// age=7 → id%50==7 → 1..1000 内 20 个; age=8 → 20 个
	if n != 40 {
		t.Fatalf("IN count 应得 40, 得到 %d", n)
	}
}