235 lines
6.9 KiB
Go
235 lines
6.9 KiB
Go
|
/*
|
||
|
|
* Copyright 2017 Dgraph Labs, Inc. and Contributors
|
||
|
|
*
|
||
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
* you may not use this file except in compliance with the License.
|
||
|
|
* You may obtain a copy of the License at
|
||
|
|
*
|
||
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
*
|
||
|
|
* Unless required by applicable law or agreed to in writing, software
|
||
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
* See the License for the specific language governing permissions and
|
||
|
|
* limitations under the License.
|
||
|
|
*/
|
||
|
|
|
||
|
|
package y
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sort"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type SimpleIterator struct {
|
||
|
|
keys [][]byte
|
||
|
|
vals [][]byte
|
||
|
|
idx int
|
||
|
|
reversed bool
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
closeCount int
|
||
|
|
)
|
||
|
|
|
||
|
|
func (s *SimpleIterator) Close() error { closeCount++; return nil }
|
||
|
|
|
||
|
|
func (s *SimpleIterator) Next() {
|
||
|
|
if !s.reversed {
|
||
|
|
s.idx++
|
||
|
|
} else {
|
||
|
|
s.idx--
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SimpleIterator) Rewind() {
|
||
|
|
if !s.reversed {
|
||
|
|
s.idx = 0
|
||
|
|
} else {
|
||
|
|
s.idx = len(s.keys) - 1
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SimpleIterator) Seek(key []byte) {
|
||
|
|
key = KeyWithTs(key, 0)
|
||
|
|
if !s.reversed {
|
||
|
|
s.idx = sort.Search(len(s.keys), func(i int) bool {
|
||
|
|
return CompareKeys(s.keys[i], key) >= 0
|
||
|
|
})
|
||
|
|
} else {
|
||
|
|
n := len(s.keys)
|
||
|
|
s.idx = n - 1 - sort.Search(n, func(i int) bool {
|
||
|
|
return CompareKeys(s.keys[n-1-i], key) <= 0
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *SimpleIterator) Key() []byte { return s.keys[s.idx] }
|
||
|
|
func (s *SimpleIterator) Value() ValueStruct {
|
||
|
|
return ValueStruct{
|
||
|
|
Value: s.vals[s.idx],
|
||
|
|
UserMeta: 55,
|
||
|
|
Meta: 0,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
func (s *SimpleIterator) Valid() bool {
|
||
|
|
return s.idx >= 0 && s.idx < len(s.keys)
|
||
|
|
}
|
||
|
|
|
||
|
|
func newSimpleIterator(keys []string, vals []string, reversed bool) *SimpleIterator {
|
||
|
|
k := make([][]byte, len(keys))
|
||
|
|
v := make([][]byte, len(vals))
|
||
|
|
AssertTrue(len(keys) == len(vals))
|
||
|
|
for i := 0; i < len(keys); i++ {
|
||
|
|
k[i] = KeyWithTs([]byte(keys[i]), 0)
|
||
|
|
v[i] = []byte(vals[i])
|
||
|
|
}
|
||
|
|
return &SimpleIterator{
|
||
|
|
keys: k,
|
||
|
|
vals: v,
|
||
|
|
idx: -1,
|
||
|
|
reversed: reversed,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func getAll(it Iterator) ([]string, []string) {
|
||
|
|
var keys, vals []string
|
||
|
|
for ; it.Valid(); it.Next() {
|
||
|
|
k := it.Key()
|
||
|
|
keys = append(keys, string(ParseKey(k)))
|
||
|
|
v := it.Value()
|
||
|
|
vals = append(vals, string(v.Value))
|
||
|
|
}
|
||
|
|
return keys, vals
|
||
|
|
}
|
||
|
|
|
||
|
|
func closeAndCheck(t *testing.T, it Iterator, expected int) {
|
||
|
|
closeCount = 0
|
||
|
|
it.Close()
|
||
|
|
require.EqualValues(t, expected, closeCount)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSimpleIterator(t *testing.T) {
|
||
|
|
keys := []string{"1", "2", "3"}
|
||
|
|
vals := []string{"v1", "v2", "v3"}
|
||
|
|
it := newSimpleIterator(keys, vals, false)
|
||
|
|
it.Rewind()
|
||
|
|
k, v := getAll(it)
|
||
|
|
require.EqualValues(t, keys, k)
|
||
|
|
require.EqualValues(t, vals, v)
|
||
|
|
|
||
|
|
closeAndCheck(t, it, 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
func reversed(a []string) []string {
|
||
|
|
var out []string
|
||
|
|
for i := len(a) - 1; i >= 0; i-- {
|
||
|
|
out = append(out, a[i])
|
||
|
|
}
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeSingle(t *testing.T) {
|
||
|
|
keys := []string{"1", "2", "3"}
|
||
|
|
vals := []string{"v1", "v2", "v3"}
|
||
|
|
it := newSimpleIterator(keys, vals, false)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it}, false)
|
||
|
|
mergeIt.Rewind()
|
||
|
|
k, v := getAll(mergeIt)
|
||
|
|
require.EqualValues(t, keys, k)
|
||
|
|
require.EqualValues(t, vals, v)
|
||
|
|
closeAndCheck(t, mergeIt, 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeSingleReversed(t *testing.T) {
|
||
|
|
keys := []string{"1", "2", "3"}
|
||
|
|
vals := []string{"v1", "v2", "v3"}
|
||
|
|
it := newSimpleIterator(keys, vals, true)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it}, true)
|
||
|
|
mergeIt.Rewind()
|
||
|
|
k, v := getAll(mergeIt)
|
||
|
|
require.EqualValues(t, reversed(keys), k)
|
||
|
|
require.EqualValues(t, reversed(vals), v)
|
||
|
|
closeAndCheck(t, mergeIt, 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeMore(t *testing.T) {
|
||
|
|
it := newSimpleIterator([]string{"1", "3", "7"}, []string{"a1", "a3", "a7"}, false)
|
||
|
|
it2 := newSimpleIterator([]string{"2", "3", "5"}, []string{"b2", "b3", "b5"}, false)
|
||
|
|
it3 := newSimpleIterator([]string{"1"}, []string{"c1"}, false)
|
||
|
|
it4 := newSimpleIterator([]string{"1", "7", "9"}, []string{"d1", "d7", "d9"}, false)
|
||
|
|
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it, it2, it3, it4}, false)
|
||
|
|
expectedKeys := []string{"1", "2", "3", "5", "7", "9"}
|
||
|
|
expectedVals := []string{"a1", "b2", "a3", "b5", "a7", "d9"}
|
||
|
|
mergeIt.Rewind()
|
||
|
|
k, v := getAll(mergeIt)
|
||
|
|
require.EqualValues(t, expectedKeys, k)
|
||
|
|
require.EqualValues(t, expectedVals, v)
|
||
|
|
closeAndCheck(t, mergeIt, 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Ensure MergeIterator satisfies the Iterator interface
|
||
|
|
func TestMergeIteratorNested(t *testing.T) {
|
||
|
|
keys := []string{"1", "2", "3"}
|
||
|
|
vals := []string{"v1", "v2", "v3"}
|
||
|
|
it := newSimpleIterator(keys, vals, false)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it}, false)
|
||
|
|
mergeIt2 := NewMergeIterator([]Iterator{mergeIt}, false)
|
||
|
|
mergeIt2.Rewind()
|
||
|
|
k, v := getAll(mergeIt2)
|
||
|
|
require.EqualValues(t, keys, k)
|
||
|
|
require.EqualValues(t, vals, v)
|
||
|
|
closeAndCheck(t, mergeIt2, 1)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeIteratorSeek(t *testing.T) {
|
||
|
|
it := newSimpleIterator([]string{"1", "3", "7"}, []string{"a1", "a3", "a7"}, false)
|
||
|
|
it2 := newSimpleIterator([]string{"2", "3", "5"}, []string{"b2", "b3", "b5"}, false)
|
||
|
|
it3 := newSimpleIterator([]string{"1"}, []string{"c1"}, false)
|
||
|
|
it4 := newSimpleIterator([]string{"1", "7", "9"}, []string{"d1", "d7", "d9"}, false)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it, it2, it3, it4}, false)
|
||
|
|
mergeIt.Seek([]byte("4"))
|
||
|
|
k, v := getAll(mergeIt)
|
||
|
|
require.EqualValues(t, []string{"5", "7", "9"}, k)
|
||
|
|
require.EqualValues(t, []string{"b5", "a7", "d9"}, v)
|
||
|
|
closeAndCheck(t, mergeIt, 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeIteratorSeekReversed(t *testing.T) {
|
||
|
|
it := newSimpleIterator([]string{"1", "3", "7"}, []string{"a1", "a3", "a7"}, true)
|
||
|
|
it2 := newSimpleIterator([]string{"2", "3", "5"}, []string{"b2", "b3", "b5"}, true)
|
||
|
|
it3 := newSimpleIterator([]string{"1"}, []string{"c1"}, true)
|
||
|
|
it4 := newSimpleIterator([]string{"1", "7", "9"}, []string{"d1", "d7", "d9"}, true)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it, it2, it3, it4}, true)
|
||
|
|
mergeIt.Seek([]byte("5"))
|
||
|
|
k, v := getAll(mergeIt)
|
||
|
|
require.EqualValues(t, []string{"5", "3", "2", "1"}, k)
|
||
|
|
require.EqualValues(t, []string{"b5", "a3", "b2", "a1"}, v)
|
||
|
|
closeAndCheck(t, mergeIt, 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeIteratorSeekInvalid(t *testing.T) {
|
||
|
|
it := newSimpleIterator([]string{"1", "3", "7"}, []string{"a1", "a3", "a7"}, false)
|
||
|
|
it2 := newSimpleIterator([]string{"2", "3", "5"}, []string{"b2", "b3", "b5"}, false)
|
||
|
|
it3 := newSimpleIterator([]string{"1"}, []string{"c1"}, false)
|
||
|
|
it4 := newSimpleIterator([]string{"1", "7", "9"}, []string{"d1", "d7", "d9"}, false)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it, it2, it3, it4}, false)
|
||
|
|
mergeIt.Seek([]byte("f"))
|
||
|
|
require.False(t, mergeIt.Valid())
|
||
|
|
closeAndCheck(t, mergeIt, 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeIteratorSeekInvalidReversed(t *testing.T) {
|
||
|
|
it := newSimpleIterator([]string{"1", "3", "7"}, []string{"a1", "a3", "a7"}, true)
|
||
|
|
it2 := newSimpleIterator([]string{"2", "3", "5"}, []string{"b2", "b3", "b5"}, true)
|
||
|
|
it3 := newSimpleIterator([]string{"1"}, []string{"c1"}, true)
|
||
|
|
it4 := newSimpleIterator([]string{"1", "7", "9"}, []string{"d1", "d7", "d9"}, true)
|
||
|
|
mergeIt := NewMergeIterator([]Iterator{it, it2, it3, it4}, true)
|
||
|
|
mergeIt.Seek([]byte("0"))
|
||
|
|
require.False(t, mergeIt.Valid())
|
||
|
|
closeAndCheck(t, mergeIt, 4)
|
||
|
|
}
|