113 lines
1.9 KiB
Go
113 lines
1.9 KiB
Go
package tsm1
|
|
|
|
import (
|
|
"bytes"
|
|
"container/heap"
|
|
)
|
|
|
|
type keyIterator struct {
|
|
f TSMFile
|
|
c int // current key index
|
|
n int // key count
|
|
key []byte
|
|
typ byte
|
|
}
|
|
|
|
func newKeyIterator(f TSMFile, seek []byte) *keyIterator {
|
|
c, n := 0, f.KeyCount()
|
|
if len(seek) > 0 {
|
|
c = f.Seek(seek)
|
|
}
|
|
|
|
if c >= n {
|
|
return nil
|
|
}
|
|
|
|
k := &keyIterator{f: f, c: c, n: n}
|
|
k.next()
|
|
|
|
return k
|
|
}
|
|
|
|
func (k *keyIterator) next() bool {
|
|
if k.c < k.n {
|
|
k.key, k.typ = k.f.KeyAt(k.c)
|
|
k.c++
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
type mergeKeyIterator struct {
|
|
itrs keyIterators
|
|
key []byte
|
|
typ byte
|
|
}
|
|
|
|
func newMergeKeyIterator(files []TSMFile, seek []byte) *mergeKeyIterator {
|
|
m := &mergeKeyIterator{}
|
|
itrs := make(keyIterators, 0, len(files))
|
|
for _, f := range files {
|
|
if ki := newKeyIterator(f, seek); ki != nil {
|
|
itrs = append(itrs, ki)
|
|
}
|
|
}
|
|
m.itrs = itrs
|
|
heap.Init(&m.itrs)
|
|
|
|
return m
|
|
}
|
|
|
|
func (m *mergeKeyIterator) Next() bool {
|
|
merging := len(m.itrs) > 1
|
|
|
|
RETRY:
|
|
if len(m.itrs) == 0 {
|
|
return false
|
|
}
|
|
|
|
key, typ := m.itrs[0].key, m.itrs[0].typ
|
|
more := m.itrs[0].next()
|
|
|
|
switch {
|
|
case len(m.itrs) > 1:
|
|
if !more {
|
|
// remove iterator from heap
|
|
heap.Pop(&m.itrs)
|
|
} else {
|
|
heap.Fix(&m.itrs, 0)
|
|
}
|
|
|
|
case len(m.itrs) == 1:
|
|
if !more {
|
|
m.itrs = nil
|
|
}
|
|
}
|
|
|
|
if merging && bytes.Equal(m.key, key) {
|
|
// same as previous key, keep iterating
|
|
goto RETRY
|
|
}
|
|
|
|
m.key, m.typ = key, typ
|
|
|
|
return true
|
|
}
|
|
|
|
func (m *mergeKeyIterator) Read() ([]byte, byte) { return m.key, m.typ }
|
|
|
|
type keyIterators []*keyIterator
|
|
|
|
func (k keyIterators) Len() int { return len(k) }
|
|
func (k keyIterators) Less(i, j int) bool { return bytes.Compare(k[i].key, k[j].key) == -1 }
|
|
func (k keyIterators) Swap(i, j int) { k[i], k[j] = k[j], k[i] }
|
|
func (k *keyIterators) Push(x interface{}) { *k = append(*k, x.(*keyIterator)) }
|
|
|
|
func (k *keyIterators) Pop() interface{} {
|
|
old := *k
|
|
n := len(old)
|
|
x := old[n-1]
|
|
*k = old[:n-1]
|
|
return x
|
|
}
|