From 00ef2380e5969e21664f53d670e16228a47d0405 Mon Sep 17 00:00:00 2001 From: Travis Turner Date: Sat, 29 Oct 2022 14:24:33 -0500 Subject: [PATCH] Batch insert via SQL (multiple tuples) (#2243) * Formatting adjustments made during code review. While reviewing the BULK INSERT logic (in order to decide how best to approach "ingest via sql" in the cloud), I made a few formatting and comment changes. I'm just adding them here as a separate commit so they don't muddy up my actual work. * Parser modifications to support mulitple tuples in INSERT INTO This commit doesn't include all of the changes required in the planner. Fow now, the planner is simply modified to continue supporting a single tuple (the first tuple in the list). * Update the planner to handle multiple INSERT INTO tuples This is part 1. It's still using the existing logic which builds an ImportRequest for every record (and every field!). The next step will involve using a client.Batch to handle the records. * Introduce client.Importer interface (used by client.Batch) Instead of the Batch having a pointer to a client, this puts an interface there instead (which the client implements). It also allows us to inject a different importer (i.e. other than a featurebase.client) into the Batch. * Decouple batch from client This commit pulls batch-specific code out of the client package and into a new batch package. It introduces the batch.Importer interface, the methods of which replace all the calls that batch was previously making directly to client methods. Finally, it contains two implementations of the batch.Importer interface: one is a wrapper around client, and the other is a wrapper around featurebase.API. * Use docker (instead of MustRunCluster) for internal batch tests Because the `batch` package tests are internal, using test.MustRunCluster() resulted in an import loop (because it eventually imports `server`, and we can't have that). So this commit replaces the use of `test.MustRunCluster()` with docker. The setup is basically the same as that used in the idk docker tests. Here we also remove all client-side references to `UseIngestAPI`, which is an experimental (json) ingest api. It's still suppored on the server, but here we remove the external usage of it. * cherry-pick fix * Use batch.Import() for sql3 INSERT INTO statements * Thread logger into sql3 * fix batch test * Fix some shadowing complaint by linter * Address some test issues related to stringsets * Exclude batch integration tests from CI * Address PR feedback - Added description to batch.README - Consolidated grep commands in .gitlab-ci.yml - Removed some debugging comments - Replaces some inadvertantly removed license headers * Add batch package to gitlab CI * Updated CI for batch package Updated CI include path Update gitlab ci Update CI Update CI Trying new include path for ci Updated gitlab ci include path Made idk race job optional for sonarcloud upload add testdata directory remove testenv from dockercompose file use GIT_STRATEGY clone in batch CI add testdata volume to dockercompose Co-authored-by: Fletcher Haynes --- .gitignore | 2 + .gitlab/.gitlab-ci.yml | 20 +- .gitlab/batch-ci.yml | 28 + Makefile | 6 +- api.go | 5 + batch/Dockerfile-test | 11 + batch/Dockerfile-wait | 8 + batch/Makefile | 54 + batch/README.md | 46 + {client => batch}/batch.go | 319 +-- batch/batch_test.go | 2359 +++++++++++++++++++++++ batch/convert.go | 145 ++ batch/docker-compose.yml | 26 + {client => batch}/egpool/egpool.go | 0 {client => batch}/egpool/egpool_test.go | 2 +- batch/error.go | 8 + batch/importer.go | 211 ++ {client => batch}/metrics.go | 3 +- batch/testdata/README.md | 3 + batch/wait.sh | 26 + client/api.go | 275 +++ client/batch_test.go | 1883 ------------------ client/client.go | 2 - client/importer.go | 219 +++ client/ingest_api_batch.go | 145 -- client/ingest_api_batch_test.go | 305 --- client/orm.go | 12 +- idk/datagen/cmd.go | 2 - idk/ingest.go | 75 +- idk/ingest_test.go | 15 +- server/server.go | 8 +- sql3/parser/ast.go | 27 +- sql3/parser/ast_test.go | 22 +- sql3/parser/parser.go | 16 +- sql3/parser/parser_test.go | 51 +- sql3/parser/walk.go | 18 +- sql3/planner/compileinsert.go | 50 +- sql3/planner/executionplanner.go | 18 +- sql3/planner/expression.go | 53 +- sql3/planner/opbulkinsert.go | 65 +- sql3/planner/opinsert.go | 578 +++--- sql3/planner/planoptimizer.go | 5 +- sql3/sql_definitions_test.go | 32 +- sql3/sql_defs_bool_test.go | 88 + sql3/sql_defs_cast_test.go | 5 +- sql3/sql_defs_in_test.go | 4 +- sql3/sql_test.go | 354 +--- 47 files changed, 4353 insertions(+), 3256 deletions(-) create mode 100644 .gitlab/batch-ci.yml create mode 100644 batch/Dockerfile-test create mode 100644 batch/Dockerfile-wait create mode 100644 batch/Makefile create mode 100644 batch/README.md rename {client => batch}/batch.go (84%) create mode 100644 batch/batch_test.go create mode 100644 batch/convert.go create mode 100644 batch/docker-compose.yml rename {client => batch}/egpool/egpool.go (100%) rename {client => batch}/egpool/egpool_test.go (91%) create mode 100644 batch/error.go create mode 100644 batch/importer.go rename {client => batch}/metrics.go (94%) create mode 100644 batch/testdata/README.md create mode 100755 batch/wait.sh create mode 100644 client/api.go delete mode 100644 client/batch_test.go create mode 100644 client/importer.go delete mode 100644 client/ingest_api_batch.go delete mode 100644 client/ingest_api_batch_test.go create mode 100644 sql3/sql_defs_bool_test.go 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