541 lines
13 KiB
Go
541 lines
13 KiB
Go
// Copyright 2021 Dolthub, Inc.
|
|
//
|
|
// 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 prolly
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/dolthub/dolt/go/store/prolly/tree"
|
|
"github.com/dolthub/dolt/go/store/val"
|
|
)
|
|
|
|
type rangeIterTest struct {
|
|
name string
|
|
testRange Range
|
|
expCount int
|
|
}
|
|
|
|
func testIterRange(t *testing.T, om testMap, tuples [][2]val.Tuple) {
|
|
ctx := context.Background()
|
|
desc := keyDescFromMap(om)
|
|
|
|
for i := 0; i < 100; i++ {
|
|
|
|
cnt := len(tuples)
|
|
a, z := testRand.Intn(cnt), testRand.Intn(cnt)
|
|
if a == z {
|
|
continue
|
|
} else if a > z {
|
|
a, z = z, a
|
|
}
|
|
start, stop := tuples[a][0], tuples[z][0]
|
|
|
|
mustRange := func(rng Range, err error) Range {
|
|
require.NoError(t, err)
|
|
return rng
|
|
}
|
|
|
|
tests := []rangeIterTest{
|
|
// two-sided ranges
|
|
{
|
|
name: "OpenRange",
|
|
testRange: mustRange(openRange(ctx, start, stop, desc)),
|
|
expCount: nonNegative((z - a) - 1),
|
|
},
|
|
{
|
|
name: "OpenStartRange",
|
|
testRange: mustRange(openStartRange(ctx, start, stop, desc)),
|
|
expCount: z - a,
|
|
},
|
|
{
|
|
name: "OpenStopRange",
|
|
testRange: mustRange(openStopRange(ctx, start, stop, desc)),
|
|
expCount: z - a,
|
|
},
|
|
{
|
|
name: "closedRange",
|
|
testRange: mustRange(closedRange(ctx, start, stop, desc)),
|
|
expCount: (z - a) + 1,
|
|
},
|
|
|
|
// one-sided ranges
|
|
{
|
|
name: "GreaterRange",
|
|
testRange: greaterRange(start, desc),
|
|
expCount: nonNegative(cnt - a - 1),
|
|
},
|
|
{
|
|
name: "GreaterOrEqualRange",
|
|
testRange: greaterOrEqualRange(start, desc),
|
|
expCount: cnt - a,
|
|
},
|
|
{
|
|
name: "LesserRange",
|
|
testRange: lesserRange(stop, desc),
|
|
expCount: z,
|
|
},
|
|
{
|
|
name: "LesserOrEqualRange",
|
|
testRange: lesserOrEqualRange(stop, desc),
|
|
expCount: z + 1,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
// s := fmt.Sprintf(test.testRange.format())
|
|
// fmt.Println(s)
|
|
|
|
iter, err := om.IterRange(ctx, test.testRange)
|
|
require.NoError(t, err)
|
|
|
|
key, _, err := iter.Next(ctx)
|
|
actCount := 0
|
|
for err != io.EOF {
|
|
actCount++
|
|
prev := key
|
|
key, _, err = iter.Next(ctx)
|
|
|
|
if key != nil {
|
|
cmp, cmpErr := desc.Compare(ctx, prev, key)
|
|
require.NoError(t, cmpErr)
|
|
assert.True(t, cmp < 0)
|
|
}
|
|
}
|
|
assert.Equal(t, io.EOF, err)
|
|
|
|
if !assert.Equal(t, test.expCount, actCount) {
|
|
fmt.Println("here")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func nonNegative(x int) int {
|
|
if x < 0 {
|
|
x = 0
|
|
}
|
|
return x
|
|
}
|
|
|
|
type prefixRangeTest struct {
|
|
name string
|
|
testRange Range
|
|
}
|
|
|
|
func testIterPrefixRange(t *testing.T, om testMap, tuples [][2]val.Tuple) {
|
|
ctx := context.Background()
|
|
prefixDesc := getDescPrefix(keyDescFromMap(om), 1)
|
|
|
|
for i := 0; i < 100; i++ {
|
|
|
|
cnt := len(tuples)
|
|
a, z := testRand.Intn(cnt), testRand.Intn(cnt)
|
|
if a > z {
|
|
a, z = z, a
|
|
}
|
|
start, err := getKeyPrefix(tuples[a][0], prefixDesc)
|
|
require.NoError(t, err)
|
|
stop, err := getKeyPrefix(tuples[z][0], prefixDesc)
|
|
require.NoError(t, err)
|
|
|
|
mustRange := func(rng Range, err error) Range {
|
|
require.NoError(t, err)
|
|
return rng
|
|
}
|
|
|
|
tests := []prefixRangeTest{
|
|
// two-sided ranges
|
|
{
|
|
name: "OpenRange",
|
|
testRange: mustRange(openRange(ctx, start, stop, prefixDesc)),
|
|
},
|
|
{
|
|
name: "OpenStartRange",
|
|
testRange: mustRange(openStartRange(ctx, start, stop, prefixDesc)),
|
|
},
|
|
{
|
|
name: "OpenStopRange",
|
|
testRange: mustRange(openStopRange(ctx, start, stop, prefixDesc)),
|
|
},
|
|
{
|
|
name: "closedRange",
|
|
testRange: mustRange(closedRange(ctx, start, stop, prefixDesc)),
|
|
},
|
|
|
|
// one-sided ranges
|
|
{
|
|
name: "GreaterRange",
|
|
testRange: greaterRange(start, prefixDesc),
|
|
},
|
|
{
|
|
name: "GreaterOrEqualRange",
|
|
testRange: greaterOrEqualRange(start, prefixDesc),
|
|
},
|
|
{
|
|
name: "LesserRange",
|
|
testRange: lesserRange(stop, prefixDesc),
|
|
},
|
|
{
|
|
name: "LesserOrEqualRange",
|
|
testRange: lesserOrEqualRange(stop, prefixDesc),
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
iter, err := om.IterRange(ctx, test.testRange)
|
|
require.NoError(t, err)
|
|
|
|
key, _, err := iter.Next(ctx)
|
|
actCount := 0
|
|
for err != io.EOF {
|
|
actCount++
|
|
prev := key
|
|
key, _, err = iter.Next(ctx)
|
|
|
|
if key != nil {
|
|
cmp, cmpErr := prefixDesc.Compare(ctx, prev, key)
|
|
require.NoError(t, cmpErr)
|
|
assert.True(t, cmp < 0)
|
|
}
|
|
}
|
|
assert.Equal(t, io.EOF, err)
|
|
|
|
expCount := getExpectedRangeSize(t, test.testRange, tuples)
|
|
assert.Equal(t, expCount, actCount)
|
|
}
|
|
}
|
|
}
|
|
|
|
func getDescPrefix(desc *val.TupleDesc, sz int) *val.TupleDesc {
|
|
return val.NewTupleDescriptor(desc.Types[:sz]...)
|
|
}
|
|
|
|
func getKeyPrefix(key val.Tuple, desc *val.TupleDesc) (partial val.Tuple, err error) {
|
|
tb := val.NewTupleBuilder(desc, ns)
|
|
for i := range desc.Types {
|
|
tb.PutRaw(i, key.GetField(i))
|
|
}
|
|
return tb.Build(context.Background(), sharedPool)
|
|
}
|
|
|
|
// computes expected range on full tuples set
|
|
func getExpectedRangeSize(t *testing.T, rng Range, tuples [][2]val.Tuple) (sz int) {
|
|
ctx := context.Background()
|
|
for i := range tuples {
|
|
k := tuples[i][0]
|
|
above, err := rng.aboveStart(ctx, k)
|
|
require.NoError(t, err)
|
|
|
|
below, err := rng.belowStop(ctx, k)
|
|
require.NoError(t, err)
|
|
|
|
if above && below {
|
|
sz++
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func TestMapIterRange(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
kd := val.NewTupleDescriptor(
|
|
val.Type{Enc: val.Int32Enc},
|
|
val.Type{Enc: val.Int32Enc},
|
|
)
|
|
vd := val.NewTupleDescriptor()
|
|
|
|
tuples := []val.Tuple{
|
|
intTuple(1, 1), intTuple(), // 0
|
|
intTuple(1, 2), intTuple(), // 2
|
|
intTuple(1, 3), intTuple(), // 4
|
|
intTuple(2, 1), intTuple(), // 6
|
|
intTuple(2, 2), intTuple(), // 8
|
|
intTuple(2, 3), intTuple(), // 10
|
|
intTuple(3, 1), intTuple(), // 12
|
|
intTuple(3, 2), intTuple(), // 14
|
|
intTuple(3, 3), intTuple(), // 16
|
|
intTuple(4, 1), intTuple(), // 18
|
|
intTuple(4, 2), intTuple(), // 20
|
|
intTuple(4, 3), intTuple(), // 22
|
|
}
|
|
require.Equal(t, 24, len(tuples))
|
|
|
|
index, err := NewMapFromTuples(ctx, ns, kd, vd, tuples...)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int(12), countOrderedMap(t, index))
|
|
|
|
partialDesc := val.NewTupleDescriptor(
|
|
val.Type{Enc: val.Int32Enc},
|
|
)
|
|
fullDesc := val.NewTupleDescriptor(
|
|
val.Type{Enc: val.Int32Enc},
|
|
val.Type{Enc: val.Int32Enc},
|
|
)
|
|
|
|
mustRange := func(rng Range, err error) Range {
|
|
require.NoError(t, err)
|
|
return rng
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
logical []int
|
|
rng Range
|
|
physical [2]int
|
|
}{
|
|
// partial-key range scan
|
|
{
|
|
name: "range [1:4]",
|
|
rng: mustRange(closedRange(ctx, intTuple(1), intTuple(4), partialDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22},
|
|
},
|
|
{
|
|
name: "range (1:4]",
|
|
rng: mustRange(openStartRange(ctx, intTuple(1), intTuple(4), partialDesc)),
|
|
physical: [2]int{6, 24},
|
|
logical: []int{6, 8, 10, 12, 14, 16, 18, 20, 22},
|
|
},
|
|
{
|
|
name: "range [1:4)",
|
|
rng: mustRange(openStopRange(ctx, intTuple(1), intTuple(4), partialDesc)),
|
|
physical: [2]int{0, 18},
|
|
logical: []int{0, 2, 4, 6, 8, 10, 12, 14, 16},
|
|
},
|
|
{
|
|
name: "range (1:4)",
|
|
rng: mustRange(openRange(ctx, intTuple(1), intTuple(4), partialDesc)),
|
|
physical: [2]int{6, 18},
|
|
logical: []int{6, 8, 10, 12, 14, 16},
|
|
},
|
|
|
|
// full-key range scan
|
|
{
|
|
name: "range (1,1:4,3)",
|
|
rng: mustRange(openRange(ctx, intTuple(1, 1), intTuple(4, 3), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{2, 8, 14, 20},
|
|
},
|
|
{
|
|
name: "range (1,1:4,3]",
|
|
rng: mustRange(openStartRange(ctx, intTuple(1, 1), intTuple(4, 3), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{2, 4, 8, 10, 14, 16, 20, 22},
|
|
},
|
|
{
|
|
name: "range [1,1:4,3)",
|
|
rng: mustRange(openStopRange(ctx, intTuple(1, 1), intTuple(4, 3), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{0, 2, 6, 8, 12, 14, 18, 20},
|
|
},
|
|
{
|
|
name: "range [1,1:4,3]",
|
|
rng: mustRange(closedRange(ctx, intTuple(1, 1), intTuple(4, 3), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22},
|
|
},
|
|
{
|
|
name: "range [1,2:4,2]",
|
|
rng: mustRange(closedRange(ctx, intTuple(1, 2), intTuple(4, 2), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{2, 8, 14, 20},
|
|
},
|
|
{
|
|
name: "range (1,2:4,2]",
|
|
rng: mustRange(openStartRange(ctx, intTuple(1, 2), intTuple(4, 2), fullDesc)),
|
|
physical: [2]int{0, 24},
|
|
logical: []int{},
|
|
},
|
|
{
|
|
name: "range [2,2:3,2]",
|
|
rng: mustRange(closedRange(ctx, intTuple(2, 2), intTuple(3, 2), fullDesc)),
|
|
physical: [2]int{6, 18},
|
|
logical: []int{8, 14},
|
|
},
|
|
{
|
|
name: "range [2,2:2,3]",
|
|
rng: mustRange(closedRange(ctx, intTuple(2, 2), intTuple(2, 3), fullDesc)),
|
|
physical: [2]int{8, 12},
|
|
logical: []int{8, 10},
|
|
},
|
|
{
|
|
name: "range [2,2:2,2]",
|
|
rng: mustRange(closedRange(ctx, intTuple(2, 2), intTuple(2, 2), fullDesc)),
|
|
physical: [2]int{8, 10},
|
|
logical: []int{8},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
// validate physical range (unfiltered iter)
|
|
iter, err := treeIterFromRange(ctx, index.Node(), ns, test.rng)
|
|
require.NoError(t, err)
|
|
|
|
var k, v val.Tuple
|
|
act := make([]val.Tuple, 0, len(test.physical))
|
|
for {
|
|
k, v, err = iter.Next(ctx)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
assert.NoError(t, err)
|
|
act = append(act, k, v)
|
|
}
|
|
assert.Error(t, io.EOF, err)
|
|
|
|
inRange := tuples[test.physical[0]:test.physical[1]]
|
|
assert.Equal(t, len(inRange), len(act))
|
|
if len(inRange) == len(act) {
|
|
for i := range inRange {
|
|
assert.Equal(t, inRange[i], act[i])
|
|
}
|
|
}
|
|
|
|
// validate logical range
|
|
iter2, err := index.IterRange(ctx, test.rng)
|
|
require.NoError(t, err)
|
|
|
|
act2 := make([]val.Tuple, 0, len(test.logical))
|
|
for {
|
|
k, _, err = iter2.Next(ctx)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
assert.NoError(t, err)
|
|
act2 = append(act2, k)
|
|
}
|
|
assert.Error(t, io.EOF, err)
|
|
|
|
exp2 := make([]val.Tuple, len(test.logical))
|
|
for i := range test.logical {
|
|
exp2[i] = tuples[test.logical[i]]
|
|
}
|
|
|
|
if !assert.Equal(t, len(exp2), len(act2)) {
|
|
t.Fail()
|
|
}
|
|
if len(exp2) == len(act2) {
|
|
for i := range exp2 {
|
|
assert.Equal(t, exp2[i], act2[i])
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func intTuple(ints ...int32) val.Tuple {
|
|
types := make([]val.Type, len(ints))
|
|
for i := range types {
|
|
types[i] = val.Type{Enc: val.Int32Enc}
|
|
}
|
|
|
|
desc := val.NewTupleDescriptor(types...)
|
|
tb := val.NewTupleBuilder(desc, ns)
|
|
for i := range ints {
|
|
tb.PutInt32(i, ints[i])
|
|
}
|
|
tup, err := tb.Build(context.Background(), sharedPool)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return tup
|
|
}
|
|
|
|
func testIterOrdinalRange(t *testing.T, om Map, tuples [][2]val.Tuple) {
|
|
cnt := len(tuples)
|
|
t.Run("test two sided bounds", func(t *testing.T) {
|
|
bounds := make([][2]int, 100)
|
|
for i := range bounds {
|
|
bounds[i] = [2]int{testRand.Intn(cnt)}
|
|
}
|
|
testIterOrdinalRangeWithBounds(t, om, tuples, bounds)
|
|
})
|
|
t.Run("test one sided bounds", func(t *testing.T) {
|
|
bounds := make([][2]int, 100)
|
|
for i := range bounds {
|
|
if i%2 == 0 {
|
|
bounds[i] = [2]int{0, testRand.Intn(cnt)}
|
|
} else {
|
|
bounds[i] = [2]int{testRand.Intn(cnt), cnt}
|
|
}
|
|
}
|
|
testIterOrdinalRangeWithBounds(t, om, tuples, bounds)
|
|
})
|
|
}
|
|
|
|
func testIterOrdinalRangeWithBounds(t *testing.T, om Map, tuples [][2]val.Tuple, bounds [][2]int) {
|
|
ctx := context.Background()
|
|
t.Run("IterOrdinalRange", func(t *testing.T) {
|
|
for _, bound := range bounds {
|
|
start, stop := bound[0], bound[1]
|
|
if start > stop {
|
|
start, stop = stop, start
|
|
} else if start == stop {
|
|
continue
|
|
}
|
|
expected := tuples[start:stop]
|
|
|
|
iter, err := om.IterOrdinalRange(ctx, uint64(start), uint64(stop))
|
|
require.NoError(t, err)
|
|
actual := iterOrdinalRange(t, ctx, iter)
|
|
assert.Equal(t, len(expected), len(actual),
|
|
"expected equal tuple slices for bounds (%d, %d)", start, stop)
|
|
assert.Equal(t, expected, actual)
|
|
}
|
|
})
|
|
t.Run("FetchOrdinalRange", func(t *testing.T) {
|
|
for _, bound := range bounds {
|
|
start, stop := bound[0], bound[1]
|
|
if start > stop {
|
|
start, stop = stop, start
|
|
} else if start == stop {
|
|
continue
|
|
}
|
|
expected := tuples[start:stop]
|
|
|
|
iter, err := om.FetchOrdinalRange(ctx, uint64(start), uint64(stop))
|
|
require.NoError(t, err)
|
|
actual := iterOrdinalRange(t, ctx, iter)
|
|
assert.Equal(t, len(expected), len(actual),
|
|
"expected equal tuple slices for bounds (%d, %d)", start, stop)
|
|
assert.Equal(t, expected, actual)
|
|
}
|
|
})
|
|
}
|
|
|
|
func iterOrdinalRange(t *testing.T, ctx context.Context, iter MapIter) (actual [][2]val.Tuple) {
|
|
for {
|
|
k, v, err := iter.Next(ctx)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, k)
|
|
assert.NotNil(t, v)
|
|
actual = append(actual, [2]val.Tuple{k, v})
|
|
}
|
|
return
|
|
}
|