address comments

Signed-off-by: Wenqi Mou <[email protected]>

Author: Tristan1900

pkg/statistics/BUILD.bazel

@@ -18,7 +18,6 @@ go_library(
         "row_sampler.go",
         "sample.go",
         "scalar.go",
-        "sort_utils.go",
         "table.go",
     ],
     importpath = "github.com/pingcap/tidb/pkg/statistics",
@@ -48,6 +47,7 @@ go_library(
         "//pkg/util/context",
         "//pkg/util/dbterror",
         "//pkg/util/fastrand",
+        "//pkg/util/generic",
         "//pkg/util/hack",
         "//pkg/util/intest",
         "//pkg/util/logutil",
@@ -78,7 +78,6 @@ go_test(
         "main_test.go",
         "sample_test.go",
         "scalar_test.go",
-        "sort_utils_test.go",
         "statistics_test.go",
         "table_test.go",
     ],

pkg/statistics/builder.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/codec"
 	"github.com/pingcap/tidb/pkg/util/collate"
+	"github.com/pingcap/tidb/pkg/util/generic"
 	"github.com/pingcap/tidb/pkg/util/memory"
 )
 
@@ -84,11 +85,11 @@ func (s *SequentialRangeChecker) IsIndexInTopNRange(idx int64) bool {
 	return false
 }
 
-// processTopNValueHeap handles the logic for a complete TopN value count using existing heap with range tracking.
-func processTopNValueHeap(topNHeap *TopNHeap[TopNWithRange], encoded []byte, curCnt float64,
+// processTopNValue handles the logic for a complete TopN value count using existing priority queue with range tracking.
+func processTopNValue(topNPriorityQueue *generic.PriorityQueue[TopNWithRange], encoded []byte, curCnt float64,
 	startIdx, endIdx int64, numTopN int, allowPruning bool, sampleFactor float64) {
 	// case 1: prune values with count of 1 under certain conditions
-	if curCnt == 1 && allowPruning && (topNHeap.Len() >= (numTopN/topNPruningThreshold) || sampleFactor > 1) {
+	if curCnt == 1 && allowPruning && (topNPriorityQueue.Len() >= (numTopN/topNPruningThreshold) || sampleFactor > 1) {
 		return
 	}
 
@@ -98,7 +99,7 @@ func processTopNValueHeap(topNHeap *TopNHeap[TopNWithRange], encoded []byte, cur
 		startIdx: startIdx,
 		endIdx:   endIdx,
 	}
-	topNHeap.Add(newItem)
+	topNPriorityQueue.Add(newItem)
 }
 
 // SortedBuilder is used to build histograms for PK and index.
@@ -430,7 +431,7 @@ func BuildHistAndTopN(
 
 	// Step1: collect topn from samples using heap and track their index ranges
 	// use heap for efficient TopN maintenance - O(log N) insertions
-	topNHeap := NewTopNHeap(numTopN, func(a, b TopNWithRange) int {
+	topNHeap := generic.NewPriorityQueue(numTopN, func(a, b TopNWithRange) int {
 		if a.Count < b.Count {
 			return -1 // min-heap: smaller counts at root
 		} else if a.Count > b.Count {
@@ -472,7 +473,7 @@ func BuildHistAndTopN(
 		// case 2, meet a different value: counting for the "current" is complete
 		sampleNDV++
 		// process the completed value using heap with range tracking
-		processTopNValueHeap(topNHeap, cur, curCnt, curStartIdx, int64(i)-1, numTopN, allowPruning, sampleFactor)
+		processTopNValue(topNHeap, cur, curCnt, curStartIdx, i-1, numTopN, allowPruning, sampleFactor)
 
 		cur, curCnt = sampleBytes, 1
 		curStartIdx = i // new value group starts at current index
@@ -485,7 +486,7 @@ func BuildHistAndTopN(
 
 	// handle the counting for the last value
 	if numTopN != 0 {
-		processTopNValueHeap(topNHeap, cur, curCnt, curStartIdx, sampleNum-1, numTopN, allowPruning, sampleFactor)
+		processTopNValue(topNHeap, cur, curCnt, curStartIdx, sampleNum-1, numTopN, allowPruning, sampleFactor)
 	}
 
 	// convert heap to sorted slice

pkg/statistics/builder_test.go

@@ -130,3 +130,37 @@ func BenchmarkBuildHistAndTopNWithLowNDV(b *testing.B) {
 		_, _, _ = BuildHistAndTopN(ctx, 256, 500, 0, collector, filedType, true, memoryTracker, false)
 	}
 }
+
+// SequentialRangeChecker tests
+func TestSequentialRangeChecker(t *testing.T) {
+	ranges := []TopNWithRange{
+		{TopNMeta: TopNMeta{Count: 10}, startIdx: 5, endIdx: 8},   // range [5,8]
+		{TopNMeta: TopNMeta{Count: 15}, startIdx: 12, endIdx: 15}, // range [12,15]
+		{TopNMeta: TopNMeta{Count: 8}, startIdx: 20, endIdx: 22},  // range [20,22]
+	}
+
+	checker := NewSequentialRangeChecker(ranges)
+
+	// test basic functionality
+	require.False(t, checker.IsIndexInTopNRange(4))  // before first range
+	require.True(t, checker.IsIndexInTopNRange(6))   // in first range
+	require.False(t, checker.IsIndexInTopNRange(10)) // between ranges
+	require.True(t, checker.IsIndexInTopNRange(13))  // in second range
+	require.True(t, checker.IsIndexInTopNRange(21))  // in third range
+	require.False(t, checker.IsIndexInTopNRange(25)) // after all ranges
+
+	// test edge cases
+	emptyChecker := NewSequentialRangeChecker([]TopNWithRange{})
+	require.False(t, emptyChecker.IsIndexInTopNRange(0))
+
+	// test unsorted input - should be sorted internally
+	unsortedRanges := []TopNWithRange{
+		{TopNMeta: TopNMeta{Count: 8}, startIdx: 20, endIdx: 22},  // third
+		{TopNMeta: TopNMeta{Count: 10}, startIdx: 5, endIdx: 8},   // first
+		{TopNMeta: TopNMeta{Count: 15}, startIdx: 12, endIdx: 15}, // second
+	}
+	unsortedChecker := NewSequentialRangeChecker(unsortedRanges)
+	require.True(t, unsortedChecker.IsIndexInTopNRange(6))
+	require.True(t, unsortedChecker.IsIndexInTopNRange(13))
+	require.True(t, unsortedChecker.IsIndexInTopNRange(21))
+}

pkg/statistics/sort_utils_test.go

@@ -2,18 +2,22 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "generic",
-    srcs = ["sync_map.go"],
+    srcs = [
+        "priority_queue.go",
+        "sync_map.go",
+    ],
     importpath = "github.com/pingcap/tidb/pkg/util/generic",
     visibility = ["//visibility:public"],
 )
 
 go_test(
     name = "generic_test",
     timeout = "short",
-    srcs = ["sync_map_test.go"],
-    flaky = True,
-    deps = [
-        ":generic",
-        "@com_github_stretchr_testify//require",
+    srcs = [
+        "priority_queue_test.go",
+        "sync_map_test.go",
     ],
+    embed = [":generic"],
+    flaky = True,
+    deps = ["@com_github_stretchr_testify//require"],
 )

pkg/util/generic/BUILD.bazel

@@ -12,65 +12,70 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package statistics
+package generic
 
 import "container/heap"
 
-// TopNHeap implements a min-heap for maintaining top N items efficiently.
+// PriorityQueue implements a min-heap for maintaining the best N items efficiently.
 // It keeps the N items with the highest values according to the comparison function.
 // The root of the heap is always the smallest item, making it easy to remove when adding better items.
-type TopNHeap[T any] struct {
+type PriorityQueue[T any] struct {
 	cmpFunc func(T, T) int
 	items   []T
 	maxSize int
 }
 
-// NewTopNHeap creates a new TopN heap with the specified maximum size and comparison function.
-func NewTopNHeap[T any](maxSize int, cmpFunc func(T, T) int) *TopNHeap[T] {
-	return &TopNHeap[T]{
+// NewPriorityQueue creates a new priority queue with the specified maximum size and comparison function.
+func NewPriorityQueue[T any](maxSize int, cmpFunc func(T, T) int) *PriorityQueue[T] {
+	return &PriorityQueue[T]{
 		items:   make([]T, 0, maxSize),
 		maxSize: maxSize,
 		cmpFunc: cmpFunc,
 	}
 }
 
 // Len returns the number of items in the heap.
-func (h *TopNHeap[T]) Len() int { return len(h.items) }
+func (h *PriorityQueue[T]) Len() int { return len(h.items) }
 
-// Less compares two items in the heap. For TopN, we use a min-heap.
-func (h *TopNHeap[T]) Less(i, j int) bool {
+// Less compares two items in the heap. We use a min-heap for efficient priority queue operations.
+func (h *PriorityQueue[T]) Less(i, j int) bool {
 	return h.cmpFunc(h.items[i], h.items[j]) < 0
 }
 
 // Swap swaps two items in the heap.
-func (h *TopNHeap[T]) Swap(i, j int) {
+func (h *PriorityQueue[T]) Swap(i, j int) {
 	h.items[i], h.items[j] = h.items[j], h.items[i]
 }
 
 // Push adds an item to the heap (used by container/heap).
-func (h *TopNHeap[T]) Push(x any) {
+func (h *PriorityQueue[T]) Push(x any) {
 	h.items = append(h.items, x.(T))
 }
 
 // Pop removes and returns the smallest item from the heap (used by container/heap).
-func (h *TopNHeap[T]) Pop() any {
+func (h *PriorityQueue[T]) Pop() any {
 	old := h.items
 	n := len(old)
 	item := old[n-1]
 	h.items = old[0 : n-1]
 	return item
 }
 
-// Add adds an item to the TopN heap. If the heap is full and the new item
+// Add adds an item to the priority queue. If the queue is full and the new item
 // is better than the worst item, it replaces the worst item.
-func (h *TopNHeap[T]) Add(item T) {
+func (h *PriorityQueue[T]) Add(item T) {
+	// handle zero capacity case
+	if h.maxSize <= 0 {
+		return
+	}
+
 	if len(h.items) < h.maxSize {
-		// heap not full, just add the item
+		// queue not full, just add the item
 		heap.Push(h, item)
 		return
 	}
 
-	// heap is full, check if new item is better than the worst (root of min-heap)
+	// queue is full, check if new item is better than the worst (root of min-heap)
 	if h.cmpFunc(item, h.items[0]) > 0 {
 		// new item is better, replace the worst
 		h.items[0] = item
@@ -79,12 +84,12 @@ func (h *TopNHeap[T]) Add(item T) {
 }
 
 // ToSortedSlice returns all items in the heap as a sorted slice (best to worst).
-func (h *TopNHeap[T]) ToSortedSlice() []T {
+func (h *PriorityQueue[T]) ToSortedSlice() []T {
 	if len(h.items) == 0 {
 		return nil
 	}
 
-	// copy items to avoid modifying the original heap
+	// copy items to avoid modifying the original queue
 	result := make([]T, len(h.items))
 	copy(result, h.items)