diff --git a/.gitignore b/.gitignore index 4d6f1a763..6b7c1cdd5 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,8 @@ builds/ *.tfstate.backup .vscode +batch/testdata/batch*.out + idk/testdata/idk*.out idk/testenv/certs/* diff --git a/.gitlab/.gitlab-ci.yml b/.gitlab/.gitlab-ci.yml index b2bb9e5d3..8bb8e2bfd 100644 --- a/.gitlab/.gitlab-ci.yml +++ b/.gitlab/.gitlab-ci.yml @@ -1,4 +1,5 @@ include: + - local: /.gitlab/batch-ci.yml - template: Security/SAST.gitlab-ci.yml - template: Security/License-Scanning.gitlab-ci.yml - template: Security/Dependency-Scanning.gitlab-ci.yml @@ -197,7 +198,7 @@ run go tests: retry: 1 script: - echo "Running featurebase unit tests..." - - go test -v -timeout=10m $(go list ./... | grep -Ev idk) + - go test -v -timeout=10m $(go list ./... | grep -Ev 'batch|idk') tags: - aws @@ -211,7 +212,7 @@ run go tests race: needs: ["smoke build"] # we do block on smoke build though bc it's pretty dumb to test stuff if it doesn't build script: - echo "Running featurebase race tests..." - - go test -race -v -timeout=10m $(go list ./... | grep -Ev idk) + - go test -race -v -timeout=10m $(go list ./... | grep -Ev 'batch|idk') tags: - aws @@ -223,7 +224,7 @@ run go tests shardwidth22: - if: '$CI_PIPELINE_SOURCE == "push" || $CI_PIPELINE_SOURCE == "schedule" || $CI_PIPELINE_SOURCE == "web"' script: - echo "Running featurebase shardwidth22 tests..." - - go test -timeout=10m -tags=shardwidth22 -v $(go list ./... | grep -Ev idk) + - go test -timeout=10m -tags=shardwidth22 -v $(go list ./... | grep -Ev 'batch|idk') tags: - aws @@ -239,8 +240,8 @@ run go tests future: retry: 1 script: - echo "Running featurebase unit tests..." - - PKG_LIST=$(go list ./... | grep -Ev 'internal/clustertests|simulacraData' | grep -Ev 'idk' | paste -s -d, -) - - go test -timeout=10m -json -coverprofile=coverage.out -covermode=atomic -coverpkg=${PKG_LIST} $(go list ./... | grep -Ev idk) | tee test-report.out + - PKG_LIST=$(go list ./... | grep -Ev 'internal/clustertests|simulacraData' | grep -Ev 'batch|idk' | paste -s -d, -) + - go test -timeout=10m -json -coverprofile=coverage.out -covermode=atomic -coverpkg=${PKG_LIST} $(go list ./... | grep -Ev 'batch|idk') | tee test-report.out artifacts: paths: - coverage.out @@ -257,8 +258,8 @@ run go tests future plg: retry: 1 script: - echo "Running featurebase plg-specific unit tests..." - - PKG_LIST=$(go list ./... | grep -Ev 'internal/clustertests|simulacraData' | grep -Ev 'idk' | paste -s -d, -) - - go test -tags=plg -timeout=10m -coverprofile=coverage-plg.out -covermode=atomic -coverpkg=${PKG_LIST} $(go list ./... | grep -Ev idk) | tee test-report-plg.out + - PKG_LIST=$(go list ./... | grep -Ev 'internal/clustertests|simulacraData' | grep -Ev 'batch|idk' | paste -s -d, -) + - go test -tags=plg -timeout=10m -coverprofile=coverage-plg.out -covermode=atomic -coverpkg=${PKG_LIST} $(go list ./... | grep -Ev 'batch|idk') | tee test-report-plg.out artifacts: paths: - coverage-plg.out @@ -399,6 +400,7 @@ run go tests idk sasl: needs: - job: build amd container fb + upload to sonarcloud: stage: integration image: sonarsource/sonar-scanner-cli:4.6 @@ -407,7 +409,7 @@ upload to sonarcloud: rules: - if: '$CI_COMMIT_TAG == null && ($CI_PIPELINE_SOURCE == "push" || $CI_PIPELINE_SOURCE == "schedule" || $CI_PIPELINE_SOURCE == "web")' script: - - sonar-scanner -Dsonar.projectKey=molecula_featurebase -Dsonar.organization=molecula -Dsonar.sources=. -Dsonar.host.url=https://sonarcloud.io -Dsonar.go.coverage.reportPaths=coverage*.out,results/coverage*out,idk/testdata/*coverage.out -Dsonar.go.tests.reportPaths=test-report*.out,idk/testdata/*report.out -Dsonar.javascript.lcov.reportPaths=lattice/coverage/lcov.info + - sonar-scanner -Dsonar.projectKey=molecula_featurebase -Dsonar.organization=molecula -Dsonar.sources=. -Dsonar.host.url=https://sonarcloud.io -Dsonar.go.coverage.reportPaths=coverage*.out,results/coverage*out,idk/testdata/*coverage.out,batch/testdata/*coverage.out -Dsonar.go.tests.reportPaths=test-report*.out,idk/testdata/*report.out,batch/testdata/*report.out -Dsonar.javascript.lcov.reportPaths=lattice/coverage/lcov.info needs: - job: run go tests future plg - job: run go tests future @@ -421,6 +423,8 @@ upload to sonarcloud: optional: true - job: run go tests idk 533 optional: true + - job: run go tests batch + optional: true package for linux amd64: stage: build diff --git a/.gitlab/batch-ci.yml b/.gitlab/batch-ci.yml new file mode 100644 index 000000000..21baf1014 --- /dev/null +++ b/.gitlab/batch-ci.yml @@ -0,0 +1,28 @@ +run go tests batch: + extends: + - .setup_ssh + variables: + USERNAME: fb-idk-access + PROJECT: batch_${CI_CONCURRENT_ID} + GIT_STRATEGY: clone + stage: test + retry: 1 + script: + - echo "Running test-all" + - cd ./batch/ + - echo $PROJECT + - echo $DOCKER_PASSWORD | docker login registry.gitlab.com --username "$USERNAME" --password-stdin + - BRANCH_NAME=${CI_COMMIT_REF_SLUG} make test-all + after_script: + - make save-pilosa-logs + - make shutdown + artifacts: + paths: + - ./batch/testdata/*_coverage.out + - ./batch/testdata/*_report.out + - ./batch/testdata/*_logs.txt + tags: + - shell + - aws + needs: + - job: build amd container fb diff --git a/Makefile b/Makefile index ca08e4646..6dad98652 100644 --- a/Makefile +++ b/Makefile @@ -45,9 +45,9 @@ vendor: go.mod version: @echo $(VERSION) -# We build a list of packages that omits the IDK packages because the IDK -# packages require fancy environment setup. -GOPACKAGES := $(shell $(GO) list ./... | grep -v "/idk") +# We build a list of packages that omits the IDK and batch packages because +# those packages require fancy environment setup. +GOPACKAGES := $(shell $(GO) list ./... | grep -v "/idk" | grep -v "/batch") # Run test suite test: diff --git a/api.go b/api.go index cecc11d3e..61a6a7060 100644 --- a/api.go +++ b/api.go @@ -3391,6 +3391,11 @@ type ComputeAPI interface { Txf() *TxFactory } +// QueryAPI is a subset of the API methods which have to do with query. +type QueryAPI interface { + Query(ctx context.Context, req *QueryRequest) (QueryResponse, error) +} + // FeatureBaseSchemaAPI is a wrapper around pilosa.API. It implements the // SchemaAPI interface with methods which are not a part of pilosa.API. type FeatureBaseSchemaAPI struct { diff --git a/batch/Dockerfile-test b/batch/Dockerfile-test new file mode 100644 index 000000000..dd6ca6962 --- /dev/null +++ b/batch/Dockerfile-test @@ -0,0 +1,11 @@ +ARG GO_VERSION=1.19 + +FROM golang:${GO_VERSION} + +WORKDIR /go/src/github.com/molecula/featurebase/ + +COPY . . + +WORKDIR /go/src/github.com/molecula/featurebase/batch/ + +CMD ["go","test","-v","-mod=vendor","-tags=odbc,dynamic","./..."] diff --git a/batch/Dockerfile-wait b/batch/Dockerfile-wait new file mode 100644 index 000000000..5b76cfb7b --- /dev/null +++ b/batch/Dockerfile-wait @@ -0,0 +1,8 @@ +FROM ubuntu:18.04 + +RUN ["apt-get", "update", "-y"] +RUN ["apt-get", "install", "-y", "curl", "netcat"] + +ADD wait.sh /wait + +ENTRYPOINT ["/wait"] diff --git a/batch/Makefile b/batch/Makefile new file mode 100644 index 000000000..dac750101 --- /dev/null +++ b/batch/Makefile @@ -0,0 +1,54 @@ +GO ?= go + +# We allow setting a custom docker-compose "project". Multiple of the +# same docker-compose environment can exist simultaneously as long as +# they use different projects (the project name is prepended to +# container names and such). This is useful in a CI environment where +# we might be running multiple instances of the tests concurrently. +PROJECT ?= batch +DOCKER_COMPOSE = docker-compose -p $(PROJECT) +BRANCH_NAME ?= "" + +.pulled: + $(DOCKER_COMPOSE) pull + touch .pulled + +vendor: ../go.mod + $(GO) mod vendor + +build-%: + $(DOCKER_COMPOSE) build $* + +pull-%: + $(DOCKER_COMPOSE) pull $* + +test-all: + $(MAKE) startup + $(MAKE) test-run + $(MAKE) shutdown + +start-all: .pulled build-wait + echo "branch name" ${BRANCH_NAME} + BRANCH_NAME=${BRANCH_NAME} $(DOCKER_COMPOSE) up -d featurebase + $(DOCKER_COMPOSE) run -T wait featurebase curl --silent --fail http://featurebase:10101/status + +startup: start-all + +shutdown: + $(DOCKER_COMPOSE) down -v --remove-orphans + rm -f .pulled + +save-%-logs: + $(DOCKER_COMPOSE) logs $* > ./testdata/$(PROJECT)_$*_logs.txt + +TCMD ?= ./... +# do "make startup", then e.g. "make test-run-local TCMD='-run=MyFavTest ./kafka'" +test-run-local: + pwd + $(DOCKER_COMPOSE) build batch-test + $(DOCKER_COMPOSE) run -T batch-test go test -mod=vendor -tags=odbc,dynamic $(TCMD) + +TPKG ?= ./... +test-run: vendor + $(DOCKER_COMPOSE) build batch-test + $(DOCKER_COMPOSE) run -T batch-test bash -c "set -o pipefail; go test -v -mod=vendor -tags=odbc,dynamic $(TPKG) -covermode=atomic -coverpkg=$(TPKG) -json -coverprofile=/testdata/$(PROJECT)_base_coverage.out | tee /testdata/$(PROJECT)_report.out" diff --git a/batch/README.md b/batch/README.md new file mode 100644 index 000000000..4ef894a83 --- /dev/null +++ b/batch/README.md @@ -0,0 +1,46 @@ +# batch + +The `batch` package provides a standard tool set for batching records in a way +that is most performant for ingesting those records into FeatureBase. The main +implementation is `Batch` (which can be initated with the `NewBatch()` +function). The `NewBatch()` function takes an `Importer` which contains all of + the methods required to interact with FeatureBase; these include methods for + doing string/id translation as well as for importing shards of data. + + IDK uses the `batch` package internally. Another example where the `batch` + package is used in the `sql3` package. When an "INSERT INTO" statement is + executed, the SQL engine uses a `Batch` to do key translation and build import + batches prior to doing the final import. +## Integration tests + +To run the tests, you will need to install the following dependencies: + +1. [Docker](https://docs.docker.com/install/) +2. [Docker Compose](https://docs.docker.com/compose/install/) + +In addition to these dependancies, you will need to be added to the molecula [Gitlab](https://registry.gitlab.com/molecula) account. + +First start the test environment. This is a docker-compose environment that includes featurebase. + + BRANCH_NAME=master make startup + +To build and run the integration tests, run: + + make test-run-local + +Then to shut down the test environment, run: + + make shutdown + +The previous command is equivalent to running the following: + + make startup + sleep 30 # wait for services to come up + make test-run + make shutdown + +To run an individual test, you can run the command directly using docker-compose. Note that you must run `docker-compose build batch-test` for docker to run the latest code. Modify the following as needed: + + make startup + docker-compose build batch-test + docker-compose run batch-test /usr/local/go/bin/go test -count=1 -mod=vendor -run=TestCmdMainOne . diff --git a/client/batch.go b/batch/batch.go similarity index 84% rename from client/batch.go rename to batch/batch.go index 0bbff77e5..90aa99c23 100644 --- a/client/batch.go +++ b/batch/batch.go @@ -1,16 +1,18 @@ -// Copyright 2021 Molecula Corp. All rights reserved. -package client +// Package batch provides tooling to prepare batches of records for ingest. +package batch import ( "bytes" + "context" "math/bits" "sort" "sync" "time" featurebase "github.com/molecula/featurebase/v3" - "github.com/molecula/featurebase/v3/client/egpool" + "github.com/molecula/featurebase/v3/batch/egpool" "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/pql" "github.com/molecula/featurebase/v3/roaring" "github.com/pkg/errors" ) @@ -18,6 +20,7 @@ import ( // Batch defaults. const ( DefaultKeyTranslateBatchSize = 100000 + existenceFieldName = "_exists" ) // TODO if using column translation, column ids might get way out of @@ -76,9 +79,9 @@ type agedTranslation struct { // Batch implements RecordBatch. // -// It supports Values of type string, uint64, int64, or nil. The -// following table describes what Pilosa field each type of value must -// map to. Fields are set up when calling "NewBatch". +// It supports Values of type string, uint64, int64, float64, or nil. The +// following table describes what Pilosa field each type of value must map to. +// Fields are set up when calling "NewBatch". // // | type | pilosa field type | options | // |--------+-------------------+-----------| @@ -91,10 +94,10 @@ type agedTranslation struct { // // nil values are ignored. type Batch struct { - client *Client - index *Index - header []*Field - headerMap map[string]*Field + importer Importer + index *featurebase.IndexInfo + header []*featurebase.FieldInfo + headerMap map[string]*featurebase.FieldInfo // prevDuration records the time that each doImport() takes. This // is used to set the timeout for transactions to a reasonable @@ -230,16 +233,25 @@ func OptUseShardTransactionalEndpoint(use bool) BatchOption { } } -// NewBatch initializes a new Batch object which will use the given -// Pilosa client, index, set of fields, and will take "size" records -// before returning ErrBatchNowFull. The positions of the Fields in -// 'fields' correspond to the positions of values in the Row's Values -// passed to Batch.Add(). -func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...BatchOption) (*Batch, error) { - if len(fields) == 0 || size == 0 { - return nil, errors.New("can't batch with no fields or batch size") +func OptImporter(i Importer) BatchOption { + return func(b *Batch) error { + b.importer = i + return nil } - headerMap := make(map[string]*Field, len(fields)) +} + +// NewBatch initializes a new Batch object which will use the given Importer, +// index, set of fields, and will take "size" records before returning +// ErrBatchNowFull. The positions of the Fields in 'fields' correspond to the +// positions of values in the Row's Values passed to Batch.Add(). +func NewBatch(importer Importer, size int, index *featurebase.IndexInfo, fields []*featurebase.FieldInfo, opts ...BatchOption) (*Batch, error) { + if len(fields) == 0 { + return nil, errors.New("can't batch with no fields") + } else if size == 0 { + return nil, errors.New("can't batch with no batch size") + } + + headerMap := make(map[string]*featurebase.FieldInfo, len(fields)) rowIDs := make(map[int][]uint64, len(fields)) values := make(map[string][]int64) boolValues := make(map[string]map[int]bool) @@ -248,35 +260,43 @@ func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...B ttSets := make(map[string]map[string][]int) hasTime := false for i, field := range fields { - headerMap[field.Name()] = field - opts := field.Opts() - switch typ := opts.Type(); typ { - case FieldTypeDefault, FieldTypeSet, FieldTypeTime: - if opts.Keys() { + headerMap[field.Name] = field + opts := field.Options + + // The client package has a FieldTypeDefault, but featurebase does not. + // When this code was moved from the client package to the batch + // package, FieldTypeDefault was no longer available. It probably isn't + // necessary, but to ensure backwards compatiblity, we continue to + // support it here with an unexported variable. + fieldTypeDefault := "" + + switch typ := opts.Type; typ { + case fieldTypeDefault, featurebase.FieldTypeSet, featurebase.FieldTypeTime: + if opts.Keys { tt[i] = make(map[string][]int) - ttSets[field.Name()] = make(map[string][]int) + ttSets[field.Name] = make(map[string][]int) } - hasTime = typ == FieldTypeTime || hasTime - case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp: + hasTime = typ == featurebase.FieldTypeTime || hasTime + case featurebase.FieldTypeInt, featurebase.FieldTypeDecimal, featurebase.FieldTypeTimestamp: // tt line only needed if int field is string foreign key tt[i] = make(map[string][]int) - values[field.Name()] = make([]int64, 0, size) - case FieldTypeMutex: + values[field.Name] = make([]int64, 0, size) + case featurebase.FieldTypeMutex: // similar to set/time fields, but no need to support sets // of values (hence no ttSets) - if opts.Keys() { + if opts.Keys { tt[i] = make(map[string][]int) } rowIDs[i] = make([]uint64, 0, size) - case FieldTypeBool: - boolValues[field.Name()] = make(map[int]bool) + case featurebase.FieldTypeBool: + boolValues[field.Name] = make(map[int]bool) default: return nil, errors.Errorf("field type '%s' is not currently supported through Batch", typ) } } b := &Batch{ - client: client, + importer: importer, header: fields, headerMap: headerMap, prevDuration: time.Minute * 11, @@ -375,7 +395,7 @@ func (qt *QuantizedTime) Reset() { // views builds the list of Pilosa views for this particular time, // given a quantum. -func (qt *QuantizedTime) views(q TimeQuantum) ([]string, error) { +func (qt *QuantizedTime) views(q featurebase.TimeQuantum) ([]string, error) { zero := QuantizedTime{} if *qt == zero { return nil, nil @@ -495,19 +515,19 @@ func (b *Batch) Add(rec Row) error { field := b.header[i] switch val := rec.Values[i].(type) { case string: - switch field.Opts().Type() { - case FieldTypeInt: + switch field.Options.Type { + case featurebase.FieldTypeInt: if val == "" { // copied from the `case nil:` section for ints and decimals - b.values[field.Name()] = append(b.values[field.Name()], 0) - nullIndices, ok := b.nullIndices[field.Name()] + b.values[field.Name] = append(b.values[field.Name], 0) + nullIndices, ok := b.nullIndices[field.Name] if !ok { nullIndices = make([]uint64, 0) } nullIndices = append(nullIndices, uint64(curPos)) - b.nullIndices[field.Name()] = nullIndices - } else if intVal, ok := b.getRowTranslation(field.Name(), val); ok { - b.values[field.Name()] = append(b.values[field.Name()], int64(intVal)) + b.nullIndices[field.Name] = nullIndices + } else if intVal, ok := b.getRowTranslation(field.Name, val); ok { + b.values[field.Name] = append(b.values[field.Name], int64(intVal)) } else { ints, ok := b.toTranslate[i][val] if !ok { @@ -515,9 +535,9 @@ func (b *Batch) Add(rec Row) error { } ints = append(ints, curPos) b.toTranslate[i][val] = ints - b.values[field.Name()] = append(b.values[field.Name()], 0) + b.values[field.Name] = append(b.values[field.Name], 0) } - case FieldTypeBool: + case featurebase.FieldTypeBool: // If we want to support bools as string values, we would do // that here. default: @@ -530,7 +550,7 @@ func (b *Batch) Add(rec Row) error { if val == "" { // b.rowIDs[i] = append(rowIDs, nilSentinel) - } else if rowID, ok := b.getRowTranslation(field.Name(), val); ok { + } else if rowID, ok := b.getRowTranslation(field.Name, val); ok { b.rowIDs[i] = append(rowIDs, rowID) } else { ints, ok := b.toTranslate[i][val] @@ -549,15 +569,15 @@ func (b *Batch) Add(rec Row) error { } b.rowIDs[i] = append(b.rowIDs[i], val) case int64: - b.values[field.Name()] = append(b.values[field.Name()], val) + b.values[field.Name] = append(b.values[field.Name], val) case []string: if len(val) == 0 { continue } - rowIDSets, ok := b.rowIDSets[field.Name()] + rowIDSets, ok := b.rowIDSets[field.Name] if !ok { rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids)) - b.rowIDSets[field.Name()] = rowIDSets + b.rowIDSets[field.Name] = rowIDSets } for len(rowIDSets) < len(b.ids)-1 { rowIDSets = append(rowIDSets, nil) // nil extend @@ -568,53 +588,53 @@ func (b *Batch) Add(rec Row) error { if k == "" { continue } - if rowID, ok := b.getRowTranslation(field.Name(), k); ok { + if rowID, ok := b.getRowTranslation(field.Name, k); ok { rowIDs = append(rowIDs, rowID) } else { - ttsets, ok := b.toTranslateSets[field.Name()] + ttsets, ok := b.toTranslateSets[field.Name] if !ok { ttsets = make(map[string][]int) - b.toTranslateSets[field.Name()] = make(map[string][]int) + b.toTranslateSets[field.Name] = make(map[string][]int) } ints, ok := ttsets[k] if !ok { ints = make([]int, 0, 1) } ints = append(ints, curPos) - b.toTranslateSets[field.Name()][k] = ints + b.toTranslateSets[field.Name][k] = ints } } - b.rowIDSets[field.Name()] = append(rowIDSets, rowIDs) + b.rowIDSets[field.Name] = append(rowIDSets, rowIDs) case []uint64: if len(val) == 0 { continue } - rowIDSets, ok := b.rowIDSets[field.Name()] + rowIDSets, ok := b.rowIDSets[field.Name] if !ok { rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids)) } for len(rowIDSets) < len(b.ids)-1 { rowIDSets = append(rowIDSets, nil) // nil extend } - b.rowIDSets[field.Name()] = append(rowIDSets, val) + b.rowIDSets[field.Name] = append(rowIDSets, val) case nil: - switch field.Opts().Type() { - case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp: - b.values[field.Name()] = append(b.values[field.Name()], 0) - nullIndices, ok := b.nullIndices[field.Name()] + switch field.Options.Type { + case featurebase.FieldTypeInt, featurebase.FieldTypeDecimal, featurebase.FieldTypeTimestamp: + b.values[field.Name] = append(b.values[field.Name], 0) + nullIndices, ok := b.nullIndices[field.Name] if !ok { nullIndices = make([]uint64, 0) } nullIndices = append(nullIndices, uint64(curPos)) - b.nullIndices[field.Name()] = nullIndices + b.nullIndices[field.Name] = nullIndices - case FieldTypeBool: - boolNulls, ok := b.boolNulls[field.Name()] + case featurebase.FieldTypeBool: + boolNulls, ok := b.boolNulls[field.Name] if !ok { boolNulls = make([]uint64, 0) } boolNulls = append(boolNulls, uint64(curPos)) - b.boolNulls[field.Name()] = boolNulls + b.boolNulls[field.Name] = boolNulls default: // only append nil to rowIDs if this field already has @@ -629,7 +649,10 @@ func (b *Batch) Add(rec Row) error { } case bool: - b.boolValues[field.Name()][curPos] = val + b.boolValues[field.Name][curPos] = val + + case pql.Decimal: + b.values[field.Name] = append(b.values[field.Name], val.ToInt64(field.Options.Scale)) default: return errors.Errorf("Val %v Type %[1]T is not currently supported. Use string, uint64 (row id), or int64 (integer value)", val) @@ -645,7 +668,7 @@ func (b *Batch) Add(rec Row) error { case string: clearRows := b.clearRowIDs[i] // translate val and add to clearRows - if rowID, ok := b.getRowTranslation(field.Name(), val); ok { + if rowID, ok := b.getRowTranslation(field.Name, val); ok { clearRows[curPos] = rowID } else { _, ok := b.toTranslateClear[i] @@ -662,14 +685,14 @@ func (b *Batch) Add(rec Row) error { case uint64: b.clearRowIDs[i][curPos] = val case nil: - if field.Opts().Type() == FieldTypeMutex { + if field.Options.Type == featurebase.FieldTypeMutex { for len(b.rowIDs[i]) <= curPos { b.rowIDs[i] = append(b.rowIDs[i], nilSentinel) } b.rowIDs[i][len(b.rowIDs[i])-1] = clearSentinel } default: - return errors.Errorf("Clearing a value '%v' Type %[1]T is not currently supported (field '%s')", val, field.Name()) + return errors.Errorf("Clearing a value '%v' Type %[1]T is not currently supported (field '%s')", val, field.Name) } // nil extend b.rowIDs so we don't run into a horrible bug // where we skip doing clears because b.rowIDs doesn't have a @@ -693,8 +716,8 @@ func (b *Batch) Add(rec Row) error { return nil } -// ErrBatchNowFull, similar to io.EOF, is a marker error to notify the -// user of a batch that it is time to call Import. +// ErrBatchNowFull — similar to io.EOF — is a marker error to notify the user of +// a batch that it is time to call Import. var ErrBatchNowFull = errors.New("batch is now full - you cannot add any more records (though the one you just added was accepted)") // ErrBatchAlreadyFull is a real error saying that Batch.Add did not @@ -714,17 +737,19 @@ var ErrBatchNowStale = errors.New("batch is stale and needs to be imported (howe // continues. split batch mode DOES NOT CURRENTLY SUPPORT MUTEX // OR INT FIELDS! func (b *Batch) Import() error { + ctx := context.Background() start := time.Now() - trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour) + trns, err := b.importer.StartTransaction(ctx, "", b.prevDuration*10, false, time.Hour) if err != nil { return errors.Wrap(err, "starting transaction") } defer func() { - trnsl, err := b.client.FinishTransaction(trns.ID) - if err != nil { - b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trnsl) + if trns != nil { + if trnsl, err := b.importer.FinishTransaction(ctx, trns.ID); err != nil { + b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trnsl) + } } - b.client.Stats.Timing(MetricBatchImportDurationSeconds, time.Since(start), 1.0) + b.importer.StatsTiming(MetricBatchImportDurationSeconds, time.Since(start), 1.0) }() size := len(b.ids) @@ -781,21 +806,23 @@ func (b *Batch) Import() error { // imports the stored data to Pilosa. Otherwise it simply returns // nil. func (b *Batch) Flush() error { + ctx := context.Background() + if !b.splitBatchMode { return nil } start := time.Now() - trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour) + trns, err := b.importer.StartTransaction(ctx, "", b.prevDuration*10, false, time.Hour) if err != nil { return errors.Wrap(err, "starting transaction") } defer func() { - trnsl, err := b.client.FinishTransaction(trns.ID) + trnsl, err := b.importer.FinishTransaction(ctx, trns.ID) if err != nil { b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trnsl) } - b.client.Stats.Timing(MetricBatchFlushDurationSeconds, time.Since(start), 1.0) + b.importer.StatsTiming(MetricBatchFlushDurationSeconds, time.Since(start), 1.0) }() importStart := time.Now() @@ -886,7 +913,7 @@ func (b *Batch) doTranslation() error { // Look up the associated field. field := b.header[i] - fieldName := field.Name() + fieldName := field.Name // Fetch the translation cache. rowCache := b.rowTranslations[fieldName] @@ -924,8 +951,8 @@ func (b *Batch) doTranslation() error { } rowCacheLock.Unlock() - switch ftype := field.Opts().Type(); ftype { - case FieldTypeSet, FieldTypeMutex, FieldTypeTime: + switch ftype := field.Options.Type; ftype { + case featurebase.FieldTypeSet, featurebase.FieldTypeMutex, featurebase.FieldTypeTime: // Fill out missing IDs in local batch records with translated IDs. rows := b.rowIDs[i] for key, idxs := range tt { @@ -952,7 +979,7 @@ func (b *Batch) doTranslation() error { } } - case FieldTypeInt: + case featurebase.FieldTypeInt: // Handle foreign key int fields — fill out b.values instead of b.rows. vals := b.values[fieldName] for key, idxs := range tt { @@ -1036,10 +1063,12 @@ func (b *Batch) doTranslation() error { return eg.Wait() } -func (b *Batch) createIndexKeys(index *Index, keys ...string) (map[string]uint64, error) { +func (b *Batch) createIndexKeys(index *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + ctx := context.Background() + batchSize := b.keyTranslateBatchSize if batchSize <= 0 || len(keys) <= batchSize { - return b.client.CreateIndexKeys(index, keys...) + return b.importer.CreateIndexKeys(ctx, index, keys...) } results := make(map[string]uint64, len(keys)) @@ -1049,7 +1078,7 @@ func (b *Batch) createIndexKeys(index *Index, keys ...string) (map[string]uint64 keySlice = keySlice[:batchSize] } - trans, err := b.client.CreateIndexKeys(index, keySlice...) + trans, err := b.importer.CreateIndexKeys(ctx, index, keySlice...) if err != nil { return nil, err } else if len(trans) != len(keySlice) { @@ -1065,10 +1094,12 @@ func (b *Batch) createIndexKeys(index *Index, keys ...string) (map[string]uint64 return results, nil } -func (b *Batch) createFieldKeys(field *Field, keys ...string) (map[string]uint64, error) { +func (b *Batch) createFieldKeys(field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + ctx := context.Background() + batchSize := b.keyTranslateBatchSize if batchSize <= 0 || len(keys) <= batchSize { - return b.client.CreateFieldKeys(field, keys...) + return b.importer.CreateFieldKeys(ctx, b.index.Name, field, keys...) } results := make(map[string]uint64, len(keys)) @@ -1078,7 +1109,7 @@ func (b *Batch) createFieldKeys(field *Field, keys ...string) (map[string]uint64 keySlice = keySlice[:batchSize] } - trans, err := b.client.CreateFieldKeys(field, keySlice...) + trans, err := b.importer.CreateFieldKeys(ctx, b.index.Name, field, keySlice...) if err != nil { return nil, err } else if len(trans) != len(keySlice) { @@ -1095,6 +1126,8 @@ func (b *Batch) createFieldKeys(field *Field, keys ...string) (map[string]uint64 } func (b *Batch) doImportShardTransactional(frags, clearFrags fragments) error { + ctx := context.Background() + start := time.Now() requests := make(map[uint64]*featurebase.ImportRoaringShardRequest) getOrCreate := func(requests map[uint64]*featurebase.ImportRoaringShardRequest, shard uint64) *featurebase.ImportRoaringShardRequest { @@ -1149,24 +1182,25 @@ func (b *Batch) doImportShardTransactional(frags, clearFrags fragments) error { } } - b.client.Stats.Timing(MetricBatchShardImportBuildRequestsSeconds, time.Since(start), 1.0) + b.importer.StatsTiming(MetricBatchShardImportBuildRequestsSeconds, time.Since(start), 1.0) start = time.Now() eg := egpool.Group{PoolSize: 20} for shard, request := range requests { shard := shard request := request eg.Go(func() error { - return b.client.ImportRoaringShard(b.index.Name(), shard, request) + return b.importer.ImportRoaringShard(ctx, b.index.Name, shard, request) }) } err := eg.Wait() dur := time.Since(start) - b.client.Stats.Timing(MetricBatchShardImportDurationSeconds, dur, 1.0) + b.importer.StatsTiming(MetricBatchShardImportDurationSeconds, dur, 1.0) b.log.Printf("import shard took: %v\n", dur) return errors.Wrap(err, "doing shard-transactional imports") } func (b *Batch) doImport(frags, clearFrags fragments) error { + ctx := context.Background() start := time.Now() eg := egpool.Group{PoolSize: 20} @@ -1188,7 +1222,7 @@ func (b *Batch) doImport(frags, clearFrags fragments) error { clearViewMap := clearFrags.GetViewMap(shard, field) if len(clearViewMap) > 0 { startx := time.Now() - err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, clearViewMap, true) + err := b.importer.ImportRoaringBitmap(ctx, b.index.Name, b.indexField(field), shard, clearViewMap, true) if err != nil { return errors.Wrapf(err, "import clearing clearing data for %s", field) } @@ -1196,7 +1230,7 @@ func (b *Batch) doImport(frags, clearFrags fragments) error { } starty := time.Now() - err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, viewMap, false) + err := b.importer.ImportRoaringBitmap(ctx, b.index.Name, b.indexField(field), shard, viewMap, false) b.log.Debugf("imp-roar %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(starty)) return errors.Wrapf(err, "importing data for %s", field) }) @@ -1215,6 +1249,22 @@ func (b *Batch) doImport(frags, clearFrags fragments) error { return nil } +// indexField is a helper function which was introduced when we switched the +// index and field types from being client types (e.g client.Index, +// client.Field) to being featurebase types (e.g. featurebase.IndexInfo, +// featurebase.FieldInfo). Unlike client.Index, featurebase.IndexInfo is not +// expected to contain the "_exists" field. So calling Field("_exists") on +// IndexInfo results in a nil field. This method creates an instance of +// FieldInfo for the "_exists" field. +func (b *Batch) indexField(field string) *featurebase.FieldInfo { + if field == existenceFieldName { + return &featurebase.FieldInfo{ + Name: existenceFieldName, + } + } + return b.index.Field(field) +} + func anyCause(cause error, errs ...error) error { if cause == nil { return nil @@ -1229,11 +1279,7 @@ func anyCause(cause error, errs ...error) error { } func (b *Batch) shardWidth() uint64 { - shardWidth := b.index.ShardWidth() - if shardWidth == 0 { - shardWidth = DefaultShardWidth - } - return shardWidth + return featurebase.ShardWidth } // this is kind of bad as it means we can never import column id @@ -1251,7 +1297,7 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments emptyClearRows := make(map[int]uint64) // create _exists fragments if needed - if b.index.Opts().TrackExistence() { + if b.index.Options.TrackExistence { var curBM *roaring.Bitmap curShard := ^uint64(0) // impossible sentinel value for shard. for _, col := range b.ids { @@ -1272,8 +1318,8 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments clearRows = emptyClearRows } field := b.header[i] - opts := field.Opts() - if opts.Type() == FieldTypeMutex { + opts := field.Options + if opts.Type == featurebase.FieldTypeMutex { continue // we handle mutex fields separately — they can't use importRoaring } curShard := ^uint64(0) // impossible sentinel value for shard. @@ -1291,8 +1337,8 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments if col/shardWidth != curShard { curShard = col / shardWidth - curBM = frags.GetOrCreate(curShard, field.Name(), "") - clearBM = clearFrags.GetOrCreate(curShard, field.Name(), "") + curBM = frags.GetOrCreate(curShard, field.Name, "") + clearBM = clearFrags.GetOrCreate(curShard, field.Name, "") } if row != nilSentinel { // TODO this is super ugly, but we want to avoid setting @@ -1300,16 +1346,16 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments // there isn't one. Should probably refactor this whole // loop to be more general w.r.t. views. Also... tests for // the NoStandardView case would be great. - if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) { + if !(opts.Type == featurebase.FieldTypeTime && opts.NoStandardView) { curBM.DirectAdd(row*shardWidth + (col % shardWidth)) } - if opts.Type() == FieldTypeTime { - views, err := b.times[j].views(opts.TimeQuantum()) + if opts.Type == featurebase.FieldTypeTime { + views, err := b.times[j].views(opts.TimeQuantum) if err != nil { return nil, nil, errors.Wrap(err, "calculating views") } for _, view := range views { - tbm := frags.GetOrCreate(curShard, field.Name(), view) + tbm := frags.GetOrCreate(curShard, field.Name, view) tbm.DirectAdd(row*shardWidth + (col % shardWidth)) } } @@ -1337,7 +1383,7 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments rowIDSets = rowIDSets[:len(b.ids)] } field := b.headerMap[fname] - opts := field.Opts() + opts := field.Options curShard := ^uint64(0) // impossible sentinel value for shard. var curBM *roaring.Bitmap for j := range b.ids { @@ -1354,13 +1400,13 @@ func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments // there isn't one. Should probably refactor this whole // loop to be more general w.r.t. views. Also... tests for // the NoStandardView case would be great. - if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) { + if !(opts.Type == featurebase.FieldTypeTime && opts.NoStandardView) { for _, row := range rowIDs { curBM.DirectAdd(row*shardWidth + (col % shardWidth)) } } - if opts.Type() == FieldTypeTime { - views, err := b.times[j].views(opts.TimeQuantum()) + if opts.Type == featurebase.FieldTypeTime { + views, err := b.times[j].views(opts.TimeQuantum) if err != nil { return nil, nil, errors.Wrap(err, "calculating views") } @@ -1409,8 +1455,8 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, sort.Stable(sc) } field := b.headerMap[fieldName] - base := field.Options().base - if field.Options().Type() == FieldTypeTimestamp { + base := field.Options.Base + if field.Options.Type == featurebase.FieldTypeTimestamp { base = 0 } @@ -1454,7 +1500,7 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, // ------------------------- for findex, rowIDs := range b.rowIDs { field := b.header[findex] - if field.Opts().Type() != FieldTypeMutex { + if field.Options.Type != featurebase.FieldTypeMutex { continue } ids = ids[:0] @@ -1483,8 +1529,8 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, } shard := ids[0] / shardWidth - bitmap := frags.GetOrCreate(shard, field.Name(), "standard") - clearBM := clearFrags.GetOrCreate(shard, field.Name(), "standard") + bitmap := frags.GetOrCreate(shard, field.Name, "standard") + clearBM := clearFrags.GetOrCreate(shard, field.Name, "standard") for i, id := range ids { if i+1 < len(ids) { // we only want the last value set for each id @@ -1495,8 +1541,8 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, row := rowIDs[i] if shard != id/shardWidth { shard = id / shardWidth - bitmap = frags.GetOrCreate(shard, field.Name(), "standard") - clearBM = clearFrags.GetOrCreate(shard, field.Name(), "standard") + bitmap = frags.GetOrCreate(shard, field.Name, "standard") + clearBM = clearFrags.GetOrCreate(shard, field.Name, "standard") } fragmentColumn := id % shardWidth clearBM.Add(fragmentColumn) // Will use this to clear columns. @@ -1521,15 +1567,16 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, // records (for all rows) to clear. for fieldname, boolNulls := range b.boolNulls { field := b.headerMap[fieldname] - if field.Opts().Type() != featurebase.FieldTypeBool { + if field.Options.Type != featurebase.FieldTypeBool { continue } for _, pos := range boolNulls { recID := b.ids[pos] shard := recID / shardWidth - clearBM := clearFrags.GetOrCreate(shard, field.Name(), "standard") + clearBM := clearFrags.GetOrCreate(shard, field.Name, "standard") fragmentColumn := recID % shardWidth + clearBM.Add(fragmentColumn) } } @@ -1540,7 +1587,7 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, // the bit in the "true" row. for fieldname, boolMap := range b.boolValues { field := b.headerMap[fieldname] - if field.Opts().Type() != featurebase.FieldTypeBool { + if field.Options.Type != featurebase.FieldTypeBool { continue } @@ -1548,8 +1595,8 @@ func (b *Batch) makeSingleValFragments(frags, clearFrags fragments) (fragments, recID := b.ids[pos] shard := recID / shardWidth - bitmap := frags.GetOrCreate(shard, field.Name(), "standard") - clearBM := clearFrags.GetOrCreate(shard, field.Name(), "standard") + bitmap := frags.GetOrCreate(shard, field.Name, "standard") + clearBM := clearFrags.GetOrCreate(shard, field.Name, "standard") fragmentColumn := recID % shardWidth clearBM.Add(fragmentColumn) @@ -1582,10 +1629,9 @@ func (v *valsByIDsSortable) Swap(i, j int) { // importValueData imports data for int fields. func (b *Batch) importValueData() error { - shardWidth := b.index.ShardWidth() - if shardWidth == 0 { - shardWidth = DefaultShardWidth - } + ctx := context.Background() + + shardWidth := uint64(featurebase.ShardWidth) eg := egpool.Group{PoolSize: 20} ids := make([]uint64, len(b.ids)) @@ -1630,15 +1676,15 @@ func (b *Batch) importValueData() error { endIdx := i shard := curShard field := b.headerMap[fieldName] - path, data, err := b.client.EncodeImportValues(field, shard, bvalues[startIdx:endIdx], ids[startIdx:endIdx], false) + path, data, err := b.importer.EncodeImportValues(ctx, b.index.Name, field, shard, bvalues[startIdx:endIdx], ids[startIdx:endIdx], false) if err != nil { return errors.Wrap(err, "encoding import values") } eg.Go(func() error { start := time.Now() - err := b.client.DoImportValues(b.index.Name(), shard, path, data) + err := b.importer.DoImport(ctx, b.index.Name, field, shard, path, data) b.log.Debugf("imp-vals %s,shard:%d,data:%d %v", field, shard, len(data), time.Since(start)) - return errors.Wrapf(err, "importing values for field = %s", field) + return errors.Wrapf(err, "importing values for field = %s", field.Name) }) startIdx = i curShard = recordID / shardWidth @@ -1673,16 +1719,15 @@ func (v *rowsByIDsSortable) Swap(i, j int) { // TODO this should work for bools as well - just need to support them // at batch creation time and when calling Add, I think. func (b *Batch) importMutexData() error { - shardWidth := b.index.ShardWidth() - if shardWidth == 0 { - shardWidth = DefaultShardWidth - } + ctx := context.Background() + + shardWidth := uint64(featurebase.ShardWidth) eg := egpool.Group{PoolSize: 20} ids := make([]uint64, 0, len(b.ids)) for findex, rowIDs := range b.rowIDs { field := b.header[findex] - if field.Opts().Type() != FieldTypeMutex { + if field.Options.Type != featurebase.FieldTypeMutex { continue } ids = ids[:0] @@ -1724,15 +1769,15 @@ func (b *Batch) importMutexData() error { endIdx := i shard := curShard field := field - path, data, err := b.client.EncodeImport(field, shard, rowIDs[startIdx:endIdx], ids[startIdx:endIdx], false) + path, data, err := b.importer.EncodeImport(ctx, b.index.Name, field, shard, rowIDs[startIdx:endIdx], ids[startIdx:endIdx], false) if err != nil { return errors.Wrap(err, "encoding mutex import") } eg.Go(func() error { start := time.Now() - err := b.client.DoImport(b.index.Name(), shard, path, data) - b.log.Debugf("imp-mux %s,shard:%d,data:%d %v", field.Name(), shard, len(data), time.Since(start)) - return errors.Wrapf(err, "importing values for field = %s", field) + err := b.importer.DoImport(ctx, b.index.Name, field, shard, path, data) + b.log.Debugf("imp-mux %s,shard:%d,data:%d %v", field.Name, shard, len(data), time.Since(start)) + return errors.Wrapf(err, "importing values for field = %s", field.Name) }) startIdx = i curShard = recordID / shardWidth diff --git a/batch/batch_test.go b/batch/batch_test.go new file mode 100644 index 000000000..77db39ab4 --- /dev/null +++ b/batch/batch_test.go @@ -0,0 +1,2359 @@ +package batch + +import ( + "context" + "fmt" + "math/rand" + "reflect" + "sort" + "strconv" + "testing" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/client" + "github.com/molecula/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) + + importer := client.NewImporter(cli) + sapi := client.NewSchemaAPI(cli) + qapi := client.NewQueryAPI(cli) + + 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 5, idx, 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)", + }) + 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 100, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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, idx, 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 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, + }, + }, + }, + Options: featurebase.IndexOptions{ + TrackExistence: true, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.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, idx, 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, idx, 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}, + }, + } + 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + // Set a bit. + _ = tq(t, ctx, qapi, &featurebase.QueryRequest{ + Index: idx.Name, + Query: "Set(1, zero='row1')", + }) + + b, err := NewBatch(importer, 1, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 10, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 6, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b1, err := NewBatch(importer, 6, idx, idx.Fields) + if err != nil { + t.Fatalf("getting batch 1: %v", err) + } + b2, err := NewBatch(importer, 6, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 6, idx, 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 beacuse 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 4, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 11, idx, idx.Fields, OptUseShardTransactionalEndpoint(true)) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + + col := uint64(0) + row := uint64(1) + 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 11, idx, idx.Fields, OptUseShardTransactionalEndpoint(true)) + if err != nil { + t.Fatalf("getting batch: %v", err) + } + + col := uint64(0) + row := uint64(1) + // 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 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, + }, + } + + createIndexAndFields(t, ctx, sapi, idx) + defer func() { + assert.NoError(t, sapi.DeleteIndex(ctx, idx.Name)) + }() + + b, err := NewBatch(importer, 3, idx, 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 createIndexAndFields(t *testing.T, ctx context.Context, sapi featurebase.SchemaAPI, idx *featurebase.IndexInfo) { + fields := make([]featurebase.CreateFieldObj, 0, len(idx.Fields)) + for _, fld := range idx.Fields { + opts := []featurebase.FieldOption{} + if fld.Options.Keys { + opts = append(opts, featurebase.OptFieldKeys()) + } + switch fld.Options.Type { + case featurebase.FieldTypeMutex: + opts = append(opts, featurebase.OptFieldTypeMutex(fld.Options.CacheType, fld.Options.CacheSize)) + case featurebase.FieldTypeSet: + opts = append(opts, featurebase.OptFieldTypeSet(fld.Options.CacheType, fld.Options.CacheSize)) + case featurebase.FieldTypeInt: + opts = append(opts, featurebase.OptFieldTypeInt(fld.Options.Min.ToInt64(0), fld.Options.Max.ToInt64(0))) + case featurebase.FieldTypeTime: + opts = append(opts, featurebase.OptFieldTypeTime(fld.Options.TimeQuantum, fld.Options.TTL.String(), fld.Options.NoStandardView)) + case featurebase.FieldTypeTimestamp: + opts = append(opts, featurebase.OptFieldTypeTimestamp(time.Unix(0, 0), fld.Options.TimeUnit)) + case featurebase.FieldTypeBool: + opts = append(opts, featurebase.OptFieldTypeBool()) + default: + t.Fatalf("unsupported field type: %s", fld.Options.Type) + } + field := featurebase.CreateFieldObj{ + Name: fld.Name, + Options: opts, + } + fields = append(fields, field) + } + + assert.NoError(t, sapi.CreateIndexAndFields(ctx, + idx.Name, + idx.Options, + fields, + )) +} diff --git a/batch/convert.go b/batch/convert.go new file mode 100644 index 000000000..85b530c8b --- /dev/null +++ b/batch/convert.go @@ -0,0 +1,145 @@ +package batch + +import ( + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/errors" +) + +var ( + MinTimestampNano = time.Unix(-1<<32, 0).UTC() // 1833-11-24T17:31:44Z + MaxTimestampNano = time.Unix(1<<32, 0).UTC() // 2106-02-07T06:28:16Z + MinTimestamp = time.Unix(-62135596799, 0).UTC() // 0001-01-01T00:00:01Z + MaxTimestamp = time.Unix(253402300799, 0).UTC() // 9999-12-31T23:59:59Z + + ErrTimestampOutOfRange = errors.New("", "value provided for timestamp field is out of range") +) + +type TimeUnit string + +const ( + TimeUnitSeconds = TimeUnit(featurebase.TimeUnitSeconds) + TimeUnitMilliseconds = TimeUnit(featurebase.TimeUnitMilliseconds) + TimeUnitMicroseconds = TimeUnit(featurebase.TimeUnitMicroseconds) + TimeUnitUSeconds = TimeUnit(featurebase.TimeUnitUSeconds) + TimeUnitNanoseconds = TimeUnit(featurebase.TimeUnitNanoseconds) +) + +// TimestampToInt64 converts the provided timestamp to an int64 as the number of +// units past the epoch. +func TimestampToInt64(unit TimeUnit, epoch time.Time, ts time.Time) (int64, error) { + var err error + + unit, err = validateTimeUnit(unit) + if err != nil { + return 0, errors.Wrap(err, "validating time unit") + } + + epoch, err = validateEpoch(epoch) + if err != nil { + return 0, errors.Wrap(err, "validating epoch") + } + + // Check if the epoch alone is out-of-range. If so, ingest should halt, + // regardless of state of the timestamp out-of-range CLI option. + if err := validateTimestamp(unit, epoch); err != nil { + return 0, errors.Wrap(err, "validating epoch") + } + + epochAsInt64 := timestampToInt(unit, epoch) + + // Check if the timestamp is out-of-range. + if err := validateTimestamp(unit, ts); err != nil { + return 0, errors.Wrapf(ErrTimestampOutOfRange, "validating timestamp: %s", ts) + } + + tsAsInt64 := timestampToInt(unit, ts) + + return tsAsInt64 - epochAsInt64, nil +} + +// validateTimeUnit checks if the time unit is supported. If the provided unit +// is blank, validateTimeUnit returns the default TimeUnit. +func validateTimeUnit(unit TimeUnit) (TimeUnit, error) { + switch unit { + case "": + return TimeUnitSeconds, nil + + case TimeUnitSeconds, + TimeUnitMilliseconds, + TimeUnitMicroseconds, + TimeUnitUSeconds, + TimeUnitNanoseconds: + return unit, nil + } + + return "", errors.Errorf("unsupported time unit: %s", unit) +} + +// validateEpoch checks if the epoch is supported. If the provided epoch +// is "zero", validateEpoch returns the default epoch value. +func validateEpoch(epoch time.Time) (time.Time, error) { + if epoch.IsZero() { + return time.Unix(0, 0), nil + } + return epoch, nil +} + +// validateTimestamp checks if the timestamp is within the range of what FB accepts. +func validateTimestamp(unit TimeUnit, ts time.Time) error { + // Min and Max timestamps that Featurebase accepts + var minStamp, maxStamp time.Time + switch unit { + case TimeUnitNanoseconds: + minStamp = MinTimestampNano + maxStamp = MaxTimestampNano + default: + minStamp = MinTimestamp + maxStamp = MaxTimestamp + } + + if ts.Before(minStamp) || ts.After(maxStamp) { + return errors.Errorf("timestamp value (%v) must be within min: %v and max: %v", ts, minStamp, maxStamp) + } + return nil +} + +// timestampToInt takes a time unit and a time.Time and converts it to an +// integer value. +func timestampToInt(unit TimeUnit, ts time.Time) int64 { + switch unit { + case TimeUnitSeconds: + return ts.Unix() + case TimeUnitMilliseconds: + return ts.UnixMilli() + case TimeUnitMicroseconds, TimeUnitUSeconds: + return ts.UnixMicro() + case TimeUnitNanoseconds: + return ts.UnixNano() + } + return 0 +} + +// intToTimestamp takes a timeunit and an integer value and converts it to +// time.Time. +func intToTimestamp(unit TimeUnit, val int64) (time.Time, error) { + switch unit { + case TimeUnitSeconds: + return time.Unix(val, 0).UTC(), nil + case TimeUnitMilliseconds: + return time.UnixMilli(val).UTC(), nil + case TimeUnitMicroseconds, TimeUnitUSeconds: + return time.UnixMicro(val).UTC(), nil + case TimeUnitNanoseconds: + return time.Unix(0, val).UTC(), nil + default: + return time.Time{}, errors.Errorf("Unknown time unit: '%v'", unit) + } +} + +// Int64ToTimestamp converts the provided int64 to a timestamp based on the time unit +// and epoch. +func Int64ToTimestamp(unit TimeUnit, epoch time.Time, val int64) (time.Time, error) { + return intToTimestamp(unit, timestampToInt(unit, epoch)+val) +} diff --git a/batch/docker-compose.yml b/batch/docker-compose.yml new file mode 100644 index 000000000..09b2c108c --- /dev/null +++ b/batch/docker-compose.yml @@ -0,0 +1,26 @@ +version: '3' + +services: + featurebase: + build: + context: ../. + dockerfile: ./Dockerfile + environment: + PILOSA_DATA_DIR: /data + PILOSA_BIND: 0.0.0.0:10101 + PILOSA_BIND_GRPC: 0.0.0.0:20101 + PILOSA_ADVERTISE: featurebase:10101 + volumes: + - ./testdata:/testdata + + batch-test: + build: + context: ../. + dockerfile: ./batch/Dockerfile-test + volumes: + - ./testdata:/testdata + + wait: + build: + context: . + dockerfile: Dockerfile-wait diff --git a/client/egpool/egpool.go b/batch/egpool/egpool.go similarity index 100% rename from client/egpool/egpool.go rename to batch/egpool/egpool.go diff --git a/client/egpool/egpool_test.go b/batch/egpool/egpool_test.go similarity index 91% rename from client/egpool/egpool_test.go rename to batch/egpool/egpool_test.go index af5132e50..c068ac245 100644 --- a/client/egpool/egpool_test.go +++ b/batch/egpool/egpool_test.go @@ -5,7 +5,7 @@ import ( "errors" "testing" - "github.com/molecula/featurebase/v3/client/egpool" + "github.com/molecula/featurebase/v3/batch/egpool" ) func TestEGPool(t *testing.T) { diff --git a/batch/error.go b/batch/error.go new file mode 100644 index 000000000..3ba31fc0f --- /dev/null +++ b/batch/error.go @@ -0,0 +1,8 @@ +package batch + +import "github.com/pkg/errors" + +// Predefined batch-related errors. +var ( + ErrPreconditionFailed = errors.New("Precondition failed") +) diff --git a/batch/importer.go b/batch/importer.go new file mode 100644 index 000000000..7f1e54478 --- /dev/null +++ b/batch/importer.go @@ -0,0 +1,211 @@ +package batch + +import ( + "context" + "time" + + "github.com/golang/protobuf/proto" //nolint:staticcheck + featurebase "github.com/molecula/featurebase/v3" + featurebaseproto "github.com/molecula/featurebase/v3/encoding/proto" + "github.com/molecula/featurebase/v3/pb" + "github.com/molecula/featurebase/v3/roaring" + "github.com/pkg/errors" +) + +type Importer interface { + StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) + FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) + CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) + CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) + ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error + ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error + EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) + EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) + DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error + + StatsTiming(name string, value time.Duration, rate float64) +} + +// Ensure type implements interface. +var _ Importer = &nopImporter{} + +// NopImporter is an implementation of the Importer interface that doesn't do +// anything. +var NopImporter Importer = &nopImporter{} + +type nopImporter struct{} + +func (n *nopImporter) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return nil, nil +} +func (n *nopImporter) FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) { + return nil, nil +} +func (n *nopImporter) CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + return nil, nil +} +func (n *nopImporter) CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + return nil, nil +} +func (n *nopImporter) ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + return nil +} +func (n *nopImporter) ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error { + return nil +} +func (n *nopImporter) EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + return "", nil, nil +} +func (n *nopImporter) EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + return "", nil, nil +} +func (n *nopImporter) DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error { + return nil +} + +func (n *nopImporter) StatsTiming(name string, value time.Duration, rate float64) {} + +// Ensure type implements interface. +var _ Importer = &FeaturebaseImporter{} + +// FeaturebaseImporter is a wrapper around featurebase.API, making it a +// batch.Importer. +type FeaturebaseImporter struct { + *featurebase.API +} + +func (f *FeaturebaseImporter) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return f.API.StartTransaction(ctx, id, timeout, exclusive, false) +} + +func (f *FeaturebaseImporter) FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) { + return f.API.FinishTransaction(ctx, id, false) +} + +func (f *FeaturebaseImporter) CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + return f.API.CreateIndexKeys(ctx, idx.Name, keys...) +} + +func (f *FeaturebaseImporter) CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + return f.API.CreateFieldKeys(ctx, index, field.Name, keys...) +} + +func (f *FeaturebaseImporter) ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + vs := make(map[string][]byte) + for k, v := range views { + data := roaring.BitmapsToRoaring([]*roaring.Bitmap{v}) + if len(data) > 0 { + vs[k] = data + } + } + req := &featurebase.ImportRoaringRequest{ + IndexCreatedAt: 0, + FieldCreatedAt: field.CreatedAt, + Clear: clear, + Views: vs, + } + return f.API.ImportRoaring(ctx, index, field.Name, shard, false, req) +} + +// ImportRoaringShard doesn't technically need to be implemented here, because +// the method on f.API has the same signature and already satisfies the +// interface. But we put this here to avoid possible confusion. +func (f *FeaturebaseImporter) ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error { + return f.API.ImportRoaringShard(ctx, index, shard, request) +} + +// EncodeImportValues is kind of weird. We're trying to mimic what the client +// does here (because the Importer interface was originally based off of the +// client methods). So we end up generating a protobuf-encode byte slice. And we +// don't really use path. +func (f *FeaturebaseImporter) EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + msg := &pb.ImportValueRequest{ + Index: index, + IndexCreatedAt: 0, + Field: field.Name, + FieldCreatedAt: field.CreatedAt, + Shard: shard, + ColumnIDs: ids, + Values: vals, + } + data, err = proto.Marshal(msg) + if err != nil { + return "", nil, errors.Wrap(err, "marshaling ImportValue to protobuf") + } + return "", data, nil +} + +func (f *FeaturebaseImporter) EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + msg := &pb.ImportRequest{ + Index: index, + IndexCreatedAt: 0, + Field: field.Name, + FieldCreatedAt: field.CreatedAt, + Shard: shard, + RowIDs: vals, + ColumnIDs: ids, + } + data, err = proto.Marshal(msg) + if err != nil { + return "", nil, errors.Wrap(err, "marshaling Import to protobuf") + } + return "", data, nil +} + +func (f *FeaturebaseImporter) DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error { + serializer := featurebaseproto.Serializer{} + + // Unmarshal request based on field type. + switch field.Options.Type { + case featurebase.FieldTypeInt, featurebase.FieldTypeDecimal, featurebase.FieldTypeTimestamp: + // Marshal into request object. + req := &featurebase.ImportValueRequest{} + if err := serializer.Unmarshal(data, req); err != nil { + return errors.Wrap(err, "unmarshaling import value request") + } + + qcx := f.API.Txf().NewQcx() + defer qcx.Abort() + + opts := []featurebase.ImportOption{ + featurebase.OptImportOptionsClear(req.Clear), + } + + if err := f.API.ImportValue(ctx, qcx, req, opts...); err != nil { + return errors.Wrap(err, "importing import value request") + } + + if err := qcx.Finish(); err != nil { + return errors.Wrap(err, "finishing qcx") + } + + default: + // Marshal into request object. + req := &featurebase.ImportRequest{} + if err := serializer.Unmarshal(data, req); err != nil { + return errors.Wrap(err, "unmarshaling import request") + } + + qcx := f.API.Txf().NewQcx() + defer qcx.Abort() + + opts := []featurebase.ImportOption{ + featurebase.OptImportOptionsClear(req.Clear), + } + if len(req.RowIDs) > 0 { + opts = append(opts, featurebase.OptImportOptionsIgnoreKeyCheck(true)) + } + + if err := f.API.Import(ctx, qcx, req, opts...); err != nil { + return errors.Wrap(err, "importing import request") + } + + if err := qcx.Finish(); err != nil { + return errors.Wrap(err, "finishing qcx") + } + } + + return nil +} + +func (f *FeaturebaseImporter) StatsTiming(name string, value time.Duration, rate float64) {} diff --git a/client/metrics.go b/batch/metrics.go similarity index 94% rename from client/metrics.go rename to batch/metrics.go index ab476051e..f46277265 100644 --- a/client/metrics.go +++ b/batch/metrics.go @@ -1,5 +1,4 @@ -// Copyright 2021 Molecula Corp. All rights reserved. -package client +package batch const ( // MetricBatchImportDurationSeconds records the full time of the diff --git a/batch/testdata/README.md b/batch/testdata/README.md new file mode 100644 index 000000000..05ad5c954 --- /dev/null +++ b/batch/testdata/README.md @@ -0,0 +1,3 @@ +# testdata + +This directory is used in CI tests. I think. diff --git a/batch/wait.sh b/batch/wait.sh new file mode 100755 index 000000000..36b9cd0e4 --- /dev/null +++ b/batch/wait.sh @@ -0,0 +1,26 @@ +#!/bin/sh + +name=$1 +shift + +_start_ts=$(date +%s) +elapsed=0 +timeout=120 +while : +do + $@ > /dev/null + _ret=$? + _end_ts=$(date +%s) + if [ $_ret -eq 0 ]; then + echo "$name is available after $((_end_ts - _start_ts)) seconds." + break + else + echo "Waiting for $name after $((_end_ts - _start_ts)) seconds." + fi + sleep 1s + elapsed=$((elapsed+1)) + if [ $elapsed -ge $timeout ]; then + exit 110 + fi +done +set -ex diff --git a/client/api.go b/client/api.go new file mode 100644 index 000000000..d01ffb97e --- /dev/null +++ b/client/api.go @@ -0,0 +1,275 @@ +package client + +import ( + "context" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/errors" +) + +var _ featurebase.SchemaAPI = &schemaAPI{} + +// schemaAPI is a featurebase client wrapper which implements the +// featurebase.SchemaAPI interface. This was introduced for use in batch tests +// when we decoupled Batch from the client package. In other words, this is only +// used for those tests, it may not be functionally complete, and should not be +// used otherwise without further testing and review of this code. +type schemaAPI struct { + *Client +} + +func NewSchemaAPI(c *Client) *schemaAPI { + return &schemaAPI{ + Client: c, + } +} + +func (s *schemaAPI) CreateIndexAndFields(ctx context.Context, indexName string, options featurebase.IndexOptions, fields []featurebase.CreateFieldObj) error { + schema, err := s.Client.Schema() + if err != nil { + return errors.Wrap(err, "getting schema") + } + + // Add the index. + idx := schema.Index(indexName, + OptIndexKeys(options.Keys), + OptIndexTrackExistence(true), + ) + if err := s.Client.CreateIndex(idx); err != nil { + return errors.Wrap(err, "creating index") + } + + // Now add fields. + for _, f := range fields { + fld, err := s.addFieldToIndex(idx, f.Name, f.Options...) + if err != nil { + return errors.Wrapf(err, "adding field to index") + } + if err := s.Client.CreateField(fld); err != nil { + return errors.Wrapf(err, "creating field") + } + } + + return nil +} + +func (s *schemaAPI) CreateField(ctx context.Context, indexName string, fieldName string, opts ...featurebase.FieldOption) (*featurebase.Field, error) { + schema, err := s.Client.Schema() + if err != nil { + return nil, errors.Wrap(err, "getting schema") + } + + if !schema.HasIndex(indexName) { + return nil, featurebase.ErrIndexNotFound + } + + idx := schema.Index(indexName) + + fld, err := s.addFieldToIndex(idx, fieldName, opts...) + if err != nil { + return nil, errors.Wrapf(err, "adding field to index") + } + + if err := s.Client.CreateField(fld); err != nil { + return nil, errors.Wrapf(err, "creating field") + } + + return nil, nil +} + +func (s *schemaAPI) addFieldToIndex(idx *Index, fieldName string, opts ...featurebase.FieldOption) (*Field, error) { + ffos := &featurebase.FieldOptions{} + for _, opt := range opts { + opt(ffos) + } + + cfos := []FieldOption{} + + switch ffos.Type { + case featurebase.FieldTypeBool: + cfos = append(cfos, OptFieldTypeBool()) + case featurebase.FieldTypeInt: + cfos = append(cfos, OptFieldTypeInt(ffos.Min.ToInt64(0), ffos.Max.ToInt64(0))) + case featurebase.FieldTypeSet: + cfos = append(cfos, + OptFieldTypeSet(CacheType(ffos.CacheType), int(ffos.CacheSize)), + OptFieldKeys(ffos.Keys), + ) + case featurebase.FieldTypeMutex: + cfos = append(cfos, + OptFieldTypeMutex(CacheType(ffos.CacheType), int(ffos.CacheSize)), + OptFieldKeys(ffos.Keys), + ) + case featurebase.FieldTypeDecimal: + cfos = append(cfos, OptFieldTypeDecimal(ffos.Scale, ffos.Min, ffos.Max)) + case featurebase.FieldTypeTime: + cfos = append(cfos, + OptFieldTypeTime(TimeQuantum(ffos.TimeQuantum), ffos.NoStandardView), + OptFieldKeys(ffos.Keys), + ) + case featurebase.FieldTypeTimestamp: + cfos = append(cfos, OptFieldTypeTimestamp(featurebase.DefaultEpoch, ffos.TimeUnit)) + default: + return nil, errors.Errorf("unsupported field type: %s", ffos.Type) + } + + return idx.Field(fieldName, cfos...), nil +} + +func (s *schemaAPI) DeleteField(ctx context.Context, indexName string, fieldName string) error { + schema, err := s.Client.Schema() + if err != nil { + return errors.Wrap(err, "getting schema") + } + + if !schema.HasIndex(indexName) { + return featurebase.ErrIndexNotFound + } + + idx := schema.Index(indexName) + + return s.Client.DeleteField(&Field{ + name: fieldName, + index: idx, + }) +} + +func (s *schemaAPI) DeleteIndex(ctx context.Context, indexName string) error { + return s.Client.DeleteIndexByName(indexName) +} + +func (s *schemaAPI) IndexInfo(ctx context.Context, indexName string) (*featurebase.IndexInfo, error) { + schema, err := s.Client.Schema() + if err != nil { + return nil, errors.Wrap(err, "getting schema") + } + + if !schema.HasIndex(indexName) { + return nil, featurebase.ErrIndexNotFound + } + + idx := schema.Index(indexName) + return FromClientIndex(idx), nil +} + +func (s *schemaAPI) Schema(ctx context.Context, withViews bool) ([]*featurebase.IndexInfo, error) { + return nil, errors.New("", "schemaAPI.Schema is not implemented") +} + +var _ featurebase.QueryAPI = &queryAPI{} + +// queryAPI is a featurebase client wrapper which implements the +// featurebase.QueryAPI interface. This was introduced for use in batch tests +// when we decoupled Batch from the client package. In other words, this is only +// used for those tests, it may not be functionally complete, and should not be +// used otherwise without further testing and review of this code. +type queryAPI struct { + *Client +} + +func NewQueryAPI(c *Client) *queryAPI { + return &queryAPI{ + Client: c, + } +} + +func (q *queryAPI) Query(ctx context.Context, req *featurebase.QueryRequest) (featurebase.QueryResponse, error) { + schema, err := q.Client.Schema() + if err != nil { + return featurebase.QueryResponse{}, errors.Wrap(err, "getting schema") + } + + if !schema.HasIndex(req.Index) { + return featurebase.QueryResponse{}, featurebase.ErrIndexNotFound + } + + idx := schema.Index(req.Index) + + qry := NewPQLBaseQuery(req.Query, idx, nil) + res, err := q.Client.Query(qry) + if err != nil { + return featurebase.QueryResponse{}, errors.Wrap(err, "querying client") + } + + var rerr error + if res.ErrorMessage != "" { + rerr = errors.New("", res.ErrorMessage) + } + + fbResults := make([]interface{}, 0) + + // Row, PairsField, GroupCounts + for _, result := range res.ResultList { + switch result.Type() { + case QueryResultTypeRow: + if len(result.Row().Keys) > 0 { + row := featurebase.NewRow() + row.Keys = result.Row().Keys + fbResults = append(fbResults, row) + } else { + fbResults = append(fbResults, featurebase.NewRow(result.Row().Columns...)) + } + + case QueryResultTypeUint64: + fbResults = append(fbResults, uint64(result.Count())) + + case QueryResultTypeBool: + fbResults = append(fbResults, result.Changed()) + + case QueryResultTypePairsField: + pairs := []featurebase.Pair{} + for _, ci := range result.CountItems() { + pairs = append(pairs, featurebase.Pair{ + ID: ci.ID, + Key: ci.Key, + Count: ci.Count, + }) + } + pf := &featurebase.PairsField{ + Pairs: pairs, + Field: "", + } + fbResults = append(fbResults, pf) + + case QueryResultTypeGroupCounts: + groups := []featurebase.GroupCount{} + for _, grpCnt := range result.GroupCounts() { + fieldRows := []featurebase.FieldRow{} + for _, fr := range grpCnt.Groups { + fieldRows = append(fieldRows, featurebase.FieldRow{ + Field: fr.FieldName, + RowID: fr.RowID, + RowKey: fr.RowKey, + Value: fr.Value, + }) + } + groups = append(groups, featurebase.GroupCount{ + Group: fieldRows, + Count: uint64(grpCnt.Count), + Agg: grpCnt.Agg, + // DecimalAgg: ??, + }) + } + + gc := featurebase.NewGroupCounts("", groups...) + fbResults = append(fbResults, gc) + + case QueryResultTypeValCount: + vc := featurebase.ValCount{ + Val: result.Value(), + Count: result.Count(), + } + fbResults = append(fbResults, vc) + + default: + return featurebase.QueryResponse{}, errors.Errorf("unsupported query result type: %d", result.Type()) + } + } + + resp := &featurebase.QueryResponse{ + Results: fbResults, + Err: rerr, + } + + return *resp, nil +} diff --git a/client/batch_test.go b/client/batch_test.go deleted file mode 100644 index caed90cfa..000000000 --- a/client/batch_test.go +++ /dev/null @@ -1,1883 +0,0 @@ -// Copyright 2021 Molecula Corp. All rights reserved. - -package client - -import ( - "fmt" - "math/rand" - "reflect" - "sort" - "strconv" - "testing" - "time" - - featurebase "github.com/molecula/featurebase/v3" - "github.com/molecula/featurebase/v3/test" - - "github.com/pkg/errors" -) - -func NewTestClient(t *testing.T, c *test.Cluster) *Client { - client, err := NewClient(c.Nodes[0].URL()) - if err != nil { - t.Fatal(err) - } - return client -} - -func TestAgainstCluster(t *testing.T) { - c := test.MustRunCluster(t, 1) - defer c.Close() - client := NewTestClient(t, c) - t.Run("string-slice-combos", func(t *testing.T) { testStringSliceCombos(t, c, client) }) - t.Run("import-batch-ints", func(t *testing.T) { testImportBatchInts(t, c, client) }) - t.Run("import-batch-bools", func(t *testing.T) { testImportBatchBools(t, c, client) }) - t.Run("import-batch-sorting", func(t *testing.T) { testImportBatchSorting(t, c, client) }) - t.Run("test-trim-null", func(t *testing.T) { testTrimNull(t, c, client) }) - t.Run("test-string-slice-empty-and-nil", func(t *testing.T) { testStringSliceEmptyAndNil(t, c, client) }) - t.Run("test-string-slice", func(t *testing.T) { testStringSlice(t, c, client) }) - t.Run("test-single-clear-batch-regression", func(t *testing.T) { testSingleClearBatchRegression(t, c, client) }) - t.Run("test-batches", func(t *testing.T) { testBatches(t, c, client) }) - t.Run("batches-strings-ids", func(t *testing.T) { testBatchesStringIDs(t, c, client) }) - t.Run("test-batch-staleness", func(t *testing.T) { testBatchStaleness(t, c, client) }) - t.Run("test-import-batch-multiple-ints", func(t *testing.T) { testImportBatchMultipleInts(t, c, client) }) - t.Run("test-import-batch-multiple-timestamps", func(t *testing.T) { testImportBatchMultipleTimestamps(t, c, client) }) - t.Run("test-import-batch-sets-clears", func(t *testing.T) { testImportBatchSetsAndClears(t, c, client) }) - t.Run("test-topn-cache-regression", func(t *testing.T) { testTopNCacheRegression(t, c, client) }) - t.Run("test-multiple-int-same-batch", func(t *testing.T) { testMultipleIntSameBatch(t, c, client) }) - t.Run("test-mutex-clearing-regression", func(t *testing.T) { mutexClearRegression(t, c, client) }) - t.Run("test-mutex-nil-clear-id", func(t *testing.T) { mutexNilClearID(t, c, client) }) - t.Run("test-mutex-nil-clear-key", func(t *testing.T) { mutexNilClearKey(t, c, client) }) -} - -func testStringSliceCombos(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-string-slice-combos") - fields := make([]*Field, 1) - fields[0] = idx.Field("a1", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - b, err := NewBatch(client, 5, 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{}}}, - } - - err = ingestRecords(records, b) - if err != nil { - t.Fatalf("importing: %v", err) - } - - a1 := fields[0] - - result := tq(t, client, a1.TopN(10)) - rez := sortableCRI(result.CountItems()) - sort.Sort(rez) - exp := sortableCRI{ - {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) - - result = tq(t, client, a1.Row("a")) - errorIfNotEqual(t, result.Row().Columns, []uint64{0, 5, 6, 11}) - result = tq(t, client, a1.Row("b")) - errorIfNotEqual(t, result.Row().Columns, []uint64{0, 5, 6, 11}) - result = tq(t, client, a1.Row("c")) - errorIfNotEqual(t, result.Row().Columns, []uint64{0, 3, 5, 6, 11}) - result = tq(t, client, a1.Row("z")) - errorIfNotEqual(t, result.Row().Columns, []uint64{1, 7}) - result = tq(t, client, a1.Row("q")) - errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) - result = tq(t, client, a1.Row("r")) - errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) - result = tq(t, client, a1.Row("s")) - errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) - result = tq(t, client, a1.Row("t")) - errorIfNotEqual(t, result.Row().Columns, []uint64{3, 9}) - - result = tq(t, client, idx.RawQuery("Count(All())")) - errorIfNotEqual(t, result.Count(), int64(14)) -} - -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) - } -} - -type sortableCRI []CountResultItem - -func (s sortableCRI) Len() int { return len(s) } -func (s sortableCRI) 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 sortableCRI) Swap(i, j int) { - s[i], s[j] = s[j], s[i] -} - -func tq(t *testing.T, client *Client, query PQLQuery) QueryResult { - resp, err := client.Query(query) - if err != nil { - t.Fatalf("querying: %v", err) - } - return resp.Results()[0] -} - -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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-ints") - field := idx.Field("anint", OptFieldTypeInt()) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 3, idx, []*Field{field}) - 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, err := client.Query(idx.BatchQuery(field.Equals(0), field.Equals(7), field.Equals(2))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - for i, result := range resp.Results() { - if !reflect.DeepEqual(result.Row().Columns, []uint64{uint64(i)}) { - t.Errorf("expected %v for %d, but got %v", []uint64{uint64(i)}, i, result.Row().Columns) - } - } -} - -func testImportBatchSorting(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-sorting") - field := idx.Field("anint", OptFieldTypeInt()) - field2 := idx.Field("amutex", OptFieldTypeMutex(CacheTypeNone, 0)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 100, idx, []*Field{field, field2}) - 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, err := client.Query(idx.RawQuery("Count(All())")) - if err != nil { - t.Fatalf("querying: %v", err) - } - if res := resp.Results()[0]; res.Count() != 100 { - t.Fatalf("unexpected result: %+v", res) - } -} - -func testTrimNull(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-trim-null") - field := idx.Field("empty", OptFieldTypeInt()) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - b, err := NewBatch(client, 3, idx, []*Field{field}) - if err != nil { - t.Fatalf("getting batch: %v", err) - } - b.nullIndices = make(map[string][]uint64, 1) - b.nullIndices[field.Name()] = []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, err := client.Query(idx.BatchQuery(field.Equals(0), field.Equals(1), field.Equals(2))) - if err != nil { - t.Fatalf("querying: %v", err) - } - for i, result := range resp.Results() { - if !reflect.DeepEqual(result.Row().Columns, []uint64(nil)) { - t.Errorf("expected %#v for %d, but got %#v", []uint64(nil), i, result.Row().Columns) - } - } - - b, err = NewBatch(client, 4, idx, []*Field{field}) - 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) - } - - resp, err = client.Query(idx.BatchQuery(field.Equals(10), field.Equals(40), field.Equals(20), field.Equals(30))) - if err != nil { - t.Fatalf("querying: %v", err) - } - for i, result := range resp.Results() { - if i == 1 { - if !reflect.DeepEqual(result.Row().Columns, []uint64(nil)) { - t.Errorf("expected %#v for %d, but got %#v", []uint64(nil), i, result.Row().Columns) - } - } else { - if !reflect.DeepEqual(result.Row().Columns, []uint64{result.Row().Columns[0]}) { - t.Errorf("expected %#v for %d, but got %#v", []uint64{result.Row().Columns[0]}, i, result.Row().Columns) - } - } - } - -} - -func testStringSliceEmptyAndNil(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-string-slice-nil") - fields := make([]*Field, 1) - fields[0] = idx.Field("strslice", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - // 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(client, 2, idx, fields) - if err != nil { - t.Fatalf("creating new batch: %v", err) - } - r := Row{Values: make([]interface{}, len(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(client, 6, idx, fields) - if err != nil { - t.Fatalf("creating new batch: %v", err) - } - r = Row{Values: make([]interface{}, len(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) - } - - rows := []interface{}{"a", "b", "c", "z"} - resp, err := client.Query(idx.BatchQuery(fields[0].Row(rows[0]), fields[0].Row(rows[1]), fields[0].Row(rows[2]), fields[0].Row(rows[3]))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - // TODO test is flaky because we can't guarantee what a,b,c map to - expectations := [][]uint64{{0, 2}, {2, 3}, {3}, {2}} - for i, re := range resp.Results() { - if !reflect.DeepEqual(re.Row().Columns, expectations[i]) { - t.Errorf("expected row %v to have columns %v, but got %v", rows[i], expectations[i], re.Row().Columns) - } - } - -} - -func testStringSlice(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-string-slice") - fields := make([]*Field, 1) - fields[0] = idx.Field("strslice", OptFieldKeys(true), OptFieldTypeSet(CacheTypeRanked, 100)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - b, err := NewBatch(client, 3, 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(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, err := client.Query(idx.BatchQuery(fields[0].Row("a"))) - if err != nil { - t.Fatalf("querying: %v", err) - } - result := resp.Result() - if !reflect.DeepEqual(result.Row().Columns, []uint64{0, 1}) { - t.Fatalf("expected a to be [0,1], got %v", result.Row().Columns) - } -} - -func testSingleClearBatchRegression(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-single-clear-batch-regression") - numFields := 1 - fields := make([]*Field, numFields) - fields[0] = idx.Field("zero", OptFieldKeys(true)) - - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - _, err = client.Query(fields[0].Set("row1", 1)) - if err != nil { - t.Fatalf("setting bit: %v", err) - } - - b, err := NewBatch(client, 1, idx, fields) - if err != nil { - t.Fatalf("getting new batch: %v", err) - } - r := Row{ID: uint64(1), Values: make([]interface{}, numFields), 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, err := client.Query(fields[0].Row("row1")) - if err != nil { - t.Fatalf("error querying: %v", err) - } - result := resp.Results()[0].Row().Columns - if len(result) != 0 { - t.Fatalf("unexpected values in row: result %+v", result) - } - -} - -func testBatches(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-batches") - numFields := 5 - fields := make([]*Field, numFields) - fields[0] = idx.Field("zero", OptFieldKeys(true)) - fields[1] = idx.Field("one", OptFieldKeys(true)) - fields[2] = idx.Field("two", OptFieldKeys(true)) - fields[3] = idx.Field("three", OptFieldTypeInt()) - fields[4] = idx.Field("four", OptFieldTypeTime(TimeQuantumYearMonthDay)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - b, err := NewBatch(client, 10, idx, fields) - if err != nil { - t.Fatalf("getting new batch: %v", err) - } - r := Row{Values: make([]interface{}, numFields), 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) - } - - resp, err := client.Query(idx.BatchQuery(fields[0].Row("a"), - fields[1].Row("b"), - fields[2].Row("c"), - fields[3].Equals(99))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - results := resp.Results() - for _, j := range []int{0, 2, 3} { - cols := results[j].Row().Columns - if !reflect.DeepEqual(cols, []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18}) { - t.Fatalf("unexpected columns for a: %v", cols) - } - } - res := results[1] - - if cols := res.Row().Columns; !reflect.DeepEqual(cols, []uint64{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}) { - t.Fatalf("unexpected columns for field 1 row b: %v", cols) - } - - resp, err = client.Query(idx.BatchQuery(fields[0].Row("d"), - fields[1].Row("e"), - fields[2].Row("f"))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - results = resp.Results() - for _, res := range results { - cols := res.Row().Columns - if !reflect.DeepEqual(cols, []uint64{20, 22, 24, 26, 28}) { - t.Fatalf("unexpected columns: %v", cols) - } - } - - resp, err = client.Query(idx.BatchQuery(fields[3].GT(-11), - fields[3].Equals(0), - fields[3].Equals(100), - fields[4].Range(1, time.Date(2019, time.January, 1, 0, 0, 0, 0, time.UTC), time.Date(2019, time.January, 29, 0, 0, 0, 0, time.UTC)), - fields[4].Range(1, time.Date(2019, time.February, 1, 0, 0, 0, 0, time.UTC), time.Date(2019, time.February, 29, 0, 0, 0, 0, time.UTC)))) - if err != nil { - t.Fatalf("querying: %v", err) - } - results = resp.Results() - - if cols := results[0].Row().Columns; !reflect.DeepEqual(cols, []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}) { - t.Fatalf("all columns (but 8) should be greater than -11, but got: %v", cols) - } - - if cols := results[1].Row().Columns; !reflect.DeepEqual(cols, []uint64{19, 21, 23, 25, 27}) { - t.Fatalf("wrong cols for ==0: %v", cols) - } - - if cols := results[2].Row().Columns; !reflect.DeepEqual(cols, []uint64{20, 22, 24, 26, 28}) { - t.Fatalf("wrong cols for ==100: %v", cols) - } - - cols := results[3].Row().Columns - exp := []uint64{0, 2, 4, 6, 10, 12, 14, 16, 18} - if !reflect.DeepEqual(cols, exp) { - t.Fatalf("wrong cols for January: got/want\n%v\n%v", cols, exp) - } - - cols = results[4].Row().Columns - exp = []uint64{1, 3, 5, 7} - if !reflect.DeepEqual(cols, exp) { - t.Fatalf("wrong cols for January: got/want\n%v\n%v", cols, exp) - } - - 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, err = client.Query(idx.BatchQuery( - fields[0].Row("a"), - fields[0].Row("x"), - fields[1].Row("b"), - )) - if err != nil { - t.Fatalf("querying after clears: %v", err) - } - if arow := resp.Results()[0].Row().Columns; arow[0] == 0 { - t.Errorf("shouldn't have id 0 in row a after clearing! %v", arow) - } - if xrow := resp.Results()[1].Row().Columns; xrow[0] != 0 { - t.Errorf("should have id 0 in row x after setting %v", xrow) - } - if brow := resp.Results()[2].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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("batches-strings-ids", OptIndexKeys(true)) - fields := make([]*Field, 3) - fields[0] = idx.Field("zero", OptFieldKeys(true)) - fields[1] = idx.Field("one", OptFieldTypeMutex(CacheTypeNone, 0), OptFieldKeys(true)) - fields[2] = idx.Field("two", OptFieldTypeTime("YMDH"), OptFieldKeys(true)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - b, err := NewBatch(client, 3, 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) - } - - resp, err := client.Query(idx.BatchQuery(fields[0].Row("a"), fields[0].Row("x"), fields[1].Row("b"), fields[1].Row("y"), fields[2].Row("c"), fields[2].Row("z"))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - results := resp.Results() - for i, res := range results { - cols := res.Row().Keys - if i%2 == 0 && !reflect.DeepEqual(cols, []string{"0", "2"}) && !reflect.DeepEqual(cols, []string{"2", "0"}) { - t.Fatalf("unexpected columns: %v", cols) - } - if i%2 == 1 && !reflect.DeepEqual(cols, []string{"1"}) { - t.Fatalf("unexpected columns: %v", cols) - } - } - - 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) - } - - resp, err = client.Query(idx.BatchQuery(fields[0].Row("a"), fields[0].Row("z"))) - if err != nil { - t.Fatalf("querying: %v", err) - } - - results = resp.Results() - for i, res := range results { - cols := res.Row().Keys - if err := isPermutationOf(cols, []string{"0", "1", "2"}); i == 0 && err != nil { - t.Fatalf("unexpected columns: %v: %v", cols, err) - } - if i == 1 && !reflect.DeepEqual(cols, []string{"3"}) { - t.Fatalf("unexpected columns: %v", cols) - } - } - -} - -func isPermutationOf(one, two []string) error { - if len(one) != len(two) { - return errors.Errorf("different lengths %d and %d", len(one), len(two)) - } -outer: - for _, vOne := range one { - for j, vTwo := range two { - if vOne == vTwo { - two = append(two[:j], two[j+1:]...) - continue outer - } - } - return errors.Errorf("%s in one but not two", vOne) - } - if len(two) != 0 { - return errors.Errorf("vals in two but not one: %v", two) - } - return nil -} - -func TestQuantizedTime(t *testing.T) { - cases := []struct { - name string - time time.Time - year string - month string - day string - hour string - quantum 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: TimeQuantumYear, - 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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-batch-staleness") - field := idx.Field("anint", OptFieldTypeInt()) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - b, err := NewBatch(client, 3, idx, []*Field{field}, 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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-multi-int") - field := idx.Field("anint", OptFieldTypeInt()) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 6, idx, []*Field{field}, 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) - } - - if resp, err := client.Query(field.Equals(4)); err != nil { - t.Fatalf("querying: %v", err) - } else if res := resp.Results()[0].Row().Columns; len(res) != 1 || res[0] != 1 { - t.Fatalf("unepxected result: %v", res) - } - -} - -// testImportBatchMultipleTimestamps tests if nils are handles correctly for TS in batch imports -func testImportBatchMultipleTimestamps(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-multi-timestamp") - field := idx.Field("ts2", OptFieldTypeTimestamp(time.Unix(0, 0), "s")) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b1, err := NewBatch(client, 6, idx, []*Field{field}) - if err != nil { - t.Fatalf("getting batch: %v", err) - } - b2, err := NewBatch(client, 6, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true)) - if err != nil { - t.Fatalf("getting batch: %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[field.name][0] != 3 { - t.Fatalf("unexpected nulls, got/want: %v/%v", b.nullIndices[field.name], []uint64{3}) - } - for i, val := range chkVals { - if b.values[field.name][i] != val { - t.Fatalf("unexpected value, got/want: %v/%v", b.values[field.name][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) - } - - qr := c.Query(t, idx.name, `Extract(All(), Rows(ts2))`) - results := qr.Results[0].(featurebase.ExtractedTable) - for k, res := range results.Columns { - if chkImport[k] != nil { - if res.Rows[0] != chkImport[k].(time.Time).UTC() { - t.Fatalf("unexpected result, got/want: %v/%v", res.Rows[0], chkImport[k].(time.Time).UTC()) - } - } else { - if res.Rows[0] != nil { - t.Fatalf("unexpected result, got/want: %v/%v", res.Rows[0], nil) - } - } - } - }) - } -} - -func testImportBatchSetsAndClears(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-set-and-clear") - field := idx.Field("aset", OptFieldTypeSet(featurebase.DefaultCacheType, featurebase.DefaultCacheSize)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 6, idx, []*Field{field}, 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) - } - - if resp, err := client.Query(field.TopN(6)); err != nil { - t.Fatalf("querying topn: %v", err) - } else if res := resp.Result().CountItems(); len(res) != 3 { - t.Fatalf("unexpected topn: %+v", res) - } - - exp := [][]uint64{ - {}, - {1}, - {}, - {}, - {}, - {2}, - {3}, - } - for row := 0; row < 7; row++ { - resp, err := client.Query(field.Row(row)) - if err != nil { - t.Fatalf("querying: %v", err) - } - res := resp.Results()[0].Row().Columns - if !reflect.DeepEqual(exp[row], res) && !(len(exp[row]) == 0 && len(res) == 0) { - t.Errorf("row: %d, exp: %v, got %v", row, exp[row], res) - } - } - -} - -// 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 beacuse 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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-topn-cache-regression") - field := idx.Field("aset", OptFieldTypeSet(featurebase.DefaultCacheType, featurebase.DefaultCacheSize)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 3, idx, []*Field{field}, 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) - } - - if resp, err := client.Query(field.TopN(6)); err != nil { - t.Fatalf("querying topn: %v", err) - } else if res := resp.Result().CountItems(); len(res) != 1 { - t.Fatalf("unexpected topn: %+v", res) - } else if res[0].ID != 1 || res[0].Count != 1 { - t.Fatalf("unexpected topn result: %v", res) - } -} - -// 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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-multiple-int-same-batch") - field := idx.Field("age", OptFieldTypeInt(0, 10000)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 4, idx, []*Field{field}, 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) - } - - if resp, err := client.Query(field.Sum(nil)); err != nil { - t.Fatalf("querying sum: %v", err) - } else if res := resp.Result().Value(); res != 2 { - t.Errorf("unexpected sum: %+v", res) - } -} - -// 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, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-multiple-mut-same-batch") - field := idx.Field("mut", OptFieldTypeMutex(CacheTypeNone, 0)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 11, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true)) - if err != nil { - t.Fatalf("getting batch: %v", err) - } - - col := uint64(0) - row := uint64(1) - 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, err := client.Query(idx.GroupBy(field.Rows(), field.Rows())) - if err != nil { - t.Fatalf("querying groupby: %v", err) - } - groupCounts := resp.Result().GroupCounts() - for j, gc := range groupCounts { - if gc.Groups[0].RowID != gc.Groups[1].RowID { - t.Errorf("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, err := client.Query(idx.GroupBy(field.Rows(), field.Rows())) - if err != nil { - t.Fatalf("querying groupby: %v", err) - } - groupCounts := resp.Result().GroupCounts() - for i, gc := range groupCounts { - if gc.Groups[0].RowID != gc.Groups[1].RowID { - t.Fatalf("bmismatched group at %d, %v", i, gc) - } - } -} - -// test clearing record with explict nil -func mutexNilClearID(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-mut-nil-clear-id") - field := idx.Field("mut", OptFieldTypeMutex(CacheTypeNone, 0)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 11, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true)) - if err != nil { - t.Fatalf("getting batch: %v", err) - } - - col := uint64(0) - row := uint64(1) - // 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, err := client.Query(idx.RawQuery("Row(mut=0)")) - if err != nil { - t.Fatalf("Fetching data: %v", err) - } - items := resp.Result().Row().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, err = client.Query(idx.RawQuery("Row(mut=0)")) - if err != nil { - t.Fatalf("Fetching data: %v", err) - } - errorIfNotEqual(t, resp.Result().Row().Columns, items) - -} - -// similar test to above but with string keys -func mutexNilClearKey(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-mut-nil-clear-key", OptIndexKeys(true)) - fields := make([]*Field, 1) - fields[0] = idx.Field("mut", OptFieldTypeMutex(CacheTypeNone, 0), OptFieldKeys(true)) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - defer func() { - err := client.DeleteIndex(idx) - if err != nil { - t.Logf("problem cleaning up from test: %v", err) - } - }() - - b, err := NewBatch(client, 3, 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, _ := client.Query(idx.RawQuery(`Row(mut="a")`)) - errorIfNotEqual(t, resp.Result().Row().Keys, []string{"0", "2"}) - - 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, err = client.Query(idx.RawQuery(`Row(mut="a")`)) - if err != nil { - t.Fatalf("importing: %v", err) - } - errorIfNotEqual(t, resp.Result().Row().Keys, []string{"0"}) -} - -func testImportBatchBools(t *testing.T, c *test.Cluster, client *Client) { - schema := NewSchema() - idx := schema.Index("test-import-batch-bools") - field := idx.Field("boolcol", OptFieldTypeBool()) - err := client.SyncSchema(schema) - if err != nil { - t.Fatalf("syncing schema: %v", err) - } - - b, err := NewBatch(client, 3, idx, []*Field{field}, OptUseShardTransactionalEndpoint(true)) - if err != nil { - t.Fatalf("getting 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, err := client.Query(idx.RawQuery("Count(All())")) - if err != nil { - t.Fatalf("querying: %v", err) - } - if res := resp.Results()[0]; res.Count() != 2 { - t.Fatalf("unexpected result: %+v", res) - } -} diff --git a/client/client.go b/client/client.go index 1ec290e0b..e41add3bb 100644 --- a/client/client.go +++ b/client/client.go @@ -39,8 +39,6 @@ const PQLVersion = "1.0" // DefaultShardWidth is used if an index doesn't have it defined. const DefaultShardWidth = pilosa.ShardWidth -const maxHosts = 10 - // Client is the HTTP client for Pilosa server. type Client struct { cluster *Cluster diff --git a/client/importer.go b/client/importer.go new file mode 100644 index 000000000..c55cc04bc --- /dev/null +++ b/client/importer.go @@ -0,0 +1,219 @@ +package client + +import ( + "context" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/roaring" + "github.com/pkg/errors" +) + +////////////////////////////////////////////////////////////////////////////// +// +// The following was introduced upon splitting batch out of the client package +// (and into its own package). During that work, we changed batch from using +// client.Index, and instead using featurebase.IndexInfo. The functions below +// convert client types from/to featurebase types in order to satisfy the batch +// requirements. +// +////////////////////////////////////////////////////////////////////////////// + +// FromClientIndex converts a client Index to a featurebase IndexInfo. +func FromClientIndex(ci *Index) *featurebase.IndexInfo { + return &featurebase.IndexInfo{ + Name: ci.Name(), + CreatedAt: ci.CreatedAt(), + Options: fromClientIndexOptions(ci.Opts()), + Fields: fromClientFieldsMap(ci.Fields()), + ShardWidth: ci.ShardWidth(), + } +} + +// fromClientIndexOptions +func fromClientIndexOptions(cio IndexOptions) featurebase.IndexOptions { + return featurebase.IndexOptions{ + Keys: cio.Keys(), + TrackExistence: cio.TrackExistence(), + PartitionN: 0, // TODO(tlt): this shouldn't be 0 once we support it. + } +} + +// toClientIndex +func toClientIndex(fi *featurebase.IndexInfo) *Index { + sch := NewSchema() + return sch.Index(fi.Name, + OptIndexKeys(fi.Options.Keys), + OptIndexTrackExistence(fi.Options.TrackExistence), + ) +} + +// fromClientField +func fromClientField(cf *Field) *featurebase.FieldInfo { + return &featurebase.FieldInfo{ + Name: cf.Name(), + CreatedAt: cf.CreatedAt(), + Options: fromClientFieldOptions(cf.Opts()), + // TODO(tlt): do we need Views? Because we can't currently get views + // from client. + // Views: ?? + } +} + +// fromClientFieldOptions +func fromClientFieldOptions(cfo FieldOptions) featurebase.FieldOptions { + return featurebase.FieldOptions{ + Base: cfo.Base(), + BitDepth: 0, // TODO(tlt): set this? + Min: cfo.Min(), + Max: cfo.Max(), + Scale: cfo.Scale(), + Keys: cfo.Keys(), + NoStandardView: cfo.NoStandardView(), + CacheSize: uint32(cfo.CacheSize()), + CacheType: string(cfo.CacheType()), + Type: string(cfo.Type()), + TimeUnit: cfo.TimeUnit(), + TimeQuantum: featurebase.TimeQuantum(cfo.TimeQuantum()), + ForeignIndex: cfo.ForeignIndex(), + TTL: cfo.TTL(), + } +} + +// toClientField +func toClientField(index string, ff *featurebase.FieldInfo) (*Field, error) { + sch := NewSchema() + idx := sch.Index(index) + + opts := []FieldOption{} + + switch ff.Options.Type { + case featurebase.FieldTypeBool: + opts = append(opts, + OptFieldTypeBool(), + ) + case featurebase.FieldTypeDecimal: + opts = append(opts, + OptFieldTypeDecimal(ff.Options.Scale, ff.Options.Min, ff.Options.Max), + ) + case featurebase.FieldTypeMutex: + opts = append(opts, + OptFieldTypeMutex(CacheType(ff.Options.CacheType), int(ff.Options.CacheSize)), + OptFieldKeys(ff.Options.Keys), + ) + case featurebase.FieldTypeSet: + opts = append(opts, + OptFieldTypeSet(CacheType(ff.Options.CacheType), int(ff.Options.CacheSize)), + OptFieldKeys(ff.Options.Keys), + ) + case featurebase.FieldTypeInt: + opts = append(opts, + OptFieldTypeInt(ff.Options.Min.ToInt64(0), ff.Options.Max.ToInt64(0)), + ) + case featurebase.FieldTypeTimestamp: + epoch, err := featurebase.ValToTimestamp(ff.Options.TimeUnit, ff.Options.Base) + if err != nil { + return nil, errors.Wrapf(err, "calculating epoch: %s, %d", ff.Options.TimeUnit, ff.Options.Base) + } + opts = append(opts, + OptFieldTypeTimestamp(epoch, ff.Options.TimeUnit), + ) + } + + return idx.Field(ff.Name, opts...), nil +} + +// FromClientFields converts a slice of client Fields to a slice of featurebase +// FieldInfo. +func FromClientFields(cf []*Field) []*featurebase.FieldInfo { + ff := make([]*featurebase.FieldInfo, len(cf)) + for i := range cf { + ff[i] = fromClientField(cf[i]) + } + return ff +} + +// fromClientFieldsMap +func fromClientFieldsMap(cf map[string]*Field) []*featurebase.FieldInfo { + ff := make([]*featurebase.FieldInfo, 0, len(cf)) + for _, v := range cf { + ff = append(ff, fromClientField(v)) + } + return ff +} + +// We can't import the batch package into client because it results in an import +// loop. That's probably an indication that this interface implementation should +// be moved somewhere else; for example, into a sub-package of the batch package +// (since it's an implementation of one of batch's interfaces). +// var _ batch.Importer = &importer{} + +// importer is a pilosa client which implements the batch.Importer interface. +// This wrapper is necessary because of the call into client.Stats.Timing(), and +// because the client takes client specific types (like Index and Field), but +// the interface takes FeatureBase specific types. +type importer struct { + *Client +} + +func NewImporter(c *Client) *importer { + return &importer{ + Client: c, + } +} + +func (i *importer) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return i.Client.StartTransaction(id, timeout, exclusive, requestTimeout) +} + +func (i *importer) FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) { + return i.Client.FinishTransaction(id) +} + +func (i *importer) CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + return i.Client.CreateIndexKeys(toClientIndex(idx), keys...) +} + +func (i *importer) CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + fld, err := toClientField(index, field) + if err != nil { + return nil, errors.Wrap(err, "converting to client field") + } + return i.Client.CreateFieldKeys(fld, keys...) +} + +func (i *importer) ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + fld, err := toClientField(index, field) + if err != nil { + return errors.Wrap(err, "converting to client field") + } + return i.Client.ImportRoaringBitmap(fld, shard, views, clear) +} + +func (i *importer) ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error { + return i.Client.ImportRoaringShard(index, shard, request) +} + +func (i *importer) EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + fld, err := toClientField(index, field) + if err != nil { + return "", nil, errors.Wrap(err, "converting to client field") + } + return i.Client.EncodeImportValues(fld, shard, vals, ids, clear) +} + +func (i *importer) EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + fld, err := toClientField(index, field) + if err != nil { + return "", nil, errors.Wrap(err, "converting to client field") + } + return i.Client.EncodeImport(fld, shard, vals, ids, clear) +} + +func (i *importer) DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error { + return i.Client.DoImport(index, shard, path, data) +} + +func (i *importer) StatsTiming(name string, value time.Duration, rate float64) { + i.Client.Stats.Timing(name, value, rate) +} diff --git a/client/ingest_api_batch.go b/client/ingest_api_batch.go deleted file mode 100644 index a1ae7a1c5..000000000 --- a/client/ingest_api_batch.go +++ /dev/null @@ -1,145 +0,0 @@ -package client - -import ( - "time" - - "github.com/molecula/featurebase/v3/logger" - "github.com/pkg/errors" -) - -// NewIngestAPIBatch creates an alternate implementation of -// RecordBatch which exists to aid in testing the new Ingest API and -// is likely far slower than the Batch. -func NewIngestAPIBatch(client *Client, size int, logger logger.Logger, fields []*Field) *ingestAPIBatch { - if len(fields) == 0 { - return nil - } - - return &ingestAPIBatch{ - client: client, - log: logger, - fields: fields, - keyed: fields[0].index.Opts().Keys(), - index: fields[0].index.Name(), - batchSize: size, - - recordsK: make(map[string]map[string]interface{}), - records: make(map[uint64]map[string]interface{}), - } -} - -type ingestAPIBatch struct { - client *Client - log logger.Logger - batchSize int - - fields []*Field - keyed bool - index string - - // map[recordKey][fieldName]value - recordsK map[string]map[string]interface{} - records map[uint64]map[string]interface{} -} - -func (b *ingestAPIBatch) Add(row Row) error { - if len(row.Clears) > 0 { - return errors.New("ingest api batch does not support clears") - } - values := make(map[string]interface{}) - for i, val := range row.Values { - field := b.fields[i] - // val can be string, uint64, int64, []string, []uint64, nil - // TODO timestamp field might need special handling - // TODO check that the Row.Clears field is only used for packed bools, and then issue a warning/error (in IDK) if the ingest API mode is used in conjunction w/ packed bools. - if val == nil { - continue - } - zero := QuantizedTime{} - if field.Options().Type() == FieldTypeTime && row.Time != zero { - timeq, err := row.Time.Time() - if err != nil { - return errors.Wrap(err, "parsing row time") - } - values[field.Name()] = map[string]interface{}{"time": timeq.Format(time.RFC3339), "values": val} - } else { - values[field.Name()] = val - } - } - - if b.keyed { - switch rowID := row.ID.(type) { - case string: - b.recordsK[rowID] = values - case []byte: - b.recordsK[string(rowID)] = values - default: - return errors.Errorf("unsupported rowID %v of type %[1]T, must be string, or []byte for keyed index", rowID) - } - if len(b.recordsK) >= b.batchSize { - return ErrBatchNowFull - } - } else { - rowID, ok := row.ID.(uint64) - if !ok { - return errors.Errorf("unsupported rowID %v of type %[1]T, must be uint64 for unkeyed index", row.ID) - } - b.records[rowID] = values - if len(b.records) >= b.batchSize { - return ErrBatchNowFull - } - } - return nil -} - -func (b *ingestAPIBatch) Import() error { - if b.keyed { - return b.importKeyed() - } - return b.importUnkeyed() -} - -func (b *ingestAPIBatch) importKeyed() error { - req := []map[string]interface{}{ - { - "action": "set", - "records": b.recordsK, - }, - } - bod, err := b.client.IngestData(b.index, req) - if err != nil { - return errors.Wrapf(err, "importKeyed, body: %s", bod) - } - - for k := range b.recordsK { - delete(b.recordsK, k) - } - return nil -} - -func (b *ingestAPIBatch) importUnkeyed() error { - req := []map[string]interface{}{ - { - "action": "set", - "records": b.records, - }, - } - bod, err := b.client.IngestData(b.index, req) - if err != nil { - return errors.Wrapf(err, "importKeyed, body: %s", bod) - } - - for v := range b.records { - delete(b.records, v) - } - return nil -} - -func (b *ingestAPIBatch) Len() int { - if b.keyed { - return len(b.recordsK) - } - return len(b.records) -} - -func (b *ingestAPIBatch) Flush() error { return nil } diff --git a/client/ingest_api_batch_test.go b/client/ingest_api_batch_test.go deleted file mode 100644 index bc6858c12..000000000 --- a/client/ingest_api_batch_test.go +++ /dev/null @@ -1,305 +0,0 @@ -package client - -import ( - "strings" - "testing" - "time" - - "github.com/molecula/featurebase/v3/logger" - "github.com/molecula/featurebase/v3/test" -) - -func TestIngestAPIBatchAdd(t *testing.T) { - t.Run("unkeyed", func(t *testing.T) { - batch := NewIngestAPIBatch(nil, 10, logger.NopLogger, []*Field{ - { - name: "a", - index: &Index{name: "idxname", options: &IndexOptions{}}, - options: &FieldOptions{ - fieldType: FieldTypeSet, - }, - }, - { - name: "b", - index: &Index{name: "idxname", options: &IndexOptions{}}, - options: &FieldOptions{ - fieldType: FieldTypeSet, - keys: true, - }, - }, - { - name: "c", - index: &Index{name: "idxname", options: &IndexOptions{}}, - options: &FieldOptions{ - fieldType: FieldTypeTime, - keys: true, - }, - }, - }) - qt := QuantizedTime{} - qt.Set(time.Date(2007, time.January, 1, 15, 0, 0, 0, time.UTC)) - err := batch.Add(Row{ - ID: uint64(1), - Values: []interface{}{uint64(2), "bkey", "ckey"}, - Time: qt, - }) - if err != nil { - t.Fatalf("adding row to batch: %v", err) - } - - if batch.records[1]["a"] != uint64(2) { - t.Fatalf("unexpected batch.records: %+v", batch.records) - } - if batch.records[1]["b"] != "bkey" { - t.Fatalf("unexpected batch.records: %+v", batch.records) - } - if batch.records[1]["c"].(map[string]interface{})["time"] != "2007-01-01T15:00:00Z" { - t.Fatalf("unexpected batch.records: %+v", batch.records) - } - if batch.records[1]["c"].(map[string]interface{})["values"] != "ckey" { - t.Fatalf("unexpected batch.records: %+v", batch.records) - } - - }) - - t.Run("keyed", func(t *testing.T) { - batch := NewIngestAPIBatch(nil, 10, logger.NopLogger, []*Field{ - { - name: "a", - index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, - options: &FieldOptions{ - fieldType: FieldTypeSet, - }, - }, - { - name: "b", - index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, - options: &FieldOptions{ - fieldType: FieldTypeSet, - keys: true, - }, - }, - { - name: "c", - index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, - options: &FieldOptions{ - fieldType: FieldTypeTime, - keys: true, - }, - }, - }) - qt := QuantizedTime{} - qt.Set(time.Date(2007, time.January, 1, 15, 0, 0, 0, time.UTC)) - err := batch.Add(Row{ - ID: "1", - Values: []interface{}{uint64(2), "bkey", "ckey"}, - Time: qt, - }) - - checkResult := func(batch *ingestAPIBatch, id string, err error) { - if err != nil { - t.Fatalf("adding row to batch: %v", err) - } - - if batch.recordsK[id]["a"] != uint64(2) { - t.Fatalf("unexpected batch.records: %+v", batch.recordsK) - } - if batch.recordsK[id]["b"] != "bkey" { - t.Fatalf("unexpected batch.records: %+v", batch.recordsK) - } - if batch.recordsK[id]["c"].(map[string]interface{})["time"] != "2007-01-01T15:00:00Z" { - t.Fatalf("unexpected batch.records: %+v", batch.recordsK) - } - if batch.recordsK[id]["c"].(map[string]interface{})["values"] != "ckey" { - t.Fatalf("unexpected batch.records: %+v", batch.recordsK) - } - } - checkResult(batch, "1", err) - - // test wrong type row ID - if err := batch.Add(Row{ID: 64.5}); !strings.Contains(err.Error(), "unsupported rowID") { - t.Fatalf("unexpected error w/ floating point rowID: %v", err) - } - - // test that byte slice ID works same as string - err = batch.Add(Row{ - ID: []byte("2"), - Values: []interface{}{uint64(2), "bkey", "ckey"}, - Time: qt, - }) - checkResult(batch, "2", err) - - }) -} - -func TestIngestAPIBatch(t *testing.T) { - c := test.MustRunCluster(t, 3) - defer c.Close() - - urls := make([]string, len(c.Nodes)) - for i, n := range c.Nodes { - urls[i] = n.URL() - } - - // Create a new client for the cluster - cli, err := newClientFromAddresses(urls, &ClientOptions{}) - if err != nil { - t.Fatalf("getting new client: %v", err) - } - defer cli.Close() - - cli.IngestSchema(map[string]interface{}{ - "index-name": "test-1", - "index-action": "create", - "primary-key-type": "uint", - "field-action": "create", - "fields": []map[string]interface{}{ - { - "field-name": "astr", - "field-type": "string", - "field-options": map[string]interface{}{}, - }, - { - "field-name": "bint", - "field-type": "int", - "field-options": map[string]interface{}{}, - }, - { - "field-name": "cid", - "field-type": "id", - "field-options": map[string]interface{}{}, - }, - { - "field-name": "dtimestamp", - "field-type": "timestamp", - "field-options": map[string]interface{}{ - "unit": "s", - }, - }, - { - "field-name": "etime", - "field-type": "string", - "field-options": map[string]interface{}{ - "time-quantum": "YMD", - }, - }, - { - "field-name": "fdecimal", - "field-type": "decimal", - "field-options": map[string]interface{}{ - "scale": 3, - }, - }, - { - "field-name": "gbool", - "field-type": "bool", - "field-options": map[string]interface{}{}, - }, - }, - }) - - schema, err := cli.Schema() - if err != nil { - t.Fatalf("getting schema: %v", err) - } - index := schema.Index("test-1") - defer cli.DeleteIndex(index) - - batch := NewIngestAPIBatch(cli, 10, logger.NopLogger, []*Field{ - { - name: "astr", - index: &Index{name: "test-1", options: &IndexOptions{}}, - options: &FieldOptions{fieldType: FieldTypeSet, keys: true}, - }, - { - name: "bint", - options: &FieldOptions{fieldType: FieldTypeInt}, - }, - { - name: "cid", - options: &FieldOptions{fieldType: FieldTypeSet, keys: false}, - }, - { - name: "dtimestamp", - options: &FieldOptions{fieldType: FieldTypeTimestamp}, - }, - { - name: "etime", - options: &FieldOptions{fieldType: FieldTypeTime, keys: true, timeQuantum: TimeQuantumYearMonthDay}, - }, - { - name: "fdecimal", - options: &FieldOptions{fieldType: FieldTypeDecimal, scale: 3}, - }, - { - name: "gbool", - options: &FieldOptions{fieldType: FieldTypeBool}, - }, - }) - - qt0 := &QuantizedTime{} - qt0.Set(time.Date(2010, time.January, 1, 0, 0, 0, 0, time.UTC)) - if err := batch.Add(Row{ - ID: uint64(7), - Values: []interface{}{"a", -2, 9, 1287367623, "e", 1.2345, true}, - Time: *qt0, - }); err != nil { - t.Fatalf("adding row: %v", err) - } - - // test nil value case - if err := batch.Add(Row{ - ID: uint64(8), - Values: []interface{}{nil, nil, nil, nil, nil, nil, nil}, - Time: QuantizedTime{}, - }); err != nil { - t.Fatalf("error adding all nil batch which should affect nothing: %v", err) - } - - if err := batch.Import(); err != nil { - t.Fatalf("importing row: %v", err) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(astr=a)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(bint==-2)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(bint==-2) result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(cid=9)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(cid=9) result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(dtimestamp=='2010-10-18T02:07:03Z')", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(dtimestamp=='2010-10-18T02:07:03Z') result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(etime=e, from='2010-01-01', to='2010-01-02')", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(etime=e, from='2010-01-01', to='2010-01-02') result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(fdecimal==1.234)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(fdecimal==1.234) result: %+v", resp.Result().Row().Columns) - } - - if resp, err := cli.Query(NewPQLBaseQuery("Row(gbool=true)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { - t.Fatalf("querying: %v", err) - } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { - t.Fatalf("unexpected Row(gbool=true) result: %+v", resp.Result().Row().Columns) - } - -} diff --git a/client/orm.go b/client/orm.go index 98a72d700..3f8538ad0 100644 --- a/client/orm.go +++ b/client/orm.go @@ -208,10 +208,10 @@ func (q PQLRowQuery) Error() error { // // Usage: // -// repo, err := NewIndex("repository") -// stargazer, err := repo.Field("stargazer") -// query := repo.BatchQuery( -// stargazer.Row(5), +// repo, err := NewIndex("repository") +// stargazer, err := repo.Field("stargazer") +// query := repo.BatchQuery( +// stargazer.Row(5), // stargazer.Row(15), // repo.Union(stargazer.Row(20), stargazer.Row(25))) type PQLBatchQuery struct { @@ -797,6 +797,10 @@ func (fo FieldOptions) TimeUnit() string { return fo.timeUnit } +func (fo FieldOptions) Base() int64 { + return fo.base +} + // NoStandardView suppresses creating the standard view for supported field types (currently, time) func (fo FieldOptions) NoStandardView() bool { return fo.noStandardView diff --git a/idk/datagen/cmd.go b/idk/datagen/cmd.go index da4376261..f3f51dccc 100644 --- a/idk/datagen/cmd.go +++ b/idk/datagen/cmd.go @@ -51,7 +51,6 @@ type Main struct { Seed int64 `short:"" help:"Seed to use for any random number generation."` TrackProgress bool `short:"" help:"Periodically print status updates on how many records have been sourced."` - UseIngestAPI bool `help:"Experimental: use new HTTP/JSON ingest API instead of low-level import API. Probably slow, does not support packed bools."` UseShardTransactionalEndpoint bool `flag:"use-shard-transactional-endpoint" help:"Use experimental transactional endpoint"` @@ -281,7 +280,6 @@ func (m *Main) Preload() error { m.idkMain.CacheLength = m.Pilosa.CacheLength m.idkMain.NewSource = m.newSource m.idkMain.TrackProgress = m.TrackProgress - m.idkMain.UseIngestAPI = m.UseIngestAPI m.idkMain.AuthToken = m.AuthToken m.idkMain.UseShardTransactionalEndpoint = m.UseShardTransactionalEndpoint if len(m.Pilosa.Hosts) > 0 { diff --git a/idk/ingest.go b/idk/ingest.go index 76e00d1c0..8a463dccd 100644 --- a/idk/ingest.go +++ b/idk/ingest.go @@ -26,6 +26,7 @@ import ( "github.com/felixge/fgprof" pilosacore "github.com/molecula/featurebase/v3" pilosagrpc "github.com/molecula/featurebase/v3/api/client" + pilosabatch "github.com/molecula/featurebase/v3/batch" pilosaclient "github.com/molecula/featurebase/v3/client" "github.com/molecula/featurebase/v3/logger" "github.com/molecula/featurebase/v3/pql" @@ -86,7 +87,6 @@ type Main struct { OffsetMode bool `short:"" help:"Set offset-mode based Autogenerated IDs, for use with a data-source that is offset-based (must be set alongside auto-generate and external-generate)."` LookupDBDSN string `flag:"lookup-db-dsn" help:"Connection string for connecting to Lookup database."` LookupBatchSize int `help:"Number of records to batch before writing them to Lookup database."` - UseIngestAPI bool `help:"Experimental: use new HTTP/JSON ingest API instead of low-level import API. Probably slow. Does not support packed bools."` AuthToken string `flag:"auth-token" help:"Authentication Token for FeatureBase"` CommitTimeout time.Duration `help:"Maximum time before canceling commit."` AllowIntOutOfRange bool `help:"Allow ingest to continue when it encounters out of range integers in IntFields. (default false)"` @@ -322,10 +322,10 @@ func (m *Main) ingest(ctx context.Context, source Source, nexter IDAllocator, so defer func() { m.log.Printf("metrics: import=%s\n", time.Duration(atomic.LoadInt64((*int64)(&m.importDuration)))) }() - var batch pilosaclient.RecordBatch + var batch pilosabatch.RecordBatch var recordizers []Recordizer var prevRec Record - var row *pilosaclient.Row + var row *pilosabatch.Row var errorCounter int // keeps track of consecuitive errors across records var anyRecordSuccessful bool if m.progress != nil { @@ -554,7 +554,7 @@ initialFetch: m.stats.Count(MetricIngesterRowsAdded, 1, 1) } - if err == pilosaclient.ErrBatchNowFull || err == pilosaclient.ErrBatchNowStale { + if err == pilosabatch.ErrBatchNowFull || err == pilosabatch.ErrBatchNowStale { batchLen := batch.Len() err = m.importBatch(batch) if err != nil { @@ -763,7 +763,7 @@ func (m *Main) Setup() (onFinishRun func(), err error) { if m.AutoGenerate { shardWidth := m.index.ShardWidth() if shardWidth == 0 { - shardWidth = pilosaclient.DefaultShardWidth + shardWidth = pilosacore.ShardWidth } m.ra = NewLocalRangeAllocator(shardWidth) } @@ -1060,7 +1060,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error { } var recordizers []Recordizer - var row *pilosaclient.Row + var row *pilosabatch.Row rec, err := source.Record() if err == nil { err = ErrSchemaChange // always need to fetch the schema the first time @@ -1328,7 +1328,7 @@ func (m *Main) findPrimary() (*url.URL, error) { } // importBatch executes batch.Import() and saves its timing. -func (m *Main) importBatch(batch pilosaclient.RecordBatch) error { +func (m *Main) importBatch(batch pilosabatch.RecordBatch) error { t := time.Now() err := batch.Import() elapsed := time.Since(t) @@ -1336,9 +1336,9 @@ func (m *Main) importBatch(batch pilosaclient.RecordBatch) error { return err } -type Recordizer func(rawRec []interface{}, rec *pilosaclient.Row) error +type Recordizer func(rawRec []interface{}, rec *pilosabatch.Row) error -func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.RecordBatch, *pilosaclient.Row, []int, error) { +func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosabatch.RecordBatch, *pilosabatch.Row, []int, error) { // Before attempting to do anything, check for duplicates in the schema. { dedup := make(map[string]struct{}) @@ -1355,10 +1355,10 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor } // From the schema, and the configuration stored on Main, we need - // to create a []pilosaclient.Field and a []Recordizer processing + // to create a []pilosacore.Field and a []Recordizer processing // functions which take a []interface{} which conforms to the // schema, and converts it to a record which conforms to the - // []pilosaclient.Field. + // []pilosacore.Field. // // The relevant config options on Main are: // 1. PrimaryKeyFields, IDField, AutoGenerate @@ -1406,7 +1406,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor } } fieldIndex := fieldIndex - rz = func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + rz = func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { id, err := field.PilosafyVal(rawRec[fieldIndex]) if err != nil { return errors.Wrapf(err, "converting %+v to ID", rawRec[fieldIndex]) @@ -1503,7 +1503,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor // TODO may need to have more sophisticated recordizer by type at some point switch idkField.(type) { case RecordTimeField: - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { tyme, err := idkField.PilosafyVal(rawRec[i]) if err != nil { @@ -1514,7 +1514,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor return nil }) case IntField, DecimalField, TimestampField: - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Clears[valIdx] = uint64(0) @@ -1525,7 +1525,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor }) case IDField, StringField: hasMutex := HasMutex(idkField) - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: if hasMutex { //need to clear the mutex @@ -1540,7 +1540,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor return errors.Wrapf(err, "converting field %d:%+v, val:%+v", i, idkField, rawRec[i]) }) case BoolField: - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Values[valIdx] = nil @@ -1550,7 +1550,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor return errors.Wrapf(err, "converting field %d:%+v, val:%+v", i, idkField, rawRec[i]) }) default: - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { rec.Values[valIdx], err = idkField.PilosafyVal(rawRec[i]) return errors.Wrapf(err, "converting field %d:%+v, val:%+v", i, idkField, rawRec[i]) }) @@ -1561,7 +1561,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor // now handle this field if it was not already found in pilosa switch fld := idkField.(type) { case RecordTimeField: - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { tyme, err := idkField.PilosafyVal(rawRec[i]) if err != nil { return errors.Wrap(err, "converting recordtimefield") @@ -1595,7 +1595,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor fields = append(fields, m.index.Field(fld.DestName(), opts...)) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: if hasMutex { //need to clear the mutex @@ -1614,7 +1614,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor if m.PackBools == "" { fields = append(fields, m.index.Field(fld.DestName(), pilosaclient.OptFieldTypeBool())) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Values[valIdx] = nil @@ -1632,7 +1632,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor } else { fields = append(fields, boolField, boolFieldExists) fieldIdx := len(fields) - 2 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Clears[fieldIdx] = idkField.DestName() // clear bools bit for this field name @@ -1672,7 +1672,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor } fields = append(fields, m.index.Field(fld.DestName(), opts...)) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Clears[valIdx] = uint64(0) @@ -1684,7 +1684,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor case DecimalField: fields = append(fields, m.index.Field(fld.DestName(), pilosaclient.OptFieldTypeDecimal(fld.Scale))) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Clears[valIdx] = uint64(0) @@ -1696,7 +1696,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor case TimestampField: fields = append(fields, m.index.Field(fld.DestName(), pilosaclient.OptFieldTypeTimestamp(fld.epoch(), string(fld.granularity())))) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { switch rawRec[i].(type) { case DeleteSentinel: rec.Clears[valIdx] = uint64(0) @@ -1708,7 +1708,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor case DateIntField: fields = append(fields, m.index.Field(fld.DestName(), pilosaclient.OptFieldTypeInt())) valIdx := len(fields) - 1 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { rec.Values[valIdx], err = idkField.PilosafyVal(rawRec[i]) return errors.Wrapf(err, "converting field %d:%+v, val:%+v", i, idkField, rawRec[i]) }) @@ -1719,7 +1719,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor m.index.Field(name+Exists, pilosaclient.OptFieldTypeSet(pilosaclient.CacheTypeRanked, pilosacore.DefaultCacheSize)), ) valIdx := len(fields) - 2 - recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizers = append(recordizers, func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { val, err := idkField.PilosafyVal(rawRec[i]) if val == nil && err == nil { rec.Values[valIdx] = nil @@ -1769,24 +1769,21 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosaclient.Recor if err != nil { return nil, nil, nil, nil, errors.Wrap(err, "creating batch") } - row := &pilosaclient.Row{ + row := &pilosabatch.Row{ Values: make([]interface{}, len(fields)), Clears: make(map[int]interface{}), } return recordizers, batch, row, lookupWriteIdxs, nil } -func (m *Main) newBatch(fields []*pilosaclient.Field) (pilosaclient.RecordBatch, error) { - if m.UseIngestAPI { - return pilosaclient.NewIngestAPIBatch(m.client, m.BatchSize, m.log, fields), nil - } - return pilosaclient.NewBatch(m.client, m.BatchSize, m.index, fields, - pilosaclient.OptLogger(m.log), - pilosaclient.OptCacheMaxAge(m.CacheLength), - pilosaclient.OptSplitBatchMode(m.ExpSplitBatchMode), - pilosaclient.OptMaxStaleness(m.BatchMaxStaleness), - pilosaclient.OptKeyTranslateBatchSize(m.KeyTranslateBatchSize), - pilosaclient.OptUseShardTransactionalEndpoint(m.UseShardTransactionalEndpoint), +func (m *Main) newBatch(fields []*pilosaclient.Field) (pilosabatch.RecordBatch, error) { + return pilosabatch.NewBatch(pilosaclient.NewImporter(m.client), m.BatchSize, pilosaclient.FromClientIndex(m.index), pilosaclient.FromClientFields(fields), + pilosabatch.OptLogger(m.log), + pilosabatch.OptCacheMaxAge(m.CacheLength), + pilosabatch.OptSplitBatchMode(m.ExpSplitBatchMode), + pilosabatch.OptMaxStaleness(m.BatchMaxStaleness), + pilosabatch.OptKeyTranslateBatchSize(m.KeyTranslateBatchSize), + pilosabatch.OptUseShardTransactionalEndpoint(m.UseShardTransactionalEndpoint), ) } @@ -2023,7 +2020,7 @@ func getPrimaryKeyRecordizer(schema []Field, pkFields []string) (recordizer Reco skipFields[fieldIndices[0]] = struct{}{} } } - recordizer = func(rawRec []interface{}, rec *pilosaclient.Row) (err error) { + recordizer = func(rawRec []interface{}, rec *pilosabatch.Row) (err error) { // first, special case for performance when there is a single // primary key field and it is a byte slice already. if len(fieldIndices) == 1 { diff --git a/idk/ingest_test.go b/idk/ingest_test.go index 63cd78553..738dd7976 100644 --- a/idk/ingest_test.go +++ b/idk/ingest_test.go @@ -16,6 +16,7 @@ import ( "github.com/golang-jwt/jwt" "github.com/molecula/featurebase/v3/authn" + "github.com/molecula/featurebase/v3/batch" pilosaclient "github.com/molecula/featurebase/v3/client" "github.com/molecula/featurebase/v3/idk/idktest" "github.com/molecula/featurebase/v3/logger" @@ -713,7 +714,7 @@ func TestGetPrimaryKeyRecordizer(t *testing.T) { t.Errorf("unmatched skips exp/got\n%+v\n%+v", test.expSkip, skips) } - row := &pilosaclient.Row{} + row := &batch.Row{} err = rdz(test.rawRec, row) if err != nil { t.Fatalf("unexpected error from recordizer: %v", err) @@ -750,11 +751,11 @@ func TestBatchFromSchema(t *testing.T) { err string batchErr string rdzErrs []string - time pilosaclient.QuantizedTime + time batch.QuantizedTime lookupWriteIdxs []int } - getQuantizedTime := func(t time.Time) pilosaclient.QuantizedTime { - qt := pilosaclient.QuantizedTime{} + getQuantizedTime := func(t time.Time) batch.QuantizedTime { + qt := batch.QuantizedTime{} qt.Set(t) return qt } @@ -844,13 +845,13 @@ func TestBatchFromSchema(t *testing.T) { { name: "empty", autogen: true, - err: "can't batch with no fields or batch size", + err: "can't batch with no fields", }, { name: "empty-w/ExtGen", autogen: true, extgen: true, - err: "can't batch with no fields or batch size", + err: "can't batch with no fields", }, { name: "no id field", @@ -1275,7 +1276,7 @@ type testSource struct { schema []Field } -func (t *testSource) Close() error { +func (s *testSource) Close() error { return nil } diff --git a/server/server.go b/server/server.go index 7517b8107..c0452a8d1 100644 --- a/server/server.go +++ b/server/server.go @@ -30,6 +30,7 @@ import ( pilosa "github.com/molecula/featurebase/v3" "github.com/molecula/featurebase/v3/authn" "github.com/molecula/featurebase/v3/authz" + "github.com/molecula/featurebase/v3/batch" "github.com/molecula/featurebase/v3/boltdb" "github.com/molecula/featurebase/v3/encoding/proto" petcd "github.com/molecula/featurebase/v3/etcd" @@ -438,9 +439,10 @@ func (m *Command) SetupServer() error { m.Config.Etcd.Id = m.Config.Name // TODO(twg) rethink this e := petcd.NewEtcd(m.Config.Etcd, m.logger, m.Config.Cluster.ReplicaN, version) - executionPlannerFn := func(e pilosa.Executor, a *pilosa.API, s string) sql3.CompilePlanner { - fapi := &pilosa.FeatureBaseSchemaAPI{API: a} - return planner.NewExecutionPlanner(e, fapi, a, s) + executionPlannerFn := func(e pilosa.Executor, api *pilosa.API, sql string) sql3.CompilePlanner { + fapi := &pilosa.FeatureBaseSchemaAPI{API: api} + fimp := &batch.FeaturebaseImporter{API: api} + return planner.NewExecutionPlanner(e, fapi, api, fimp, m.logger, sql) } serverOptions := []pilosa.ServerOption{ diff --git a/sql3/parser/ast.go b/sql3/parser/ast.go index 39bb3141d..9e63ab5da 100644 --- a/sql3/parser/ast.go +++ b/sql3/parser/ast.go @@ -2018,7 +2018,7 @@ func (l *ExprList) Clone() *ExprList { return &other } -/*func cloneExprLists(a []*ExprList) []*ExprList { +func cloneExprLists(a []*ExprList) []*ExprList { if a == nil { return nil } @@ -2027,7 +2027,7 @@ func (l *ExprList) Clone() *ExprList { other[i] = a[i].Clone() } return other -}*/ +} // String returns the string representation of the expression. func (l *ExprList) String() string { @@ -2775,8 +2775,8 @@ type InsertStatement struct { Columns []*Ident // optional column list ColumnsRparen Pos // position of column list right paren - Values Pos // position of VALUES keyword - ValueList *ExprList // list of values + Values Pos // position of VALUES keyword + TupleList []*ExprList // multiple tuples // Select *SelectStatement // SELECT statement @@ -2796,7 +2796,7 @@ func (s *InsertStatement) Clone() *InsertStatement { other.Table = s.Table.Clone() other.Alias = s.Alias.Clone() other.Columns = cloneIdents(s.Columns) - other.ValueList = s.ValueList.Clone() + other.TupleList = cloneExprLists(s.TupleList) //other.Select = s.Select.Clone() //other.UpsertClause = s.UpsertClause.Clone() return &other @@ -2848,14 +2848,19 @@ func (s *InsertStatement) String() string { // fmt.Fprintf(&buf, " %s", s.Select.String()) //} else { buf.WriteString(" VALUES") - buf.WriteString(" (") - for j, expr := range s.ValueList.Exprs { - if j != 0 { - buf.WriteString(", ") + for i, tuple := range s.TupleList { + if i != 0 { + buf.WriteString(",") } - buf.WriteString(expr.String()) + buf.WriteString(" (") + for j, expr := range tuple.Exprs { + if j != 0 { + buf.WriteString(", ") + } + buf.WriteString(expr.String()) + } + buf.WriteString(")") } - buf.WriteString(")") //} //if s.UpsertClause != nil { diff --git a/sql3/parser/ast_test.go b/sql3/parser/ast_test.go index a5b87843a..bc4466171 100644 --- a/sql3/parser/ast_test.go +++ b/sql3/parser/ast_test.go @@ -592,11 +592,29 @@ func TestInsertStatement_String(t *testing.T) { {Name: "x"}, {Name: "y"}, }, - ValueList: &parser.ExprList{ - Exprs: []parser.Expr{&parser.NullLit{}, &parser.NullLit{}}, + TupleList: []*parser.ExprList{ + { + Exprs: []parser.Expr{&parser.NullLit{}, &parser.NullLit{}}, + }, }, }, `INSERT INTO "tbl" ("x", "y") VALUES (NULL, NULL)`) + AssertStatementStringer(t, &parser.InsertStatement{ + Table: &parser.Ident{Name: "tbl"}, + Columns: []*parser.Ident{ + {Name: "x"}, + {Name: "y"}, + }, + TupleList: []*parser.ExprList{ + { + Exprs: []parser.Expr{&parser.IntegerLit{Value: "1"}, &parser.IntegerLit{Value: "2"}}, + }, + { + Exprs: []parser.Expr{&parser.IntegerLit{Value: "3"}, &parser.IntegerLit{Value: "4"}}, + }, + }, + }, `INSERT INTO "tbl" ("x", "y") VALUES (1, 2), (3, 4)`) + // AssertStatementStringer(t, &sql.InsertStatement{ // WithClause: &sql.WithClause{ // CTEs: []*sql.CTE{ diff --git a/sql3/parser/parser.go b/sql3/parser/parser.go index d8227894d..f7ad8daba 100644 --- a/sql3/parser/parser.go +++ b/sql3/parser/parser.go @@ -1584,19 +1584,23 @@ func (p *Parser) parseInsertStatement(withClause *WithClause) (_ *InsertStatemen switch p.peek() { case VALUES: stmt.Values, _, _ = p.scan() + + // Parse out the value tuples. + stmt.TupleList = make([]*ExprList, 0) + for { - var list ExprList + var tuple ExprList if p.peek() != LP { return &stmt, p.errorExpected(p.pos, p.tok, "left paren") } - list.Lparen, _, _ = p.scan() + tuple.Lparen, _, _ = p.scan() for { expr, err := p.ParseExpr() if err != nil { return &stmt, err } - list.Exprs = append(list.Exprs, expr) + tuple.Exprs = append(tuple.Exprs, expr) if p.peek() == RP { break @@ -1605,14 +1609,16 @@ func (p *Parser) parseInsertStatement(withClause *WithClause) (_ *InsertStatemen } p.scan() } - list.Rparen, _, _ = p.scan() - stmt.ValueList = &list + tuple.Rparen, _, _ = p.scan() + + stmt.TupleList = append(stmt.TupleList, &tuple) if p.peek() != COMMA { break } p.scan() } + //case SELECT: // if stmt.Select, err = p.parseSelectStatement(false, nil); err != nil { // return &stmt, err diff --git a/sql3/parser/parser_test.go b/sql3/parser/parser_test.go index c96c9676a..43faf760c 100644 --- a/sql3/parser/parser_test.go +++ b/sql3/parser/parser_test.go @@ -2269,15 +2269,50 @@ func TestParser_ParseStatement(t *testing.T) { }, ColumnsRparen: pos(21), Values: pos(23), - ValueList: &parser.ExprList{ - Lparen: pos(30), - Exprs: []parser.Expr{ - &parser.IntegerLit{ValuePos: pos(31), Value: "1"}, - &parser.IntegerLit{ValuePos: pos(34), Value: "2"}, + TupleList: []*parser.ExprList{ + { + Lparen: pos(30), + Exprs: []parser.Expr{ + &parser.IntegerLit{ValuePos: pos(31), Value: "1"}, + &parser.IntegerLit{ValuePos: pos(34), Value: "2"}, + }, + Rparen: pos(35), }, - Rparen: pos(35), }, }) + + // Ensure we can parse multiple tuple values in an INSERT INTO statement. + AssertParseStatement(t, `INSERT INTO tbl (x, y) VALUES (1, 2), (3, 4)`, &parser.InsertStatement{ + Insert: pos(0), + Into: pos(7), + Table: &parser.Ident{NamePos: pos(12), Name: "tbl"}, + ColumnsLparen: pos(16), + Columns: []*parser.Ident{ + {NamePos: pos(17), Name: "x"}, + {NamePos: pos(20), Name: "y"}, + }, + ColumnsRparen: pos(21), + Values: pos(23), + TupleList: []*parser.ExprList{ + { + Lparen: pos(30), + Exprs: []parser.Expr{ + &parser.IntegerLit{ValuePos: pos(31), Value: "1"}, + &parser.IntegerLit{ValuePos: pos(34), Value: "2"}, + }, + Rparen: pos(35), + }, + { + Lparen: pos(38), + Exprs: []parser.Expr{ + &parser.IntegerLit{ValuePos: pos(39), Value: "3"}, + &parser.IntegerLit{ValuePos: pos(42), Value: "4"}, + }, + Rparen: pos(43), + }, + }, + }) + /*AssertParseStatement(t, `REPLACE INTO tbl (x, y) VALUES (1, 2), (3, 4)`, &parser.InsertStatement{ Replace: pos(0), Into: pos(8), @@ -3604,7 +3639,3 @@ func AssertParseExprError(tb testing.TB, s string, want string) { func pos(offset int) parser.Pos { return parser.Pos{Offset: offset, Line: 1, Column: offset + 1} } - -func deepEqual(a, b interface{}) string { - return strings.Join(deep.Equal(a, b), "\n") -} diff --git a/sql3/parser/walk.go b/sql3/parser/walk.go index 2104bdcf1..8433961c1 100644 --- a/sql3/parser/walk.go +++ b/sql3/parser/walk.go @@ -275,15 +275,17 @@ func walk(v Visitor, node Node) (_ Node, err error) { if err := walkIdentList(v, n.Columns); err != nil { return node, err } - //for i := range n.ValueLists { - if list, err := walk(v, n.ValueList); err != nil { - return node, err - } else if list != nil { - n.ValueList = list.(*ExprList) - } else { - n.ValueList = nil + + for i, tuple := range n.TupleList { + if list, err := walk(v, tuple); err != nil { + return node, err + } else if list != nil { + n.TupleList[i] = list.(*ExprList) + } else { + n.TupleList[i] = nil + } } - //} + /*if n.Select != nil { if sel, err := walk(v, n.Select); err != nil { return node, err diff --git a/sql3/planner/compileinsert.go b/sql3/planner/compileinsert.go index e76d3d96a..6a6b80bac 100644 --- a/sql3/planner/compileinsert.go +++ b/sql3/planner/compileinsert.go @@ -18,7 +18,7 @@ func (p *ExecutionPlanner) compileInsertStatement(stmt *parser.InsertStatement) tableName := parser.IdentName(stmt.Table) targetColumns := []*qualifiedRefPlanExpression{} - insertValues := []types.PlanExpression{} + insertValues := [][]types.PlanExpression{} table, err := p.schemaAPI.IndexInfo(context.Background(), tableName) if err != nil { @@ -54,12 +54,16 @@ func (p *ExecutionPlanner) compileInsertStatement(stmt *parser.InsertStatement) } //add expressions from values list - for _, expr := range stmt.ValueList.Exprs { - e, err := p.compileExpr(expr) - if err != nil { - return nil, err + for _, tuple := range stmt.TupleList { + tupleValues := []types.PlanExpression{} + for _, expr := range tuple.Exprs { + e, err := p.compileExpr(expr) + if err != nil { + return nil, err + } + tupleValues = append(tupleValues, e) } - insertValues = append(insertValues, e) + insertValues = append(insertValues, tupleValues) } return NewPlanOpInsert(p, tableName, targetColumns, insertValues), nil @@ -92,8 +96,10 @@ func (p *ExecutionPlanner) analyzeInsertStatement(stmt *parser.InsertStatement) } // Make sure (implicit) insert list and expression list have the same // number of items. - if len(typeNames) != len(stmt.ValueList.Exprs) { - return sql3.NewErrInsertExprTargetCountMismatch(stmt.ValueList.Lparen.Line, stmt.ValueList.Lparen.Column) + for _, tuple := range stmt.TupleList { + if len(typeNames) != len(tuple.Exprs) { + return sql3.NewErrInsertExprTargetCountMismatch(tuple.Lparen.Line, tuple.Lparen.Column) + } } } else { // Check column list refers to actual columns, and that there are no @@ -153,24 +159,28 @@ func (p *ExecutionPlanner) analyzeInsertStatement(stmt *parser.InsertStatement) } // Make sure insert list and expression list have the same number of items. - if len(stmt.Columns) != len(stmt.ValueList.Exprs) { - return sql3.NewErrInsertExprTargetCountMismatch(stmt.ValueList.Lparen.Line, stmt.ValueList.Lparen.Column) + for _, tuple := range stmt.TupleList { + if len(stmt.Columns) != len(tuple.Exprs) { + return sql3.NewErrInsertExprTargetCountMismatch(tuple.Lparen.Line, tuple.Lparen.Column) + } } } // Check each of the expressions. - for i, expr := range stmt.ValueList.Exprs { - e, err := p.analyzeExpression(expr, stmt) - if err != nil { - return err - } + for _, tuple := range stmt.TupleList { + for i, expr := range tuple.Exprs { + e, err := p.analyzeExpression(expr, stmt) + if err != nil { + return err + } - // Type check against same ordinal position in column type list. - if !typesAreAssignmentCompatible(typeNames[i], e.DataType()) { - return sql3.NewErrTypeAssignmentIncompatible(expr.Pos().Line, expr.Pos().Column, e.DataType().TypeName(), typeNames[i].TypeName()) - } + // Type check against same ordinal position in column type list. + if !typesAreAssignmentCompatible(typeNames[i], e.DataType()) { + return sql3.NewErrTypeAssignmentIncompatible(expr.Pos().Line, expr.Pos().Column, e.DataType().TypeName(), typeNames[i].TypeName()) + } - stmt.ValueList.Exprs[i] = e + tuple.Exprs[i] = e + } } return nil diff --git a/sql3/planner/executionplanner.go b/sql3/planner/executionplanner.go index 63978095d..77f50fc11 100644 --- a/sql3/planner/executionplanner.go +++ b/sql3/planner/executionplanner.go @@ -5,9 +5,10 @@ package planner import ( "context" "encoding/json" - "log" pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/batch" + "github.com/molecula/featurebase/v3/logger" "github.com/molecula/featurebase/v3/sql3" "github.com/molecula/featurebase/v3/sql3/parser" "github.com/molecula/featurebase/v3/sql3/planner/types" @@ -26,15 +27,19 @@ type ExecutionPlanner struct { executor pilosa.Executor schemaAPI pilosa.SchemaAPI computeAPI pilosa.ComputeAPI + importer batch.Importer + logger logger.Logger sql string scopeStack *scopeStack } -func NewExecutionPlanner(executor pilosa.Executor, schemaAPI pilosa.SchemaAPI, computeAPI pilosa.ComputeAPI, sql string) *ExecutionPlanner { +func NewExecutionPlanner(executor pilosa.Executor, schemaAPI pilosa.SchemaAPI, computeAPI pilosa.ComputeAPI, importer batch.Importer, logger logger.Logger, sql string) *ExecutionPlanner { return &ExecutionPlanner{ executor: executor, schemaAPI: schemaAPI, computeAPI: computeAPI, + importer: importer, + logger: logger, sql: sql, scopeStack: newScopeStack(), } @@ -80,10 +85,15 @@ func (p *ExecutionPlanner) CompilePlan(ctx context.Context, stmt parser.Statemen } // Log the plan. This happens even if an error occurred. - if rootOperator != nil { + switch rootOperator.(type) { + case *PlanOpInsert: + // Don't log the insert plan since it can be very large. + case nil: + // pass + default: plan := rootOperator.Plan() a, _ := json.MarshalIndent(plan, "", " ") - log.Println(string(a)) + p.logger.Debugf(string(a)) } return rootOperator, err diff --git a/sql3/planner/expression.go b/sql3/planner/expression.go index d3f436c0c..7cf55be7b 100644 --- a/sql3/planner/expression.go +++ b/sql3/planner/expression.go @@ -7,6 +7,7 @@ import ( "fmt" "math" "regexp" + "sort" "strconv" "strings" "time" @@ -1912,8 +1913,27 @@ func (n *castPlanExpression) Evaluate(currentRow []interface{}) (interface{}, er case *parser.DataTypeStringSet: return nl, nil case *parser.DataTypeString: - //TODO(pok) come up with a better string representation of string set - return fmt.Sprintf("%v", nl), nil + sort.Strings(nl) + + var ret strings.Builder + + // open bracket + ret.WriteString("[") + + // elements + var afterFirst bool + for i := range nl { + if afterFirst { + ret.WriteString(",") + } + ret.WriteString(`"` + strings.ReplaceAll(nl[i], `"`, `\"`) + `"`) + afterFirst = true + } + + // close braket + ret.WriteString("]") + + return ret.String(), nil } case *parser.DataTypeTimestamp: @@ -2097,17 +2117,12 @@ func (n *exprTupleLiteralPlanExpression) Evaluate(currentRow []interface{}) (int return nil, err } - //if it is a string, do a coercion - val, ok := timestampEval.(string) - if ok { - if tm, err := time.ParseInLocation(time.RFC3339Nano, val, time.UTC); err == nil { - timestampEval = tm - } else if tm, err := time.ParseInLocation(time.RFC3339, val, time.UTC); err == nil { - timestampEval = tm - } else if tm, err := time.ParseInLocation("2006-01-02", val, time.UTC); err == nil { - timestampEval = tm - } else { + // if it is a string, do a coercion + if val, ok := timestampEval.(string); ok { + if tm, err := timestampFromString(val); err != nil { return nil, sql3.NewErrInvalidTypeCoercion(0, 0, val, n.members[0].Type().TypeName()) + } else { + timestampEval = tm } } @@ -2482,3 +2497,17 @@ func wildCardToRegexp(pattern string) string { return result.String() } + +// timeFromString attempts to parse the string to a time.Time using a series of +// time formats. +func timestampFromString(s string) (time.Time, error) { + if tm, err := time.ParseInLocation(time.RFC3339Nano, s, time.UTC); err == nil { + return tm, nil + } else if tm, err := time.ParseInLocation(time.RFC3339, s, time.UTC); err == nil { + return tm, nil + } else if tm, err := time.ParseInLocation("2006-01-02", s, time.UTC); err == nil { + return tm, nil + } + + return time.Time{}, sql3.NewErrInvalidTypeCoercion(0, 0, s, "time.Time") +} diff --git a/sql3/planner/opbulkinsert.go b/sql3/planner/opbulkinsert.go index 4522bb25b..19df38cf0 100644 --- a/sql3/planner/opbulkinsert.go +++ b/sql3/planner/opbulkinsert.go @@ -138,7 +138,11 @@ type bulkInsertCSVRowIter struct { isKeyed bool options *bulkInsertOptions - latch *struct{} + // latch is used to indicate if the CSV has been processed. It will + // be set to a non-nil value upon processing. After that, the file + // should not be processed again. + latch *struct{} + currentBatch []interface{} lastKeyValue uint64 } @@ -146,39 +150,46 @@ type bulkInsertCSVRowIter struct { var _ types.RowIterator = (*bulkInsertCSVRowIter)(nil) func (i *bulkInsertCSVRowIter) Next(ctx context.Context) (types.Row, error) { - if i.latch == nil { - i.latch = &struct{}{} - i.lastKeyValue = 0 + // If Next has already been called, return early. We only want to process + // the file once. + if i.latch != nil { + return nil, types.ErrNoMoreRows + } - f, err := os.Open(i.options.fileName) - if err != nil { + // Set latch to indicate that Next() has been called. + i.latch = &struct{}{} + + i.lastKeyValue = 0 + + f, err := os.Open(i.options.fileName) + if err != nil { + return nil, err + } + + defer f.Close() + + linesRead := 0 + csvReader := csv.NewReader(f) + for { + rec, err := csvReader.Read() + if err == io.EOF { + break + } else if err != nil { return nil, err } - defer f.Close() + // do something with read line + if err = i.processCSVLine(ctx, rec); err != nil { + return nil, err + } + linesRead += 1 - linesRead := 0 - csvReader := csv.NewReader(f) - for { - rec, err := csvReader.Read() - if err == io.EOF { - break - } - if err != nil { - return nil, err - } - // do something with read line - err = i.processCSVLine(ctx, rec) - if err != nil { - return nil, err - } - linesRead += 1 - // bail if we have a rows limit and we've hit it - if i.options.rowsLimit > 0 && linesRead >= i.options.rowsLimit { - break - } + // bail if we have a rows limit and we've hit it + if i.options.rowsLimit > 0 && linesRead >= i.options.rowsLimit { + break } } + return nil, types.ErrNoMoreRows } diff --git a/sql3/planner/opinsert.go b/sql3/planner/opinsert.go index 00c5c1888..5969eb0d4 100644 --- a/sql3/planner/opinsert.go +++ b/sql3/planner/opinsert.go @@ -9,10 +9,10 @@ import ( "time" pilosa "github.com/molecula/featurebase/v3" - "github.com/molecula/featurebase/v3/pql" + fbbatch "github.com/molecula/featurebase/v3/batch" "github.com/molecula/featurebase/v3/sql3" - "github.com/molecula/featurebase/v3/sql3/parser" "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/pkg/errors" ) // PlanOpInsert plan operator to handle INSERT. @@ -20,11 +20,11 @@ type PlanOpInsert struct { planner *ExecutionPlanner tableName string targetColumns []*qualifiedRefPlanExpression - insertValues []types.PlanExpression + insertValues [][]types.PlanExpression warnings []string } -func NewPlanOpInsert(p *ExecutionPlanner, tableName string, targetColumns []*qualifiedRefPlanExpression, insertValues []types.PlanExpression) *PlanOpInsert { +func NewPlanOpInsert(p *ExecutionPlanner, tableName string, targetColumns []*qualifiedRefPlanExpression, insertValues [][]types.PlanExpression) *PlanOpInsert { return &PlanOpInsert{ planner: p, tableName: tableName, @@ -48,11 +48,15 @@ func (p *PlanOpInsert) Plan() map[string]interface{} { ps = append(ps, e.Plan()) } result["targetColumns"] = ps - ps = make([]interface{}, 0) - for _, e := range p.insertValues { - ps = append(ps, e.Plan()) + pps := make([]interface{}, 0) + for _, tuple := range p.insertValues { + ps := make([]interface{}, 0) + for _, e := range tuple { + ps = append(ps, e.Plan()) + } + pps = append(pps, ps) } - result["insertValues"] = ps + result["insertValues"] = pps return result } @@ -93,362 +97,236 @@ type insertRowIter struct { planner *ExecutionPlanner tableName string targetColumns []*qualifiedRefPlanExpression - insertValues []types.PlanExpression + insertValues [][]types.PlanExpression } var _ types.RowIterator = (*insertRowIter)(nil) func (i *insertRowIter) Next(ctx context.Context) (types.Row, error) { - qcx := i.planner.computeAPI.Txf().NewQcx() + // posID is the position of the "_id" column in both the targetColumns and + // values lists. + var posID int - colIDs := make([]uint64, 0) - colKeys := make([]string, 0) + // posVals maps the position in the tuple to the position in the row.Values. + // It essentially takes the _id column into account and skips it. + // + // So for example, if we have sql: + // INSERT INTO (a, _id, b, c) VALUES ('aa', 1, 'bb', 'cc'); + // + // Then we want row.ID = 1 and row.Values to contain {'aa', 'bb', 'cc'}. + // This means posVals would contain the map []int{0, 1*, 1, 2}, which maps + // VALUES positions (0,2,3) to row.Values (0, 1, 2). Note, the _id position + // (shown as 1* in the example above) isn't used because we handle it + // separately. + posVals := make([]int, len(i.targetColumns)) - addColID := func(v interface{}) error { - switch id := v.(type) { - case int64: - colIDs = append(colIDs, uint64(id)) - case uint64: - colIDs = append(colIDs, id) - case string: - colKeys = append(colKeys, id) - default: - return sql3.NewErrInternalf("unhandled _id data type '%T'", id) + var foundPosID bool + for j := range i.targetColumns { + if foundPosID { + posVals[j] = j - 1 + continue } - return nil + if strings.EqualFold(i.targetColumns[j].columnName, "_id") { + posID = j + foundPosID = true + } + posVals[j] = j } - //find the _id column and evaluate - var err error - var columnID interface{} - for idx, iv := range i.insertValues { - targetColumn := i.targetColumns[idx] - if strings.EqualFold(targetColumn.columnName, "_id") { - columnID, err = iv.Evaluate(nil) - if err != nil { - return nil, err + // batchSize is currently set to the size of the entire + // VALUES list. In the future we may want to break this up into smaller + // batches. + batchSize := len(i.insertValues) + + // idxInfoBase is the full IndexInfo stored in the schema. The instance of + // IndexInfo used in the import (and created below) will be based on the + // information from idxInfoBase, but the fields may be a limited subset, and + // may be in a different order. + idxInfoBase, err := i.planner.schemaAPI.IndexInfo(ctx, i.tableName) + if err != nil { + return nil, sql3.NewErrTableNotFound(0, 0, i.tableName) + } + + // idxInfo is a subset of idxInfoBase, containing only those fields included + // in the INSERT INTO statement (i.e. only i.targetcolumns), and in the + // order specified. + idxInfo := &pilosa.IndexInfo{ + Name: idxInfoBase.Name, + CreatedAt: idxInfoBase.CreatedAt, + Options: idxInfoBase.Options, + Fields: make([]*pilosa.FieldInfo, len(i.targetColumns)-1), + ShardWidth: idxInfoBase.ShardWidth, + } + + // Set up Fields based on i.targetColumns. + var counter int + for ii, targetColumn := range i.targetColumns { + // Skip the "_id" column. + if ii == posID { + continue + } + idxInfo.Fields[counter] = idxInfoBase.Field(targetColumn.columnName) + counter++ + } + + batch, err := fbbatch.NewBatch(i.planner.importer, batchSize, idxInfo, idxInfo.Fields, + fbbatch.OptUseShardTransactionalEndpoint(true), + ) + if err != nil { + return nil, errors.Wrap(err, "setting up batch") + } + + // row is the single instance of batch.Row allocated. It is re-used + // throughout the for loop to minimize memory allocation. + var row fbbatch.Row + + // Initialize row.Values to the size of the target columns, but exclude the + // record ID ("_id") since that's stored in row.ID. + row.Values = make([]interface{}, len(i.targetColumns)-1) + + for _, tuple := range i.insertValues { + // Evaluate and set the record ID. + if eval, err := tuple[posID].Evaluate(nil); err != nil { + return nil, errors.Wrapf(err, "evaluating record id: %v", tuple[posID]) + } else { + // These value types correspond to the types supported in batch.Add(). + switch recid := eval.(type) { + case string, uint64, []byte: + row.ID = recid + case int64: + if recid < 0 { + return nil, sql3.NewErrInternalf("_id value cannot be negative: %d", recid) + } + row.ID = uint64(recid) + default: + // If we get to here, it's likey that the id type is unsupported + // and will cause an error in batch.Add(). So there's no need to + // return an error here in this default. + row.ID = eval } - break + } + + // Loop over the values in the tuple and populate the Row values. + for idx, iv := range tuple { + // Skip the record ID because that was already handled above + // (prior to this loop). + if idx == posID { + continue + } + + eval, err := iv.Evaluate(nil) + if err != nil { + return nil, errors.Wrapf(err, "evaluating tuple value: %v", iv) + } + + // batch.Add does not typically look at field type to determine how + // to handle a particular value in a row. Instead, it uses value + // type. As an example, if the value type is int64, then batch.Add + // assumes that it should be handled as if going into an `int` + // field. Therefore, we need to look at the field type here and make + // sure that the value types being sent through batch.Add align with + // the field type assumptions that batch.Add is making. + switch opts := idxInfo.Fields[posVals[idx]].Options; opts.Type { + + // By the time we get here, we assume that the planner has already + // determined the value type such that the Evaluate() method called + // above results in the correct value types. There is one exception: + // sql3 treats all integer values as int64. This means that an ID + // field, which expects a uint64 value, would get treated as an int + // field. In order to avoid that, we cast int64 values to uint64 + // when the field type is set or mutex. + case pilosa.FieldTypeSet, pilosa.FieldTypeMutex: + switch v := eval.(type) { + case int64: + if v < 0 { + return nil, sql3.NewErrInternalf("converting negative value to uint64: %d", v) + } + row.Values[posVals[idx]] = uint64(v) + case []int64: + uint64s := make([]uint64, len(v)) + for i := range v { + if v[i] < 0 { + return nil, sql3.NewErrInternalf("converting negative slice value to uint64: %d", v[i]) + } + uint64s[i] = uint64(v[i]) + } + row.Values[posVals[idx]] = uint64s + default: + row.Values[posVals[idx]] = eval + } + + case pilosa.FieldTypeTimestamp: + switch v := eval.(type) { + + // time.Time is used for date literals generated in the parser. + // For example, if using `current_time`, the type received here + // will be a time.Time. + case time.Time: + // Convert Base, which is the epoch for Timestamp fields, to + // a time.Time value. + unit := fbbatch.TimeUnit(opts.TimeUnit) + epoch, err := fbbatch.Int64ToTimestamp(unit, time.Time{}, opts.Base) + if err != nil { + return nil, errors.Wrapf(err, "converting base to epoch: %d", opts.Base) + } + + i64, err := fbbatch.TimestampToInt64(unit, epoch, v) + if err != nil { + return nil, errors.Wrapf(err, "converting timestamp to int64: %s", v) + } + row.Values[posVals[idx]] = i64 + + // string is the normal case for dates; used when the date is + // provided as a string in the INSERT INTO statement. + case string: + ts, err := timestampFromString(v) + if err != nil { + return nil, errors.Wrapf(err, "parsing timestamp: %s", v) + } + + // Convert Base, which is the epoch for Timestamp fields, to + // a time.Time value. + unit := fbbatch.TimeUnit(opts.TimeUnit) + epoch, err := fbbatch.Int64ToTimestamp(unit, time.Time{}, opts.Base) + if err != nil { + return nil, errors.Wrapf(err, "converting base to epoch: %d", opts.Base) + } + + i64, err := fbbatch.TimestampToInt64(unit, epoch, ts) + if err != nil { + return nil, errors.Wrapf(err, "converting timestamp to int64: %s", v) + } + row.Values[posVals[idx]] = i64 + + // nil is to support `null` values. + case nil: + row.Values[posVals[idx]] = eval + + default: + return nil, sql3.NewErrInternalf("unsupported timestamp type: %T", eval) + } + + default: + row.Values[posVals[idx]] = eval + } + } + + if err := batch.Add(row); err != nil { + // Breaking here on ErrBatchNowFull is only valid because we are + // explicity setting the batch size to the number of tuples in the + // INSERT INTO statement. Which means we're only handling a single + // batch. If this evolves to handle multiple batches, this will need + // to instead call batch.Import() and continue looping over tuples. + // We may also need to handle ErrBatchNowStale. + if err == fbbatch.ErrBatchNowFull { + break + } + return nil, errors.Wrap(err, "adding record") } } - //eval all the expressions and do the insert - for idx, iv := range i.insertValues { - colIDs = make([]uint64, 0) - colKeys = make([]string, 0) - - targetColumn := i.targetColumns[idx] - - if strings.EqualFold(targetColumn.columnName, "_id") { - continue - } - - eval, err := iv.Evaluate(nil) - if err != nil { - return nil, err - } - - //nothing to do if a value is null - if eval == nil { - continue - } - - sourceType := iv.Type() - switch targetType := i.targetColumns[idx].dataType.(type) { - case *parser.DataTypeInt: - - err = addColID(columnID) - if err != nil { - return nil, err - } - - vals := make([]int64, 1) - vals[0] = eval.(int64) - - req := &pilosa.ImportValueRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - Values: vals, - } - - err = i.planner.computeAPI.ImportValue(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeBool: - err = addColID(columnID) - if err != nil { - return nil, err - } - - val := eval.(bool) - vals := make([]uint64, 1) - if val { - vals[0] = 1 - } else { - vals[0] = 0 - } - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: vals, - } - - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeDecimal: - err = addColID(columnID) - if err != nil { - return nil, err - } - - vals := make([]float64, 1) - vals[0] = eval.(pql.Decimal).Float64() - - req := &pilosa.ImportValueRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - FloatValues: vals, - } - - err = i.planner.computeAPI.ImportValue(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeID: - err = addColID(columnID) - if err != nil { - return nil, err - } - - coercedVal, err := coerceValue(sourceType, targetType, eval, parser.Pos{Line: 0, Column: 0}) - if err != nil { - return nil, err - } - - vals := make([]uint64, 1) - vals[0] = uint64(coercedVal.(int64)) - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: vals, - } - - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeIDSet: - rowIDs := make([]uint64, 0) - rowSet := eval.([]int64) - for k := range rowSet { - err = addColID(columnID) - if err != nil { - return nil, err - } - rowIDs = append(rowIDs, uint64(rowSet[k])) - } - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: rowIDs, - } - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeIDSetQuantum: - rowIDs := make([]uint64, 0) - timestamps := make([]int64, 0) - - coercedVal, err := coerceValue(sourceType, targetType, eval, parser.Pos{Line: 0, Column: 0}) - if err != nil { - return nil, err - } - - record := coercedVal.([]interface{}) - rowSet := record[1].([]int64) - for k := range rowSet { - err = addColID(columnID) - if err != nil { - return nil, err - } - rowIDs = append(rowIDs, uint64(rowSet[k])) - } - - if record[0] == nil { - timestamps = nil - } else { - timestamp, ok := record[0].(time.Time) - if !ok { - return nil, sql3.NewErrInternalf("unexpected type '%T'", record[0]) - } - for range rowSet { - timestamps = append(timestamps, timestamp.Unix()) - } - } - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: rowIDs, - Timestamps: timestamps, - } - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeString: - err = addColID(columnID) - if err != nil { - return nil, err - } - - rowKeys := make([]string, 1) - rowKeys[0] = eval.(string) - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowKeys: rowKeys, - } - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeStringSet: - rowKeys := make([]string, 0) - rowSet := eval.([]string) - for k := range rowSet { - err = addColID(columnID) - if err != nil { - return nil, err - } - rowKeys = append(rowKeys, rowSet[k]) - } - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowKeys: rowKeys, - } - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeStringSetQuantum: - rowKeys := make([]string, 0) - timestamps := make([]int64, 0) - - coercedVal, err := coerceValue(sourceType, targetType, eval, parser.Pos{Line: 0, Column: 0}) - if err != nil { - return nil, err - } - - record := coercedVal.([]interface{}) - rowSet := record[1].([]string) - for k := range rowSet { - err = addColID(columnID) - if err != nil { - return nil, err - } - rowKeys = append(rowKeys, rowSet[k]) - } - - if record[0] == nil { - timestamps = nil - } else { - timestamp, ok := record[0].(time.Time) - if !ok { - return nil, sql3.NewErrInternalf("unexpected type '%T'", record[0]) - } - for range rowSet { - timestamps = append(timestamps, timestamp.Unix()) - } - } - - req := &pilosa.ImportRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowKeys: rowKeys, - Timestamps: timestamps, - } - err = i.planner.computeAPI.Import(ctx, qcx, req) - if err != nil { - return nil, err - } - - case *parser.DataTypeTimestamp: - err = addColID(columnID) - if err != nil { - return nil, err - } - - coercedVal, err := coerceValue(sourceType, targetType, eval, parser.Pos{Line: 0, Column: 0}) - if err != nil { - return nil, err - } - - vals := make([]time.Time, 1) - vals[0] = coercedVal.(time.Time) - - req := &pilosa.ImportValueRequest{ - Index: i.tableName, - Field: targetColumn.columnName, - Shard: 0, //TODO: handle non-0 shards - ColumnIDs: colIDs, - ColumnKeys: colKeys, - TimestampValues: vals, - } - - err = i.planner.computeAPI.ImportValue(ctx, qcx, req) - if err != nil { - return nil, err - } - - default: - return nil, sql3.NewErrInternalf("unhandled data type '%T'", targetType) - } + if err := batch.Import(); err != nil { + return nil, errors.Wrap(err, "importing batch") } return nil, types.ErrNoMoreRows diff --git a/sql3/planner/planoptimizer.go b/sql3/planner/planoptimizer.go index 065dc3484..24aa425fc 100644 --- a/sql3/planner/planoptimizer.go +++ b/sql3/planner/planoptimizer.go @@ -5,7 +5,6 @@ package planner import ( "context" "fmt" - "log" "reflect" "strings" @@ -676,11 +675,11 @@ func tryToRewriteSubtableJoins(ctx context.Context, a *ExecutionPlanner, n types // for each of the projection operators, for each of the projections // transform each of the referenced values with a the first arg - log.Printf("%T", projections) + a.logger.Debugf("%T", projections) } // there is a join condition, make sure it is one that is permissible (range queries only?) - log.Printf("%T", tvf) + a.logger.Debugf("%T", tvf) return nl, true, nil default: diff --git a/sql3/sql_definitions_test.go b/sql3/sql_definitions_test.go index 0039e0b99..3bde76a4f 100644 --- a/sql3/sql_definitions_test.go +++ b/sql3/sql_definitions_test.go @@ -71,7 +71,8 @@ var tableTests []tableTest = []tableTest{ row(int64(3), int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}, pql.NewDecimal(34567, 2)), row(int64(4), int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}, pql.NewDecimal(45678, 2)), ), - compare: compareExactUnordered, + compare: compareExactUnordered, + sortStringKeys: true, }, { // Select all with top. @@ -92,7 +93,8 @@ var tableTests []tableTest = []tableTest{ row(int64(1), int64(11), []int64{11, 12, 13}, int64(101), "str1", []string{"a1", "b1", "c1"}, pql.NewDecimal(12345, 2)), row(int64(2), int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}, pql.NewDecimal(23456, 2)), ), - compare: compareExactUnordered, + compare: compareExactUnordered, + sortStringKeys: true, }, { // Select all with where on each field. @@ -114,7 +116,8 @@ var tableTests []tableTest = []tableTest{ expRows: rows( row(int64(2), int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}, pql.NewDecimal(23456, 2)), ), - compare: compareExactOrdered, + compare: compareExactOrdered, + sortStringKeys: true, }, }, }, @@ -158,7 +161,8 @@ var tableTests []tableTest = []tableTest{ row("three", int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), row("four", int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}), ), - compare: compareExactUnordered, + compare: compareExactUnordered, + sortStringKeys: true, }, { // Select all with top. @@ -180,8 +184,9 @@ var tableTests []tableTest = []tableTest{ row("three", int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), row("four", int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}), ), - compare: compareIncludedIn, - expRowCount: 2, + compare: compareIncludedIn, + sortStringKeys: true, + expRowCount: 2, }, { // Select all with where on int field. @@ -202,7 +207,8 @@ var tableTests []tableTest = []tableTest{ expRows: rows( row("two", int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), ), - compare: compareExactUnordered, + compare: compareExactUnordered, + sortStringKeys: true, }, }, }, @@ -346,6 +352,9 @@ var tableTests []tableTest = []tableTest{ joinTestsOrders, joinTests, + //bool (batch logic) + boolTests, + //time quantums // Skip for now - timeQuantumInsertTest, } @@ -384,6 +393,15 @@ var insertTest = tableTest{ expRows: rows(), compare: compareExactUnordered, }, + { + // Insert multiple tuples + sqls: sqls( + "insert into testinsert (_id, a, b, s, bl, d, event, ievent) values (4, 40, 400, 'foo', false, 10.12, ['A', 'B', 'C'], [1, 2, 3]), (5, 50, 500, 'var', true, 20.24, ['X', 'Y', 'Z'], [4, 5, 6])", + ), + expHdrs: hdrs(), + expRows: rows(), + compare: compareExactUnordered, + }, { // Insert with nulls sqls: sqls( diff --git a/sql3/sql_defs_bool_test.go b/sql3/sql_defs_bool_test.go new file mode 100644 index 000000000..de58620be --- /dev/null +++ b/sql3/sql_defs_bool_test.go @@ -0,0 +1,88 @@ +package sql3_test + +// BOOL tests +var boolTests = tableTest{ + name: "single-bool-field", + table: tbl( + "singleboolfield", + srcHdrs( + srcHdr("_id", fldTypeID), + srcHdr("a_bool", fldTypeBool), + ), + srcRows(), + ), + sqlTests: []sqlTest{ + { + // Insert, step 1. + name: "insert1", + sqls: sqls( + `insert into singleboolfield (_id, a_bool) values + (1, true), + (2, true), + (3, false), + (4, false), + (5, null), + (6, null)`, + ), + expHdrs: hdrs(), + expRows: rows(), + compare: compareExactOrdered, + }, + { + // Select all, step 1. + name: "select-all1", + sqls: sqls( + "select * from singleboolfield", + ), + expHdrs: hdrs( + hdr("_id", fldTypeID), + hdr("a_bool", fldTypeBool), + ), + expRows: rows( + row(int64(1), true), + row(int64(2), true), + row(int64(3), false), + row(int64(4), false), + row(int64(5), nil), + row(int64(6), nil), + ), + compare: compareExactOrdered, + }, + { + // Insert, step 2. Change bool values to all other combinations. + name: "insert2", + sqls: sqls( + `insert into singleboolfield (_id, a_bool) values + (1, false), + (2, null), + (3, true), + (4, null), + (5, false), + (6, true)`, + ), + expHdrs: hdrs(), + expRows: rows(), + compare: compareExactOrdered, + }, + { + // Select all, step 2. + name: "select-all2", + sqls: sqls( + "select * from singleboolfield", + ), + expHdrs: hdrs( + hdr("_id", fldTypeID), + hdr("a_bool", fldTypeBool), + ), + expRows: rows( + row(int64(1), false), + row(int64(2), nil), + row(int64(3), true), + row(int64(4), nil), + row(int64(5), false), + row(int64(6), true), + ), + compare: compareExactOrdered, + }, + }, +} diff --git a/sql3/sql_defs_cast_test.go b/sql3/sql_defs_cast_test.go index d144c9f41..ea2b1c2d4 100644 --- a/sql3/sql_defs_cast_test.go +++ b/sql3/sql_defs_cast_test.go @@ -785,7 +785,7 @@ var castStringSet = tableTest{ hdr("", fldTypeString), ), expRows: rows( - row(int64(1), string("[101 102]")), + row(int64(1), string(`["101","102"]`)), ), compare: compareExactUnordered, }, @@ -800,7 +800,8 @@ var castStringSet = tableTest{ expRows: rows( row(int64(1), []string{"101", "102"}), ), - compare: compareExactUnordered, + compare: compareExactUnordered, + sortStringKeys: true, }, { sqls: sqls( diff --git a/sql3/sql_defs_in_test.go b/sql3/sql_defs_in_test.go index e07f05f4c..76467b7b5 100644 --- a/sql3/sql_defs_in_test.go +++ b/sql3/sql_defs_in_test.go @@ -1,6 +1,6 @@ package sql3_test -//IN tests +// IN tests var inTests = tableTest{ table: tbl( "in_all_types", @@ -131,7 +131,7 @@ var inTests = tableTest{ }, } -//NOT IN tests +// NOT IN tests var notInTests = tableTest{ table: tbl( "not_in_all_types", diff --git a/sql3/sql_test.go b/sql3/sql_test.go index f10718fc0..429a3861c 100644 --- a/sql3/sql_test.go +++ b/sql3/sql_test.go @@ -2,14 +2,13 @@ package sql3_test import ( - "context" "fmt" "log" + "sort" "strings" "testing" "time" - pilosa "github.com/molecula/featurebase/v3" "github.com/molecula/featurebase/v3/sql3/parser" planner_types "github.com/molecula/featurebase/v3/sql3/planner/types" sql_test "github.com/molecula/featurebase/v3/sql3/test" @@ -22,18 +21,17 @@ func TestSQL_Execute(t *testing.T) { c := test.MustRunCluster(t, 1) defer c.Close() - ctx := context.Background() - api := c.GetNode(0).API svr := c.GetNode(0).Server for i, test := range tableTests { tableTestName := fmt.Sprintf("table-%d", i) if test.name != "" { tableTestName = test.name + } else if test.table.name != "" { + tableTestName = test.table.name } t.Run(tableTestName, func(t *testing.T) { - var err error // Create a table with all field types. if test.table.columns != nil { _, _, err := sql_test.MustQueryRows(t, svr, test.table.createTable()) @@ -41,241 +39,9 @@ func TestSQL_Execute(t *testing.T) { } if len(test.table.rows) > 0 { - // Populate fields with data. - qcx := api.Txf().NewQcx() - - // idIdx is the index position of the _id column. If a source provides the - // _id somewhere other than column 0, then we need to add logic here to find - // its index. - idIdx := 0 - for i, col := range test.table.columns { - if col.name == "_id" { - continue - } - - colIDs := make([]uint64, 0) - colKeys := make([]string, 0) - - addColID := func(v interface{}) { - switch id := v.(type) { - case uint64: - colIDs = append(colIDs, id) - case int64: - colIDs = append(colIDs, uint64(id)) - case string: - colKeys = append(colKeys, id) - default: - t.Fatalf("unexpected type for colid '%T'", v) - } - } - - switch col.typ.(type) { - case *parser.DataTypeInt: - vals := make([]int64, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - vals = append(vals, row[i].(int64)) - } - if len(vals) == 0 { - continue - } - req := &pilosa.ImportValueRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - Values: vals, - } - - err = api.ImportValue(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeBool: - vals := make([]uint64, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - if row[i].(bool) { - vals = append(vals, 1) - } else { - vals = append(vals, 0) - } - } - if len(vals) == 0 { - continue - } - req := &pilosa.ImportRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: vals, - } - - err = api.Import(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeDecimal: - vals := make([]float64, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - vals = append(vals, row[i].(float64)) - } - if len(vals) == 0 { - continue - } - req := &pilosa.ImportValueRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - FloatValues: vals, - } - - err = api.ImportValue(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeIDSet: - rowIDs := make([]uint64, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - rowSet := row[i].([]int64) - for k := range rowSet { - addColID(row[idIdx]) - rowIDs = append(rowIDs, uint64(rowSet[k])) - } - } - if len(rowIDs) == 0 { - continue - } - req := &pilosa.ImportRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: rowIDs, - } - err = api.Import(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeID: - rowIDs := make([]uint64, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - rowIDs = append(rowIDs, uint64(row[i].(int64))) - } - - if len(rowIDs) == 0 { - continue - } - req := &pilosa.ImportRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowIDs: rowIDs, - } - err = api.Import(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeString: - rowKeys := make([]string, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - rowKeys = append(rowKeys, row[i].(string)) - } - - if len(rowKeys) == 0 { - continue - } - req := &pilosa.ImportRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowKeys: rowKeys, - } - err = api.Import(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeStringSet: - rowKeys := make([]string, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - rowSet := row[i].([]string) - for k := range rowSet { - addColID(row[idIdx]) - rowKeys = append(rowKeys, rowSet[k]) - } - } - - if len(rowKeys) == 0 { - continue - } - req := &pilosa.ImportRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - RowKeys: rowKeys, - } - err = api.Import(ctx, qcx, req) - assert.NoError(t, err) - - case *parser.DataTypeTimestamp: - vals := make([]time.Time, 0) - for _, row := range test.table.rows { - if row[i] == nil { - continue - } - addColID(row[idIdx]) - vals = append(vals, row[i].(time.Time)) - } - if len(vals) == 0 { - continue - } - req := &pilosa.ImportValueRequest{ - Index: test.table.name, - Field: col.name, - Shard: 0, - ColumnIDs: colIDs, - ColumnKeys: colKeys, - TimestampValues: vals, - } - - err = api.ImportValue(ctx, qcx, req) - assert.NoError(t, err) - - default: - t.Fatalf("column type not supported: %s", col.typ) - } - } + _, _, err := sql_test.MustQueryRows(t, svr, test.table.insertInto(t)) + assert.NoError(t, err) } for i, sqltest := range test.sqlTests { @@ -315,22 +81,25 @@ func TestSQL_Execute(t *testing.T) { exp[i] = make([]interface{}, len(headers)) for j := range sqltest.expHdrs { targetIdx := m[sqltest.expHdrs[j].ColumnName] - if !assert.GreaterOrEqual(t, len(sqltest.expRows[i]), len(headers)) { - t.Fatalf("expected row set has fewer columns than returned headers") - } + assert.GreaterOrEqual(t, len(sqltest.expRows[i]), len(headers), + "expected row set has fewer columns than returned headers") exp[i][targetIdx] = sqltest.expRows[i][j] } } + if sqltest.sortStringKeys { + sortStringKeys(rows) + } + switch sqltest.compare { case compareExactOrdered: - assert.EqualValues(t, len(sqltest.expRows), len(rows)) + assert.Equal(t, len(sqltest.expRows), len(rows)) assert.EqualValues(t, exp, rows) case compareExactUnordered: - assert.EqualValues(t, len(sqltest.expRows), len(rows)) + assert.Equal(t, len(sqltest.expRows), len(rows)) assert.ElementsMatch(t, exp, rows) case compareIncludedIn: - assert.EqualValues(t, sqltest.expRowCount, len(rows)) + assert.Equal(t, sqltest.expRowCount, len(rows)) for _, row := range rows { assert.Contains(t, exp, row) } @@ -343,6 +112,23 @@ func TestSQL_Execute(t *testing.T) { } } +// sortStringKeys goes through an entire set of rows, and for any []string it +// finds, it orders the elements. This is obviously only useful in tests, and +// only in cases where we expect the elements to match, but we don't care what +// order they're in. It's basically the equivalent of assert.ElementsMatch(), +// but the way we use that on rows doesn't recurse down into the field values +// within each row. +func sortStringKeys(in [][]interface{}) { + for i := range in { + for j := range in[i] { + switch v := in[i][j].(type) { + case []string: + sort.Strings(v) + } + } + } +} + ////////////////////////////////////////////////////////////////////// type fldType parser.ExprDataType @@ -375,13 +161,14 @@ type tableTest struct { } type sqlTest struct { - name string - sqls []string - expHdrs []*planner_types.PlannerColumn - expRows [][]interface{} - expErr string - compare compareMethod - expRowCount int + name string + sqls []string + expHdrs []*planner_types.PlannerColumn + expRows [][]interface{} + expErr string + compare compareMethod + sortStringKeys bool + expRowCount int } // The following "source" types are helpers for creating a test table. @@ -420,6 +207,63 @@ func srcRow(cells ...interface{}) sourceRow { type sourceRow []interface{} +type sourceRows []sourceRow + +// insertTuples returns the list of tuples (as a single string) to use as the +// VALUES value in an INSERT INTO statement. +func (sr sourceRows) insertTuples(t *testing.T) string { + var afterFirstRow bool + var sb strings.Builder + for _, row := range sr { + if afterFirstRow { + sb.WriteString(",") + } + + var afterFirstCell bool + sb.WriteString("(") + + for _, cell := range row { + if afterFirstCell { + sb.WriteString(",") + } + switch v := cell.(type) { + case string: + sb.WriteString("'" + v + "'") + case int64: + sb.WriteString(fmt.Sprintf("%d", v)) + case float64: + sb.WriteString(fmt.Sprintf("%.2f", v)) + case []int64: + strs := make([]string, len(v)) + for i := range v { + strs[i] = fmt.Sprintf("%d", v[i]) + } + sb.WriteString("[" + strings.Join(strs, ",") + "]") + case []string: + if len(v) == 0 { + sb.WriteString("[]") + } else { + sb.WriteString("['" + strings.Join(v, "','") + "']") + } + case bool: + sb.WriteString(fmt.Sprintf("%v", v)) + case nil: + sb.WriteString("null") + case time.Time: + sb.WriteString("'" + v.Format(time.RFC3339) + "'") + + default: + t.Fatalf("unsupported cell type: %T", cell) + } + afterFirstCell = true + } + + sb.WriteString(")") + afterFirstRow = true + } + return sb.String() +} + type source struct { name string columns []sourceColumn @@ -445,6 +289,12 @@ func (s source) createTable() string { return ct } +func (s source) insertInto(t *testing.T) string { + ii := "INSERT INTO " + s.name + " VALUES " + ii += sourceRows(s.rows).insertTuples(t) + return ii +} + // hdrs is just a helper function to make the test definition look cleaner. func hdrs(hdrs ...*planner_types.PlannerColumn) []*planner_types.PlannerColumn { return hdrs