dolthub--dolt
wehub-resource-sync
883 行
30 KiB
Go
883 行
30 KiB
Go
// Copyright 2024 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"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"math/rand"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
"github.com/dolthub/go-mysql-server/sql/expression/function/vector"
|
|
"github.com/dolthub/go-mysql-server/sql/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/dolthub/dolt/go/store/hash"
|
|
"github.com/dolthub/dolt/go/store/pool"
|
|
"github.com/dolthub/dolt/go/store/prolly/tree"
|
|
"github.com/dolthub/dolt/go/store/val"
|
|
)
|
|
|
|
func newJsonValue(t *testing.T, ctx context.Context, v interface{}) sql.JSONWrapper {
|
|
doc, _, err := types.JSON.Convert(ctx, v)
|
|
require.NoError(t, err)
|
|
return doc.(sql.JSONWrapper)
|
|
}
|
|
|
|
// newJsonDocument creates a JSON value from a provided value.
|
|
func newJsonDocument(t *testing.T, ctx context.Context, ns tree.NodeStore, v interface{}) hash.Hash {
|
|
doc := newJsonValue(t, ctx, v)
|
|
root, err := tree.SerializeJsonToAddr(ctx, ns, doc)
|
|
require.NoError(t, err)
|
|
return root.HashOf()
|
|
}
|
|
|
|
var jsonTestKeyDesc = val.NewTupleDescriptor(
|
|
val.Type{Enc: val.JSONAddrEnc, Nullable: true},
|
|
)
|
|
|
|
var vectorTestKeyDesc = val.NewTupleDescriptor(
|
|
val.Type{Enc: val.BytesAdaptiveEnc, Nullable: true},
|
|
)
|
|
|
|
var testValDesc = val.NewTupleDescriptor(
|
|
val.Type{Enc: val.Int64Enc, Nullable: true},
|
|
)
|
|
|
|
func buildTuple(t *testing.T, ctx context.Context, ns tree.NodeStore, pool pool.BuffPool, desc *val.TupleDesc, row []interface{}) val.Tuple {
|
|
builder := val.NewTupleBuilder(desc, ns)
|
|
for i, column := range row {
|
|
err := tree.PutField(ctx, ns, builder, i, column)
|
|
require.NoError(t, err)
|
|
}
|
|
tup, err := builder.Build(context.Background(), pool)
|
|
require.NoError(t, err)
|
|
return tup
|
|
}
|
|
|
|
func buildTuples(t *testing.T, ctx context.Context, ns tree.NodeStore, pool pool.BuffPool, desc *val.TupleDesc, rows [][]interface{}) [][]byte {
|
|
result := make([][]byte, len(rows))
|
|
for i, row := range rows {
|
|
result[i] = buildTuple(t, ctx, ns, pool, desc, row)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func createAndValidateProximityMap(t *testing.T, ctx context.Context, ns tree.NodeStore, keyDesc *val.TupleDesc, keyBytes [][]byte, valueDesc *val.TupleDesc, valueBytes [][]byte, logChunkSize uint8) ProximityMap {
|
|
m := createProximityMap(t, ctx, ns, keyDesc, keyBytes, valueDesc, valueBytes, logChunkSize)
|
|
validateProximityMapSkipHistoryIndependenceCheck(t, ctx, ns, &m, keyDesc, valueDesc, keyBytes, valueBytes)
|
|
return m
|
|
}
|
|
|
|
func createProximityMap(t *testing.T, ctx context.Context, ns tree.NodeStore, keyDesc *val.TupleDesc, keyBytes [][]byte, valueDesc *val.TupleDesc, valueBytes [][]byte, logChunkSize uint8) ProximityMap {
|
|
count := len(keyBytes)
|
|
require.Equal(t, count, len(valueBytes))
|
|
|
|
distanceType := vector.DistanceL2Squared{}
|
|
|
|
builder, err := NewProximityMapBuilder(ctx, ns, distanceType, keyDesc, valueDesc, logChunkSize)
|
|
require.NoError(t, err)
|
|
|
|
for i, key := range keyBytes {
|
|
value := valueBytes[i]
|
|
err = builder.Insert(ctx, key, value)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
m, err := builder.Flush(ctx)
|
|
require.NoError(t, err)
|
|
|
|
mapCount, err := m.Count()
|
|
require.NoError(t, err)
|
|
require.Equal(t, count, mapCount)
|
|
|
|
return m
|
|
}
|
|
|
|
func validateProximityMap(t *testing.T, ctx context.Context, ns tree.NodeStore, m *ProximityMap, keyDesc, valDesc *val.TupleDesc, keys, values [][]byte, logChunkSize uint8) {
|
|
validateProximityMapSkipHistoryIndependenceCheck(t, ctx, ns, m, keyDesc, valDesc, keys, values)
|
|
validateHistoryIndependence(t, ctx, ns, m, keyDesc, keys, valDesc, values, logChunkSize)
|
|
}
|
|
|
|
func validateProximityMapSkipHistoryIndependenceCheck(t *testing.T, ctx context.Context, ns tree.NodeStore, m *ProximityMap, keyDesc, valDesc *val.TupleDesc, keys, values [][]byte) {
|
|
expectedSize := len(keys)
|
|
actualSize, err := m.Count()
|
|
require.NoError(t, err)
|
|
require.Equal(t, expectedSize, actualSize)
|
|
// Check that every key and value appears in the map exactly once.
|
|
matches := 0
|
|
for i := 0; i < actualSize; i++ {
|
|
err = m.Get(ctx, keys[i], func(foundKey val.Tuple, foundValue val.Tuple) error {
|
|
require.Equal(t, val.Tuple(keys[i]), foundKey)
|
|
require.Equal(t, val.Tuple(values[i]), foundValue)
|
|
matches++
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
require.Equal(t, expectedSize, matches)
|
|
|
|
// Check that the invariant holds: each vector is closer to its parent than any of its uncles.
|
|
err = tree.WalkNodes(ctx, m.tuples.Root, ns, func(ctx context.Context, nd *tree.Node) error {
|
|
validateProximityMapNode(t, ctx, ns, nd, vector.DistanceL2Squared{}, keyDesc, valDesc)
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Finally, build a new map with the supplied keys and values and confirm that it has the same root hash.
|
|
}
|
|
|
|
func validateHistoryIndependence(t *testing.T, ctx context.Context, ns tree.NodeStore, m *ProximityMap, keyDesc *val.TupleDesc, keyBytes [][]byte, valueDesc *val.TupleDesc, valueBytes [][]byte, logChunkSize uint8) {
|
|
// Build a new map with the supplied keys and values and confirm that it has the same root hash.
|
|
other := createProximityMap(t, ctx, ns, keyDesc, keyBytes, valueDesc, valueBytes, logChunkSize)
|
|
require.Equal(t, other.HashOf(), m.HashOf())
|
|
}
|
|
|
|
func vectorFromKey(t *testing.T, keyDesc *val.TupleDesc, key []byte) []float32 {
|
|
encodedVector := keyDesc.GetField(0, key)
|
|
return decodeVector(t, keyDesc, encodedVector)
|
|
}
|
|
|
|
func validateProximityMapNode(t *testing.T, ctx context.Context, ns tree.NodeStore, nd *tree.Node, distanceType vector.DistanceType, keyDesc *val.TupleDesc, desc *val.TupleDesc) {
|
|
// For each node, the node's grandchildren should be closer to their parent than the other children.
|
|
if nd.Level() == 0 {
|
|
// Leaf node
|
|
return
|
|
}
|
|
if nd.Count() <= 1 {
|
|
// A node with only one child is trivially valid.
|
|
return
|
|
}
|
|
// Get the vector in each key
|
|
vectors := make([][]float32, nd.Count())
|
|
for vectorIdx := 0; vectorIdx < nd.Count(); vectorIdx++ {
|
|
vectorKey := nd.GetKey(vectorIdx)
|
|
vectors[vectorIdx] = vectorFromKey(t, keyDesc, vectorKey)
|
|
}
|
|
for childIdx := 0; childIdx < nd.Count(); childIdx++ {
|
|
// Get the child node
|
|
childHash := hash.New(nd.GetValue(childIdx))
|
|
childNode, err := ns.Read(ctx, childHash)
|
|
require.NoError(t, err)
|
|
for childKeyIdx := 0; childKeyIdx < childNode.Count(); childKeyIdx++ {
|
|
childVectorKey := childNode.GetKey(childKeyIdx)
|
|
childVector := vectorFromKey(t, keyDesc, childVectorKey)
|
|
minDistance := math.MaxFloat64
|
|
closestKeyIdx := -1
|
|
for otherChildIdx := 0; otherChildIdx < nd.Count(); otherChildIdx++ {
|
|
distance, err := distanceType.Eval(childVector, vectors[otherChildIdx])
|
|
require.NoError(t, err)
|
|
if distance < minDistance {
|
|
minDistance = distance
|
|
closestKeyIdx = otherChildIdx
|
|
}
|
|
}
|
|
require.Equal(t, closestKeyIdx, childIdx)
|
|
}
|
|
}
|
|
}
|
|
|
|
func encodeVector(t *testing.T, keyDesc *val.TupleDesc, vec ...float32) []byte {
|
|
enc := keyDesc.Types[0].Enc
|
|
switch enc {
|
|
case val.JSONAddrEnc:
|
|
res, err := json.Marshal(vec)
|
|
require.NoError(t, err)
|
|
return res
|
|
case val.BytesAdaptiveEnc:
|
|
return sql.EncodeVector(vec)
|
|
default:
|
|
panic("unexpected encoding")
|
|
}
|
|
}
|
|
|
|
func decodeVector(t *testing.T, keyDesc *val.TupleDesc, valBytes []byte) []float32 {
|
|
ctx := context.Background()
|
|
enc := keyDesc.Types[0].Enc
|
|
var vectorValue interface{}
|
|
var err error
|
|
switch enc {
|
|
case val.JSONAddrEnc:
|
|
vectorValue, err = tree.NewJSONDoc(hash.New(valBytes), ns).ToIndexedJSONDocument(ctx)
|
|
require.NoError(t, err)
|
|
case val.BytesAdaptiveEnc:
|
|
var ok bool
|
|
vectorValue, ok, err = val.GetBytesAdaptiveValue(ctx, ns, valBytes)
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
default:
|
|
panic("unexpected encoding")
|
|
}
|
|
res, err := sql.ConvertToVector(ctx, vectorValue)
|
|
require.NoError(t, err)
|
|
return res
|
|
}
|
|
|
|
func putVector(t *testing.T, keyBuilder *val.TupleBuilder, v []byte) {
|
|
ctx := context.Background()
|
|
enc := keyBuilder.Desc.Types[0].Enc
|
|
switch enc {
|
|
case val.JSONAddrEnc:
|
|
keyBuilder.PutJSONAddr(0, newJsonDocument(t, ctx, ns, v))
|
|
case val.BytesAdaptiveEnc:
|
|
err := keyBuilder.PutAdaptiveBytesFromInline(ctx, 0, v)
|
|
require.NoError(t, err)
|
|
default:
|
|
panic("unexpected encoding")
|
|
}
|
|
}
|
|
|
|
func TestProximityMap(t *testing.T) {
|
|
t.Run("JSON vector encoding", func(t *testing.T) {
|
|
testProximityMapWithEncoding(t, jsonTestKeyDesc)
|
|
})
|
|
t.Run("VECTOR vector encoding", func(t *testing.T) {
|
|
testProximityMapWithEncoding(t, vectorTestKeyDesc)
|
|
})
|
|
}
|
|
|
|
func testProximityMapWithEncoding(t *testing.T, keyDesc *val.TupleDesc) {
|
|
testEmptyProximityMap(t, keyDesc)
|
|
testSingleEntryProximityMap(t, keyDesc)
|
|
testDoubleEntryProximityMapGetExact(t, keyDesc)
|
|
testDoubleEntryProximityMapGetClosest(t, keyDesc)
|
|
testProximityMapGetManyClosest(t, keyDesc)
|
|
testProximityMapWithOverflowNode(t, keyDesc)
|
|
testMultilevelProximityMap(t, keyDesc)
|
|
testLargerMultilevelProximityMap(t, keyDesc)
|
|
testInsertOrderIndependence(t, keyDesc)
|
|
testIncrementalInserts(t, keyDesc)
|
|
testIncrementalUpdates(t, keyDesc)
|
|
testIncrementalDeletes(t, keyDesc)
|
|
testNonlexographicKey(t, keyDesc)
|
|
testManyDimensions(t, keyDesc)
|
|
}
|
|
|
|
func testEmptyProximityMap(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("empty map", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
createAndValidateProximityMap(t, ctx, ns, keyDesc, nil, testValDesc, nil, 10)
|
|
})
|
|
}
|
|
|
|
func testSingleEntryProximityMap(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("single entry map", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, [][]interface{}{{encodeVector(t, keyDesc, 1.0)}})
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, [][]interface{}{{int64(1)}})
|
|
createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 10)
|
|
})
|
|
}
|
|
|
|
func testDoubleEntryProximityMapGetExact(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("double entry map get exact", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows := [][]interface{}{{encodeVector(t, keyDesc, 0.0, 6.0)}, {encodeVector(t, keyDesc, 3.0, 4.0)}}
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{int64(1)}, {int64(2)}}
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 10)
|
|
matches := 0
|
|
for i, key := range keys {
|
|
err := m.Get(ctx, key, func(foundKey val.Tuple, foundValue val.Tuple) error {
|
|
require.Equal(t, val.Tuple(key), foundKey)
|
|
require.Equal(t, val.Tuple(values[i]), foundValue)
|
|
matches++
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
require.Equal(t, matches, len(keys))
|
|
})
|
|
}
|
|
|
|
func testDoubleEntryProximityMapGetClosest(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("double entry map get closest", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{int64(1)}, {int64(2)}}
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 10)
|
|
|
|
matches := 0
|
|
|
|
mapIter, err := m.GetClosest(ctx, sql.EncodeVector([]float32{0.0, 0.0}), 1)
|
|
require.NoError(t, err)
|
|
for {
|
|
k, v, err := mapIter.Next(ctx)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
require.NoError(t, err)
|
|
require.Equal(t, val.Tuple(keys[1]), k)
|
|
require.Equal(t, val.Tuple(values[1]), v)
|
|
matches++
|
|
}
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, matches, 1)
|
|
})
|
|
}
|
|
|
|
func testProximityMapGetManyClosest(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("get many closest", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 0.0)},
|
|
{encodeVector(t, keyDesc, 0.0, 10.0)},
|
|
{encodeVector(t, keyDesc, 10.0, 10.0)},
|
|
{encodeVector(t, keyDesc, 10.0, 0.0)},
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 10)
|
|
|
|
queryVector := sql.EncodeVector([]float32{3.0, 1.0})
|
|
sortOrder := []int{0, 3, 1, 2} // indexes in sorted order: [0.0, 0.0], [10.0, 0.0], [0.0, 10.0], [10.0, 10.0]
|
|
|
|
for limit := 0; limit <= 4; limit++ {
|
|
t.Run(fmt.Sprintf("limit %d", limit), func(t *testing.T) {
|
|
matches := 0
|
|
|
|
mapIter, err := m.GetClosest(ctx, queryVector, limit)
|
|
require.NoError(t, err)
|
|
for {
|
|
k, v, err := mapIter.Next(ctx)
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
require.NoError(t, err)
|
|
require.Equal(t, val.Tuple(keys[sortOrder[matches]]), k)
|
|
require.Equal(t, val.Tuple(values[sortOrder[matches]]), v)
|
|
matches++
|
|
}
|
|
require.NoError(t, err)
|
|
require.Equal(t, limit, matches)
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func testProximityMapWithOverflowNode(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("node too large to fit in a single physical chunk", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
// Create an index with enough rows that it can't fit in a single physical chunk
|
|
keyRows := make([][]interface{}, 0, 4000)
|
|
valueRows := make([][]interface{}, 0, 4000)
|
|
|
|
for i := int64(0); i < 4000; i++ {
|
|
keyRows = append(keyRows, []interface{}{encodeVector(t, keyDesc, float32(i))})
|
|
valueRows = append(valueRows, []interface{}{i})
|
|
}
|
|
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
// Set logChunkSize to a high enough value that everything goes in a single chunk
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 16)
|
|
|
|
count, err := m.Count()
|
|
require.NoError(t, err)
|
|
require.Equal(t, 4000, count)
|
|
})
|
|
}
|
|
|
|
func testMultilevelProximityMap(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("map with multiple levels", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 1)
|
|
matches := 0
|
|
for i, key := range keys {
|
|
err := m.Get(ctx, key, func(foundKey val.Tuple, foundValue val.Tuple) error {
|
|
require.Equal(t, val.Tuple(key), foundKey)
|
|
require.Equal(t, val.Tuple(values[i]), foundValue)
|
|
matches++
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
require.Equal(t, matches, len(keys))
|
|
})
|
|
}
|
|
|
|
func testLargerMultilevelProximityMap(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("larger map with multiple levels", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
{encodeVector(t, keyDesc, 9.0, 10.0)},
|
|
{encodeVector(t, keyDesc, 11.0, 12.0)},
|
|
{encodeVector(t, keyDesc, 13.0, 14.0)},
|
|
{encodeVector(t, keyDesc, 15.0, 16.0)},
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, keyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}, {int64(5)}, {int64(6)}, {int64(7)}, {int64(8)}}
|
|
values := buildTuples(t, ctx, ns, pb, testValDesc, valueRows)
|
|
|
|
m := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys, testValDesc, values, 1)
|
|
matches := 0
|
|
for i, key := range keys {
|
|
err := m.Get(ctx, key, func(foundKey val.Tuple, foundValue val.Tuple) error {
|
|
require.Equal(t, val.Tuple(key), foundKey)
|
|
require.Equal(t, val.Tuple(values[i]), foundValue)
|
|
matches++
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
require.Equal(t, matches, len(keys))
|
|
})
|
|
}
|
|
|
|
func testInsertOrderIndependence(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("insert order independence", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
keyRows1 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
keys1 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows1)
|
|
|
|
valueRows1 := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values1 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows1)
|
|
|
|
keyRows2 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
}
|
|
keys2 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows2)
|
|
|
|
valueRows2 := [][]interface{}{{int64(4)}, {int64(3)}, {int64(2)}, {int64(1)}}
|
|
values2 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows2)
|
|
|
|
m1 := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys1, testValDesc, values1, 1)
|
|
m2 := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys2, testValDesc, values2, 1)
|
|
|
|
if !assert.Equal(t, m1.tuples.Root.HashOf(), m2.tuples.Root.HashOf(), "trees have different hashes") {
|
|
require.NoError(t, tree.OutputProllyNodeBytes(os.Stdout, m1.tuples.Root))
|
|
require.NoError(t, tree.OutputProllyNodeBytes(os.Stdout, m2.tuples.Root))
|
|
}
|
|
})
|
|
}
|
|
|
|
func testIncrementalInserts(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("incremental inserts", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns, keyDesc := keyDescWithNodeStore(keyDesc)
|
|
pb := pool.NewBuffPool()
|
|
logChunkSize := uint8(1)
|
|
distanceType := vector.DistanceL2Squared{}
|
|
flusher := ProximityFlusher{logChunkSize: logChunkSize, distanceType: distanceType}
|
|
keyRows1 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
keys1 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows1)
|
|
|
|
valueRows1 := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values1 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows1)
|
|
|
|
m1 := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys1, testValDesc, values1, logChunkSize)
|
|
|
|
l1 := m1.tuples.Root.Level()
|
|
_ = l1
|
|
mutableMap := newProximityMutableMap(m1)
|
|
|
|
keyRows2 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 9.0, 10.0)},
|
|
{encodeVector(t, keyDesc, 11.0, 12.0)},
|
|
{encodeVector(t, keyDesc, 13.0, 14.0)},
|
|
{encodeVector(t, keyDesc, 15.0, 16.0)},
|
|
}
|
|
keys2 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows2)
|
|
|
|
valueRows2 := [][]interface{}{{int64(5)}, {int64(6)}, {int64(7)}, {int64(8)}}
|
|
values2 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows2)
|
|
|
|
for i, key := range keys2 {
|
|
err := mutableMap.Put(ctx, key, values2[i])
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Check that map looks how we expect.
|
|
newMap, err := flusher.Map(ctx, mutableMap)
|
|
require.NoError(t, err)
|
|
|
|
l2 := m1.tuples.Root.Level()
|
|
_ = l2
|
|
|
|
combinedKeyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
{encodeVector(t, keyDesc, 9.0, 10.0)},
|
|
{encodeVector(t, keyDesc, 11.0, 12.0)},
|
|
{encodeVector(t, keyDesc, 13.0, 14.0)},
|
|
{encodeVector(t, keyDesc, 15.0, 16.0)},
|
|
}
|
|
combinedKeys := buildTuples(t, ctx, ns, pb, keyDesc, combinedKeyRows)
|
|
|
|
combinedValueRows := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}, {int64(5)}, {int64(6)}, {int64(7)}, {int64(8)}}
|
|
combinedValues := buildTuples(t, ctx, ns, pb, testValDesc, combinedValueRows)
|
|
|
|
validateProximityMap(t, ctx, ns, &newMap, keyDesc, testValDesc, combinedKeys, combinedValues, logChunkSize)
|
|
})
|
|
}
|
|
|
|
// keyDescWithNodeStore returns a new NodeStore and a TupleDesc based on the input that uses it.
|
|
//
|
|
// TODO: this is necessary for VECTOR encoding because the map mutator needs to be able to call Compare() on the
|
|
// list of keys under edit, which requires the NodeStore for `BytesAdaptiveEnc`. For this use case, this comparison
|
|
// is expensive and unnecessary. We should change the map mutator used by vector indexes to use a custom comparator.
|
|
func keyDescWithNodeStore(keyDesc *val.TupleDesc) (tree.NodeStore, *val.TupleDesc) {
|
|
ns := tree.NewTestNodeStore()
|
|
keyDesc = val.NewTupleDescriptorWithArgs(val.TupleDescriptorArgs{
|
|
ValueStore: ns,
|
|
}, keyDesc.Types...)
|
|
return ns, keyDesc
|
|
}
|
|
|
|
func testIncrementalUpdates(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("incremental updates", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns, keyDesc := keyDescWithNodeStore(keyDesc)
|
|
pb := pool.NewBuffPool()
|
|
logChunkSize := uint8(1)
|
|
distanceType := vector.DistanceL2Squared{}
|
|
flusher := ProximityFlusher{logChunkSize: logChunkSize, distanceType: distanceType}
|
|
keyRows1 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
keys1 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows1)
|
|
|
|
valueRows1 := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values1 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows1)
|
|
|
|
m1 := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys1, testValDesc, values1, logChunkSize)
|
|
|
|
mutableMap := newProximityMutableMap(m1)
|
|
|
|
bp := pool.NewBuffPool()
|
|
|
|
keyBuilder := val.NewTupleBuilder(keyDesc, ns)
|
|
valueBuilder := val.NewTupleBuilder(testValDesc, ns)
|
|
|
|
// update leaf node
|
|
{
|
|
putVector(t, keyBuilder, encodeVector(t, keyDesc, 0.0, 1.0))
|
|
nextKey, err := keyBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
valueBuilder.PutInt64(0, 5)
|
|
nextValue, err := valueBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
err = mutableMap.Put(ctx, nextKey, nextValue)
|
|
require.NoError(t, err)
|
|
|
|
newMap, err := flusher.Map(ctx, mutableMap)
|
|
require.NoError(t, err)
|
|
|
|
newCount, err := newMap.Count()
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, 4, newCount)
|
|
|
|
// validate
|
|
|
|
combinedKeyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
combinedKeys := buildTuples(t, ctx, ns, pb, keyDesc, combinedKeyRows)
|
|
combinedValueRows := [][]interface{}{{int64(5)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
combinedValues := buildTuples(t, ctx, ns, pb, testValDesc, combinedValueRows)
|
|
|
|
validateProximityMap(t, ctx, ns, &newMap, keyDesc, testValDesc, combinedKeys, combinedValues, logChunkSize)
|
|
}
|
|
|
|
// update root node
|
|
{
|
|
putVector(t, keyBuilder, encodeVector(t, keyDesc, 5.0, 6.0))
|
|
nextKey, err := keyBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
valueBuilder.PutInt64(0, 6)
|
|
nextValue, err := valueBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
err = mutableMap.Put(ctx, nextKey, nextValue)
|
|
require.NoError(t, err)
|
|
|
|
newMap, err := flusher.Map(ctx, mutableMap)
|
|
require.NoError(t, err)
|
|
|
|
combinedKeyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
combinedKeys := buildTuples(t, ctx, ns, pb, keyDesc, combinedKeyRows)
|
|
combinedValueRows := [][]interface{}{{int64(5)}, {int64(2)}, {int64(6)}, {int64(4)}}
|
|
combinedValues := buildTuples(t, ctx, ns, pb, testValDesc, combinedValueRows)
|
|
|
|
validateProximityMap(t, ctx, ns, &newMap, keyDesc, testValDesc, combinedKeys, combinedValues, logChunkSize)
|
|
|
|
}
|
|
})
|
|
}
|
|
|
|
func testIncrementalDeletes(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("incremental deletes", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns, keyDesc := keyDescWithNodeStore(keyDesc)
|
|
pb := pool.NewBuffPool()
|
|
logChunkSize := uint8(1)
|
|
distanceType := vector.DistanceL2Squared{}
|
|
flusher := ProximityFlusher{logChunkSize: logChunkSize, distanceType: distanceType}
|
|
keyRows1 := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 1.0)},
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
keys1 := buildTuples(t, ctx, ns, pb, keyDesc, keyRows1)
|
|
|
|
valueRows1 := [][]interface{}{{int64(1)}, {int64(2)}, {int64(3)}, {int64(4)}}
|
|
values1 := buildTuples(t, ctx, ns, pb, testValDesc, valueRows1)
|
|
|
|
m1 := createAndValidateProximityMap(t, ctx, ns, keyDesc, keys1, testValDesc, values1, logChunkSize)
|
|
|
|
mutableMap := newProximityMutableMap(m1)
|
|
|
|
bp := pool.NewBuffPool()
|
|
|
|
keyBuilder := val.NewTupleBuilder(keyDesc, ns)
|
|
|
|
// delete leaf node
|
|
{
|
|
putVector(t, keyBuilder, encodeVector(t, keyDesc, 0.0, 1.0))
|
|
nextKey, err := keyBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
err = mutableMap.Put(ctx, nextKey, nil)
|
|
require.NoError(t, err)
|
|
|
|
newMap, err := flusher.Map(ctx, mutableMap)
|
|
require.NoError(t, err)
|
|
|
|
combinedKeyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 5.0, 6.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
combinedKeys := buildTuples(t, ctx, ns, pb, keyDesc, combinedKeyRows)
|
|
combinedValueRows := [][]interface{}{{int64(2)}, {int64(3)}, {int64(4)}}
|
|
combinedValues := buildTuples(t, ctx, ns, pb, testValDesc, combinedValueRows)
|
|
|
|
validateProximityMap(t, ctx, ns, &newMap, keyDesc, testValDesc, combinedKeys, combinedValues, logChunkSize)
|
|
|
|
}
|
|
|
|
// delete root node
|
|
{
|
|
putVector(t, keyBuilder, encodeVector(t, keyDesc, 5.0, 6.0))
|
|
nextKey, err := keyBuilder.Build(context.Background(), bp)
|
|
require.NoError(t, err)
|
|
|
|
err = mutableMap.Put(ctx, nextKey, nil)
|
|
require.NoError(t, err)
|
|
|
|
newMap, err := flusher.Map(ctx, mutableMap)
|
|
require.NoError(t, err)
|
|
|
|
combinedKeyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 3.0, 4.0)},
|
|
{encodeVector(t, keyDesc, 7.0, 8.0)},
|
|
}
|
|
combinedKeys := buildTuples(t, ctx, ns, pb, keyDesc, combinedKeyRows)
|
|
combinedValueRows := [][]interface{}{{int64(2)}, {int64(4)}}
|
|
combinedValues := buildTuples(t, ctx, ns, pb, testValDesc, combinedValueRows)
|
|
|
|
validateProximityMap(t, ctx, ns, &newMap, keyDesc, testValDesc, combinedKeys, combinedValues, logChunkSize)
|
|
|
|
}
|
|
})
|
|
}
|
|
|
|
// As part of the algorithm for building proximity maps, we store the map keys as bytestrings in a temporary table.
|
|
// The sorting order of a key is not always the same as the lexographic ordering of these bytestrings.
|
|
// This test makes sure that even when this is not the case we still generate correct output.
|
|
func testNonlexographicKey(t *testing.T, keyDesc *val.TupleDesc) {
|
|
t.Run("non-lexographic key", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
pb := pool.NewBuffPool()
|
|
|
|
testKeyDesc := val.NewTupleDescriptor(
|
|
keyDesc.Types[0],
|
|
val.Type{Enc: val.Int64Enc, Nullable: true},
|
|
)
|
|
|
|
valDesc := val.NewTupleDescriptor()
|
|
|
|
keyRows := [][]interface{}{
|
|
{encodeVector(t, keyDesc, 0.0, 0.0), int64(4 + 0*256)},
|
|
{encodeVector(t, keyDesc, 0.0, 0.0), int64(3 + 1*256)},
|
|
{encodeVector(t, keyDesc, 0.0, 0.0), int64(2 + 2*256)},
|
|
{encodeVector(t, keyDesc, 0.0, 0.0), int64(1 + 3*256)},
|
|
{encodeVector(t, keyDesc, 0.0, 0.0), int64(0 + 4*256)},
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, testKeyDesc, keyRows)
|
|
|
|
valueRows := [][]interface{}{{}, {}, {}, {}, {}}
|
|
values := buildTuples(t, ctx, ns, pb, valDesc, valueRows)
|
|
|
|
// The way the validation test is currently written it assumes that all vectors are unique, but this is not a
|
|
// requirement. Skip validation for now.
|
|
_ = createProximityMap(t, ctx, ns, testKeyDesc, keys, valDesc, values, 1)
|
|
})
|
|
}
|
|
|
|
func testManyDimensions(t *testing.T, keyDesc *val.TupleDesc) {
|
|
ctx := context.Background()
|
|
ns := tree.NewTestNodeStore()
|
|
numRows := 50
|
|
dimensions := 50
|
|
testManyDimensionsHelper(ctx, t, keyDesc, ns, numRows, dimensions)
|
|
}
|
|
|
|
func testManyDimensionsHelper(ctx context.Context, t *testing.T, keyDesc *val.TupleDesc, ns tree.NodeStore, numRows int, dimensions int) {
|
|
pb := pool.NewBuffPool()
|
|
testKeyDesc := val.NewTupleDescriptor(
|
|
keyDesc.Types[0],
|
|
val.Type{Enc: val.Int64Enc, Nullable: true},
|
|
)
|
|
|
|
valDesc := val.NewTupleDescriptor()
|
|
|
|
t.Run(fmt.Sprintf("numRows = %d, dimensions = %d", numRows, dimensions), func(t *testing.T) {
|
|
keyRows := make([][]interface{}, numRows)
|
|
valueRows := make([][]interface{}, numRows)
|
|
for i := 0; i < numRows; i++ {
|
|
keyRows[i] = []interface{}{makeManyDimensionalVector(keyDesc.Types[0].Enc, dimensions, int64(i)), i}
|
|
valueRows[i] = []interface{}{}
|
|
}
|
|
keys := buildTuples(t, ctx, ns, pb, testKeyDesc, keyRows)
|
|
values := buildTuples(t, ctx, ns, pb, testKeyDesc, valueRows)
|
|
|
|
_ = createAndValidateProximityMap(t, ctx, ns, testKeyDesc, keys, valDesc, values, 3)
|
|
})
|
|
}
|
|
|
|
func makeManyDimensionalVector(encoding val.Encoding, dimensions int, seed int64) interface{} {
|
|
rng := rand.New(rand.NewSource(seed))
|
|
switch encoding {
|
|
case val.JSONAddrEnc:
|
|
var builder strings.Builder
|
|
|
|
builder.WriteRune('[')
|
|
if dimensions > 0 {
|
|
|
|
builder.WriteString(strconv.Itoa(rng.Int()))
|
|
for d := 1; d < dimensions; d++ {
|
|
builder.WriteRune(',')
|
|
builder.WriteString(strconv.Itoa(rng.Int()))
|
|
}
|
|
}
|
|
builder.WriteRune(']')
|
|
return builder.String()
|
|
case val.BytesAdaptiveEnc:
|
|
result := make([]float32, dimensions)
|
|
for i := 0; i < dimensions; i++ {
|
|
result[i] = rng.Float32()
|
|
}
|
|
return sql.EncodeVector(result)
|
|
default:
|
|
panic("unexpected encoding")
|
|
}
|
|
}
|