github.com/dolthub/go-mysql-server@v0.18.0/sql/stats/sql_heap.go (about)

     1  package stats
     2  
     3  import (
     4  	"container/heap"
     5  
     6  	"github.com/dolthub/go-mysql-server/sql"
     7  )
     8  
     9  // An SqlHeap is a min-heap of ints.
    10  type SqlHeap struct {
    11  	vals []sql.Row
    12  	cnts []uint64
    13  	k    int
    14  }
    15  
    16  type HeapRow struct {
    17  	Row   sql.Row
    18  	Count int
    19  }
    20  
    21  func NewHeapRow(r sql.Row, cnt int) HeapRow {
    22  	return HeapRow{Row: r, Count: cnt}
    23  }
    24  
    25  func NewSqlHeap(k int) *SqlHeap {
    26  	return &SqlHeap{vals: make([]sql.Row, 0), cnts: make([]uint64, 0), k: k}
    27  }
    28  
    29  func (h SqlHeap) Array() []sql.Row {
    30  	return h.vals
    31  }
    32  
    33  func (h SqlHeap) Counts() []uint64 {
    34  	return h.cnts
    35  }
    36  
    37  func (h SqlHeap) Len() int { return len(h.vals) }
    38  func (h SqlHeap) Less(i, j int) bool {
    39  	return h.cnts[i] < h.cnts[j]
    40  }
    41  func (h SqlHeap) Swap(i, j int) {
    42  	h.vals[i], h.vals[j] = h.vals[j], h.vals[i]
    43  	h.cnts[i], h.cnts[j] = h.cnts[j], h.cnts[i]
    44  }
    45  
    46  func (h *SqlHeap) Push(x any) {
    47  	// Push and Pop use pointer receivers because they modify the slice's length,
    48  	// not just its contents.
    49  	hr := x.(HeapRow)
    50  	h.vals = append(h.vals, hr.Row)
    51  	h.cnts = append(h.cnts, uint64(hr.Count))
    52  	if len(h.vals) > h.k {
    53  		heap.Pop(h)
    54  	}
    55  }
    56  
    57  func (h *SqlHeap) Pop() any {
    58  	n := len(h.vals)
    59  	r := h.vals[n-1]
    60  	h.vals = h.vals[0 : n-1]
    61  	c := h.cnts[n-1]
    62  	h.cnts = h.cnts[0 : n-1]
    63  	return HeapRow{Row: r, Count: int(c)}
    64  }