feat: updates to btree/index/query/sstable/table

2025-10-11 13:19:26 +08:00
parent c8cbe4178f
commit 03ec262ca5
7 changed files with 1229 additions and 54 deletions
--- a/btree_test.go
+++ b/btree_test.go
@@ -132,6 +132,386 @@ func TestBTreeSerialization(t *testing.T) {
 	t.Log("Serialization test passed!")
 }

+// TestBTreeForEach 测试升序迭代
+func TestBTreeForEach(t *testing.T) {
+	// 创建测试文件
+	file, err := os.Create("test_foreach.sst")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove("test_foreach.sst")
+
+	// 构建 B+Tree
+	builder := NewBTreeBuilder(file, 256)
+	for i := int64(1); i <= 100; i++ {
+		err := builder.Add(i, i*100, int32(i*10))
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	rootOffset, err := builder.Build()
+	if err != nil {
+		t.Fatal(err)
+	}
+	file.Close()
+
+	// 打开并 mmap
+	file, _ = os.Open("test_foreach.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	// 测试 1: 完整升序迭代
+	t.Run("Complete", func(t *testing.T) {
+		var keys []int64
+		var offsets []int64
+		var sizes []int32
+
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			keys = append(keys, key)
+			offsets = append(offsets, offset)
+			sizes = append(sizes, size)
+			return true
+		})
+
+		// 验证数量
+		if len(keys) != 100 {
+			t.Errorf("Expected 100 keys, got %d", len(keys))
+		}
+
+		// 验证顺序（升序）
+		for i := 0; i < len(keys)-1; i++ {
+			if keys[i] >= keys[i+1] {
+				t.Errorf("Keys not in ascending order: keys[%d]=%d, keys[%d]=%d",
+					i, keys[i], i+1, keys[i+1])
+			}
+		}
+
+		// 验证第一个和最后一个
+		if keys[0] != 1 {
+			t.Errorf("Expected first key=1, got %d", keys[0])
+		}
+		if keys[99] != 100 {
+			t.Errorf("Expected last key=100, got %d", keys[99])
+		}
+
+		// 验证 offset 和 size
+		for i, key := range keys {
+			expectedOffset := key * 100
+			expectedSize := int32(key * 10)
+			if offsets[i] != expectedOffset {
+				t.Errorf("Key %d: expected offset %d, got %d", key, expectedOffset, offsets[i])
+			}
+			if sizes[i] != expectedSize {
+				t.Errorf("Key %d: expected size %d, got %d", key, expectedSize, sizes[i])
+			}
+		}
+	})
+
+	// 测试 2: 提前终止
+	t.Run("EarlyTermination", func(t *testing.T) {
+		var keys []int64
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			keys = append(keys, key)
+			return len(keys) < 5 // 只收集 5 个
+		})
+
+		if len(keys) != 5 {
+			t.Errorf("Expected 5 keys, got %d", len(keys))
+		}
+		if keys[0] != 1 || keys[4] != 5 {
+			t.Errorf("Expected keys [1,2,3,4,5], got %v", keys)
+		}
+	})
+
+	// 测试 3: 条件过滤
+	t.Run("ConditionalFilter", func(t *testing.T) {
+		var evenKeys []int64
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			if key%2 == 0 {
+				evenKeys = append(evenKeys, key)
+			}
+			return true
+		})
+
+		if len(evenKeys) != 50 {
+			t.Errorf("Expected 50 even keys, got %d", len(evenKeys))
+		}
+
+		// 验证都是偶数
+		for _, key := range evenKeys {
+			if key%2 != 0 {
+				t.Errorf("Key %d is not even", key)
+			}
+		}
+	})
+
+	// 测试 4: 查找第一个满足条件的
+	t.Run("FindFirst", func(t *testing.T) {
+		var foundKey int64
+		count := 0
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			count++
+			if key > 50 {
+				foundKey = key
+				return false // 找到后停止
+			}
+			return true
+		})
+
+		if foundKey != 51 {
+			t.Errorf("Expected to find key 51, got %d", foundKey)
+		}
+		if count != 51 {
+			t.Errorf("Expected to iterate 51 times, got %d", count)
+		}
+	})
+
+	// 测试 5: 与 GetAllKeys 结果一致性
+	t.Run("ConsistencyWithGetAllKeys", func(t *testing.T) {
+		var iterKeys []int64
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			iterKeys = append(iterKeys, key)
+			return true
+		})
+
+		allKeys := reader.GetAllKeys()
+
+		if len(iterKeys) != len(allKeys) {
+			t.Errorf("Length mismatch: ForEach=%d, GetAllKeys=%d", len(iterKeys), len(allKeys))
+		}
+
+		for i := range iterKeys {
+			if iterKeys[i] != allKeys[i] {
+				t.Errorf("Key mismatch at index %d: ForEach=%d, GetAllKeys=%d",
+					i, iterKeys[i], allKeys[i])
+			}
+		}
+	})
+}
+
+// TestBTreeForEachDesc 测试降序迭代
+func TestBTreeForEachDesc(t *testing.T) {
+	// 创建测试文件
+	file, err := os.Create("test_foreach_desc.sst")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove("test_foreach_desc.sst")
+
+	// 构建 B+Tree
+	builder := NewBTreeBuilder(file, 256)
+	for i := int64(1); i <= 100; i++ {
+		err := builder.Add(i, i*100, int32(i*10))
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	rootOffset, err := builder.Build()
+	if err != nil {
+		t.Fatal(err)
+	}
+	file.Close()
+
+	// 打开并 mmap
+	file, _ = os.Open("test_foreach_desc.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	// 测试 1: 完整降序迭代
+	t.Run("Complete", func(t *testing.T) {
+		var keys []int64
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			keys = append(keys, key)
+			return true
+		})
+
+		// 验证数量
+		if len(keys) != 100 {
+			t.Errorf("Expected 100 keys, got %d", len(keys))
+		}
+
+		// 验证顺序（降序）
+		for i := 0; i < len(keys)-1; i++ {
+			if keys[i] <= keys[i+1] {
+				t.Errorf("Keys not in descending order: keys[%d]=%d, keys[%d]=%d",
+					i, keys[i], i+1, keys[i+1])
+			}
+		}
+
+		// 验证第一个和最后一个
+		if keys[0] != 100 {
+			t.Errorf("Expected first key=100, got %d", keys[0])
+		}
+		if keys[99] != 1 {
+			t.Errorf("Expected last key=1, got %d", keys[99])
+		}
+	})
+
+	// 测试 2: 获取最新的 N 条记录（时序数据库常见需求）
+	t.Run("GetLatestN", func(t *testing.T) {
+		var latestKeys []int64
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			latestKeys = append(latestKeys, key)
+			return len(latestKeys) < 10 // 只取最新的 10 条
+		})
+
+		if len(latestKeys) != 10 {
+			t.Errorf("Expected 10 keys, got %d", len(latestKeys))
+		}
+
+		// 验证是最新的 10 条（100, 99, 98, ..., 91）
+		for i, key := range latestKeys {
+			expected := int64(100 - i)
+			if key != expected {
+				t.Errorf("latestKeys[%d]: expected %d, got %d", i, expected, key)
+			}
+		}
+	})
+
+	// 测试 3: 与 GetAllKeysDesc 结果一致性
+	t.Run("ConsistencyWithGetAllKeysDesc", func(t *testing.T) {
+		var iterKeys []int64
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			iterKeys = append(iterKeys, key)
+			return true
+		})
+
+		allKeys := reader.GetAllKeysDesc()
+
+		if len(iterKeys) != len(allKeys) {
+			t.Errorf("Length mismatch: ForEachDesc=%d, GetAllKeysDesc=%d", len(iterKeys), len(allKeys))
+		}
+
+		for i := range iterKeys {
+			if iterKeys[i] != allKeys[i] {
+				t.Errorf("Key mismatch at index %d: ForEachDesc=%d, GetAllKeysDesc=%d",
+					i, iterKeys[i], allKeys[i])
+			}
+		}
+	})
+
+	// 测试 4: 降序查找第一个满足条件的
+	t.Run("FindFirstDesc", func(t *testing.T) {
+		var foundKey int64
+		count := 0
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			count++
+			if key < 50 {
+				foundKey = key
+				return false // 找到后停止
+			}
+			return true
+		})
+
+		if foundKey != 49 {
+			t.Errorf("Expected to find key 49, got %d", foundKey)
+		}
+		if count != 52 { // 100, 99, ..., 50, 49
+			t.Errorf("Expected to iterate 52 times, got %d", count)
+		}
+	})
+}
+
+// TestBTreeForEachEmpty 测试空树的迭代
+func TestBTreeForEachEmpty(t *testing.T) {
+	// 创建空的 B+Tree
+	file, _ := os.Create("test_empty.sst")
+	defer os.Remove("test_empty.sst")
+
+	builder := NewBTreeBuilder(file, 256)
+	rootOffset, _ := builder.Build()
+	file.Close()
+
+	file, _ = os.Open("test_empty.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	// 测试升序迭代
+	t.Run("ForEach", func(t *testing.T) {
+		called := false
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			called = true
+			return true
+		})
+
+		if called {
+			t.Error("Callback should not be called on empty tree")
+		}
+	})
+
+	// 测试降序迭代
+	t.Run("ForEachDesc", func(t *testing.T) {
+		called := false
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			called = true
+			return true
+		})
+
+		if called {
+			t.Error("Callback should not be called on empty tree")
+		}
+	})
+}
+
+// TestBTreeForEachSingle 测试单个元素的迭代
+func TestBTreeForEachSingle(t *testing.T) {
+	// 创建只有一个元素的 B+Tree
+	file, _ := os.Create("test_single.sst")
+	defer os.Remove("test_single.sst")
+
+	builder := NewBTreeBuilder(file, 256)
+	builder.Add(42, 4200, 420)
+	rootOffset, _ := builder.Build()
+	file.Close()
+
+	file, _ = os.Open("test_single.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	// 测试升序迭代
+	t.Run("ForEach", func(t *testing.T) {
+		var keys []int64
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			keys = append(keys, key)
+			if offset != 4200 || size != 420 {
+				t.Errorf("Unexpected data: offset=%d, size=%d", offset, size)
+			}
+			return true
+		})
+
+		if len(keys) != 1 || keys[0] != 42 {
+			t.Errorf("Expected single key 42, got %v", keys)
+		}
+	})
+
+	// 测试降序迭代
+	t.Run("ForEachDesc", func(t *testing.T) {
+		var keys []int64
+		reader.ForEachDesc(func(key int64, offset int64, size int32) bool {
+			keys = append(keys, key)
+			return true
+		})
+
+		if len(keys) != 1 || keys[0] != 42 {
+			t.Errorf("Expected single key 42, got %v", keys)
+		}
+	})
+}
+
 func BenchmarkBTreeGet(b *testing.B) {
 	// 构建测试数据
 	file, _ := os.Create("bench.sst")
@@ -159,3 +539,86 @@ func BenchmarkBTreeGet(b *testing.B) {
 		reader.Get(key)
 	}
 }
+
+// BenchmarkBTreeForEach 性能测试：完整迭代
+func BenchmarkBTreeForEach(b *testing.B) {
+	file, _ := os.Create("bench_foreach.sst")
+	defer os.Remove("bench_foreach.sst")
+
+	builder := NewBTreeBuilder(file, 256)
+	for i := int64(1); i <= 10000; i++ {
+		builder.Add(i, i*100, 100)
+	}
+	rootOffset, _ := builder.Build()
+	file.Close()
+
+	file, _ = os.Open("bench_foreach.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	b.ResetTimer()
+	for b.Loop() {
+		count := 0
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			count++
+			return true
+		})
+	}
+}
+
+// BenchmarkBTreeForEachEarlyTermination 性能测试：提前终止
+func BenchmarkBTreeForEachEarlyTermination(b *testing.B) {
+	file, _ := os.Create("bench_foreach_early.sst")
+	defer os.Remove("bench_foreach_early.sst")
+
+	builder := NewBTreeBuilder(file, 256)
+	for i := int64(1); i <= 100000; i++ {
+		builder.Add(i, i*100, 100)
+	}
+	rootOffset, _ := builder.Build()
+	file.Close()
+
+	file, _ = os.Open("bench_foreach_early.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	b.ResetTimer()
+	for b.Loop() {
+		count := 0
+		reader.ForEach(func(key int64, offset int64, size int32) bool {
+			count++
+			return count < 10 // 只读取前 10 个
+		})
+	}
+}
+
+// BenchmarkBTreeGetAllKeys vs ForEach 对比
+func BenchmarkBTreeGetAllKeys(b *testing.B) {
+	file, _ := os.Create("bench_getall.sst")
+	defer os.Remove("bench_getall.sst")
+
+	builder := NewBTreeBuilder(file, 256)
+	for i := int64(1); i <= 10000; i++ {
+		builder.Add(i, i*100, 100)
+	}
+	rootOffset, _ := builder.Build()
+	file.Close()
+
+	file, _ = os.Open("bench_getall.sst")
+	defer file.Close()
+	mmapData, _ := mmap.Map(file, mmap.RDONLY, 0)
+	defer mmapData.Unmap()
+
+	reader := NewBTreeReader(mmapData, rootOffset)
+
+	b.ResetTimer()
+	for b.Loop() {
+		_ = reader.GetAllKeys()
+	}
+}