scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-08-28 10:54:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
featurebase/batch/batch_test.go
Seebs 54dbeec1af support null/non-null tests for non-BSI fields 
There's a lot going on here. First, we were treating "the test is
a Condition" as implying BSI, which it doesn't anymore. Second, the
behavior of conditions was weird and BSI-specific. Third, we had
to propagate these changes and features throughout a bunch of code,
including both the core featurebase code and the DAX replacements/copies
of it, plus the SQL3 layer.

We refactor this so that tests for equality and inequality work for
non-BSI fields, so now if you accidentally use `==` in a Row call
on a non-BSI field, it still works; that's not specific to BSI
fields anymore.

We add a TrackExistence flag to fields, and propagate it through
things like our protobuf code, etcetera, so that we can successfully
create fields. Newly-created fields get this by default, because
we add it unconditionally to them, but the paths that are being
called with existing fields don't add it. So, when we "create"
(really, just load the definition of) a field from something stored
in the schema, we don't add TrackExistence to it, but any path to
creating a new field should.

A time quantum field with NoStandardView will *effectively*
lack TrackExistence.

For sets, mutexes, and time quantums with a standard view, anything
that sets bits will also set a corresponding bit for the record in
a new "existence" view. This allows us to distinguish between an
empty set and a null, and also allows null checks to be constant-time.

When clearing bits, we don't clear existence bits EXCEPT that if
you clear a bit in a mutex, *and the bit actually existed*, we clear
the existence bit. For sets and time quantums, clearing bits never
clears the existence bit.

Deleting records clears the existence bit.

We also add code to the `batch` subpackage to generate suitable
existence field bitmaps and import them. This logic correctly handles
empty sets and nils. The `batch` package does not allow specification
of anything equivalent to clearing a single bit from an existing
record, so we don't have to deal with the mutex complexity in that
case, which is good because it would be impossible.

This requires a number of other subtle changes, such as allowing
new fields to have more than one FieldOption specified for them.

We also drop the handful of implementation bits relating to the
"fullySorted" internal-use-only import flag, which existed only to
support the JSON ingest API, which we've removed.

The most dangerous part of this is that the mutex semantics are
impossible to implement on top of our existing API, because they
require us to know, not how *many* bits we cleared, but which
*specific* bits we cleared. I've implemented this as a new Tx method,
which is almost certainly going to be tech debt one day; if we some
day drop the Import API, we should remove that.

The testing for this is only currently covering the Set/Clear
behavior of PQL, and the Import API. The batch tests haven't been
written yet.

Fields that don't have existence tracking enabled refuse to perform
null/not-null tests. They should also report themselves as having
no null values -- if a record exists, sets in it are considered
empty rather than null.

The SQL3 support requires a number of subtle modifications to both
featurebase and some addon tooling. The essential thing is dropping
the unconditional translation of nil slices to non-nil empty slices
in translateResult, both in the executor and the orchestrator. We
also modify the logic that handles generating results from Extract
calls, to ensure that non-null sets get an empty slice created for
them even if they never have any values assigned.

The expected results for some tests are different now; we expect to
get nil slices, rather than 0-length non-nil slices, for fields which
were never written for a given record. Most tests were not changed.
(In every case, if a test was failing, I actually checked the logic
before changing expected results. This required a lot of tracking down
of edge cases.)

The batch package now rejects as an error attempts to clear single
bits from mutex fields, because so far as I can tell it's simply
impossible to have a roaring import that specifies the correct semantics
there; you can't tell whether to clear an existence bit without
access to the currently-set bits, which the batch API doesn't have.
We already supported the special case of specifying a clear value
of nil for clearing a mutex field; now that is the only allowed
value for a mutex field to have in row.Clears.

We change the logic for fixing up incoming view names (in two places)
to stop assuming that any view in a time field other than "" that does
not have viewStandard as a prefix is a partial time quantum name that
should have "standard_" prepended to it. This allows us to submit
bitmaps for "existence" to time quantum fields and not have them
silently transformed into "standard_existence" because that's what we'd
do with "202203".

We drop the field ClearBits method, which was totally unused.

We drop the sliceDifference function, which was used in a previous
mutex implementation and hasn't been used in ages, and the test
case for it, and the helper function used only by that test case.
2023-03-24 16:01:09 -05:00

2413 lines
61 KiB
Go
Raw Permalink Blame History

// Copyright 2021 Molecula Corp. All rights reserved.
package batch
import (
"context"
"fmt"
"math/rand"
"reflect"
"sort"
"strconv"
"testing"
"time"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/client"
"github.com/featurebasedb/featurebase/v3/pql"
"github.com/stretchr/testify/assert"
"github.com/pkg/errors"
)
func TestAgainstCluster(t *testing.T) {
cli, err := client.NewClient("featurebase:10101")
assert.NoError(t, err)
sapi := client.NewSchemaAPI(cli)
qapi := client.NewQueryAPI(cli)
importer := client.NewImporter(cli, sapi)
t.Run("string-slice-combos", func(t *testing.T) { testStringSliceCombos(t, importer, sapi, qapi) })
t.Run("import-batch-ints", func(t *testing.T) { testImportBatchInts(t, importer, sapi, qapi) })
t.Run("import-batch-bools", func(t *testing.T) { testImportBatchBools(t, importer, sapi, qapi) })
t.Run("import-batch-sorting", func(t *testing.T) { testImportBatchSorting(t, importer, sapi, qapi) })
t.Run("test-trim-null", func(t *testing.T) { testTrimNull(t, importer, sapi, qapi) })
t.Run("test-string-slice-empty-and-nil", func(t *testing.T) { testStringSliceEmptyAndNil(t, importer, sapi, qapi) })
t.Run("test-string-slice", func(t *testing.T) { testStringSlice(t, importer, sapi, qapi) })
t.Run("test-single-clear-batch-regression", func(t *testing.T) { testSingleClearBatchRegression(t, importer, sapi, qapi) })
t.Run("test-batches", func(t *testing.T) { testBatches(t, importer, sapi, qapi) })
t.Run("batches-strings-ids", func(t *testing.T) { testBatchesStringIDs(t, importer, sapi, qapi) })
t.Run("test-batch-staleness", func(t *testing.T) { testBatchStaleness(t, importer, sapi, qapi) })
t.Run("test-import-batch-multiple-ints", func(t *testing.T) { testImportBatchMultipleInts(t, importer, sapi, qapi) })
t.Run("test-import-batch-multiple-timestamps", func(t *testing.T) { testImportBatchMultipleTimestamps(t, importer, sapi, qapi) })
t.Run("test-import-batch-sets-clears", func(t *testing.T) { testImportBatchSetsAndClears(t, importer, sapi, qapi) })
t.Run("test-topn-cache-regression", func(t *testing.T) { testTopNCacheRegression(t, importer, sapi, qapi) })
t.Run("test-multiple-int-same-batch", func(t *testing.T) { testMultipleIntSameBatch(t, importer, sapi, qapi) })
t.Run("test-mutex-clearing-regression", func(t *testing.T) { mutexClearRegression(t, importer, sapi, qapi) })
t.Run("test-mutex-nil-clear-id", func(t *testing.T) { mutexNilClearID(t, importer, sapi, qapi) })
t.Run("test-mutex-nil-clear-key", func(t *testing.T) { mutexNilClearKey(t, importer, sapi, qapi) })
}
func testStringSliceCombos(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-string-slice-combos",
Fields: []*featurebase.FieldInfo{
{
Name: "a1",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 5, tbl, idx.Fields)
if err != nil {
t.Fatalf("creating new batch: %v", err)
}
records := []Row{
{ID: uint64(0), Values: []interface{}{[]string{"a", "b", "c"}}},
{ID: uint64(1), Values: []interface{}{[]string{"z"}}},
{ID: uint64(2), Values: []interface{}{[]string{}}},
{ID: uint64(3), Values: []interface{}{[]string{"q", "r", "s", "t", "c"}}},
{ID: uint64(4), Values: []interface{}{nil}},
{ID: uint64(5), Values: []interface{}{[]string{"a", "b", "c"}}},
{ID: uint64(6), Values: []interface{}{[]string{"a", "b", "c"}}},
{ID: uint64(7), Values: []interface{}{[]string{"z"}}},
{ID: uint64(8), Values: []interface{}{[]string{}}},
{ID: uint64(9), Values: []interface{}{[]string{"q", "r", "s", "t"}}},
{ID: uint64(10), Values: []interface{}{nil}},
{ID: uint64(11), Values: []interface{}{[]string{"a", "b", "c"}}},
{ID: uint64(12), Values: []interface{}{[]string{}}},
{ID: uint64(13), Values: []interface{}{[]string{}}},
}
assert.NoError(t, ingestRecords(records, b))
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "TopN(a1, n=10)",
})
if resp.Err != nil {
t.Fatalf("unexpected error from TopN query: %v", resp.Err)
}
if len(resp.Results) < 1 {
t.Fatalf("expected non-empty result set, got empty results")
}
pairsField, ok := resp.Results[0].(*featurebase.PairsField)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
rez := sortablePairs(pairsField.Pairs)
sort.Sort(rez)
exp := sortablePairs{
{Key: "a", Count: 4},
{Key: "b", Count: 4},
{Key: "c", Count: 5},
{Key: "q", Count: 2},
{Key: "r", Count: 2},
{Key: "s", Count: 2},
{Key: "t", Count: 2},
{Key: "z", Count: 2},
}
sort.Sort(exp)
errorIfNotEqual(t, exp, rez)
tests := []struct {
pql string
exp interface{}
}{
{
pql: "Row(a1='a')",
exp: []uint64{0, 5, 6, 11},
},
{
pql: "Row(a1='b')",
exp: []uint64{0, 5, 6, 11},
},
{
pql: "Row(a1='c')",
exp: []uint64{0, 3, 5, 6, 11},
},
{
pql: "Row(a1='z')",
exp: []uint64{1, 7},
},
{
pql: "Row(a1='q')",
exp: []uint64{3, 9},
},
{
pql: "Row(a1='r')",
exp: []uint64{3, 9},
},
{
pql: "Row(a1='s')",
exp: []uint64{3, 9},
},
{
pql: "Row(a1='t')",
exp: []uint64{3, 9},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
tresp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, tok := tresp.Results[0].(*featurebase.Row)
assert.True(t, tok, "wrong return type: %T", tresp.Results[0])
assert.Equal(t, test.exp, row.Columns())
})
}
resp = tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Count(All())",
})
count, ok := resp.Results[0].(uint64)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, uint64(14), count)
}
func errorIfNotEqual(t *testing.T, exp, got interface{}) {
t.Helper()
if !reflect.DeepEqual(exp, got) {
t.Errorf("unequal exp/got:\n%v\n%v", exp, got)
}
}
// sortablePairs is a sortable slice of featurebase.Pair.
type sortablePairs []featurebase.Pair
func (s sortablePairs) Len() int { return len(s) }
func (s sortablePairs) Less(i, j int) bool {
if s[i].Count != s[j].Count {
return s[i].Count > s[j].Count
}
if s[i].ID != s[j].ID {
return s[i].ID < s[j].ID
}
if s[i].Key != s[j].Key {
return s[i].Key < s[j].Key
}
return true
}
func (s sortablePairs) Swap(i, j int) {
s[i], s[j] = s[j], s[i]
}
func tq(t *testing.T, ctx context.Context, qapi featurebase.QueryAPI, query *featurebase.QueryRequest) featurebase.QueryResponse {
resp, err := qapi.Query(ctx, query)
assert.NoError(t, err)
return resp
}
func ingestRecords(records []Row, batch *Batch) error {
for _, rec := range records {
err := batch.Add(rec)
if err == ErrBatchNowFull {
err = batch.Import()
if err != nil {
return errors.Wrap(err, "importing batch")
}
} else if err != nil {
return errors.Wrap(err, "while adding record")
}
}
if batch.Len() > 0 {
err := batch.Import()
if err != nil {
return errors.Wrap(err, "importing batch")
}
}
return nil
}
func testImportBatchInts(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-import-batch-ints",
Fields: []*featurebase.FieldInfo{
{
Name: "anint",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Max: pql.NewDecimal(1000, 0),
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r := Row{Values: make([]interface{}, 1)}
for i := uint64(0); i < 3; i++ {
r.ID = i
r.Values[0] = int64(i)
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
r.ID = uint64(0)
r.Values[0] = nil
err = b.Add(r)
if err != nil {
t.Fatalf("adding after import: %v", err)
}
r.ID = uint64(1)
r.Values[0] = int64(7)
err = b.Add(r)
if err != nil {
t.Fatalf("adding second after import: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("second import: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(anint=0) Row(anint=7) Row(anint=2)",
})
assert.Equal(t, 3, len(resp.Results))
for i, result := range resp.Results {
row, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []uint64{uint64(i)}, row.Columns())
}
}
func testImportBatchSorting(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-import-batch-sorting",
Fields: []*featurebase.FieldInfo{
{
Name: "anint",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Max: pql.NewDecimal(10_000_000, 0),
},
},
{
Name: "amutex",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeMutex,
CacheType: featurebase.CacheTypeNone,
CacheSize: 0,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 100, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r := Row{Values: make([]interface{}, 2)}
rnd := rand.New(rand.NewSource(7))
// generate 100 records randomly spread/ordered across multiple
// shards to test sorting on int/mutex fields
for i := 0; i < 100; i++ {
id := rnd.Intn(10_000_000)
r.ID = uint64(id)
r.Values[0] = int64(id)
r.Values[1] = uint64(id)
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("second import: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Count(All())",
})
count, ok := resp.Results[0].(uint64)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, uint64(100), count)
}
func testTrimNull(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
fieldName := "empty"
idx := &featurebase.IndexInfo{
Name: "test-trim-null",
Fields: []*featurebase.FieldInfo{
{
Name: fieldName,
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Max: pql.NewDecimal(1000, 0),
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting batch: %v", err)
}
b.nullIndices = make(map[string][]uint64, 1)
b.nullIndices[fieldName] = []uint64{0, 1, 2}
r := Row{Values: make([]interface{}, 1)}
for i := 0; i < 3; i++ {
r.ID = uint64(i)
r.Values[0] = int64(i)
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(empty=0) Row(empty=1) Row(empty=2)",
})
assert.Equal(t, 3, len(resp.Results))
for _, result := range resp.Results {
row, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []uint64{}, row.Columns())
}
b, err = NewBatch(importer, 4, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r = Row{Values: make([]interface{}, 1)}
for i := 10; i < 40; i += 10 {
r.ID = uint64(i)
r.Values[0] = int64(i)
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
r.ID = uint64(40)
r.Values[0] = nil
err = b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
tests := []struct {
pql string
exp interface{}
}{
{
pql: "Row(empty=10)",
exp: []uint64{10},
},
{
pql: "Row(empty=40)",
exp: []uint64{},
},
{
pql: "Row(empty=20)",
exp: []uint64{20},
},
{
pql: "Row(empty=30)",
exp: []uint64{30},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, test.exp, row.Columns())
})
}
}
func testStringSliceEmptyAndNil(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-string-slice-nil",
Fields: []*featurebase.FieldInfo{
{
Name: "strslice",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
TrackExistence: true,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
// first create a batch and test adding a single value with empty
// string - this failed with a translation error at one point, and
// how we catch it and treat it like a nil.
b, err := NewBatch(importer, 2, tbl, idx.Fields)
if err != nil {
t.Fatalf("creating new batch: %v", err)
}
r := Row{Values: make([]interface{}, len(idx.Fields))}
r.ID = uint64(1)
r.Values[0] = ""
err = b.Add(r)
if err != nil {
t.Fatalf("adding: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
// now create a batch and add a mixture of string slice values
b, err = NewBatch(importer, 6, tbl, idx.Fields)
if err != nil {
t.Fatalf("creating new batch: %v", err)
}
r = Row{Values: make([]interface{}, len(idx.Fields))}
r.ID = uint64(0)
r.Values[0] = []string{"a"}
err = b.Add(r)
if err != nil {
t.Fatalf("adding to batch: %v", err)
}
r.ID = uint64(1)
r.Values[0] = nil
err = b.Add(r)
if err != nil {
t.Fatalf("adding batch with nil stringslice to r: %v", err)
}
r.ID = uint64(2)
r.Values[0] = []string{"a", "b", "z"}
err = b.Add(r)
if err != nil {
t.Fatalf("adding batch with idslice to r: %v", err)
}
r.ID = uint64(3)
r.Values[0] = []string{"b", "c"}
err = b.Add(r)
if err != nil {
t.Fatalf("adding batch with stringslice to r: %v", err)
}
r.ID = uint64(4)
r.Values[0] = []string{}
err = b.Add(r)
if err != nil {
t.Fatalf("adding batch with stringslice to r: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
tests := []struct {
pql string
exp interface{}
}{
{
pql: "Row(strslice='a')",
exp: []uint64{0, 2},
},
{
pql: "Row(strslice='b')",
exp: []uint64{2, 3},
},
{
pql: "Row(strslice='c')",
exp: []uint64{3},
},
{
pql: "Row(strslice='z')",
exp: []uint64{2},
},
{
pql: "Row(strslice==null)",
exp: []uint64{1},
},
{
pql: "Row(strslice!=null)",
exp: []uint64{0, 2, 3, 4},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, test.exp, row.Columns())
})
}
}
func testStringSlice(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-string-slice",
Fields: []*featurebase.FieldInfo{
{
Name: "strslice",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields)
if err != nil {
t.Fatalf("creating new batch: %v", err)
}
rowmap := map[string]uint64{
"c": 9,
"d": 10,
"f": 13,
}
b.rowTranslations["strslice"] = make(map[string]agedTranslation)
for k, id := range rowmap {
b.rowTranslations["strslice"][k] = agedTranslation{
id: id,
}
}
r := Row{Values: make([]interface{}, len(idx.Fields))}
r.ID = uint64(0)
r.Values[0] = []string{"a"}
err = b.Add(r)
if err != nil {
t.Fatalf("adding to batch: %v", err)
}
if got := b.toTranslateSets["strslice"]["a"]; !reflect.DeepEqual(got, []int{0}) {
t.Fatalf("expected []int{0}, got: %v", got)
}
r.ID = uint64(1)
r.Values[0] = []string{"a", "b", "c"}
err = b.Add(r)
if err != nil {
t.Fatalf("adding to batch: %v", err)
}
if got := b.toTranslateSets["strslice"]["a"]; !reflect.DeepEqual(got, []int{0, 1}) {
t.Fatalf("expected []int{0,1}, got: %v", got)
}
if got := b.toTranslateSets["strslice"]["b"]; !reflect.DeepEqual(got, []int{1}) {
t.Fatalf("expected []int{1}, got: %v", got)
}
if got, ok := b.toTranslateSets["strslice"]["c"]; ok {
t.Fatalf("should be nothing at c, got: %v", got)
}
if got := b.rowIDSets["strslice"][1]; !reflect.DeepEqual(got, []uint64{9}) {
t.Fatalf("expected c to map to rowID 9 but got %v", got)
}
r.ID = uint64(2)
r.Values[0] = []string{"d", "e", "f"}
err = b.Add(r)
if err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
if got, ok := b.toTranslateSets["strslice"]["d"]; ok {
t.Fatalf("should be nothing at d, got: %v", got)
}
if got, ok := b.toTranslateSets["strslice"]["f"]; ok {
t.Fatalf("should be nothing at f, got: %v", got)
}
if got := b.toTranslateSets["strslice"]["e"]; !reflect.DeepEqual(got, []int{2}) {
t.Fatalf("expected []int{2}, got: %v", got)
}
if got := b.rowIDSets["strslice"][2]; !reflect.DeepEqual(got, []uint64{10, 13}) {
t.Fatalf("expected c to map to rowID 9 but got %v", got)
}
err = b.doTranslation()
if err != nil {
t.Fatalf("translating: %v", err)
}
if got0 := b.rowIDSets["strslice"][0]; len(got0) != 1 {
t.Errorf("after translation, rec 0, wrong len: %v", got0)
} else if got1 := b.rowIDSets["strslice"][1]; len(got1) != 3 || got1[0] != 9 || (got1[1] != got0[0] && got1[2] != got0[0]) {
t.Errorf("after translation, rec 1: %v, rec 0: %v", got1, got0)
} else if got2 := b.rowIDSets["strslice"][2]; len(got2) != 3 || got2[0] != 10 || got2[1] != 13 || got2[2] == got1[2] || got2[2] == got0[0] {
t.Errorf("after translation, rec 2: %v", got2)
}
frags, clearFrags, err := b.makeFragments(make(fragments), make(fragments))
if err != nil {
t.Errorf("making fragments: %v", err)
}
err = b.doImport(frags, clearFrags)
if err != nil {
t.Fatalf("doing import: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(strslice='a')",
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, []uint64{0, 1}, row.Columns())
}
func testSingleClearBatchRegression(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-single-clear-batch-regression",
Fields: []*featurebase.FieldInfo{
{
Name: "zero",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
// Set a bit.
_ = tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Set(1, zero='row1')",
})
b, err := NewBatch(importer, 1, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting new batch: %v", err)
}
r := Row{ID: uint64(1), Values: make([]interface{}, len(idx.Fields)), Clears: make(map[int]interface{})}
r.Values[0] = nil
r.Clears[0] = "row1"
err = b.Add(r)
if err != ErrBatchNowFull {
t.Fatalf("wrong error from batch add: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("error importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(zero='row1')",
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, []uint64{}, row.Columns())
}
func testBatches(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-batches",
Fields: []*featurebase.FieldInfo{
{
Name: "zero",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
{
Name: "one",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
{
Name: "two",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
{
Name: "three",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Min: pql.NewDecimal(-1_000_000, 0),
Max: pql.NewDecimal(1_000_000, 0),
},
},
{
Name: "four",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeTime,
TimeQuantum: "YMD",
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 10, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting new batch: %v", err)
}
r := Row{Values: make([]interface{}, len(idx.Fields)), Clears: make(map[int]interface{})}
r.Time.Set(time.Date(2019, time.January, 2, 15, 45, 0, 0, time.UTC))
for i := 0; i < 9; i++ {
r.ID = uint64(i)
if i%2 == 0 {
r.Values[0] = "a"
r.Values[1] = "b"
r.Values[2] = "c"
r.Values[3] = int64(99)
r.Values[4] = uint64(1)
r.Time.SetMonth("01")
} else {
r.Values[0] = "x"
r.Values[1] = "y"
r.Values[2] = "z"
r.Values[3] = int64(-10)
r.Values[4] = uint64(1)
r.Time.SetMonth("02")
}
if i == 8 {
r.Values[0] = nil
r.Clears[1] = uint64(97)
r.Clears[2] = "c"
r.Values[3] = nil
r.Values[4] = nil
}
err := b.Add(r)
if err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
}
if len(b.toTranslate[0]) != 2 {
t.Fatalf("wrong number of keys in toTranslate[0]")
}
for k, ints := range b.toTranslate[0] {
if k == "a" {
if !reflect.DeepEqual(ints, []int{0, 2, 4, 6}) {
t.Fatalf("wrong ints for key a in field zero: %v", ints)
}
} else if k == "x" {
if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) {
t.Fatalf("wrong ints for key x in field zero: %v", ints)
}
} else {
t.Fatalf("unexpected key %s", k)
}
}
if !reflect.DeepEqual(b.toTranslateClear, map[int]map[string][]int{2: {"c": {8}}}) {
t.Errorf("unexpected toTranslateClear: %+v", b.toTranslateClear)
}
if !reflect.DeepEqual(b.clearRowIDs, map[int]map[int]uint64{1: {8: 97}, 2: {}}) {
t.Errorf("unexpected clearRowIDs: %+v", b.clearRowIDs)
}
if !reflect.DeepEqual(b.values["three"], []int64{99, -10, 99, -10, 99, -10, 99, -10, 0}) {
t.Fatalf("unexpected values: %v", b.values["three"])
}
if !reflect.DeepEqual(b.nullIndices["three"], []uint64{8}) {
t.Fatalf("unexpected nullIndices: %v", b.nullIndices["three"])
}
if len(b.toTranslate[1]) != 2 {
t.Fatalf("wrong number of keys in toTranslate[1]")
}
for k, ints := range b.toTranslate[1] {
if k == "b" {
if !reflect.DeepEqual(ints, []int{0, 2, 4, 6, 8}) {
t.Fatalf("wrong ints for key b in field one: %v", ints)
}
} else if k == "y" {
if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) {
t.Fatalf("wrong ints for key y in field one: %v", ints)
}
} else {
t.Fatalf("unexpected key %s", k)
}
}
if len(b.toTranslate[2]) != 2 {
t.Fatalf("wrong number of keys in toTranslate[2]")
}
for k, ints := range b.toTranslate[2] {
if k == "c" {
if !reflect.DeepEqual(ints, []int{0, 2, 4, 6, 8}) {
t.Fatalf("wrong ints for key c in field two: %v", ints)
}
} else if k == "z" {
if !reflect.DeepEqual(ints, []int{1, 3, 5, 7}) {
t.Fatalf("wrong ints for key z in field two: %v", ints)
}
} else {
t.Fatalf("unexpected key %s", k)
}
}
err = b.Add(r)
if err != ErrBatchNowFull {
t.Fatalf("should have gotten full batch error, but got %v", err)
}
err = b.Add(r)
if err != ErrBatchAlreadyFull {
t.Fatalf("should have gotten already full batch error, but got %v", err)
}
if !reflect.DeepEqual(b.values["three"], []int64{99, -10, 99, -10, 99, -10, 99, -10, 0, 0}) {
t.Fatalf("unexpected values: %v", b.values["three"])
}
err = b.doTranslation()
if err != nil {
t.Fatalf("doing translation: %v", err)
}
for fidx, rowIDs := range b.rowIDs {
// we don't know which key will get translated first, but we do know the pattern
if fidx == 0 {
if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, nilSentinel, nilSentinel}) &&
!reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, nilSentinel, nilSentinel}) {
t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs)
}
} else if fidx == 4 {
if !reflect.DeepEqual(rowIDs, []uint64{1, 1, 1, 1, 1, 1, 1, 1, nilSentinel, nilSentinel}) {
t.Fatalf("unexpected rowids for time field")
}
} else if fidx == 3 {
if len(rowIDs) != 0 {
t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs)
}
} else {
if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, 1, nilSentinel}) && !reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, 2, nilSentinel}) {
t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs)
}
}
}
if !reflect.DeepEqual(b.clearRowIDs[1], map[int]uint64{8: 97}) {
t.Errorf("unexpected clearRowIDs after translation: %+v", b.clearRowIDs[1])
}
if !reflect.DeepEqual(b.clearRowIDs[2], map[int]uint64{8: 2}) && !reflect.DeepEqual(b.clearRowIDs[2], map[int]uint64{8: 1}) {
t.Errorf("unexpected clearRowIDs: after translation%+v", b.clearRowIDs[2])
}
frags, clearFrags, err := b.makeFragments(make(fragments), make(fragments))
if err != nil {
t.Errorf("making fragments: %v", err)
}
err = b.doImport(frags, clearFrags)
if err != nil {
t.Fatalf("doing import: %v", err)
}
b.reset()
for i := 9; i < 19; i++ {
r.ID = uint64(i)
if i%2 == 0 {
r.Values[0] = "a"
r.Values[1] = "b"
r.Values[2] = "c"
r.Values[3] = int64(99)
r.Values[4] = uint64(1)
} else {
r.Values[0] = "x"
r.Values[1] = "y"
r.Values[2] = "z"
r.Values[3] = int64(-10)
r.Values[4] = uint64(2)
}
err := b.Add(r)
if i != 18 && err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
if i == 18 && err != ErrBatchNowFull {
t.Fatalf("unexpected err: %v", err)
}
}
// should do nothing
err = b.doTranslation()
if err != nil {
t.Fatalf("doing translation: %v", err)
}
frags, clearFrags, err = b.makeFragments(make(fragments), make(fragments))
if err != nil {
t.Errorf("making fragments: %v", err)
}
err = b.doImport(frags, clearFrags)
if err != nil {
t.Fatalf("doing import: %v", err)
}
for fidx, rowIDs := range b.rowIDs {
if fidx == 3 {
if len(rowIDs) != 0 {
t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs)
}
continue
}
// we don't know which key will get translated first, but we do know the pattern
if !reflect.DeepEqual(rowIDs, []uint64{1, 2, 1, 2, 1, 2, 1, 2, 1, 2}) && !reflect.DeepEqual(rowIDs, []uint64{2, 1, 2, 1, 2, 1, 2, 1, 2, 1}) {
t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs)
}
}
b.reset()
for i := 19; i < 29; i++ {
r.ID = uint64(i)
if i%2 == 0 {
r.Values[0] = "d"
r.Values[1] = "e"
r.Values[2] = "f"
r.Values[3] = int64(100)
r.Values[4] = uint64(3)
} else {
r.Values[0] = "u"
r.Values[1] = "v"
r.Values[2] = "w"
r.Values[3] = int64(0)
r.Values[4] = uint64(4)
}
err := b.Add(r)
if i != 28 && err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
if i == 28 && err != ErrBatchNowFull {
t.Fatalf("unexpected err: %v", err)
}
}
err = b.doTranslation()
if err != nil {
t.Fatalf("doing translation: %v", err)
}
frags, clearFrags, err = b.makeFragments(make(fragments), make(fragments))
if err != nil {
t.Errorf("making fragments: %v", err)
}
err = b.doImport(frags, clearFrags)
if err != nil {
t.Fatalf("doing import: %v", err)
}
for fidx, rowIDs := range b.rowIDs {
// we don't know which key will get translated first, but we do know the pattern
if fidx == 3 {
if len(rowIDs) != 0 {
t.Fatalf("expected no rowIDs for int field, but got: %v", rowIDs)
}
continue
}
if !reflect.DeepEqual(rowIDs, []uint64{3, 4, 3, 4, 3, 4, 3, 4, 3, 4}) && !reflect.DeepEqual(rowIDs, []uint64{4, 3, 4, 3, 4, 3, 4, 3, 4, 3}) {
t.Fatalf("unexpected row ids for field %d: %v", fidx, rowIDs)
}
}
frags, _, err = b.makeFragments(make(fragments), make(fragments))
if err != nil {
t.Fatalf("making fragments: %v", err)
}
var n int
for key := range frags {
if key.shard == 0 {
n++
}
}
if n != 5 { // zero, one, two, four (three is an int field so not in fragments) + _exists
t.Fatalf("there should be 5 views, but have %d", n)
}
tests := []struct {
pql string
exp interface{}
}{
{
pql: "Row(zero='a')",
exp: []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18},
},
{
pql: "Row(one='b')",
exp: []uint64{0, 2, 4, 6, 8, 10, 12, 14, 16, 18},
},
{
pql: "Row(two='c')",
exp: []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18},
},
{
pql: "Row(three=99)",
exp: []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18},
},
{
pql: "Row(zero='d')",
exp: []uint64{20, 22, 24, 26, 28},
},
{
pql: "Row(one='e')",
exp: []uint64{20, 22, 24, 26, 28},
},
{
pql: "Row(two='f')",
exp: []uint64{20, 22, 24, 26, 28},
},
{
pql: "Row(three > -11)",
exp: []uint64{0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28},
},
{
pql: "Row(three=0)",
exp: []uint64{19, 21, 23, 25, 27},
},
{
pql: "Row(three=100)",
exp: []uint64{20, 22, 24, 26, 28},
},
{
pql: "Row(four=1, from=2019-01-01T00:00, to=2019-01-29T00:00)",
exp: []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18},
},
{
pql: "Row(four=1, from=2019-02-01T00:00, to=2019-02-29T00:00)",
exp: []uint64{1, 3, 5, 7},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, test.exp, row.Columns())
})
}
b.reset()
r.ID = uint64(0)
r.Values[0] = "x"
r.Values[1] = "b"
r.Clears[0] = "a"
r.Clears[1] = "b" // b should get cleared
err = b.Add(r)
if err != nil {
t.Fatalf("adding with clears: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing w/clears: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(zero='a') Row(zero='x') Row(one='b')",
})
assert.Equal(t, 3, len(resp.Results))
for i, result := range resp.Results {
row, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
switch i {
case 0:
if arow := row.Columns(); arow[0] == 0 {
t.Errorf("shouldn't have id 0 in row a after clearing! %v", arow)
}
case 1:
if xrow := row.Columns(); xrow[0] != 0 {
t.Errorf("should have id 0 in row x after setting %v", xrow)
}
case 2:
if brow := row.Columns(); brow[0] == 0 {
t.Errorf("shouldn't have id 0 in row b after clearing! %v", brow)
}
}
}
// TODO test importing across multiple shards
}
func testBatchesStringIDs(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "batches-strings-ids",
Fields: []*featurebase.FieldInfo{
{
Name: "zero",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
Keys: true,
CacheType: featurebase.CacheTypeRanked,
CacheSize: 100,
},
},
{
Name: "one",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeMutex,
Keys: true,
CacheType: featurebase.CacheTypeNone,
CacheSize: 0,
},
},
{
Name: "two",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeTime,
Keys: true,
TimeQuantum: "YMDH",
},
},
},
Options: featurebase.IndexOptions{
Keys: true,
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting new batch: %v", err)
}
r := Row{Values: make([]interface{}, 3)}
r.Time.Set(time.Date(2019, time.January, 2, 15, 45, 0, 0, time.UTC))
for i := 0; i < 3; i++ {
r.ID = strconv.Itoa(i)
if i%2 == 0 {
r.Values[0] = "a"
r.Values[1] = "b"
r.Values[2] = "c"
r.Time.SetMonth("01")
} else {
r.Values[0] = "x"
r.Values[1] = "y"
r.Values[2] = "z"
r.Time.SetMonth("02")
}
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("unexpected err adding record: %v", err)
}
}
if len(b.toTranslateID) != 3 {
t.Fatalf("id translation table unexpected size: %v", b.toTranslateID)
}
for i, k := range b.toTranslateID {
if ik, err := strconv.Atoi(k); err != nil || ik != i {
t.Errorf("unexpected toTranslateID key %s at index %d", k, i)
}
}
err = b.doTranslation()
if err != nil {
t.Fatalf("translating: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
tests := []struct {
pql string
exp interface{}
}{
{
pql: "Row(zero='a')",
exp: []string{"0", "2"},
},
{
pql: "Row(zero='x')",
exp: []string{"1"},
},
{
pql: "Row(one='b')",
exp: []string{"0", "2"},
},
{
pql: "Row(one='y')",
exp: []string{"1"},
},
{
pql: "Row(two='c')",
exp: []string{"0", "2"},
},
{
pql: "Row(two='z')",
exp: []string{"1"},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-part1-%d", i), func(t *testing.T) {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.ElementsMatch(t, test.exp, row.Keys)
})
}
b.reset()
r.ID = "1"
r.Values[0] = "a"
err = b.Add(r)
if err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
r.ID = "3"
r.Values[0] = "z"
err = b.Add(r)
if err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
tests = []struct {
pql string
exp interface{}
}{
{
pql: "Row(zero='a')",
exp: []string{"0", "1", "2"},
},
{
pql: "Row(zero='z')",
exp: []string{"3"},
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-part2-%d", i), func(t *testing.T) {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: test.pql,
})
row, ok := resp.Results[0].(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.ElementsMatch(t, test.exp, row.Keys)
})
}
}
func TestQuantizedTime(t *testing.T) {
cases := []struct {
name string
time time.Time
year string
month string
day string
hour string
quantum featurebase.TimeQuantum
reset bool
exp []string
expErr string
}{
{
name: "no time quantum",
expErr: "",
},
{
name: "no time quantum with data",
year: "2017",
exp: []string{},
expErr: "",
},
{
name: "no data",
quantum: "Y",
exp: nil,
expErr: "",
},
{
name: "timestamp",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
quantum: "YMDH",
exp: []string{"2013", "201310", "20131016", "2013101617"},
},
{
name: "timestamp-less-granular",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
quantum: "YM",
exp: []string{"2013", "201310"},
},
{
name: "timestamp-mid-granular",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
quantum: "MD",
exp: []string{"201310", "20131016"},
},
{
name: "justyear",
year: "2013",
quantum: "Y",
exp: []string{"2013"},
},
{
name: "justyear-wantmonth",
year: "2013",
quantum: "YM",
expErr: "no data set for month",
},
{
name: "timestamp-changeyear",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
year: "2019",
quantum: "YMDH",
exp: []string{"2019", "201910", "20191016", "2019101617"},
},
{
name: "yearmonthdayhour",
year: "2013",
month: "10",
day: "16",
hour: "17",
quantum: "YMDH",
exp: []string{"2013", "201310", "20131016", "2013101617"},
},
{
name: "timestamp-changehour",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
hour: "05",
quantum: "MDH",
exp: []string{"201310", "20131016", "2013101605"},
},
{
name: "timestamp",
time: time.Date(2013, time.October, 16, 17, 34, 43, 0, time.FixedZone("UTC-5", -5*60*60)),
quantum: "YMDH",
reset: true,
exp: nil,
},
}
for i, test := range cases {
t.Run(test.name+strconv.Itoa(i), func(t *testing.T) {
tq := QuantizedTime{}
var zt time.Time
if zt != test.time {
tq.Set(test.time)
}
if test.year != "" {
tq.SetYear(test.year)
}
if test.month != "" {
tq.SetMonth(test.month)
}
if test.day != "" {
tq.SetDay(test.day)
}
if test.hour != "" {
tq.SetHour(test.hour)
}
if test.reset {
tq.Reset()
}
views, err := tq.views(test.quantum)
if !reflect.DeepEqual(views, test.exp) {
t.Errorf("unexpected views, got/want:\n%v\n%v\n", views, test.exp)
}
if (err != nil && err.Error() != test.expErr) || (err == nil && test.expErr != "") {
t.Errorf("unexpected error, got/want:\n%v\n%s\n", err, test.expErr)
}
})
}
}
func testBatchStaleness(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-batch-staleness",
Fields: []*featurebase.FieldInfo{
{
Name: "anint",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Min: pql.NewDecimal(-1_000_000, 0),
Max: pql.NewDecimal(1_000_000, 0),
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields, OptMaxStaleness(time.Millisecond))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r := Row{ID: uint64(0), Values: []interface{}{int64(0)}}
err = b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
// sleep so batch becomes stale
time.Sleep(time.Millisecond)
r = Row{ID: uint64(1), Values: []interface{}{int64(0)}}
err = b.Add(r)
if err != ErrBatchNowStale {
t.Fatal("batch expected to be stale")
}
}
func testImportBatchMultipleInts(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-import-batch-multi-int",
Fields: []*featurebase.FieldInfo{
{
Name: "anint",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Min: pql.NewDecimal(-1_000_000, 0),
Max: pql.NewDecimal(1_000_000, 0),
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 6, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r := Row{Values: make([]interface{}, 1)}
vals := []int64{16, 8, 32, 1, 2, 4}
for i := uint64(0); i < 6; i++ {
r.ID = uint64(1)
r.Values[0] = vals[i]
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(anint=4)",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
row, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []uint64{1}, row.Columns())
}
// testImportBatchMultipleTimestamps tests if nils are handles correctly for TS
// in batch imports
func testImportBatchMultipleTimestamps(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
fieldName := "ts2"
idx := &featurebase.IndexInfo{
Name: "test-import-batch-multi-timestamp",
Fields: []*featurebase.FieldInfo{
{
Name: fieldName,
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeTimestamp,
TimeUnit: "s",
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b1, err := NewBatch(importer, 6, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting batch 1: %v", err)
}
b2, err := NewBatch(importer, 6, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch 2: %v", err)
}
batches := []*Batch{b1, b2}
for j := 0; j < 2; j++ {
t.Run(fmt.Sprintf("batch %d", j), func(t *testing.T) {
b := batches[j]
r := Row{Values: make([]interface{}, 1)}
rawVals := []interface{}{int64(16), int64(8), int64(32), nil, int64(2), int64(4)}
chkVals := []int64{16, 8, 32, 0, 2, 4}
chkImport := []interface{}{time.Unix(16, 0), time.Unix(8, 0), time.Unix(32, 0), nil, time.Unix(2, 0), time.Unix(4, 0)}
cols := []uint64{0, 1, 2, 3, 4, 5}
for i := range cols {
r.ID = cols[i]
r.Values[0] = rawVals[i]
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
if b.nullIndices[fieldName][0] != 3 {
t.Fatalf("unexpected nulls, got/want: %v/%v", b.nullIndices[fieldName], []uint64{3})
}
for i, val := range chkVals {
if b.values[fieldName][i] != val {
t.Fatalf("unexpected value, got/want: %v/%v", b.values[fieldName][i], val)
}
if b.ids[i] != cols[i] {
t.Fatalf("unexpected id, got/want: %v/%v", b.ids[i], cols[i])
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
for i := range chkImport {
var timeStr string
if chkImport[i] == nil {
timeStr = "null"
} else {
ttime, ok := chkImport[i].(time.Time)
assert.True(t, ok)
timeStr = fmt.Sprintf(`"%s"`, ttime.Format(time.RFC3339))
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: fmt.Sprintf("Row(ts2 == %s)", timeStr),
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
row, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []uint64{uint64(i)}, row.Columns())
}
})
}
}
func testImportBatchSetsAndClears(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-import-batch-set-and-clear",
Fields: []*featurebase.FieldInfo{
{
Name: "aset",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
CacheType: featurebase.DefaultCacheType,
CacheSize: featurebase.DefaultCacheSize,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 6, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
r := Row{
Values: make([]interface{}, 1),
Clears: make(map[int]interface{}),
}
vals := []uint64{1, 2, 3, 1, 5, 6}
clears := []interface{}{nil, uint64(1), uint64(3), nil, uint64(2), uint64(4)}
for i := uint64(0); i < 6; i++ {
r.ID = i%3 + 1
r.Values[0] = vals[i]
if clears[i] != nil {
r.Clears[0] = clears[i]
}
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("adding to batch: %v", err)
}
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "TopN(aset, n=6)",
})
pairsField, ok := resp.Results[0].(*featurebase.PairsField)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, 3, len(pairsField.Pairs))
exp := [][]uint64{
{},
{1},
{},
{},
{},
{2},
{3},
}
for row := 0; row < 7; row++ {
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: fmt.Sprintf("Row(aset=%d)", row),
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
fRow, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, exp[row], fRow.Columns())
}
}
// testTopNCacheRegression recreates an issue we saw in an IDK test
// where if a value is completely removed (all bits unset from a row),
// it didn't get removed from the cache because a full recalculation
// had no way to clear the cache, it would just reset existing
// values. We added Clear on the cache interface to fix this.
func testTopNCacheRegression(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-topn-cache-regression",
Fields: []*featurebase.FieldInfo{
{
Name: "aset",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeSet,
CacheType: featurebase.DefaultCacheType,
CacheSize: featurebase.DefaultCacheSize,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
records := []struct {
ID uint64
Set interface{}
Clear interface{}
}{
{0, 1, nil},
{featurebase.ShardWidth, 1, nil},
{featurebase.ShardWidth * 2, nil, 1},
{featurebase.ShardWidth * 2, nil, 1},
{0, nil, 1},
{featurebase.ShardWidth, nil, 1},
{featurebase.ShardWidth, 1, nil},
{featurebase.ShardWidth, nil, nil},
}
for _, rec := range records {
if rec.Set != nil {
rec.Set = uint64(rec.Set.(int))
}
row := Row{
ID: rec.ID,
Values: []interface{}{rec.Set},
}
if rec.Clear != nil {
row.Clears = map[int]interface{}{0: uint64(rec.Clear.(int))}
}
err := b.Add(row)
if err == ErrBatchNowFull {
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
}
}
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "TopN(aset, n=6)",
})
pairsField, ok := resp.Results[0].(*featurebase.PairsField)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, 1, len(pairsField.Pairs))
assert.Equal(t, uint64(1), pairsField.Pairs[0].ID)
assert.Equal(t, uint64(1), pairsField.Pairs[0].Count)
}
// testMultipleIntSameBatch checks that if the same ID is added multiple times
// with different values that only the last value is set and the bits aren't
// mixed together. It adds a different ID in between the two same ones which
// triggered a bug because we were sorting by shard rather than ID.
func testMultipleIntSameBatch(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-multiple-int-same-batch",
Fields: []*featurebase.FieldInfo{
{
Name: "age",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeInt,
Max: pql.NewDecimal(10_000, 0),
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 4, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
if err := b.Add(Row{
ID: uint64(1),
Values: []interface{}{int64(1)},
}); err != nil {
t.Fatalf("adding to batch: %v", err)
}
if err := b.Add(Row{
ID: uint64(2),
Values: []interface{}{int64(0)},
}); err != nil {
t.Fatalf("adding to batch: %v", err)
}
if err := b.Add(Row{
ID: uint64(1),
Values: []interface{}{int64(2)},
}); err != nil {
t.Fatalf("adding to batch: %v", err)
}
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Sum(field=age)",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
row, ok := result.(featurebase.ValCount)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, int64(2), row.Val)
}
// mutexClearRegression checks for a bug where shards beyond the first
// one in a batch did not get any bits set in their clear bitmap, and
// in fact, all the bits were set in the clear bitmap for the first
// shard.
func mutexClearRegression(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-multiple-mut-same-batch",
Fields: []*featurebase.FieldInfo{
{
Name: "mut",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeMutex,
CacheType: featurebase.CacheTypeNone,
CacheSize: 0,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 11, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
col := uint64(0)
row := uint64(0)
for i := uint64(0); i <= 21; i++ {
col = (i%2+1)*featurebase.ShardWidth + i%5
row = i % 3
if err := b.Add(Row{
ID: col,
Values: []interface{}{row},
}); err == ErrBatchNowFull {
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "GroupBy(Rows(field=mut), Rows(field=mut))",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
groupCounts, ok := result.(*featurebase.GroupCounts)
assert.True(t, ok, "wrong return type: %T", result)
for j, gc := range groupCounts.Groups() {
assert.Equal(t, gc.Group[0].RowID, gc.Group[1].RowID, "zmismatched group at after %d batch: %d, %v", j, i, gc)
}
} else if err != nil {
t.Fatalf("adding to batch: %v", err)
}
}
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "GroupBy(Rows(field=mut), Rows(field=mut))",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
groupCounts, ok := result.(*featurebase.GroupCounts)
assert.True(t, ok, "wrong return type: %T", result)
for i, gc := range groupCounts.Groups() {
assert.Equal(t, gc.Group[0].RowID, gc.Group[1].RowID, "bmismatched group at %d, %v", i, gc)
}
}
// test clearing record with explict nil
func mutexNilClearID(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-mut-nil-clear-id",
Fields: []*featurebase.FieldInfo{
{
Name: "mut",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeMutex,
CacheType: featurebase.CacheTypeNone,
CacheSize: 0,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 11, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting batch: %v", err)
}
col := uint64(0)
row := uint64(0)
// populate mutex with some data
for i := uint64(0); i < 11; i++ {
col = (i%2+1)*featurebase.ShardWidth + i%5
row = i % 3
if err := b.Add(Row{
ID: col,
Values: []interface{}{row},
}); err == ErrBatchNowFull {
if err := b.Import(); err != nil {
t.Fatalf("importing: %v", err)
}
} else if err != nil {
t.Fatalf("adding to batch: %v", err)
}
}
// example data just copyied from test above
// confirm expected data
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(mut=0)",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
fRow, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
items := fRow.Columns()
// delete item 0
b.Add(
Row{
ID: items[0],
Values: []interface{}{nil},
Clears: map[int]interface{}{0: nil},
},
)
b.Import()
items = items[1:]
// confirm record removed
resp = tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(mut=0)",
})
assert.Equal(t, 1, len(resp.Results))
result = resp.Results[0]
fRow, ok = result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, items, fRow.Columns())
}
// similar test to above but with string keys
func mutexNilClearKey(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-mut-nil-clear-key",
Fields: []*featurebase.FieldInfo{
{
Name: "mut",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeMutex,
CacheType: featurebase.CacheTypeNone,
CacheSize: 0,
Keys: true,
},
},
},
Options: featurebase.IndexOptions{
Keys: true,
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields)
if err != nil {
t.Fatalf("getting new batch: %v", err)
}
r := Row{Values: make([]interface{}, 1)}
for i := 0; i < 3; i++ {
r.ID = strconv.Itoa(i)
if i%2 == 0 {
r.Values[0] = "a"
} else {
r.Values[0] = "x"
}
err := b.Add(r)
if err != nil && err != ErrBatchNowFull {
t.Fatalf("unexpected err adding record: %v", err)
}
}
if len(b.toTranslateID) != 3 {
t.Fatalf("id translation table unexpected size: %v", b.toTranslateID)
}
for i, k := range b.toTranslateID {
if ik, err := strconv.Atoi(k); err != nil || ik != i {
t.Errorf("unexpected toTranslateID key %s at index %d", k, i)
}
}
err = b.doTranslation()
if err != nil {
t.Fatalf("translating: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(mut='a')",
})
assert.Equal(t, 1, len(resp.Results))
result := resp.Results[0]
fRow, ok := result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []string{"0", "2"}, fRow.Keys)
r.ID = "2"
r.Values[0] = nil
r.Clears = map[int]interface{}{0: nil}
err = b.Add(r)
if err != nil {
t.Fatalf("unexpected err adding record: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("importing: %v", err)
}
resp = tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Row(mut='a')",
})
assert.Equal(t, 1, len(resp.Results))
result = resp.Results[0]
fRow, ok = result.(*featurebase.Row)
assert.True(t, ok, "wrong return type: %T", result)
assert.Equal(t, []string{"0"}, fRow.Keys)
}
func testImportBatchBools(t *testing.T, importer featurebase.Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) {
ctx := context.Background()
idx := &featurebase.IndexInfo{
Name: "test-import-batch-bools",
Fields: []*featurebase.FieldInfo{
{
Name: "boolcol",
Options: featurebase.FieldOptions{
Type: featurebase.FieldTypeBool,
},
},
},
Options: featurebase.IndexOptions{
TrackExistence: true,
},
}
tbl := featurebase.IndexInfoToTable(idx)
assert.NoError(t, sapi.CreateTable(ctx, tbl))
defer func() {
assert.NoError(t, sapi.DeleteTable(ctx, tbl.Name))
}()
b, err := NewBatch(importer, 3, tbl, idx.Fields, OptUseShardTransactionalEndpoint(true))
if err != nil {
t.Fatalf("getting new batch: %v", err)
}
r := Row{Values: make([]interface{}, 1)}
r.ID = uint64(0)
r.Values[0] = bool(false)
err = b.Add(r)
if err != nil {
t.Fatalf("adding after import: %v", err)
}
r.ID = uint64(1)
r.Values[0] = bool(true)
err = b.Add(r)
if err != nil {
t.Fatalf("adding second after import: %v", err)
}
err = b.Import()
if err != nil {
t.Fatalf("second import: %v", err)
}
resp := tq(t, ctx, qapi, &featurebase.QueryRequest{
Index: idx.Name,
Query: "Count(All())",
})
count, ok := resp.Results[0].(uint64)
assert.True(t, ok, "wrong return type: %T", resp.Results[0])
assert.Equal(t, uint64(2), count)
}
func TestConvert(t *testing.T) {
t.Run("timestampToInt", func(t *testing.T) {
tests := []struct {
unit TimeUnit
ts string
exp int64
}{
{unit: "s", ts: "2022-01-01T00:00:00Z", exp: 1640995200},
{unit: "ms", ts: "2022-01-01T00:00:00Z", exp: 1640995200000},
{unit: "us", ts: "2022-01-01T00:00:00Z", exp: 1640995200000000},
{unit: "ns", ts: "2022-01-01T00:00:00Z", exp: 1640995200000000000},
{unit: "x", ts: "2022-01-01T00:00:00Z", exp: 0},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
ts, err := time.Parse(time.RFC3339, test.ts)
assert.NoError(t, err)
v := timestampToInt(test.unit, ts)
assert.Equal(t, test.exp, v)
})
}
})
t.Run("Int64ToTimestamp", func(t *testing.T) {
tests := []struct {
unit TimeUnit
epoch string
val int64
exp time.Time
}{
{
unit: "ms",
epoch: "2022-01-01T00:00:00Z",
val: 0,
exp: time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC),
},
{
unit: "s",
epoch: "2022-01-01T00:00:00Z",
val: 86400,
exp: time.Date(2022, 1, 2, 0, 0, 0, 0, time.UTC),
},
}
for i, test := range tests {
t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
epoch, err := time.Parse(time.RFC3339, test.epoch)
assert.NoError(t, err)
ts, err := Int64ToTimestamp(test.unit, epoch, test.val)
assert.NoError(t, err)
assert.Equal(t, test.exp, ts)
})
}
})
}
Reference in a new issue View git blame Copy permalink