influxdb/tsdb/tsm1/file_store_key_iterator.go

116 lines
2.0 KiB
Go

package tsm1
import (
"bytes"
"container/heap"
)
type keyIterator struct {
iter TSMIterator
key []byte
typ byte
}
func newKeyIterator(f TSMFile, seek []byte) *keyIterator {
k := &keyIterator{iter: f.Iterator(seek)}
k.next()
return k
}
func (k *keyIterator) next() bool {
if k.iter.Next() {
k.key, k.typ = k.iter.Key(), k.iter.Type()
return true
}
return false
}
func (k *keyIterator) Err() error { return k.iter.Err() }
type mergeKeyIterator struct {
itrs keyIterators
key []byte
typ byte
err error
}
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 {
if m.err != nil {
return false
}
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()
if !more {
if err := m.itrs[0].Err(); err != nil {
m.err = err
return false
}
}
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) Err() error { return m.err }
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
}