featurebase/dax/queryer/orchestrator.go
Travis 04aa8b18fc Add DAX - full list of squashed commits below
In this commit, the Directive is mocked; it doesn't actually reach out
to a controller.

Limits key translation to only those partitions (per index) specified in
the Directive. Attempting to create or find a key (or ID) for a
partition which is not handled by this node will result in an error;
translation requests are no longer forwarded to other nodes.

Limits import into only those shards specified, per index, by the
Directive. Attempting to import into a shard which is not handled by
this node will result in an error.

Stub out /directive endpoint

The `applyDirective()` method still needs to be implemented.

Update mds references to use the new /mds/types structure

In mds, we moved the shared types to mds/types. FeatureBase needs to
reference those instead.

This also bumps the mds version in go.mod.

Implement the Add/Remove Index part of Holder.ApplyDirective()

This adds functionality to `Holder.ApplyDirective()` which adds or
removes indexes (tables) based on those provided in the Directive. Still
to be implemented here: shards and partitions.

WIP: remove client from Batch

Move Batch into its own package: batch

Also, in order to avoid import loops, this introduces packages:
/batch/types
/client/types

Reorganize the Importer-related code

Moved the Importer interface to package: batch
Move the "pilosa client" implementation of the Importer interface to
package: client

Modify batch.NewBatch to take an Importer (not client)

This commit modifies the batch.NewBatch() function to use a functinal
option on Batch to inject an Importer into the Batch. Prior to this,
NewBatch() took a pointer to a client, which was a little too
restrictive. Now, MDS can implement an Importer which uses information
from MDS to determine to which node(s) the import calls should be directed.

Add client.SetAuthToken() method to satisfy SchemaManager interface

Update ApplyDirective logic to include fields.

This needs more work, but it was enough to get a basic test passing.

Move Transaction type into /types package.

Add interface check on batch.Importer no-op implementation

Updated ApplyDirective to create all currently support Field types

There are still the following TODOs:
- [ ] impolement field options (ex: decimal scale, int min/max, etc).
- [ ] `time` fields

Added support for Decimal.Scale in ApplyDirective

Update mds dependency

Add /health endpoint

Update to use dax (dax/mds) instead of mds.

After moving the mds repository into the dax repository as a
sub-package, this commit changes everything in FeatureBase to use the
dax repo instead of the now abandoned mds repo.

Introduce and implment the WriteLogger interfaces.

This adds both a `WriteLogReader` and `WriteLogWriter` interface. They
are both implemented by the implementation: `fileWriteLogger`. The
`fileWriteLogger` uses the dax/writelogger API to append log messages to
files on disk.

Add WriteLogWriter.ImportRoaring method to interface

This commit adds the `ImportRoaring` method to the `WriteLogWriter`
interface. Still to implement are the `Import` and `ImportValue`
methods.

Reorganize the ApplyDirective code

The primary goal was to cache the incoming Directive on the Holder prior
to applying all of the changes in the directive (i.e. loading data from
the WriteLogger) because applying those changes often validated against
the accepted state of the node.

Implement all of the WriteLogger read/write methods

Implement the HTTP WriteLogger implementation

WIP: Introduce shard.Version. Implement snapshotter.

Add HTTP Snapshotter implementation

This also recofigures server to use the HTTPSnapshotter instead of the
FileSnapshotter.

Implement snapshotter: TableKeys

Implement snapshotter: FieldKeys

Dependency dance

last of the dependency dance

Add support for prototype

This adds the Makefile targets to build the docker container and push it
to ECR.

SQL3 changes which break with dax changes

Missed TODO: implement FieldVersion version to WriteLogger

Address bug causing missing TranslateStores to error

Originally, we tried to limit the TranslateStores which get allocated to
only those for which the node is responsible. This works when adding a
new table. But if a table already exists, there's no logic to start
missing TranslateStores.

This reverts back to the old FeatureBase logic which brutishly allocates
a TranslateStore for every partition, even if one is not needed.

We need to address this by allowing the ApplyDirective logic to
initialize TranslateStores when they don't yet exist.

Move the ImportRoaringShardRequest type to the types package

Since the ImportRoaringShardRequest object is part of the Importer
interface, we need to move it to a non-root (i.e. pilosa) package. All
the other interface types are either concrete types or part of a
sub-package (such as roaring). We do this to prevent an implementer of
the interface from having to import the entire pilosa package and risk
circular imports.

buncha changes to support latest dax stuff

Move dax related types to /dax sub-package

This commit moves all the common "dax" types into the /dax sub-package.
The idea is to ensure that featurebase does not import dax at all.
It's ok if dax imports featurebase.
In the future, we might need to split the dax sub-package (common data
types used by muliple molecula data-plan services) into it's own repo.

Add type: dax.Schema

This isn't currently being used; I started to use is as a replacement
for pilosa_client.Schema, but then deferred that. But we'll need to do
it eventually, so it doesn't hurt to have this type in place.

Export RowIDs.Merge() method for use in orchestrator.

Add CreateSQL method to dax.Table type

The CreateSQL() method will return the "CREATE TABLE" statement required
to create the dax.Table.

Comment out confusing writelogger log message.

We need to revisit this, but for now, this log message is confusing.

Also, rename daxSharder to versionStore.

Remove hard-coded AWS account

Implement more FieldOptions such as Epoch

Some of the FieldOption logic was stubbed out in the dax package. This
commit fills that out more; specifically, it adds the
dax.Field.Options.Epoch parameter.

export stuff needed for TopK in orchestrator

export ValCount stuff to implement Percentile in orchestrator

export more stuff to support less code in orchestrator, shared objs

Port dax repo over to featurebase/dax (run all as sub-services)

This commit does ALOT. Sorry.
It introduces a `featurebase dax` sub-command which can be configured to
run the various dax services as sub-services within the same process, or
individually as the lone service in process.

It also changes all the URL paths to be prefixed with the service name.
So for example, instead of calling localhost:8080/status, you would now
call localhost:8080/featurebase/status.

Also, note that all services provide a /health endpoint to confirm they
are running in process.

Clean up integration tests. Remove PILOSA_ config prefix.

Remove duplicate clients (mistake from porting dax to featurebase)

Rename sub-service "featurebase" to "computer"

In the places where we have hard-coded the sub-service name into a URI
path, I've tried to tag the line with a comment containing:
`// #SERVICEPATHPREFIX`

Update copilot manifest files to reference "computer"

Port dax/README.md from dax repository

Separate (toml) Queryer Config from Injections

We needed to separate the toml config from the configuration required to
inject sub-services into the Queryer. I'm not sure this is the best
solution, but it's *a* solution. So here we are.

Clean up (i.e. remove) the queryer "implementations" package

Remove old test file

Run WriteLogger and Snapshotter as local sub-services.

Prior to this commit, the writelogger and snapshotter services only
worked when run as separate services. This allows them to be run in the
same process as all the other dax services.

There is still some naming issues that we should address, but it's
functional for now.

Clean up (i.e. organize) the intra-service interfaces.

Implement alpha Director for local messages from MDS to Computer

Prior to this commit, messages from MDS to the computer service were
still going over http. This commit introduces an interface
implementation which registers the local computer command, and use that
command's API to directly reference methods used by the Director.

Clean up a few more interface names

Add Queryer OpenAPI document.

Update copilot manifests to reflect latest changes

Add OpenAPI documents for WriteLogger and Snapshotter

Add OpenAPI document for MDS service

Add OpenAPI document for Computer service

Consolidate errors to use fb/errors package.

This commit is a first pass at trying to ensure that all of the DAX code
uses:
"github.com/molecula/featurebase/v3/errors"

This package is a wrapper for "github.com/pkg/errors", so going forward
we want to avoid importing that package.

The only method which isn't backward-compatible is `New()`; the
New() method in the featurebase/errors package takes an errors.Code. If
this becomes a problem, we could change this by reverting New() and then
introducing something like NewCoded(). But for now I think it might
actually discourage someone from just creating a New() error without
thinking about how it should be coded.

Introduce VersionStore interface

Move the existing VersionStore code to the `inmem` package as the
in-memory implementation of the new dax.VersionStore interface.

Introduce NodeService interface

With this, the Controller can maintain a registry of nodes by using this
NodeService interface as opposed to an in-memory map of nodes on the
Controller struct.

This also adds an inmem implementation of the NodeService interface.

Introduce controller.Balancer interface

This moves the existing balancer package to controller/naive package.
The idea is to allow us to add a different Balancer implementation in
the future.

Introduce DirectiveVersion interface

This commit also includes *A LOT* of refactoring to use dax.Worker and
dax.Job types everywhere instead of strings.

Introduce Schemar interface

The previous `Schemar` struct was moved to the `schemar/inmem` package,
and `Schemar` is now an interface implemented by that inmem package.

Remove unused type `nUnit`

Add boltdb implementation of VersionStore interface.

This removed the previous sqlite implementation; we decided not to use
sqlite for now (as a basic, local disk implementation) because it
requires CGO.

--------------------------------------------
No longer applicable:

Add sqlite implementation of VersionStore interface.

This commit implements the VersionStore interface using sqlite. Sqlite
requires CGO, so this may not be something we want to include, but it's
implemented here to get a feel for how an external implementation might
be used; the next step will be to determine how the user configured
FeatureBase to run using sqlite as a backing store for services like
MDS.

Add boltdb implementation of NodeService and DirectiveVersion interfaces.

Add boltdb implementation of naive Balancer interfaces.

This includes the two interfaces defined in `naive/balancer.go`:
- WorkerJobService
- FreeJobService

Add boltdb implementation of Schemar interface.

clean up a linter issue

Thread context.Context through all the interfaces.

Some of the interface implementations are going to use context, so we
need to make that part of the interface. The boltdb implementations, for
example, take a context. This is probably so we can do things like
cancel or timeout operations.

Update interfaces to return error; remove `panic(err)` everywhere.

Down-rev grpc version to 1.38.0

Later versions (after 1.42.0?) cause MustRunCluster.Close() in tests to
deadlock.

This commit also adds an `isComputeNode` feature flag around some of the
write log and shard/partition check functionality so that it doesn't run
under normal conditions (this is excercised by running the sql3 tests
for example).

Add MDS_Persistence test to cover meta data persistence

This adds a basic test which configures the MDS container to use boltdb
as its persistence storage, saved on a docker volume. Then, the mds
container is stopped/replaced, and we confirm that the data stored on
the volume is availble to the new MDS container.

Fix a few things after rebase with sql-experiment branch

The lastest version of sql-experiment contains a fairly significan
refactor of the way query iteration works. This commit adjusts for those
changes.

pull dax IDK changes in to FB IDK (#2177)

* pull dax IDK changes in to FB IDK

* Move docker-related IDK build stuff to featurebase root

Building the docker image required the root level go.mod and vendor
directory. This change moves the make targets to the root level
Makefile, and the Dockerfiles now copy the root level vendor directory
(and everything else in the root for that matter).

* Fix batch- and client-related tests

* InitializePoller on MDS restart/replacement

Prior to this change, if MDS was restarted, its internal poller (which
maintains an in-memory list of nodes to poll) is empty. This is bad,
because it doesn't know about nodes that it should be polling.

This change fixes that. Upon MDS startup, it intializes the poller with
the list of nodes that MDS keeps in persistent storage (currently:
boltdb).

* Add EFS volume to MDS Copilot manifest

This allows us to use MDS's persistent storage (via boltdb) in the
Copilot demo by saving metadata in a boltdb file on EFS.

* Thread logger.Logger through all dax components

* Revert some of the breaking changes from DAX development.

When we first started prototyping DAX, we made changes to the
featurebase core code which would have broken the existing featurebase
functionality. This commit reverts some of those changes. Anywhere that
we need to modify core featurebase functionilty, we put it behind some
kind of feature flag. This flag is typically determined by whether the
running node is a "compute" node (i.e. DAX.COMPUTER.RUN = true).

Co-authored-by: Travis Turner <travis@molecula.com>

add packaging for DAX

need cgo for datagen build

bind to 0.0.0.0, pass GOOS and GOARCH explicitly

not sure if the explicit GOOS/GOARCH is actually necessary...

Get INSERT INTO (aka ingest) working through SQL3

This commit does a few things which I'll try do describe here.

- Introduces a Qctx interface. The existing Qcx is an implementation of
  this interface, and can be used exactly how it has been. But this
  allows us to abstract away the notion of Qcx in the Queryer (which is
  handling SQL3) until we're ready to address that. As an example, the
  Qcx has a notion of a featurebase Holder, but that doesn't make sense
  when we're at the Queryer layer. For now, the Qctx used in the Queryer
  is a no-op.

- Adds a ComputeAPI interface implementation for the Queryer. This is
  effectively the Import() and ImportValues() methods used for ingest.
  The logic here handles the incoming ImportRequest by first doing any
  necessary column and row translation for the entire request, then it
  splits the records by shard, and generates a new ImportRequest per
  shard with only the shard-appropriate records.

- Changes the mds.Importer to take an MDS interface implementation
  (which can be an mds client) instead of an mdsAddress. This allows us
  to use a localy MDS implementation rather than assuming we need a
  client to make calls over a network.

Add queryer.Importer interface to handle ingest via SQL (#2203)

* Add queryer.Importer interface to handle ingest via SQL

This is meant to support ingest through SQL when the queryer and the
compute services are running in the same process, or when they are on
seperate processes and need to talk via http client.

* remove datagen from RPM

was originally added as a convenience to generate test data, but is
unused and annoying because datagen doesn't easily cross-compile due
to cgo

* add marshalUnmarshal to controller to avoid passing pointers

passing pointers across API boundaries can cause unpredictable things
in local vs remote configurations.

Co-authored-by: Matthew Jaffee <jaffee@pilosa.com>

"fix" a few issues with wrong default partition numbers

these still need to be properly fixed and actually get the correct
data from MDS

go mod tidy

Introduce TableQualifier (OrganizationID/DatabaseID) (#2220)

* add check in ApplyDirective that version is increasing

fix TestAPIDirective to make version always increasing

* fix docker image build and break out dax test in CI

We have to run the DAX integration tests separately as they call out
to Docker, and so it isn't easy to run them in a Docker container as
the other tests do. So we run them directly on the CI runner which has
Docker and Go installed.

We also explicitly exclude these tests from running during the other
tests.

Also my editor was automatically reformatting some comments badly
which is why I added the "data" thing in those two places

* add timeout to poller

* give Poller a default Logger

apparently we can NPE sometimes, seen in CI: https://gitlab.com/molecula/featurebase/-/jobs/3028286364

* bunch of testing fixes, mostly IDK/DAX related

make MDS error if sendDirectives errors, don't just
log. sendDirectives can error if computer nodes disagree about the
validity of a schema (for example), in which case it might need to get
deleted and user notified somehow. very messy, needs more thought.

re-introduce old env prefix to maintain compatibility with master
branch

make self-contained dax container for IDK testing

build IDK images from source (now that all the source is available
since it's in the same repo)

catch errors in DoExtractQuery in idktest.go

fix IDK bug where prefix path was hardcoded in all cases rather than
only when useMDS was true

fix TestBatchTargetMDS... needed to add field options and catch error
when creating table. also needed an _id field

* fix env prefix in tests

* WIP getting tests to pass, wanna see CI

* don't error if we get a zero version directive and we don't have a

directive yet

* cleanup debugging junk

* "fix" future.rename thing, run IDK tests

* Introduce TableQualifier (OrganizationID/DatabaseID)

This commit introduces a lot of new types (in dax/table.go) related to
TableQualifer (which is made up of OrganizationID and DatabaseID), as
well as things like TableID and TableKey.

For the most part, we try to thread a QualifiedTableID through the
entirety of DAX. There are some places (for example in the Balancers,
which are just aware of string keys) which use a string TableKey
(tbl__org__db__tableid).

* Remove some debugging comments

* Add Org/DB support to CLI.

This commit adds support for special commands:

SET
SET ORG acme
SET DB db1
USE db1

* remove ".pulled" from IDK Makefile

I don't think we need it any more as most things can be built
locally. I think it was only there to refresh the FeatureBase images
that were tagged as master, but we don't need to do that any more.

* Change DAX json tags to kebab-case (i.e. hyphenated)

This commit also renames some struct arguments to more accurately
reflect their type: for example, renaming `Table` to `TableKey` when the
type is TableKey.

* Return DAX TableName in SHOW TABLES (instead of Index.Name)

There are cases where SchemaAPI is used to return DAX friendly table
names (as opposed to featurebase index names, which are DAX TableKey).

This is an attempt to do that. With that said, it's not ideal because
anything could call those API methods and expect the other type.

* Fix a bug which wasn't completely dropping a table.

When using boltdb as a backend, DROP TABLE wasn't removing the
reverse-lookup key for the table in boltdb.

* Remove idk/testenv/certs which got accidentally committed.

also update .gitignore to include those.

* Fix IDK ingest tests to be TableQualifier aware.

* Add example Table types to dax/table.com godoc.

* ignore idk.Main fields for flags, upgrade commandeer

* go mod tidy

* Fix DAX integration tests: ingester using wrong ENV VARs

We change from ORGANIZATION_ID to ORG_ID
and from DATABASE_ID to DB_ID

* Clarify things around idk (docker) tests

* Stop running TestKafkaSourceIntegration with t.Parallel()

This test can't be run in parallel as it's currently written. Doing so
allows for interleaving of messages to the same kafka topic between
tests.

I didn't attempt to modify the test so it could be run in parallel. That
could be done, but left for someone more ambitious.

Co-authored-by: Matthew Jaffee <jaffee@pilosa.com>

Require Directive.Version be a non-zero value. (#2227)

Because the directive cached on the holder is not a pointer, its default
version is 0. In order to avoid having to compare against that, we just
require that Directive.Version start at 1.

General, non-invasive code cleanup and comment adjustment.

Move ImportRoaringShardRequest out of the types package

Early on in the DAX development, I moved ImportRoaringShardRequest into
a types package. There must have been some import loop going on, but
since that is not longer the case, it's safe to move this back into the
core featurebase (er... pilosa) package.

Move Transaction struct back into the pilosa package (from types)

Revert some name changes (cli -> client)

Add DAX Handler CloseTimeout

This was implemented in htt_handler.go, but it had been commented out in
the DAX handler. This just uncomments that and finishes the
implementation.

Remove Qcx from queryer.Importer interface

This sets us up to revert the Qctx interface that was initially
introduced to allow us to abstract away the need for a Qcx when calling
the ComputeAPI from a remote service (i.e. the queryer).

Add some go-doc comments and remove unused code.

Move SchemaManager setup from datagen to idk.Main (#2233)

The set for idk.SchemaManager (for dax implementations) was previously
in datagen. This may have been because of some import loop problem
during development, but that's no longer an issue.

The setup for this should be in idk.Main so anything using that can
leverage the MDS-specific SchemaManager setup.

Fix issues around nil TxFactory

First, don't return a nil. Rather return a new *TxFactory (with no
holder).

Second, don't call `f.holder` in the testhook outside of checking if
`f.holder` is nil.

Wrap all bare errors

Make service prefixes constants

Instead of having `"computer"` throughout the code, use instead a
constant: `dax.ServicePrefixComputer`.

MDS skip errors when sending empty directives

also add in the docker-login and ecr-push changes for serverless DAX

Fix the logic in Directive.IsEmpty() (#2236)

Update the cached value for Index.translatePartitions

In the case where a node already knows about an index, but its
assignment of partitions for that index changes (for example, when
another node goes down and the node in question is now responsible for
more partitions than it previously was), then we need to update the
cached value of Index.translatePartitions because that's used in
translation checks.

minor fixes for IDK-related bugs

WIP: tokenize CLI to access cloud

FB CLI cloud support with automatic token refresh

Also adds support for a GET command which allows making HTTP GET
queries to cloud CP which can be handy for debugging stuff. E.g. GET /v2/databases

buncha little fixes working on writelogger stuff

fix writelogger/snapshotter setup bugs

implement writelogging for importRoaringShard

add debug endpoint to MDS

use shard transactional endpoint in MDS datagen

add debugging to API related to writelogger

revert handleroption change

clean up big PR

remove "GET" command from CLI for making arbitrary HTTP request to
cloud control plane (was a messy hack and not that useful)

remove json tags from FB objects where we had to duplicate the object
elsewhere due to import loops and weren't actually json encoding it

unexport handlerOption which was exported to try to avoid doing
certain things if we're in DAX mode, but I didn't end up merging that code.

remove (hopefully) unecessary extra call to api.indexField

fix some formatting, unexport some vars, godoc, etc

oops, fix build failure

Update FeatureBase CLI to support a standard deployment

The standard deployment uses a different endpoint and request payload.
This commit tries to detect is the standard deployment is being used,
and if so, it uses a standard-specific FBQueryer.

It also modifies the auto-detection logic to try standard featurebase
and dax ports in the case where a port was not provided.

MDS API refactor (#2259)

* MDS API refactor

table IDs are exposed but only created server side

also cleaned up dax Makefile

* clean up review feedback

Co-authored-by: Travis Turner <travis@pilosa.com>

* remove TablesByName

* rip out inmem implementations and use boltdb everywhere

* remove inmem balancer, create bolt tempfile by default on startup

* WIP on snapshot table impl and test

* Minor comment and code layout adjustments.

This commit also adds the `Equals` method to `QualifiedTableID` for
equality comparisons. It's no longer safe to compare struct (two structs
might still be equal even if one of the structs doesn't have a `Name`
value.

* Use a unique docker network for each dax test

Ocassionally we would see some test failures due to a network already
existing. This shouldn't happen, but to avoid that, this commit
generates a unique name for each sub test (which gets deleted at the end
of every test).

* Fix one instance of NewQualifiedTableID losing Name

We should probably check the other instances and see if Name is getting
lost.

* simplify unique network stuff and fix api directive tests

* Remove TableIDRequest and TableIDResponse types for /table-id (#2267)

For the mds/table-id http requests, just use dax.QualifiedTableID as
both the request and response types.

* remove lattice from dax, no error on node re-reg, dax docker-compose

* various updates

* WIP: mds-refactor branch review

* no-op on SnapshotTableKeys if table is not keyed

* Makefile helpers

* add doWeCare so controller doesn't fail unnecessarily

* clean up table creation (#2272)

* Strip underscores from TableID stub name

* fix boltdb versionstore tests: generate unique, sorted tables

* fix controller test related to reregistering a node

* JobSet -> generic Set

Co-authored-by: Travis Turner <travis@pilosa.com>
Co-authored-by: Travis Turner <travis@molecula.com>

Cleanup after rebase on master

The latest rebase on master entailed all the client/batch changes as
well as some of the qcx refactoring. It made for a hairy rebase. This
commit fixes some of the tests that were failing after that rebase.

Fix batch/client import loop missed during rebase (#2280)

It's not surprising that `batch` can't import `client`. It was doing
that here (importing an error type from the `client` package). What is
surprising is that it's okay for `batch_test.go` to import `client` even
though `batch_test.go` is an internal test and therefore part of the
`batch` package.

different boltDB's for schemar/controller, explicit balancers

nice helpers for dax docker-compose, make build really fast

build FB binary outside of docker, then create Docker image with its
working dir in an empty subdirectory so it doesn't send a GB of
context to the daemon.

error on unassigned jobs and use client with timeout

fix CR feedback

deregister batch of nodes

also make removal faster via director dial timeout

implement WorkersForJobPrefix so orchestrator doesn't make up shards

also fix some godocs and remove unused method

Run sub-tasks of a Directive concurrently in a worker pool. (#2275)

* Run sub-tasks of a Directive concurrently in a worker pool.

This allows the compute node to concurrently load shapshot and writelog
data concurrently, instead of one keyset/partition/shard at a time.

It introduces a config parameter called `DirectiveWorkerPoolSize`.

* code review cleanup

* Use unique container names in DAX integration tests

We were seeing "container already exists" errors in CI, so just to be
safe, this commit constructs a unique container name for every container
in the DAX integration test run.

Stub in SystemAPI to Queryer (note: will not work if used)

This just makes is so that dax can compile. Actually implementing
system-table functionality for dax will take some planning.

Tlt/dax merge prep (#2282)

* Remove copilot directory

* Remove Dockerfile-datagen-long

* Remove orphaned RegisterNodeRequest

This type is not defined in the dax/mds/http package.

* implement TIMEQUANTUM and TTL in Table.Field type

* Remove the "service" misdirection in queryer/writelogger/snapshotter.

We had originally used an additional layer, er.. package, for a "service".
The main distinction was that the Config differed in that it was
internal, unlike the Config that we need to provide for the top-level
server config (i.e. toml). Having that additional layer just to support
a different Config seemed premature at best. So I'm removing it.

* Remove dax docker containers no longer used in tests

Since we run everything as "featurebase", we don't have multiple
container types anymore.

* Some minor comment updates

* Remove nfpm stuff related to dax

* Fix linter issues

Fix "duplicate" issues raised by sonarcloud.

run docker components of dax integration tests with coverage

trying to get dax integration coverage

add coverate volume mounts throughout dax integration tests

add a lock, tweak dax Makefile, remote flag on query handler

remove some unused code

convert batch tests to use clustertests to get coverage

maybe fix clustertests

more authclustertests fixes, test is failing locally

but also seems to have been silently failing in CI prior to these
changes... let's see if it's still silent

fix some lint to kick CI

just re-running the job wasn't working... strange behavior

remove RetryLogic test and pipe which don't work

RetryLogic test removed due to etcd changes. Seebs thinks we shouldn't
test this here.

Pipe was being ignored since we're no longer using "bash -c" to
execute the command. If we need to generate that output file we'll
either have to reintroduce bash -c and set -o pipefail so that it
actually fails properly, or figure out some other solution.

shooting into the dark...

first cut at bulk node registration

remove unused stuff from batch tests, set coverpkg to ../...

batch registration timeout and fix tests

disable most tests and don't run fb background batch test

debuggin!!!!!!!!!

and then he tried this....

Implement importer (for INSERT INTO) in the Queryer

Prior to this, we we passing a nil value in for the importer to the
planner.NewExecutionPlanner in the Queryer. This meant that INSERT INTO
statements didn't work. Now they should.

It uses the importer that we build for IDK in /idk/mds/importer.go, and
wrapps that with a type that can determine if the provided string
"index" is of the form indexName or TableKey.

turn off debug mode, fix log saving

Run sql3 test definitions in a dax integration test

There are currently 22 tests which are not passing. They are skipped in
the "skips" slice.

WIP, not working, pql queries to tests

Add TableQualifier to PQL query logic in the Queryer

Add more PQL tests to the keyed table

Allow instant node registration if registration-batch-timeout=0

When running dax services in process, we don't want to wait 3s for the
compute node to register; we know it's there because it's in the same
process.

Fixes related to IncludesColumn PQL test.

Tests for ConstRow and FieldValue

cleanup

add UnionRows and Options, better error reporting on bad queries

delete unused schemar client.go, clean up unused in batch test CI

move test timeouts into more reasonable territory

apparently this had already been done, but got merge-stommped at some point

move dax bolt test helpers into dax package

Add computer CheckIn routine (#2296)

* Add computer CheckIn routine

This adds a background routine which sends a "check-in" request to MDS
every <interval>. This is to address the case where the poller has
removed a computer node from the node list (due to a network fault, for
example), but the node is still healthy and becomes available again. In
that case, the node needs to "check-in" to tell MDS it is still there.
MDS will likely send the node a new directive with Method=reset telling
the node to delete all of its data an apply the latest directive.

* Don't send directives to Deregistered (i.e. removed) nodes

We have an issue where we're locking on sendDirective in the
controller, and when the node is unavailable, the send hangs and never
releases the lock. This is a temporary fix for that until we address the
real problem.

Fix .gitlab-ci.yml after rebase

fix some indentation shenanigans

(cherry picked from commit 20a8b5713a)
2022-12-12 09:01:20 -08:00

3494 lines
104 KiB
Go

// Copyright 2021 Molecula Corp. All rights reserved.
package queryer
import (
"context"
"fmt"
"math"
"sort"
"strings"
"time"
featurebase "github.com/molecula/featurebase/v3"
"github.com/molecula/featurebase/v3/dax"
"github.com/molecula/featurebase/v3/dax/mds/controller"
"github.com/molecula/featurebase/v3/dax/mds/schemar"
"github.com/molecula/featurebase/v3/errors"
"github.com/molecula/featurebase/v3/logger"
"github.com/molecula/featurebase/v3/net"
"github.com/molecula/featurebase/v3/pql"
"github.com/molecula/featurebase/v3/stats"
"github.com/molecula/featurebase/v3/tracing"
"golang.org/x/sync/errgroup"
)
// Field types.
const (
FieldTypeSet = "set"
FieldTypeInt = "int"
FieldTypeTime = "time"
FieldTypeMutex = "mutex"
FieldTypeBool = "bool"
FieldTypeDecimal = "decimal"
FieldTypeTimestamp = "timestamp"
// Row ids used for boolean fields.
falseRowID = uint64(0)
trueRowID = uint64(1)
)
var ErrFieldNotFound error = dax.NewErrFieldDoesNotExist("")
const (
errConnectionRefused = "connect: connection refused"
)
type Topologer interface {
ComputeNodes(ctx context.Context, index string, shards []uint64) ([]controller.ComputeNode, error)
}
type MDSTopology struct {
mds MDS
}
func (m *MDSTopology) ComputeNodes(ctx context.Context, index string, shards []uint64) ([]controller.ComputeNode, error) {
var daxShards = make(dax.ShardNums, len(shards))
for i, s := range shards {
daxShards[i] = dax.ShardNum(s)
}
qtid := dax.TableKey(index).QualifiedTableID()
return m.mds.ComputeNodes(ctx, qtid, daxShards...)
}
// TODO(jaffee) we need version info in here ASAP. whenever schema or topo
// changes, version gets bumped and nodes know to reject queries
// and update their info from the MDS instead of querying it every
// time.
type Translator interface {
CreateIndexKeys(ctx context.Context, index string, keys []string) (map[string]uint64, error)
CreateFieldKeys(ctx context.Context, index string, field string, keys []string) (map[string]uint64, error)
FindIndexKeys(ctx context.Context, index string, keys []string) (map[string]uint64, error)
FindFieldKeys(ctx context.Context, index, field string, keys []string) (map[string]uint64, error)
// TODO(jaffee) the naming here is a cluster. TranslateIndexIDs takes a list, but TranslateFieldIDs takes a set, both have alternate methods that take the other thing. :facepalm:
TranslateIndexIDs(ctx context.Context, index string, ids []uint64) ([]string, error)
TranslateIndexIDSet(ctx context.Context, index string, ids map[uint64]struct{}) (map[uint64]string, error)
TranslateFieldIDs(ctx context.Context, index, field string, ids map[uint64]struct{}) (map[uint64]string, error)
TranslateFieldListIDs(ctx context.Context, index, field string, ids []uint64) ([]string, error)
}
// executor recursively executes calls in a PQL query across all shards.
type orchestrator struct {
schema featurebase.SchemaInfoAPI
topology Topologer
trans Translator
// Client used for remote requests.
client *featurebase.InternalClient
stats stats.StatsClient
logger logger.Logger
}
func emptyResult(c *pql.Call) interface{} {
switch c.Name {
case "Clear", "ClearRow":
return false
case "Row":
return &featurebase.Row{Keys: []string{}}
case "Rows":
return featurebase.RowIdentifiers{Keys: []string{}}
case "IncludesColumn":
return false
}
return nil
}
// Execute executes a PQL query.
func (o *orchestrator) Execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) (featurebase.QueryResponse, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "orchestrator.Execute")
span.LogKV("pql", q.String())
defer span.Finish()
resp := featurebase.QueryResponse{}
// Check for query cancellation.
if err := validateQueryContext(ctx); err != nil {
return resp, err
}
// Verify that an index is set.
if index == "" {
return resp, featurebase.ErrIndexRequired
}
idx, err := o.schema.IndexInfo(ctx, index)
if err != nil {
return resp, errors.Wrap(err, "getting index")
}
// Default options.
if opt == nil {
opt = &featurebase.ExecOptions{}
}
results, err := o.execute(ctx, index, q, shards, opt)
if err != nil {
return resp, err
} else if err := validateQueryContext(ctx); err != nil {
return resp, err
}
resp.Results = results
if err := o.translateResults(ctx, index, idx, q.Calls, results, opt.MaxMemory); err != nil {
if errors.Cause(err) == featurebase.ErrTranslatingKeyNotFound {
// No error - return empty result
resp.Results = make([]interface{}, len(q.Calls))
for i, c := range q.Calls {
resp.Results[i] = emptyResult(c)
}
return resp, nil
}
return resp, err
} else if err := validateQueryContext(ctx); err != nil {
return resp, err
}
return resp, nil
}
func (o *orchestrator) execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) ([]interface{}, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.execute")
defer span.Finish()
// Apply translations if necessary.
var colTranslations map[string]map[string]uint64 // colID := colTranslations[index][key]
var rowTranslations map[string]map[string]map[string]uint64 // rowID := rowTranslations[index][field][key]
if !opt.Remote {
cols, rows, err := o.preTranslate(ctx, index, q.Calls...)
if err != nil {
return nil, err
}
colTranslations, rowTranslations = cols, rows
}
// Execute each call serially.
results := make([]interface{}, 0, len(q.Calls))
for i, call := range q.Calls {
if err := validateQueryContext(ctx); err != nil {
return nil, err
}
// Apply call translation.
if !opt.Remote && !opt.PreTranslated {
translated, err := o.translateCall(ctx, call, index, colTranslations, rowTranslations)
if err != nil {
return nil, errors.Wrap(err, "translating call")
}
if translated == nil {
results = append(results, emptyResult(call))
continue
}
call = translated
}
// If you actually make a top-level Distinct call, you
// want a featurebase.SignedRow back. Otherwise, it's something else
// that will be using it as a row, and we only care
// about the positive values, because only positive values
// are valid column IDs. So we don't actually eat top-level
// pre calls.
if call.Name == "Count" {
// Handle count specially, skipping the level directly underneath it.
for _, child := range call.Children {
err := o.handlePreCallChildren(ctx, index, child, shards, opt)
if err != nil {
return nil, err
}
}
} else {
err := o.handlePreCallChildren(ctx, index, call, shards, opt)
if err != nil {
return nil, err
}
}
var v interface{}
var err error
// Top-level calls don't need to precompute cross-index things,
// because we can just pick whatever index we want, but we
// still need to handle them. Since everything else was
// already precomputed by handlePreCallChildren, though,
// we don't need this logic in executeCall.
newIndex := call.CallIndex()
if newIndex != "" && newIndex != index {
v, err = o.executeCall(ctx, newIndex, call, nil, opt)
} else {
v, err = o.executeCall(ctx, index, call, shards, opt)
}
if err != nil {
return nil, err
}
if vc, ok := v.(featurebase.ValCount); ok {
vc.Cleanup()
v = vc
}
results = append(results, v)
// Some Calls can have significant data associated with them
// that gets generated during processing, such as Precomputed
// values. Dumping the precomputed data, if any, lets the GC
// free the memory before we get there.
o.dumpPrecomputedCalls(ctx, q.Calls[i])
}
return results, nil
}
// handlePreCalls traverses the call tree looking for calls that need
// precomputed values (e.g. Distinct, UnionRows, ConstRow...).
func (o *orchestrator) handlePreCalls(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) error {
if c.Name == "Precomputed" {
idx := c.Args["valueidx"].(int64)
if idx >= 0 && idx < int64(len(opt.EmbeddedData)) {
row := opt.EmbeddedData[idx]
c.Precomputed = make(map[uint64]interface{}, len(row.Segments))
for _, segment := range row.Segments {
c.Precomputed[segment.Shard()] = &featurebase.Row{Segments: []featurebase.RowSegment{segment}}
}
} else {
return fmt.Errorf("no precomputed data! index %d, len %d", idx, len(opt.EmbeddedData))
}
return nil
}
newIndex := c.CallIndex()
// A cross-index query is handled by precall. This is inefficient,
// but we have to do it for now because shards might be different and
// we haven't implemented the local precalls that would be enough
// in some cases.
//
// This makes simple cross-index queries noticably inefficient.
//
// If you're here because of that: We should be using PrecallLocal
// in cases where the call isn't already PrecallGlobal, and
// PrecallLocal should wait until we're running on a specific node
// to do the farming-out of just the sub-queries it has to run
// for its local shards.
//
// As is, we have one node querying every node, then sending out
// all the data to every node, including the data that node already
// has. We could reduce the actual copying around dramatically,
// but only in the cases where local is good enough -- not something
// like Distinct, where you can't predict output shard for a result
// from the shard being queried.
if newIndex != "" && newIndex != index {
c.Type = pql.PrecallGlobal
index = newIndex
// we need to recompute shards, then
shards = nil
}
if err := o.handlePreCallChildren(ctx, index, c, shards, opt); err != nil {
return err
}
// child calls already handled, no precall for this, so we're done
if c.Type == pql.PrecallNone {
return nil
}
// We don't try to handle sub-calls from here. I'm not 100%
// sure that's right, but I think the fact that they're happening
// inside a precomputed call may mean they need different
// handling. In any event, the sub-calls will get handled by
// the executeCall when it gets to them...
// We set c to look like a normal call, and actually execute it:
c.Type = pql.PrecallNone
// possibly override call index.
v, err := o.executeCall(ctx, index, c, shards, opt)
if err != nil {
return err
}
var row *featurebase.Row
switch r := v.(type) {
case *featurebase.Row:
row = r
case featurebase.SignedRow:
row = r.Pos
default:
return fmt.Errorf("precomputed call %s returned unexpected non-Row data: %T", c.Name, v)
}
if err := ctx.Err(); err != nil {
return err
}
c.Children = []*pql.Call{}
c.Name = "Precomputed"
c.Args = map[string]interface{}{"valueidx": len(opt.EmbeddedData)}
// stash a copy of the full results, which can be forwarded to other
// shards if the query has to go to them
opt.EmbeddedData = append(opt.EmbeddedData, row)
// and stash a copy locally, so local calls can use it
if row != nil {
c.Precomputed = make(map[uint64]interface{}, len(row.Segments))
for _, segment := range row.Segments {
c.Precomputed[segment.Shard()] = &featurebase.Row{Segments: []featurebase.RowSegment{segment}}
}
}
return nil
}
// dumpPrecomputedCalls throws away precomputed call data. this is used so we
// can drop any large data associated with a call once we've processed
// the call.
func (o *orchestrator) dumpPrecomputedCalls(ctx context.Context, c *pql.Call) {
for _, call := range c.Children {
o.dumpPrecomputedCalls(ctx, call)
}
c.Precomputed = nil
}
// handlePreCallChildren handles any pre-calls in the children of a given call.
func (o *orchestrator) handlePreCallChildren(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) error {
for i := range c.Children {
if err := ctx.Err(); err != nil {
return err
}
if err := o.handlePreCalls(ctx, index, c.Children[i], shards, opt); err != nil {
return err
}
}
for key, val := range c.Args {
// Do not precompute GroupBy aggregates
if key == "aggregate" {
continue
}
// Handle Call() operations which exist inside named arguments, too.
if call, ok := val.(*pql.Call); ok {
if err := ctx.Err(); err != nil {
return err
}
if err := o.handlePreCalls(ctx, index, call, shards, opt); err != nil {
return err
}
}
}
return nil
}
// preprocessQuery expands any calls that need preprocessing.
func (o *orchestrator) preprocessQuery(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*pql.Call, error) {
switch c.Name {
case "All":
_, hasLimit, err := c.UintArg("limit")
if err != nil {
return nil, err
}
_, hasOffset, err := c.UintArg("offset")
if err != nil {
return nil, err
}
if !hasLimit && !hasOffset {
return c, nil
}
// Rewrite the All() w/ limit to Limit(All()).
c.Children = []*pql.Call{
{
Name: "All",
},
}
c.Name = "Limit"
return c, nil
default:
// Recurse through child calls.
out := make([]*pql.Call, len(c.Children))
var changed bool
for i, child := range c.Children {
res, err := o.preprocessQuery(ctx, index, child, shards, opt)
if err != nil {
return nil, err
}
if res != child {
changed = true
}
out[i] = res
}
if changed {
c = c.Clone()
c.Children = out
}
return c, nil
}
}
// executeCall executes a call.
func (o *orchestrator) executeCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCall")
defer span.Finish()
if err := validateQueryContext(ctx); err != nil {
return nil, err
} else if err := o.validateCallArgs(c); err != nil {
return nil, errors.Wrap(err, "validating args")
}
indexTag := "index:" + index
metricName := "query_" + strings.ToLower(c.Name) + "_total"
statFn := func() {
if !opt.Remote {
o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
}
}
// Preprocess the query.
c, err := o.preprocessQuery(ctx, index, c, shards, opt)
if err != nil {
return nil, err
}
switch c.Name {
case "Sum":
statFn()
res, err := o.executeSum(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeSum")
case "Min":
statFn()
res, err := o.executeMin(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMin")
case "Max":
statFn()
res, err := o.executeMax(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMax")
case "MinRow":
statFn()
res, err := o.executeMinRow(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMinRow")
case "MaxRow":
statFn()
res, err := o.executeMaxRow(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeMaxRow")
// case "Clear":
// statFn()
// res, err := o.executeClearBit(ctx, index, c, opt)
// return res, errors.Wrap(err, "executeClearBit")
// case "ClearRow":
// statFn()
// res, err := o.executeClearRow(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeClearRow")
case "Distinct":
statFn()
res, err := o.executeDistinct(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeDistinct")
// case "Store":
// statFn()
// res, err := o.executeSetRow(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeSetRow")
case "Count":
statFn()
res, err := o.executeCount(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeCount")
// case "Set":
// statFn()
// res, err := o.executeSet(ctx, index, c, opt)
// return res, errors.Wrap(err, "executeSet")
case "TopK":
statFn()
res, err := o.executeTopK(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeTopK")
case "TopN":
statFn()
res, err := o.executeTopN(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeTopN")
case "Rows":
statFn()
res, err := o.executeRows(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeRows")
case "Extract":
statFn()
res, err := o.executeExtract(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeExtract")
case "GroupBy":
statFn()
res, err := o.executeGroupBy(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeGroupBy")
case "Options":
statFn()
res, err := o.executeOptionsCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeOptionsCall")
case "IncludesColumn":
res, err := o.executeIncludesColumnCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeIncludesColumnCall")
case "FieldValue":
statFn()
res, err := o.executeFieldValueCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeFieldValueCall")
case "Precomputed":
res, err := o.executePrecomputedCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executePrecomputedCall")
case "UnionRows":
res, err := o.executeUnionRows(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeUnionRows")
case "ConstRow":
res, err := o.executeConstRow(ctx, index, c)
return res, errors.Wrap(err, "executeConstRow")
case "Limit":
res, err := o.executeLimitCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeLimitCall")
case "Percentile":
res, err := o.executePercentile(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executePercentile")
// case "Delete":
// statFn() //TODO(twg) need this?
// res, err := o.executeDeleteRecords(ctx, index, c, shards, opt)
// return res, errors.Wrap(err, "executeDelete")
default: // o.g. "Row", "Union", "Intersect" or anything that returns a bitmap.
statFn()
res, err := o.executeBitmapCall(ctx, index, c, shards, opt)
return res, errors.Wrap(err, "executeBitmapCall")
}
}
// validateCallArgs ensures that the value types in call.Args are expected.
func (o *orchestrator) validateCallArgs(c *pql.Call) error {
if _, ok := c.Args["ids"]; ok {
switch v := c.Args["ids"].(type) {
case []int64, []uint64:
// noop
case []interface{}:
b := make([]int64, len(v))
for i := range v {
b[i] = v[i].(int64)
}
c.Args["ids"] = b
default:
return fmt.Errorf("invalid call.Args[ids]: %s", v)
}
}
return nil
}
func (o *orchestrator) executeOptionsCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeOptionsCall")
defer span.Finish()
optCopy := &featurebase.ExecOptions{}
*optCopy = *opt
if arg, ok := c.Args["shards"]; ok {
if optShards, ok := arg.([]interface{}); ok {
shards = []uint64{}
for _, s := range optShards {
if shard, ok := s.(int64); ok {
shards = append(shards, uint64(shard))
} else {
return nil, errors.New(errors.ErrUncoded, "Query(): shards must be a list of unsigned integers")
}
}
} else {
return nil, errors.New(errors.ErrUncoded, "Query(): shards must be a list of unsigned integers")
}
}
return o.executeCall(ctx, index, c.Children[0], shards, optCopy)
}
// executeIncludesColumnCall executes an IncludesColumn() call.
func (o *orchestrator) executeIncludesColumnCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (bool, error) {
// Get the shard containing the column, since that's the only
// shard that needs to execute this query.
var shard uint64
col, ok, err := c.UintArg("column")
if err != nil {
return false, errors.Wrap(err, "getting column from args")
} else if !ok {
return false, errors.New(errors.ErrUncoded, "IncludesColumn call must specify a column")
}
shard = col / featurebase.ShardWidth
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(bool)
return other || v.(bool)
}
result, err := o.mapReduce(ctx, index, []uint64{shard}, c, opt, reduceFn)
if err != nil {
return false, err
}
return result.(bool), nil
}
// executeFieldValueCall executes a FieldValue() call.
func (o *orchestrator) executeFieldValueCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) {
fieldName, ok := c.Args["field"].(string)
if !ok || fieldName == "" {
return featurebase.ValCount{}, featurebase.ErrFieldRequired
}
colKey, ok := c.Args["column"]
if !ok || colKey == "" {
return featurebase.ValCount{}, featurebase.ErrColumnRequired
}
colID, ok, err := c.UintArg("column")
if !ok || err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "getting column argument")
}
shard := colID / featurebase.ShardWidth
// Select single returned result at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.ValCount)
if other.Count == 1 {
return other
}
return v
}
result, err := o.mapReduce(ctx, index, []uint64{shard}, c, opt, reduceFn)
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "map reduce")
}
other, _ := result.(featurebase.ValCount)
return other, nil
}
// executeLimitCall executes a Limit() call.
func (o *orchestrator) executeLimitCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.Row, error) {
bitmapCall := c.Children[0]
limit, hasLimit, err := c.UintArg("limit")
if err != nil {
return nil, errors.Wrap(err, "getting limit")
}
offset, _, err := c.UintArg("offset")
if err != nil {
return nil, errors.Wrap(err, "getting offset")
}
if !hasLimit {
limit = math.MaxUint64
}
// Execute bitmap call, storing the full result on this node.
res, err := o.executeCall(ctx, index, bitmapCall, shards, opt)
if err != nil {
return nil, errors.Wrap(err, "limit map reduce")
}
if res == nil {
res = featurebase.NewRow()
}
result, ok := res.(*featurebase.Row)
if !ok {
return nil, errors.Errorf("expected Row but got %T", result)
}
if offset != 0 {
i := 0
var leadingBits []uint64
for i < len(result.Segments) && offset > 0 {
seg := result.Segments[i]
count := seg.Count()
if count > offset {
data := seg.Columns()
data = data[offset:]
leadingBits = data
i++
break
}
offset -= count
i++
}
row := featurebase.NewRow(leadingBits...)
row.Merge(&featurebase.Row{Segments: result.Segments[i:]})
result = row
}
if limit < result.Count() {
i := 0
var trailingBits []uint64
for i < len(result.Segments) && limit > 0 {
seg := result.Segments[i]
count := seg.Count()
if count > limit {
data := seg.Columns()
data = data[:limit]
trailingBits = data
break
}
limit -= count
i++
}
row := featurebase.NewRow(trailingBits...)
row.Merge(&featurebase.Row{Segments: result.Segments[:i]})
result = row
}
return result, nil
}
// executeSum executes a Sum() call.
func (o *orchestrator) executeSum(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeSum")
defer span.Finish()
fieldName, err := c.FirstStringArg("field", "_field")
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "Sum(): field required")
}
if len(c.Children) > 1 {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Sum() only accepts a single bitmap input")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.ValCount)
return other.Add(v.(featurebase.ValCount))
}
result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return featurebase.ValCount{}, err
}
other, _ := result.(featurebase.ValCount)
if other.Count == 0 {
return featurebase.ValCount{}, nil
}
// scale summed response if it's a decimal field and this is
// not a remote query (we're about to return to original client).
if !opt.Remote {
field, err := o.schema.FieldInfo(ctx, index, fieldName)
if field == nil {
return featurebase.ValCount{}, errors.Wrapf(err, "%q", fieldName)
}
if field.Options.Type == FieldTypeDecimal {
dec := pql.NewDecimal(other.Val, field.Options.Scale)
other.DecimalVal = &dec
other.FloatVal = 0
other.Val = 0
}
}
return other, nil
}
// executeDistinct executes a Distinct call on a field. It returns a
// SignedRow for int fields and a *Row for set/mutex/time fields.
func (o *orchestrator) executeDistinct(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeDistinct")
defer span.Finish()
field, hasField, err := c.StringArg("field")
if err != nil {
return featurebase.SignedRow{}, errors.Wrap(err, "loading field option in Distinct query")
} else if !hasField {
return featurebase.SignedRow{}, fmt.Errorf("missing field option in Distinct query")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
if err := ctx.Err(); err != nil {
return err
}
switch other := prev.(type) {
case featurebase.SignedRow:
return other.Union(v.(featurebase.SignedRow))
case *featurebase.Row:
if other == nil {
return v
} else if v.(*featurebase.Row) == nil {
return other
}
return other.Union(v.(*featurebase.Row))
case nil:
return v
case featurebase.DistinctTimestamp:
return other.Union(v.(featurebase.DistinctTimestamp))
default:
return errors.Errorf("unexpected return type from executeDistinctShard: %+v %T", other, other)
}
}
result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, errors.Wrap(err, "mapReduce")
}
if other, ok := result.(featurebase.SignedRow); ok {
other.Field = field
}
return result, nil
}
// executeMin executes a Min() call.
func (o *orchestrator) executeMin(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMin")
defer span.Finish()
if _, err := c.FirstStringArg("field", "_field"); err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "Min(): field required")
}
if len(c.Children) > 1 {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Min() only accepts a single bitmap input")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.ValCount)
return other.Smaller(v.(featurebase.ValCount))
}
result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return featurebase.ValCount{}, err
}
other, _ := result.(featurebase.ValCount)
if other.Count == 0 {
return featurebase.ValCount{}, nil
}
return other, nil
}
// executeMax executes a Max() call.
func (o *orchestrator) executeMax(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMax")
defer span.Finish()
if _, err := c.FirstStringArg("field", "_field"); err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "Max(): field required")
}
if len(c.Children) > 1 {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Max() only accepts a single bitmap input")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.ValCount)
return other.Larger(v.(featurebase.ValCount))
}
result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return featurebase.ValCount{}, err
}
other, _ := result.(featurebase.ValCount)
if other.Count == 0 {
return featurebase.ValCount{}, nil
}
return other, nil
}
// TODO(jaffee) fix this... valcountize assumes access to field details like base
// executePercentile executes a Percentile() call.
func (o *orchestrator) executePercentile(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executePercentile")
defer span.Finish()
// get nth
var nthFloat float64
nthArg, ok := c.Args["nth"]
if !ok {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Percentile(): nth required")
}
switch nthArg := nthArg.(type) {
case pql.Decimal:
nthFloat = nthArg.Float64()
case int64:
nthFloat = float64(nthArg)
default:
return featurebase.ValCount{}, errors.Errorf("Percentile(): invalid nth='%v' of type (%[1]T), should be a number between 0 and 100 inclusive", c.Args["nth"])
}
if nthFloat < 0 || nthFloat > 100.0 {
return featurebase.ValCount{}, errors.Errorf("Percentile(): invalid nth value (%f), should be a number between 0 and 100 inclusive", nthFloat)
}
// get field
fieldName, err := c.FirstStringArg("field", "_field")
if err != nil {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Percentile(): field required")
}
field, err := o.schema.FieldInfo(ctx, index, fieldName)
if err != nil {
return featurebase.ValCount{}, ErrFieldNotFound
}
// filter call for min & max
var filterCall *pql.Call
// check if filter provided
if filterArg, ok := c.Args["filter"].(*pql.Call); ok && filterArg != nil {
filterCall = filterArg
}
// get min
q, _ := pql.ParseString(fmt.Sprintf(`Min(field="%s")`, fieldName))
minCall := q.Calls[0]
if filterCall != nil {
minCall.Children = append(minCall.Children, filterCall)
}
minVal, err := o.executeMin(ctx, index, minCall, shards, opt)
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "executing Min call for Percentile")
}
if nthFloat == 0.0 {
return minVal, nil
}
// get max
q, _ = pql.ParseString(fmt.Sprintf(`Max(field="%s")`, fieldName))
maxCall := q.Calls[0]
if filterCall != nil {
maxCall.Children = append(maxCall.Children, filterCall)
}
maxVal, err := o.executeMax(ctx, index, maxCall, shards, opt)
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "executing Max call for Percentile")
}
// set up reusables
var countCall, rangeCall *pql.Call
if filterCall == nil {
countQuery, _ := pql.ParseString(fmt.Sprintf("Count(Row(%s < 0))", fieldName))
countCall = countQuery.Calls[0]
rangeCall = countCall.Children[0]
} else {
countQuery, _ := pql.ParseString(fmt.Sprintf(`Count(Intersect(Row(%s < 0)))`, fieldName))
countCall = countQuery.Calls[0]
intersectCall := countCall.Children[0]
intersectCall.Children = append(intersectCall.Children, filterCall)
rangeCall = intersectCall.Children[0]
}
k := (100 - nthFloat) / nthFloat
min, max := minVal.Val, maxVal.Val
// estimate nth val, eg median when nth=0.5
for min < max {
// compute average without integer overflow, then correct for division of
// odd numbers by 2
possibleNthVal := ((max / 2) + (min / 2)) + (((max % 2) + (min % 2)) / 2)
// possibleNthVal = (max + min) / 2
// get left count
rangeCall.Args[fieldName] = &pql.Condition{
Op: pql.Token(pql.LT),
Value: possibleNthVal,
}
leftCountUint64, err := o.executeCount(ctx, index, countCall, shards, opt)
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "executing Count call L for Percentile")
}
leftCount := int64(leftCountUint64)
// get right count
rangeCall.Args[fieldName] = &pql.Condition{
Op: pql.Token(pql.GT),
Value: possibleNthVal,
}
rightCountUint64, err := o.executeCount(ctx, index, countCall, shards, opt)
if err != nil {
return featurebase.ValCount{}, errors.Wrap(err, "executing Count call R for Percentile")
}
rightCount := int64(rightCountUint64)
// 'weight' the left count as per k
leftCountWeighted := int64(math.Round(k * float64(leftCount)))
// binary search
if leftCountWeighted > rightCount {
max = possibleNthVal - 1
} else if leftCountWeighted < rightCount {
min = possibleNthVal + 1
} else {
return cookValCount(possibleNthVal, 1, field), nil
}
}
return cookValCount(min, 1, field), nil
}
func cookValCount(val int64, cnt uint64, field *featurebase.FieldInfo) featurebase.ValCount {
valCount := featurebase.ValCount{Count: int64(cnt)}
base := field.Options.Base
switch field.Options.Type {
case featurebase.FieldTypeDecimal:
dec := pql.NewDecimal(val+base, field.Options.Scale)
valCount.DecimalVal = &dec
case FieldTypeTimestamp:
valCount.TimestampVal = time.Unix(0, (val+base)*featurebase.TimeUnitNanos(field.Options.TimeUnit)).UTC()
}
valCount.Val = val + base
return valCount
}
// executeMinRow executes a MinRow() call.
func (o *orchestrator) executeMinRow(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ interface{}, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMinRow")
defer span.Finish()
if field := c.Args["field"]; field == "" {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "MinRow(): field required")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
// if minRowID exists, and if it is smaller than the other one return it.
// otherwise return the minRowID of the one which exists.
if prev == nil {
return v
} else if v == nil {
return prev
}
prevp, _ := prev.(featurebase.PairField)
vp, _ := v.(featurebase.PairField)
if prevp.Pair.Count > 0 && vp.Pair.Count > 0 {
if prevp.Pair.ID < vp.Pair.ID {
return prevp
}
return vp
} else if prevp.Pair.Count > 0 {
return prevp
}
return vp
}
return o.mapReduce(ctx, index, shards, c, opt, reduceFn)
}
// executeMaxRow executes a MaxRow() call.
func (o *orchestrator) executeMaxRow(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ interface{}, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMaxRow")
defer span.Finish()
if field := c.Args["field"]; field == "" {
return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "MaxRow(): field required")
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
// if minRowID exists, and if it is smaller than the other one return it.
// otherwise return the minRowID of the one which exists.
if prev == nil {
return v
} else if v == nil {
return prev
}
prevp, _ := prev.(featurebase.PairField)
vp, _ := v.(featurebase.PairField)
if prevp.Pair.Count > 0 && vp.Pair.Count > 0 {
if prevp.Pair.ID > vp.Pair.ID {
return prevp
}
return vp
} else if prevp.Pair.Count > 0 {
return prevp
}
return vp
}
return o.mapReduce(ctx, index, shards, c, opt, reduceFn)
}
// executePrecomputedCall pretends to execute a call that we have a precomputed value for.
func (o *orchestrator) executePrecomputedCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ *featurebase.Row, err error) {
span, _ := tracing.StartSpanFromContext(ctx, "Executor.executePrecomputedCall")
defer span.Finish()
result := featurebase.NewRow()
for _, row := range c.Precomputed {
result.Merge(row.(*featurebase.Row))
}
return result, nil
}
// executeBitmapCall executes a call that returns a bitmap.
func (o *orchestrator) executeBitmapCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ *featurebase.Row, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeBitmapCall")
span.LogKV("pqlCallName", c.Name)
defer span.Finish()
indexTag := "index:" + index
metricName := "query_" + strings.ToLower(c.Name) + "_total"
if c.Name == "Row" && c.HasConditionArg() {
metricName = "query_row_bsi_total"
}
if !opt.Remote {
o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag})
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(*featurebase.Row)
if other == nil {
// TODO... what's going on on the following line
other = featurebase.NewRow() // bug! this row ends up containing Badger Txn data that should be accessed outside the Txn.
}
if err := ctx.Err(); err != nil {
return err
}
other.Merge(v.(*featurebase.Row))
return other
}
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, errors.Wrap(err, "map reduce")
}
row, _ := other.(*featurebase.Row)
return row, nil
}
type Error string // TODO(jaffee) convert to standard error package
func (e Error) Error() string { return string(e) }
const ViewNotFound = Error("view not found")
const FragmentNotFound = Error("fragment not found")
func (o *orchestrator) executeTopK(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopK")
defer span.Finish()
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
x, _ := prev.([]*featurebase.Row)
y, _ := v.([]*featurebase.Row)
return ([]*featurebase.Row)(featurebase.AddBSI(x, y))
}
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, err
}
results, _ := other.([]*featurebase.Row)
if opt.Remote {
return results, nil
}
k, hasK, err := c.UintArg("k")
if err != nil {
return nil, errors.Wrap(err, "fetching k")
}
var limit *uint64
if hasK {
limit = &k
}
var dst []featurebase.Pair
featurebase.BSIData(results).PivotDescending(featurebase.NewRow().Union(results...), 0, limit, nil, func(count uint64, ids ...uint64) {
for _, id := range ids {
dst = append(dst, featurebase.Pair{
ID: id,
Count: count,
})
}
})
fieldName, hasFieldName, err := c.StringArg("_field")
if err != nil {
return nil, errors.Wrap(err, "fetching TopK field")
} else if !hasFieldName {
return nil, errors.New(errors.ErrUncoded, "missing field in TopK")
}
return &featurebase.PairsField{
Pairs: dst,
Field: fieldName,
}, nil
}
// uint64Slice represents a sortable slice of uint64 numbers.
type uint64Slice []uint64
func (p uint64Slice) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p uint64Slice) Len() int { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }
// executeTopN executes a TopN() call.
// This first performs the TopN() to determine the top results and then
// requeries to retrieve the full counts for each of the top results.
func (o *orchestrator) executeTopN(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.PairsField, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopN")
defer span.Finish()
idsArg, _, err := c.UintSliceArg("ids")
if err != nil {
return nil, fmt.Errorf("executeTopN: %v", err)
}
fieldName, _ := c.Args["_field"].(string)
n, _, err := c.UintArg("n")
if err != nil {
return nil, fmt.Errorf("executeTopN: %v", err)
}
// Execute original query.
pairs, err := o.executeTopNShards(ctx, index, c, shards, opt)
if err != nil {
return nil, errors.Wrap(err, "finding top results")
}
// If this call is against specific ids, or we didn't get results,
// or we are part of a larger distributed query then don't refetch.
if len(pairs.Pairs) == 0 || len(idsArg) > 0 || opt.Remote {
return &featurebase.PairsField{
Pairs: pairs.Pairs,
Field: fieldName,
}, nil
}
// Only the original caller should refetch the full counts.
// TODO(@kuba--): ...but do we really need `Clone` here?
other := c.Clone()
ids := featurebase.Pairs(pairs.Pairs).Keys()
sort.Sort(uint64Slice(ids))
other.Args["ids"] = ids
trimmedList, err := o.executeTopNShards(ctx, index, other, shards, opt)
if err != nil {
return nil, errors.Wrap(err, "retrieving full counts")
}
if n != 0 && int(n) < len(trimmedList.Pairs) {
trimmedList.Pairs = trimmedList.Pairs[0:n]
}
return &featurebase.PairsField{
Pairs: trimmedList.Pairs,
Field: fieldName,
}, nil
}
func (o *orchestrator) executeTopNShards(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.PairsField, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopNShards")
defer span.Finish()
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(*featurebase.PairsField)
vpf, _ := v.(*featurebase.PairsField)
if other == nil {
return vpf
} else if vpf == nil {
return other
}
if err := ctx.Err(); err != nil {
return err
}
other.Pairs = featurebase.Pairs(other.Pairs).Add(vpf.Pairs)
return other
}
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, err
}
results, _ := other.(*featurebase.PairsField)
// Sort final merged results.
sort.Sort(featurebase.Pairs(results.Pairs))
return results, nil
}
// order denotes sort order—can be asc or desc (see constants below).
type order bool
const (
asc order = true
desc order = false
)
// groupCountSorter sorts the output of a GroupBy request (a
// []GroupCount) according to sorting instructions encoded in "fields"
// and "order".
//
// Each field in "fields" is an integer which can be -1 to denote
// sorting on the Count and -2 to denote sorting on the
// sum/aggregate. Currently nothing else is supported, but the idea
// was that if there were positive integers they would be indexes into
// GroupCount.FieldRow and allowing sorting on the values of different
// fields in the group. Each item in "order" corresponds to the same
// index in "fields" and denotes the order of the sort.
type groupCountSorter struct {
fields []int
order []order
data []featurebase.GroupCount
}
func (g *groupCountSorter) Len() int { return len(g.data) }
func (g *groupCountSorter) Swap(i, j int) { g.data[i], g.data[j] = g.data[j], g.data[i] }
func (g *groupCountSorter) Less(i, j int) bool {
gci, gcj := g.data[i], g.data[j]
for idx, fieldIndex := range g.fields {
fieldOrder := g.order[idx]
switch fieldIndex {
case -1: // Count
if gci.Count < gcj.Count {
return fieldOrder == asc
} else if gci.Count > gcj.Count {
return fieldOrder == desc
}
case -2: // Aggregate
if gci.Agg < gcj.Agg {
return fieldOrder == asc
} else if gci.Agg > gcj.Agg {
return fieldOrder == desc
}
default:
panic("impossible")
}
}
return false
}
// getSorter hackily parses the sortSpec and figures out how to sort
// the GroupBy results.
func getSorter(sortSpec string) (*groupCountSorter, error) {
gcs := &groupCountSorter{
fields: []int{},
order: []order{},
}
sortOn := strings.Split(sortSpec, ",")
for _, sortField := range sortOn {
sortField = strings.TrimSpace(sortField)
fieldDir := strings.Fields(sortField)
if len(fieldDir) == 0 {
return nil, errors.Errorf("invalid sorting directive: '%s'", sortField)
} else if fieldDir[0] == "count" {
gcs.fields = append(gcs.fields, -1)
} else if fieldDir[0] == "aggregate" || fieldDir[0] == "sum" {
gcs.fields = append(gcs.fields, -2)
} else {
return nil, errors.Errorf("sorting is only supported on count, aggregate, or sum, not '%s'", fieldDir[0])
}
if len(fieldDir) == 1 {
gcs.order = append(gcs.order, desc)
} else if len(fieldDir) > 2 {
return nil, errors.Errorf("parsing sort directive: '%s': too many elements", sortField)
} else if fieldDir[1] == "asc" {
gcs.order = append(gcs.order, asc)
} else if fieldDir[1] == "desc" {
gcs.order = append(gcs.order, desc)
} else {
return nil, errors.Errorf("unknown sort direction '%s'", fieldDir[1])
}
}
return gcs, nil
}
// findGroupCounts gets a safe-to-use but possibly empty []GroupCount from
// an interface which might be a *GroupCounts or a []GroupCount.
func findGroupCounts(v interface{}) []featurebase.GroupCount {
switch gc := v.(type) {
case []featurebase.GroupCount:
return gc
case *featurebase.GroupCounts:
return gc.Groups()
}
return nil
}
func (o *orchestrator) executeGroupBy(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.GroupCounts, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeGroupBy")
defer span.Finish()
// validate call
if len(c.Children) == 0 {
return nil, errors.New(errors.ErrUncoded, "need at least one child call")
}
limit := int(^uint(0) >> 1)
if lim, hasLimit, err := c.UintArg("limit"); err != nil {
return nil, err
} else if hasLimit {
limit = int(lim)
}
filter, _, err := c.CallArg("filter")
if err != nil {
return nil, err
}
var sorter *groupCountSorter
if sortSpec, found, err := c.StringArg("sort"); err != nil {
return nil, errors.Wrap(err, "getting sort arg")
} else if found {
sorter, err = getSorter(sortSpec)
if err != nil {
return nil, errors.Wrap(err, "parsing sort spec")
}
// don't want to prematurely limit the results if we're sorting
limit = int(^uint(0) >> 1)
}
having, hasHaving, err := c.CallArg("having")
if err != nil {
return nil, errors.Wrap(err, "getting 'having' argument")
} else if hasHaving {
// don't want to prematurely limit the results if we're filtering some out
limit = int(^uint(0) >> 1)
}
// perform necessary Rows queries (any that have limit or columns args) -
// TODO, call async? would only help if multiple Rows queries had a column
// or limit arg.
// TODO support TopN in here would be really cool - and pretty easy I think.
childRows := make([]featurebase.RowIDs, len(c.Children))
for i, child := range c.Children {
// Check "field" first for backwards compatibility, then set _field.
// TODO: remove at Pilosa 2.0
if fieldName, ok := child.Args["field"].(string); ok {
child.Args["_field"] = fieldName
}
if child.Name != "Rows" {
return nil, errors.Errorf("'%s' is not a valid child query for GroupBy, must be 'Rows'", child.Name)
}
_, hasLimit, err := child.UintArg("limit")
if err != nil {
return nil, errors.Wrap(err, "getting limit")
}
_, hasCol, err := child.UintArg("column")
if err != nil {
return nil, errors.Wrap(err, "getting column")
}
_, hasLike, err := child.StringArg("like")
if err != nil {
return nil, errors.Wrap(err, "getting like")
}
_, hasIn, err := child.UintSliceArg("in")
if err != nil {
return nil, errors.Wrap(err, "getting 'in'")
}
if hasLimit || hasCol || hasLike || hasIn { // we need to perform this query cluster-wide ahead of executeGroupByShard
if idx, ok := child.Args["valueidx"].(int64); ok {
// The rows query was already completed on the initiating node.
childRows[i] = opt.EmbeddedData[idx].Columns()
continue
}
r, er := o.executeRows(ctx, index, child, shards, opt)
if er != nil {
return nil, errors.Wrap(er, "getting rows for ")
}
// need to sort because filters assume ordering
sort.Slice(r, func(x, y int) bool { return r[x] < r[y] })
childRows[i] = r
if len(childRows[i]) == 0 { // there are no results because this field has no values.
return &featurebase.GroupCounts{}, nil
}
// Stuff the result into opt.EmbeddedData so that it gets sent to other nodes in the map-reduce.
// This is flagged as "NoSplit" to ensure that the entire row gets sent out.
rowsRow := featurebase.NewRow(childRows[i]...)
rowsRow.NoSplit = true
child.Args["valueidx"] = int64(len(opt.EmbeddedData))
opt.EmbeddedData = append(opt.EmbeddedData, rowsRow)
}
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other := findGroupCounts(prev)
if err := ctx.Err(); err != nil {
return err
}
return mergeGroupCounts(other, findGroupCounts(v), limit)
}
// Get full result set.
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, errors.Wrap(err, "mapReduce")
}
results, _ := other.([]featurebase.GroupCount)
// If there's no sorting, we want to apply limits before
// calculating the Distinct aggregate which is expensive on a
// per-result basis.
if sorter == nil && !hasHaving {
results, err = applyLimitAndOffsetToGroupByResult(c, results)
if err != nil {
return nil, errors.Wrap(err, "applying limit/offset")
}
}
// TODO as an optimization, we could apply some "having"
// conditions here long as they aren't on the Count(Distinct)
// aggregate
// Calculate Count(Distinct) aggregate if requested.
aggregate, _, err := c.CallArg("aggregate")
if err == nil && aggregate != nil && aggregate.Name == "Count" && len(aggregate.Children) > 0 && aggregate.Children[0].Name == "Distinct" && !opt.Remote {
for n, gc := range results {
intersectRows := make([]*pql.Call, 0, len(gc.Group))
for _, fr := range gc.Group {
var value interface{} = fr.RowID
// use fr.Value instead of fr.RowID if set (from int fields)
if fr.Value != nil {
value = &pql.Condition{Op: pql.EQ, Value: *fr.Value}
}
intersectRows = append(intersectRows, &pql.Call{Name: "Row", Args: map[string]interface{}{fr.Field: value}})
}
// apply any filter, if present
if filter != nil {
intersectRows = append(intersectRows, filter)
}
// also intersect with any children of Distinct
if len(aggregate.Children[0].Children) > 0 {
intersectRows = append(intersectRows, aggregate.Children[0].Children[0])
}
countDistinctIntersect := &pql.Call{
Name: "Count",
Children: []*pql.Call{
{
Name: "Distinct",
Children: []*pql.Call{
{
Name: "Intersect",
Children: intersectRows,
},
},
Args: aggregate.Children[0].Args,
Type: pql.PrecallGlobal,
},
},
}
opt.PreTranslated = true
aggregateCount, err := o.execute(ctx, index, &pql.Query{Calls: []*pql.Call{countDistinctIntersect}}, []uint64{}, opt)
if err != nil {
return nil, err
}
results[n].Agg = int64(aggregateCount[0].(uint64))
}
}
// Apply having.
if hasHaving && !opt.Remote {
// parse the condition as PQL
if having.Name != "Condition" {
return nil, errors.New(errors.ErrUncoded, "the only supported having call is Condition()")
}
if len(having.Args) != 1 {
return nil, errors.New(errors.ErrUncoded, "Condition() must contain a single condition")
}
for subj, cond := range having.Args {
switch subj {
case "count", "sum":
results = featurebase.ApplyConditionToGroupCounts(results, subj, cond.(*pql.Condition))
default:
return nil, errors.New(errors.ErrUncoded, "Condition() only supports count or sum")
}
}
}
if sorter != nil && !opt.Remote {
sorter.data = results
sort.Stable(sorter)
results, err = applyLimitAndOffsetToGroupByResult(c, results)
if err != nil {
return nil, errors.Wrap(err, "applying limit/offset")
}
} else if hasHaving && !opt.Remote {
results, err = applyLimitAndOffsetToGroupByResult(c, results)
if err != nil {
return nil, errors.Wrap(err, "applying limit/offset")
}
}
aggType := ""
if aggregate != nil {
switch aggregate.Name {
case "Sum":
aggType = "sum"
case "Count":
aggType = "aggregate"
}
}
for _, res := range results {
if res.DecimalAgg != nil && aggType == "sum" {
aggType = "decimalSum"
break
}
}
return featurebase.NewGroupCounts(aggType, results...), nil
}
func applyLimitAndOffsetToGroupByResult(c *pql.Call, results []featurebase.GroupCount) ([]featurebase.GroupCount, error) {
// Apply offset.
if offset, hasOffset, err := c.UintArg("offset"); err != nil {
return nil, err
} else if hasOffset {
if int(offset) < len(results) {
results = results[offset:]
}
}
// Apply limit.
if limit, hasLimit, err := c.UintArg("limit"); err != nil {
return nil, err
} else if hasLimit {
if int(limit) < len(results) {
results = results[:limit]
}
}
return results, nil
}
// mergeGroupCounts merges two slices of GroupCounts throwing away any that go
// beyond the limit. It assume that the two slices are sorted by the row ids in
// the fields of the group counts. It may modify its arguments.
func mergeGroupCounts(a, b []featurebase.GroupCount, limit int) []featurebase.GroupCount {
if limit > len(a)+len(b) {
limit = len(a) + len(b)
}
ret := make([]featurebase.GroupCount, 0, limit)
i, j := 0, 0
for i < len(a) && j < len(b) && len(ret) < limit {
switch a[i].Compare(b[j]) {
case -1:
ret = append(ret, a[i])
i++
case 0:
a[i].Count += b[j].Count
a[i].Agg += b[j].Agg
if a[i].DecimalAgg != nil && b[j].DecimalAgg != nil {
sum := pql.AddDecimal(*a[i].DecimalAgg, *b[j].DecimalAgg)
a[i].DecimalAgg = &sum
}
ret = append(ret, a[i])
i++
j++
case 1:
ret = append(ret, b[j])
j++
}
}
for ; i < len(a) && len(ret) < limit; i++ {
ret = append(ret, a[i])
}
for ; j < len(b) && len(ret) < limit; j++ {
ret = append(ret, b[j])
}
return ret
}
func (o *orchestrator) executeRows(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (featurebase.RowIDs, error) {
// Fetch field name from argument.
// Check "field" first for backwards compatibility.
// TODO: remove at Pilosa 2.0
var fieldName string
var ok bool
if fieldName, ok = c.Args["field"].(string); ok {
c.Args["_field"] = fieldName
}
if fieldName, ok = c.Args["_field"].(string); !ok {
return nil, errors.New(errors.ErrUncoded, "Rows() field required")
}
// TODO(tlt): this is here to prevent the linter from complaining.
// Presumably this fieldName is/was used in code which is no longer here or
// is currently commented out.
_ = fieldName
if columnID, ok, err := c.UintArg("column"); err != nil {
return nil, errors.Wrap(err, "getting column")
} else if ok {
shards = []uint64{columnID / featurebase.ShardWidth}
}
// TODO, support "in" in conjunction w/ other args... or at least error if they're present together
if ids, found, err := c.UintSliceArg("in"); err != nil {
return nil, errors.Wrapf(err, "'in' argument of Rows must be a slice")
} else if found {
// "in" not supported with other args, so check here
for arg := range c.Args {
if arg != "field" && arg != "_field" && arg != "in" {
return nil, errors.Errorf("Rows call with 'in' does not support other arguments, but found '%s'", arg)
}
}
return ids, nil
}
// Determine limit so we can use it when reducing.
limit := int(^uint(0) >> 1)
if lim, hasLimit, err := c.UintArg("limit"); err != nil {
return nil, err
} else if hasLimit {
limit = int(lim)
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.RowIDs)
if err := ctx.Err(); err != nil {
return err
}
return other.Merge(v.(featurebase.RowIDs), limit)
}
// Get full result set.
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return nil, err
}
results, _ := other.(featurebase.RowIDs)
// TODO(jaffee) enable "like" support
// if !opt.Remote {
// if like, hasLike, err := c.StringArg("like"); err != nil {
// return nil, errors.Wrap(err, "getting like pattern")
// } else if hasLike {
// matches, err := e.Cluster.matchField(ctx, e.Holder.Field(index, fieldName), like)
// if err != nil {
// return nil, errors.Wrap(err, "matching like pattern")
// }
// i, j, k := 0, 0, 0
// for i < len(results) && j < len(matches) {
// x, y := results[i], matches[j]
// switch {
// case x < y:
// i++
// case y < x:
// j++
// default:
// results[k] = x
// i++
// j++
// k++
// }
// }
// results = results[:k]
// }
// }
return results, nil
}
func (o *orchestrator) executeExtract(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (featurebase.ExtractedIDMatrix, error) {
// Extract the column filter call.
if len(c.Children) < 1 {
return featurebase.ExtractedIDMatrix{}, errors.New(errors.ErrUncoded, "missing column filter in Extract")
}
// Extract fields from rows calls.
fields := make([]string, len(c.Children)-1)
for i, rows := range c.Children[1:] {
if rows.Name != "Rows" {
return featurebase.ExtractedIDMatrix{}, errors.Errorf("child call of Extract is %q but expected Rows", rows.Name)
}
var fieldName string
var ok bool
for k, v := range rows.Args {
switch k {
case "field", "_field":
fieldName = v.(string)
ok = true
default:
return featurebase.ExtractedIDMatrix{}, errors.Errorf("unsupported Rows argument for Extract: %q", k)
}
}
if !ok {
return featurebase.ExtractedIDMatrix{}, errors.New(errors.ErrUncoded, "missing field specification in Rows")
}
fields[i] = fieldName
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(featurebase.ExtractedIDMatrix)
if err := ctx.Err(); err != nil {
return err
}
other.Append(v.(featurebase.ExtractedIDMatrix))
return other
}
// Get full result set.
other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return featurebase.ExtractedIDMatrix{}, err
}
results, _ := other.(featurebase.ExtractedIDMatrix)
sort.Slice(results.Columns, func(i, j int) bool {
return results.Columns[i].ColumnID < results.Columns[j].ColumnID
})
return results, nil
}
func (o *orchestrator) executeConstRow(ctx context.Context, index string, c *pql.Call) (res *featurebase.Row, err error) {
// Fetch user-provided columns list.
ids, ok := c.Args["columns"].([]uint64)
if !ok {
return nil, errors.New(errors.ErrUncoded, "missing columns list")
}
return featurebase.NewRow(ids...), nil
}
func (o *orchestrator) executeUnionRows(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.Row, error) {
// Turn UnionRows(Rows(...)) into Union(Row(...), ...).
var rows []*pql.Call
for _, child := range c.Children {
// Check that we can use the call.
switch child.Name {
case "Rows":
case "TopN":
default:
return nil, errors.Errorf("cannot use %v as a rows query", child)
}
// Execute the call.
rowsResult, err := o.executeCall(ctx, index, child, shards, opt)
if err != nil {
return nil, err
}
// Turn the results into rows calls.
var resultRows []*pql.Call
switch rowsResult := rowsResult.(type) {
case *featurebase.PairsField:
// Translate pairs into rows calls.
for _, p := range rowsResult.Pairs {
var val interface{}
switch {
case p.Key != "":
val = p.Key
default:
val = p.ID
}
resultRows = append(resultRows, &pql.Call{
Name: "Row",
Args: map[string]interface{}{
rowsResult.Field: val,
},
})
}
case featurebase.RowIDs:
// Translate Row IDs into Row calls.
for _, id := range rowsResult {
resultRows = append(resultRows, &pql.Call{
Name: "Row",
Args: map[string]interface{}{
child.Args["_field"].(string): id,
},
})
}
default:
return nil, errors.Errorf("unexpected Rows type %T", rowsResult)
}
// Propogate any special properties of the call.
switch child.Name {
case "Rows":
// Propogate "from" time, if set.
if v, ok := child.Args["from"]; ok {
for _, rowCall := range resultRows {
rowCall.Args["from"] = v
}
}
// Propogate "to" time, if set.
if v, ok := child.Args["to"]; ok {
for _, rowCall := range resultRows {
rowCall.Args["to"] = v
}
}
}
rows = append(rows, resultRows...)
}
// Generate a Union call over the rows.
c = &pql.Call{
Name: "Union",
Children: rows,
}
// Execute the generated Union() call.
return o.executeBitmapCall(ctx, index, c, shards, opt)
}
// executeCount executes a count() call.
func (o *orchestrator) executeCount(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (uint64, error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCount")
defer span.Finish()
if len(c.Children) == 0 {
return 0, errors.New(errors.ErrUncoded, "Count() requires an input bitmap")
} else if len(c.Children) > 1 {
return 0, errors.New(errors.ErrUncoded, "Count() only accepts a single bitmap input")
}
child := c.Children[0]
// If the child is distinct/similar, execute it directly here and count the result.
if child.Type == pql.PrecallGlobal {
result, err := o.executeCall(ctx, index, child, shards, opt)
if err != nil {
return 0, err
}
switch row := result.(type) {
case *featurebase.Row:
return row.Count(), nil
case featurebase.SignedRow:
return row.Pos.Count() + row.Neg.Count(), nil
case featurebase.DistinctTimestamp:
return uint64(len(row.Values)), nil
default:
return 0, errors.Errorf("cannot count result of type %T from call %q", row, child.String())
}
}
// Merge returned results at coordinating node.
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
other, _ := prev.(uint64)
return other + v.(uint64)
}
result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn)
if err != nil {
return 0, err
}
n, _ := result.(uint64)
return n, nil
}
// remoteExec executes a PQL query remotely for a set of shards on a node.
func (o *orchestrator) remoteExec(ctx context.Context, node dax.Address, index string, q *pql.Query, shards []uint64, embed []*featurebase.Row) (results []interface{}, err error) { // nolint: interfacer
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeExec")
defer span.Finish()
// Encode request object.
pbreq := &featurebase.QueryRequest{
Query: q.String(),
Shards: shards,
Remote: true,
EmbeddedData: embed,
}
scheme := node.Scheme()
if scheme == "" {
scheme = "http"
}
resp, err := o.client.QueryNode(ctx, &net.URI{
Scheme: scheme,
Host: node.Host(),
Port: node.Port(),
}, index, pbreq)
if err != nil {
return nil, err
}
return resp.Results, resp.Err
}
// mapReduce maps and reduces data across the cluster.
//
// If a mapping of shards to a node fails then the shards are resplit across
// secondary nodes and retried. This continues to occur until all nodes are exhausted.
//
// mapReduce has to ensure that it never returns before any work it spawned has
// terminated. It's not enough to cancel the jobs; we have to wait for them to be
// done, or we can unmap resources they're still using.
func (o *orchestrator) mapReduce(ctx context.Context, index string, shards []uint64, c *pql.Call, opt *featurebase.ExecOptions, reduceFn reduceFunc) (result interface{}, err error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapReduce")
defer span.Finish()
ch := make(chan mapResponse)
// Wrap context with a cancel to kill goroutines on exit.
ctx, cancel := context.WithCancel(ctx)
// Create an errgroup so we can wait for all the goroutines to exit
eg, ctx := errgroup.WithContext(ctx)
// After we're done processing, we have to wait for any outstanding
// functions in the ErrGroup to complete. If we didn't have an error
// already at that point, we'll report any errors from the ErrGroup
// instead.
defer func() {
cancel()
errWait := eg.Wait()
if err == nil {
err = errWait
}
}()
nodes, err := o.topology.ComputeNodes(ctx, index, shards)
if err != nil {
return nil, errors.Wrapf(err, "getting nodes/shards for index '%q'", index)
}
// Start mapping across all primary owners.
if err = o.mapper(ctx, eg, ch, index, nodes, c, opt, reduceFn); err != nil {
return nil, errors.Wrap(err, "starting mapper")
}
// Iterate over all map responses and reduce.
expected := 0
for _, n := range nodes {
expected += len(n.Shards)
}
done := ctx.Done()
for expected > 0 {
select {
case <-done:
return nil, ctx.Err()
case resp := <-ch:
if resp.err != nil {
cancel() // TODO(jaffee) I added this... seems right, but wasn't there before
return nil, errors.Wrap(resp.err, "mapping on primary node")
}
// if we got a response that we aren't discarding
// because it's an error, subtract it from our count...
expected -= len(resp.shards)
// Reduce value.
result = reduceFn(ctx, result, resp.result)
var ok bool
// note *not* shadowed.
if err, ok = result.(error); ok {
cancel()
return nil, err
}
}
}
// note the deferred Wait above which might override this nil.
return result, nil
}
// makeEmbeddedDataForShards produces new rows containing the RowSegments
// that would correspond to a given set of shards.
func makeEmbeddedDataForShards(allRows []*featurebase.Row, shards []uint64) []*featurebase.Row {
if len(allRows) == 0 || len(shards) == 0 {
return nil
}
newRows := make([]*featurebase.Row, len(allRows))
for i, row := range allRows {
if row == nil || len(row.Segments) == 0 {
continue
}
if row.NoSplit {
newRows[i] = row
continue
}
segments := row.Segments
segmentIndex := 0
newRows[i] = &featurebase.Row{
Index: row.Index,
Field: row.Field,
}
for _, shard := range shards {
for segmentIndex < len(segments) && segments[segmentIndex].Shard() < shard {
segmentIndex++
}
// no more segments in this row
if segmentIndex >= len(segments) {
break
}
if segments[segmentIndex].Shard() == shard {
newRows[i].Segments = append(newRows[i].Segments, segments[segmentIndex])
segmentIndex++
if segmentIndex >= len(segments) {
// no more segments, we're done
break
}
}
// if we got here, segments[segmentIndex].shard exists
// but is greater than the current shard, so we continue.
}
}
return newRows
}
func (o *orchestrator) mapper(ctx context.Context, eg *errgroup.Group, ch chan mapResponse, index string, nodes []controller.ComputeNode, c *pql.Call, opt *featurebase.ExecOptions, reduceFn reduceFunc) (reterr error) {
span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapper")
defer span.Finish()
// Group shards together by nodes.
done := ctx.Done()
// Execute each node in a separate goroutine.
for _, node := range nodes {
node := node
shards := make([]uint64, len(node.Shards))
for i, dshard := range node.Shards {
shards[i] = uint64(dshard)
}
eg.Go(func() error {
resp := mapResponse{node: node.Address, shards: shards}
var embeddedRowsForNode []*featurebase.Row
if opt.EmbeddedData != nil {
embeddedRowsForNode = makeEmbeddedDataForShards(opt.EmbeddedData, shards)
}
attempts := 0
for ; attempts == 0 || (resp.err != nil && strings.Contains(resp.err.Error(), errConnectionRefused) && attempts < 3); attempts++ {
// On error retry against remaining nodes. If an error returns then
// the context will cancel and cause all open goroutines to return.
//
// We distinguish here between an error which indicates that the
// node is not available (and therefore we need to failover to a
// replica) and a valid error from a healthy node. In the case of
// the latter, there's no need to retry a replica, we should trust
// the error from the healthy node and return that immediately.
// TODO(jaffee) retries should contact MDS and find out who is up and has access to shards needed
results, err := o.remoteExec(ctx, node.Address, index, &pql.Query{Calls: []*pql.Call{c}}, shards, embeddedRowsForNode)
if len(results) > 0 {
resp.result = results[0]
}
resp.err = err
}
// Return response to the channel.
select {
case <-done:
// If someone just canceled the context
// arbitrarily, we could end up here with this
// being the first non-nil error handed to
// the ErrGroup, in which case, it's the best
// explanation we have for why everything's
// stopping.
return ctx.Err()
case ch <- resp:
return nil
}
})
if reterr != nil {
return reterr // exit early if error occurs when running serially
}
}
return nil
}
func (o *orchestrator) preTranslate(ctx context.Context, index string, calls ...*pql.Call) (cols map[string]map[string]uint64, rows map[string]map[string]map[string]uint64, err error) {
// Collect all of the required keys.
collector := keyCollector{
createCols: make(map[string][]string),
findCols: make(map[string][]string),
createRows: make(map[string]map[string][]string),
findRows: make(map[string]map[string][]string),
}
for _, call := range calls {
err := o.collectCallKeys(&collector, call, index)
if err != nil {
return nil, nil, err
}
}
// Create keys.
// Both rows and columns need to be created first because of foreign index keys.
cols = make(map[string]map[string]uint64)
rows = make(map[string]map[string]map[string]uint64)
for index, keys := range collector.createCols {
translations, err := o.trans.CreateIndexKeys(ctx, index, keys)
if err != nil {
return nil, nil, errors.Wrap(err, "creating query column keys")
}
cols[index] = translations
}
for index, fields := range collector.createRows {
idxRows := make(map[string]map[string]uint64)
for field, keys := range fields {
translations, err := o.trans.CreateFieldKeys(ctx, index, field, keys)
if err != nil {
return nil, nil, errors.Wrap(err, "creating query row keys")
}
idxRows[field] = translations
}
rows[index] = idxRows
}
// Find other keys.
for index, keys := range collector.findCols {
translations, err := o.trans.FindIndexKeys(ctx, index, keys)
if err != nil {
return nil, nil, errors.Wrap(err, "finding query column keys")
}
if prev := cols[index]; prev != nil {
for key, id := range translations {
prev[key] = id
}
} else {
cols[index] = translations
}
}
for index, fields := range collector.findRows {
idxRows := rows[index]
if idxRows == nil {
idxRows = make(map[string]map[string]uint64)
rows[index] = idxRows
}
for field, keys := range fields {
translations, err := o.trans.FindFieldKeys(ctx, index, field, keys)
if err != nil {
return nil, nil, errors.Wrap(err, "finding query row keys")
}
if prev := idxRows[field]; prev != nil {
for key, id := range translations {
prev[key] = id
}
} else {
idxRows[field] = translations
}
}
}
return cols, rows, nil
}
func (o *orchestrator) collectCallKeys(dst *keyCollector, c *pql.Call, index string) error {
// Check for an overriding 'index' argument.
// This also applies to all child calls.
if callIndex := c.CallIndex(); callIndex != "" {
index = callIndex
}
// Handle the field arg.
switch c.Name {
case "Set":
if field, err := c.FieldArg(); err == nil {
if arg, ok := c.Args[field].(string); ok {
dst.CreateRows(index, field, arg)
}
}
// TODO: will have to consider how to support Store... creating the field if it doesn't exist will not be a thing though.
case "Store":
return errors.New(errors.ErrUncoded, "Store query currently unsupported")
case "Clear", "Row", "Range", "ClearRow":
if field, err := c.FieldArg(); err == nil {
switch arg := c.Args[field].(type) {
case string:
dst.FindRows(index, field, arg)
case *pql.Condition:
// This is a workaround to allow `==` and `!=` to work on foreign index fields.
if key, ok := arg.Value.(string); ok {
switch arg.Op {
case pql.EQ, pql.NEQ:
dst.FindRows(index, field, key)
default:
return errors.Errorf("operator %v not defined on strings", arg.Op)
}
}
}
}
}
// Handle _col.
if col, ok := c.Args["_col"].(string); ok {
switch c.Name {
case "Set":
dst.CreateColumns(index, col)
default:
dst.FindColumns(index, col)
}
}
// Handle _row.
if row, ok := c.Args["_row"].(string); ok {
// Find the field.
field, ok, err := c.StringArg("_field")
if err != nil {
return errors.Wrap(err, "finding field")
}
if !ok {
return errors.Wrap(ErrFieldNotFound, "finding field for _row argument")
}
dst.FindRows(index, field, row)
}
// Handle queries that need a "column" argument.
switch c.Name {
case "Rows", "GroupBy", "FieldValue", "IncludesColumn":
if col, ok := c.Args["column"].(string); ok {
dst.FindColumns(index, col)
}
}
// Handle special per-query arguments.
switch c.Name {
case "ConstRow":
// Translate the columns list.
if cols, ok := c.Args["columns"].([]interface{}); ok {
keys := make([]string, 0, len(cols))
for _, v := range cols {
switch v := v.(type) {
case string:
keys = append(keys, v)
case uint64:
case int64:
default:
return errors.Errorf("invalid column identifier %v of type %T", c, c)
}
}
dst.FindColumns(index, keys...)
}
case "Rows":
// Find the field.
var field string
if f, ok1, err := c.StringArg("_field"); err != nil {
return errors.Wrap(err, "finding _field for Rows previous translation")
} else if ok1 {
field = f
} else if f, ok2, err := c.StringArg("field"); err != nil {
return errors.Wrap(err, "finding field for Rows previous translation")
} else if ok2 {
field = f
} else {
return errors.New(errors.ErrUncoded, "missing field in Rows call")
}
if prev, ok := c.Args["previous"].(string); ok {
dst.FindRows(index, field, prev)
}
if in, ok := c.Args["in"]; ok {
inIn, ok := in.([]interface{})
if !ok {
return errors.Errorf("unexpected type for argument 'in' %v of %[1]T", inIn)
}
inStrs := make([]string, 0)
for _, v := range inIn {
if vstr, ok := v.(string); ok {
inStrs = append(inStrs, vstr)
}
}
dst.FindRows(index, field, inStrs...)
}
}
// Collect keys from child calls.
for _, child := range c.Children {
err := o.collectCallKeys(dst, child, index)
if err != nil {
return err
}
}
// Collect keys from argument calls.
for _, arg := range c.Args {
argCall, ok := arg.(*pql.Call)
if !ok {
continue
}
err := o.collectCallKeys(dst, argCall, index)
if err != nil {
return err
}
}
return nil
}
type keyCollector struct {
createCols, findCols map[string][]string // map[index] -> column keys
createRows, findRows map[string]map[string][]string // map[index]map[field] -> row keys
}
func (c *keyCollector) CreateColumns(index string, columns ...string) {
if len(columns) == 0 {
return
}
c.createCols[index] = append(c.createCols[index], columns...)
}
func (c *keyCollector) FindColumns(index string, columns ...string) {
if len(columns) == 0 {
return
}
c.findCols[index] = append(c.findCols[index], columns...)
}
func (c *keyCollector) CreateRows(index string, field string, columns ...string) {
if len(columns) == 0 {
return
}
idx := c.createRows[index]
if idx == nil {
idx = make(map[string][]string)
c.createRows[index] = idx
}
idx[field] = append(idx[field], columns...)
}
func (c *keyCollector) FindRows(index string, field string, columns ...string) {
if len(columns) == 0 {
return
}
idx := c.findRows[index]
if idx == nil {
idx = make(map[string][]string)
c.findRows[index] = idx
}
idx[field] = append(idx[field], columns...)
}
func fieldValidateValue(f *featurebase.FieldInfo, val interface{}) error {
if val == nil {
return nil
}
// Validate special types.
switch val := val.(type) {
case string:
if !f.Options.Keys {
return errors.Errorf("string value on unkeyed field %q", f.Name)
}
return nil
case *pql.Condition:
switch v := val.Value.(type) {
case nil:
case string:
case uint64:
case int64:
case float64:
case pql.Decimal:
case time.Time:
case []interface{}:
for _, v := range v {
if err := fieldValidateValue(f, v); err != nil {
return err
}
}
return nil
default:
return errors.Errorf("invalid value %v in condition %q", v, val.String())
}
return fieldValidateValue(f, val.Value)
}
switch f.Options.Type {
case FieldTypeSet, FieldTypeMutex, FieldTypeTime:
switch v := val.(type) {
case uint64:
case int64:
if v < 0 {
return errors.Errorf("negative ID %d for set field %q", v, f.Name)
}
default:
return errors.Errorf("invalid value %v for field %q of type %s", v, f.Name, f.Options.Type)
}
if f.Options.Keys {
return errors.Errorf("found integer ID %d on keyed field %q", val, f.Name)
}
case FieldTypeBool:
switch v := val.(type) {
case bool:
default:
return errors.Errorf("invalid value %v for bool field %q", v, f.Name)
}
case FieldTypeInt:
switch v := val.(type) {
case uint64:
if v > 1<<63 {
return errors.Errorf("oversized integer %d for int field %q (range: -2^63 to 2^63-1)", v, f.Name)
}
case int64:
default:
return errors.Errorf("invalid value %v for int field %q", v, f.Name)
}
case FieldTypeDecimal:
switch v := val.(type) {
case uint64:
case int64:
case float64:
case pql.Decimal:
default:
return errors.Errorf("invalid value %v for decimal field %q", v, f.Name)
}
case FieldTypeTimestamp:
switch v := val.(type) {
case time.Time:
default:
return errors.Errorf("invalid value %v for timestamp field %q", v, f.Name)
}
default:
return errors.Errorf("unsupported type %s of field %q", f.Options.Type, f.Name)
}
return nil
}
func (o *orchestrator) translateCall(ctx context.Context, c *pql.Call, index string, columnKeys map[string]map[string]uint64, rowKeys map[string]map[string]map[string]uint64) (*pql.Call, error) {
// Check for an overriding 'index' argument.
// This also applies to all child calls.
if callIndex := c.CallIndex(); callIndex != "" {
index = callIndex
}
idx, err := o.schema.IndexInfo(ctx, index)
if err != nil {
return nil, errors.Wrapf(err, "translating query on index %q", index)
}
// Fetch the column keys list for this index.
indexCols, indexRows := columnKeys[index], rowKeys[index]
// Handle the field arg.
switch c.Name {
case "Set", "Store":
if field, err := c.FieldArg(); err == nil {
f, err := o.schema.FieldInfo(ctx, index, field)
if err != nil {
return nil, errors.Wrapf(err, "validating value for field %q", field)
}
arg := c.Args[field]
if err := fieldValidateValue(f, arg); err != nil {
return nil, errors.Wrap(err, "validating store value")
}
switch arg := arg.(type) {
case string:
if translation, ok := indexRows[field][arg]; ok {
c.Args[field] = translation
} else {
return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "destination key not found %q in %q in index %q", arg, field, index)
}
case bool:
if arg {
c.Args[field] = trueRowID
} else {
c.Args[field] = falseRowID
}
}
}
case "Clear", "Row", "Range", "ClearRow":
if field, err := c.FieldArg(); err == nil {
f, err := o.schema.FieldInfo(ctx, index, field)
if err != nil {
return nil, errors.Wrapf(err, "validating value for field %q", field)
}
arg := c.Args[field]
if err := fieldValidateValue(f, arg); err != nil {
return nil, errors.Wrap(err, "validating field parameter value")
}
if c.Name == "Row" {
switch f.Options.Type {
case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp:
if _, ok := arg.(*pql.Condition); !ok {
// This is workaround to support pql.ASSIGN ('=') as condition ('==') for BSI fields.
arg = &pql.Condition{
Op: pql.EQ,
Value: arg,
}
c.Args[field] = arg
}
}
}
switch arg := arg.(type) {
case string:
if translation, ok := indexRows[field][arg]; ok {
c.Args[field] = translation
} else {
// Rewrite the call into a zero value call.
return o.callZero(c), nil
}
case bool:
if arg {
c.Args[field] = trueRowID
} else {
c.Args[field] = falseRowID
}
case *pql.Condition:
// This is a workaround to allow `==` and `!=` to work on foreign index fields.
if key, ok := arg.Value.(string); ok {
switch arg.Op {
case pql.EQ, pql.NEQ:
if translation, ok := indexRows[field][key]; ok {
arg.Value = translation
} else {
// Rewrite the call into a zero value call.
return o.callZero(c), nil
}
default:
return nil, errors.Errorf("operator %v not defined on strings", arg.Op)
}
}
}
}
}
// Handle _col.
if col, ok := c.Args["_col"].(string); ok {
if !idx.Options.Keys {
return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "translating column on unkeyed index %q", index)
}
if id, ok := indexCols[col]; ok {
c.Args["_col"] = id
} else {
switch c.Name {
case "Set":
return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "destination key not found %q in index %q", col, index)
default:
return o.callZero(c), nil
}
}
}
// Handle _row.
if row, ok := c.Args["_row"]; ok {
// Find the field.
var field string
if f, ok1, err := c.StringArg("_field"); err != nil {
return nil, errors.Wrap(err, "finding _field")
} else if ok1 {
field = f
} else if f, ok2, err := c.StringArg("field"); err != nil {
return nil, errors.Wrap(err, "finding field")
} else if ok2 {
field = f
} else {
return nil, errors.New(errors.ErrUncoded, "missing field")
}
f, err := o.schema.FieldInfo(ctx, index, field)
if err != nil {
return nil, errors.Wrapf(err, "validating value for field %q", field)
}
if err := fieldValidateValue(f, row); err != nil {
return nil, errors.Wrap(err, "validating row value")
}
switch row := row.(type) {
case string:
if translation, ok := indexRows[field][row]; ok {
c.Args["_row"] = translation
} else {
return o.callZero(c), nil
}
}
}
// Handle queries that need a "column" argument.
switch c.Name {
case "Rows", "GroupBy", "FieldValue", "IncludesColumn":
if col, ok := c.Args["column"].(string); ok {
if translation, ok := indexCols[col]; ok {
c.Args["column"] = translation
} else {
// Rewrite the call into a zero value call.
return o.callZero(c), nil
}
}
}
// Handle special per-query arguments.
switch c.Name {
case "ConstRow":
// Translate the columns list.
if cols, ok := c.Args["columns"].([]interface{}); ok {
out := make([]uint64, 0, len(cols))
for _, v := range cols {
switch v := v.(type) {
case string:
if id, ok := indexCols[v]; ok {
out = append(out, id)
}
case uint64:
out = append(out, v)
case int64:
out = append(out, uint64(v))
default:
return nil, errors.Errorf("invalid column identifier %v of type %T", c, c)
}
}
c.Args["columns"] = out
}
case "Rows":
// Find the field.
var field string
if f, ok1, err := c.StringArg("_field"); err != nil {
return nil, errors.Wrap(err, "finding _field for Rows previous translation")
} else if ok1 {
field = f
} else if f, ok2, err := c.StringArg("field"); err != nil {
return nil, errors.Wrap(err, "finding field for Rows previous translation")
} else if ok2 {
field = f
} else {
return nil, errors.New(errors.ErrUncoded, "missing field in Rows call")
}
// Translate the previous row key.
if prev, ok := c.Args["previous"]; ok {
// Validate the type.
f, err := o.schema.FieldInfo(ctx, index, field)
if err != nil {
return nil, errors.Wrapf(err, "validating value for field %q", field)
}
if err := fieldValidateValue(f, prev); err != nil {
return nil, errors.Wrap(err, "validating prev value")
}
switch prev := prev.(type) {
case string:
// Look up a translation for the previous row key.
if translation, ok := indexRows[field][prev]; ok {
c.Args["previous"] = translation
} else {
return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "translating previous key %q from field %q in index %q in Rows call", prev, field, index)
}
case bool:
if prev {
c.Args["previous"] = trueRowID
} else {
c.Args["previous"] = falseRowID
}
}
}
// Check if "like" argument is applied to keyed fields.
if _, found := c.Args["like"].(string); found {
fieldName, err := c.FirstStringArg("_field", "field")
if err != nil || fieldName == "" {
return nil, fmt.Errorf("cannot read field name for Rows call")
}
if f, err := o.schema.FieldInfo(ctx, index, fieldName); err != nil {
return nil, errors.Wrapf(err, "getting field %q", fieldName)
} else if !f.Options.Keys {
return nil, fmt.Errorf("'%s' is not a set/mutex/time field with a string key", fieldName)
}
}
if in, ok := c.Args["in"]; ok {
inIn, ok := in.([]interface{})
if !ok {
return nil, errors.Errorf("unexpected type for argument 'in' %v of %[1]T", in)
}
inIDs := make([]interface{}, 0, len(inIn))
for _, inVal := range inIn {
if inStr, ok := inVal.(string); ok {
id, found := rowKeys[index][field][inStr]
if found {
inIDs = append(inIDs, id)
}
} else {
inIDs = append(inIDs, inVal)
}
}
c.Args["in"] = inIDs
}
}
// Translate child calls.
for i, child := range c.Children {
translated, err := o.translateCall(ctx, child, index, columnKeys, rowKeys)
if err != nil {
return nil, err
}
c.Children[i] = translated
}
// Translate argument calls.
for k, arg := range c.Args {
argCall, ok := arg.(*pql.Call)
if !ok {
continue
}
translated, err := o.translateCall(ctx, argCall, index, columnKeys, rowKeys)
if err != nil {
return nil, err
}
c.Args[k] = translated
}
return c, nil
}
func (o *orchestrator) callZero(c *pql.Call) *pql.Call {
switch c.Name {
case "Row", "Range":
if field, err := c.FieldArg(); err == nil {
if cond, ok := c.Args[field].(*pql.Condition); ok {
if cond.Op == pql.NEQ {
// Turn not nothing into everything.
return &pql.Call{Name: "All"}
}
}
}
// Use an empty union as a placeholder.
return &pql.Call{Name: "Union"}
default:
return nil
}
}
func (o *orchestrator) translateResults(ctx context.Context, index string, idx *featurebase.IndexInfo, calls []*pql.Call, results []interface{}, memoryAvailable int64) (err error) {
span, _ := tracing.StartSpanFromContext(ctx, "Executor.translateResults")
defer span.Finish()
idMap := make(map[uint64]string)
if idx.Options.Keys {
// Collect all index ids.
idSet := make(map[uint64]struct{})
for i := range calls {
if err := o.collectResultIDs(ctx, idx, calls[i], results[i], idSet); err != nil {
return err
}
}
if idMap, err = o.trans.TranslateIndexIDSet(ctx, index, idSet); err != nil {
return err
}
}
for i := range results {
results[i], err = o.translateResult(ctx, idx, calls[i], results[i], idMap)
if err != nil {
return err
}
}
return nil
}
// translationStrategy denotes the several different ways the bits in
// a *Row could be translated to string keys.
type translationStrategy int
const (
// byCurrentIndex means to interpret the bits as IDs in "top
// level" index for this query (e.g. the index specified in the
// path of the HTTP request).
byCurrentIndex translationStrategy = iota + 1
// byRowField means that the bits in this *Row are row IDs which
// should be translated using the field's (*Row.Field) translation store.
byRowField
// byRowFieldForeignIndex means that the bits in this *Row should
// be interpreted as IDs in the foreign index of the *Row.Field.
byRowFieldForeignIndex
// byRowIndex means the bits in this *Row should be translated
// according to the index named by *Row.Index
byRowIndex
// noTranslation means the bits should not be translated to string
// keys.
noTranslation
)
// howToTranslate determines how a *Row object's bits should be
// translated to keys (if at all). There are several different options
// detailed by the various const values of translationStrategy. In
// order to do this it has to figure out the row's index and field
// which it also returns as the caller may need them to actually
// execute the translation or do whatever else it's doing with the
// translationStrategy information.
func (o *orchestrator) howToTranslate(ctx context.Context, idx *featurebase.IndexInfo, row *featurebase.Row) (rowIdx *featurebase.IndexInfo, rowField *featurebase.FieldInfo, strat translationStrategy, err error) {
// First get the index and field the row specifies (if any).
rowIdx = idx
if row.Index != "" && row.Index != idx.Name {
rowIdx, err = o.schema.IndexInfo(ctx, row.Index)
if err != nil {
return nil, nil, 0, errors.Wrapf(err, "got a row with unknown index: %s", row.Index)
}
}
if row.Field != "" {
rowField, err = o.schema.FieldInfo(ctx, row.Index, row.Field)
if err != nil {
return nil, nil, 0, errors.Wrapf(err, "got a row with unknown index/field %s/%s", idx.Name, row.Field)
}
}
// Handle the case where the Row has specified a field.
if rowField != nil {
// Handle the case where field has a foreign index.
if rowField.Options.ForeignIndex != "" {
fidx, err := o.schema.IndexInfo(ctx, rowField.Options.ForeignIndex)
if err != nil {
return nil, nil, 0, errors.Errorf("foreign index %s not found for field %s in index %s", rowField.Options.ForeignIndex, rowField.Name, rowIdx.Name)
}
if fidx.Options.Keys {
return rowIdx, rowField, byRowFieldForeignIndex, nil
}
} else if rowField.Options.Keys {
return rowIdx, rowField, byRowField, nil
}
return rowIdx, rowField, noTranslation, nil
}
// In this case, the row has specified an index, but not a field,
// so we translate according to that index.
if rowIdx != idx && rowIdx.Options.Keys {
return rowIdx, rowField, byRowIndex, nil
}
// Handle the normal case (row represents a set of records in
// the top level index, Row has not specifed a different index
// or field).
if rowIdx == idx && idx.Options.Keys && rowField == nil {
return rowIdx, rowField, byCurrentIndex, nil
}
return rowIdx, rowField, noTranslation, nil
}
func (o *orchestrator) collectResultIDs(ctx context.Context, idx *featurebase.IndexInfo, call *pql.Call, result interface{}, idSet map[uint64]struct{}) error {
switch result := result.(type) {
case *featurebase.Row:
// Only collect result IDs if they are in the current index.
_, _, strategy, err := o.howToTranslate(ctx, idx, result)
if err != nil {
return errors.Wrap(err, "determining how to translate")
}
if strategy == byCurrentIndex {
for _, segment := range result.Segments {
for _, col := range segment.Columns() {
idSet[col] = struct{}{}
}
}
}
case featurebase.ExtractedIDMatrix:
for _, col := range result.Columns {
idSet[col.ColumnID] = struct{}{}
}
}
return nil
}
// preTranslateMatrixSet translates the IDs of a set field in an extracted matrix.
func (o *orchestrator) preTranslateMatrixSet(ctx context.Context, mat featurebase.ExtractedIDMatrix, fieldIdx uint, index, field string) (map[uint64]string, error) {
ids := make(map[uint64]struct{}, len(mat.Columns))
for _, col := range mat.Columns {
for _, v := range col.Rows[fieldIdx] {
ids[v] = struct{}{}
}
}
return o.trans.TranslateFieldIDs(ctx, index, field, ids)
}
func (o *orchestrator) translateResult(ctx context.Context, idx *featurebase.IndexInfo, call *pql.Call, result interface{}, idSet map[uint64]string) (_ interface{}, err error) {
switch result := result.(type) {
case *featurebase.Row:
rowIdx, rowField, strategy, err := o.howToTranslate(ctx, idx, result)
if err != nil {
return nil, errors.Wrap(err, "determining translation strategy")
}
switch strategy {
case byCurrentIndex:
other := &featurebase.Row{}
for _, segment := range result.Segments {
for _, col := range segment.Columns() {
other.Keys = append(other.Keys, idSet[col])
}
}
return other, nil
case byRowField:
keys, err := o.trans.TranslateFieldListIDs(ctx, rowIdx.Name, rowField.Name, result.Columns())
if err != nil {
return nil, errors.Wrap(err, "translating Row to field keys")
}
result.Keys = keys
case byRowFieldForeignIndex:
idx, err = o.schema.IndexInfo(ctx, rowField.Options.ForeignIndex)
if err != nil {
return nil, errors.Wrapf(err, "foreign index %s not found for field %s in index %s", rowField.Options.ForeignIndex, rowField.Name, rowIdx.Name)
}
for _, segment := range result.Segments {
keys, err := o.trans.TranslateIndexIDs(ctx, rowField.Options.ForeignIndex, segment.Columns())
if err != nil {
return nil, errors.Wrap(err, "translating index ids")
}
result.Keys = append(result.Keys, keys...)
}
case byRowIndex:
for _, segment := range result.Segments {
keys, err := o.trans.TranslateIndexIDs(ctx, rowIdx.Name, segment.Columns())
if err != nil {
return nil, errors.Wrap(err, "translating index ids")
}
result.Keys = append(result.Keys, keys...)
}
return result, nil
case noTranslation:
return result, nil
default:
return nil, errors.Errorf("unknown translation strategy %d", strategy)
}
case featurebase.SignedRow:
sr, err := func() (*featurebase.SignedRow, error) {
fieldName := callArgString(call, "field")
if fieldName == "" {
return nil, nil
}
field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName)
if err != nil {
return nil, nil
}
if field.Options.Keys {
rslt := result.Pos
if rslt == nil {
return &featurebase.SignedRow{Pos: &featurebase.Row{}}, nil
}
other := &featurebase.Row{}
for _, segment := range rslt.Segments {
keys, err := o.trans.TranslateIndexIDs(ctx, field.Options.ForeignIndex, segment.Columns())
if err != nil {
return nil, errors.Wrap(err, "translating index ids")
}
other.Keys = append(other.Keys, keys...)
}
return &featurebase.SignedRow{Pos: other}, nil
}
return nil, nil
}()
if err != nil {
return nil, err
} else if sr != nil {
return *sr, nil
}
case featurebase.PairField:
if fieldName := callArgString(call, "field"); fieldName != "" {
field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName)
if err != nil {
return nil, fmt.Errorf("field %q not found", fieldName)
}
if field.Options.Keys {
// TODO(jaffee) get index name from call? CallIndex? (not just here)
keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, fieldName, []uint64{result.Pair.ID})
if err != nil {
return nil, err
}
key := keys[0]
if call.Name == "MinRow" || call.Name == "MaxRow" {
result.Pair.Key = key
return result, nil
}
return featurebase.PairField{
Pair: featurebase.Pair{Key: key, Count: result.Pair.Count},
Field: fieldName,
}, nil
}
}
case *featurebase.PairsField:
if fieldName := callArgString(call, "_field"); fieldName != "" {
field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName)
if err != nil {
return nil, errors.Wrapf(err, "field '%q'", fieldName)
}
if field.Options.Keys {
ids := make([]uint64, len(result.Pairs))
for i := range result.Pairs {
ids[i] = result.Pairs[i].ID
}
keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, fieldName, ids)
if err != nil {
return nil, err
}
other := make([]featurebase.Pair, len(result.Pairs))
for i := range result.Pairs {
other[i] = featurebase.Pair{Key: keys[i], Count: result.Pairs[i].Count}
}
return &featurebase.PairsField{
Pairs: other,
Field: fieldName,
}, nil
}
}
case *featurebase.GroupCounts:
fieldIDs := make(map[*featurebase.FieldInfo]map[uint64]struct{})
foreignIDs := make(map[*featurebase.FieldInfo]map[uint64]struct{})
groups := result.Groups()
for _, gl := range groups {
for _, g := range gl.Group {
field, err := o.schema.FieldInfo(ctx, idx.Name, g.Field)
if err != nil {
return nil, errors.Wrapf(err, "getting field '%q", g.Field)
}
if field.Options.Keys {
if g.Value != nil {
if fi := field.Options.ForeignIndex; fi != "" {
m, ok := foreignIDs[field]
if !ok {
m = make(map[uint64]struct{}, len(groups))
foreignIDs[field] = m
}
m[uint64(*g.Value)] = struct{}{}
continue
}
}
m, ok := fieldIDs[field]
if !ok {
m = make(map[uint64]struct{}, len(groups))
fieldIDs[field] = m
}
m[g.RowID] = struct{}{}
}
}
}
fieldTranslations := make(map[string]map[uint64]string)
for field, ids := range fieldIDs {
trans, err := o.trans.TranslateFieldIDs(ctx, idx.Name, field.Name, ids)
if err != nil {
return nil, errors.Wrapf(err, "translating IDs in field '%q'", field.Name)
}
fieldTranslations[field.Name] = trans
}
foreignTranslations := make(map[string]map[uint64]string)
for field, ids := range foreignIDs {
trans, err := o.trans.TranslateIndexIDSet(ctx, field.Options.ForeignIndex, ids)
if err != nil {
return nil, errors.Wrapf(err, "translating foreign IDs from index %q", field.Options.ForeignIndex)
}
foreignTranslations[field.Name] = trans
}
// We are reluctant to smash result, and I'm not sure we need
// to be but I'm not sure we don't need to be.
newGroups := make([]featurebase.GroupCount, len(groups))
copy(newGroups, groups)
for gi, gl := range groups {
group := make([]featurebase.FieldRow, len(gl.Group))
for i, g := range gl.Group {
if ft, ok := fieldTranslations[g.Field]; ok {
g.RowKey = ft[g.RowID]
} else if ft, ok := foreignTranslations[g.Field]; ok && g.Value != nil {
g.RowKey = ft[uint64(*g.Value)]
g.Value = nil
}
group[i] = g
}
// Replace with translated group.
newGroups[gi].Group = group
}
if result != nil {
return featurebase.NewGroupCounts(result.AggregateColumn(), newGroups...), nil
}
return &featurebase.GroupCounts{}, nil
case featurebase.RowIDs:
fieldName := callArgString(call, "_field")
if fieldName == "" {
return nil, ErrFieldNotFound
}
other := featurebase.RowIdentifiers{
Field: fieldName,
}
if field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName); err != nil {
return nil, errors.Wrapf(err, "'%q'", fieldName)
} else if field.Options.Keys {
keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, field.Name, result)
if err != nil {
return nil, errors.Wrap(err, "translating row IDs")
}
other.Keys = keys
} else {
other.Rows = result
}
return other, nil
case featurebase.ExtractedIDMatrix:
type fieldMapper = func([]uint64) (_ interface{}, err error)
fields := make([]featurebase.ExtractedTableField, len(result.Fields))
mappers := make([]fieldMapper, len(result.Fields))
for i, v := range result.Fields {
field, err := o.schema.FieldInfo(ctx, idx.Name, v)
if err != nil {
return nil, errors.Wrapf(err, "'%q'", v)
}
var mapper fieldMapper
var datatype string
switch typ := field.Options.Type; typ {
case FieldTypeBool:
datatype = "bool"
mapper = func(ids []uint64) (_ interface{}, err error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
switch ids[0] {
case 0:
return false, nil
case 1:
return true, nil
default:
return nil, errors.Errorf("invalid ID for boolean %q: %d", field.Name, ids[0])
}
default:
return nil, errors.Errorf("boolean %q has too many values: %v", field.Name, ids)
}
}
case FieldTypeSet, FieldTypeTime:
if field.Options.Keys {
datatype = "[]string"
translations, err := o.preTranslateMatrixSet(ctx, result, uint(i), idx.Name, field.Name)
if err != nil {
return nil, errors.Wrapf(err, "translating IDs of field %q", v)
}
mapper = func(ids []uint64) (interface{}, error) {
keys := make([]string, len(ids))
for i, id := range ids {
keys[i] = translations[id]
}
return keys, nil
}
} else {
datatype = "[]uint64"
mapper = func(ids []uint64) (interface{}, error) {
if ids == nil {
ids = []uint64{}
}
return ids, nil
}
}
case FieldTypeMutex:
if field.Options.Keys {
datatype = "string"
translations, err := o.preTranslateMatrixSet(ctx, result, uint(i), idx.Name, field.Name)
if err != nil {
return nil, errors.Wrapf(err, "translating IDs of field %q", v)
}
mapper = func(ids []uint64) (interface{}, error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return translations[ids[0]], nil
default:
return nil, errors.Errorf("mutex %q has too many values: %v", field.Name, ids)
}
}
} else {
datatype = "uint64"
mapper = func(ids []uint64) (_ interface{}, err error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return ids[0], nil
default:
return nil, errors.Errorf("mutex %q has too many values: %v", field.Name, ids)
}
}
}
case FieldTypeInt:
if fi := field.Options.ForeignIndex; fi != "" {
if field.Options.Keys {
datatype = "string"
ids := make(map[uint64]struct{}, len(result.Columns))
for _, col := range result.Columns {
for _, v := range col.Rows[i] {
ids[v] = struct{}{}
}
}
trans, err := o.trans.TranslateIndexIDSet(ctx, field.Options.ForeignIndex, ids)
if err != nil {
return nil, errors.Wrapf(err, "translating foreign IDs from index %q", field.Options.ForeignIndex)
}
mapper = func(ids []uint64) (interface{}, error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return trans[ids[0]], nil
default:
return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids)
}
}
} else {
datatype = "uint64"
mapper = func(ids []uint64) (interface{}, error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return ids[0], nil
default:
return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids)
}
}
}
} else {
datatype = "int64"
mapper = func(ids []uint64) (interface{}, error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return int64(ids[0]), nil
default:
return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids)
}
}
}
case FieldTypeDecimal:
datatype = "decimal"
scale := field.Options.Scale
mapper = func(ids []uint64) (_ interface{}, err error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return pql.NewDecimal(int64(ids[0]), scale), nil
default:
return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids)
}
}
case FieldTypeTimestamp:
datatype = "timestamp"
mapper = func(ids []uint64) (_ interface{}, err error) {
switch len(ids) {
case 0:
return nil, nil
case 1:
return time.Unix(0, int64(ids[0])*int64(featurebase.TimeUnitNanos(field.Options.TimeUnit))).UTC(), nil
default:
return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids)
}
}
default:
return nil, errors.Errorf("field type %q not yet supported", typ)
}
mappers[i] = mapper
fields[i] = featurebase.ExtractedTableField{
Name: v,
Type: datatype,
}
}
var translateCol func(uint64) (featurebase.KeyOrID, error)
if idx.Options.Keys {
translateCol = func(id uint64) (featurebase.KeyOrID, error) {
return featurebase.KeyOrID{Keyed: true, Key: idSet[id]}, nil
}
} else {
translateCol = func(id uint64) (featurebase.KeyOrID, error) {
return featurebase.KeyOrID{ID: id}, nil
}
}
cols := make([]featurebase.ExtractedTableColumn, len(result.Columns))
colData := make([]interface{}, len(cols)*len(result.Fields))
for i, col := range result.Columns {
data := colData[i*len(result.Fields) : (i+1)*len(result.Fields) : (i+1)*len(result.Fields)]
for j, rows := range col.Rows {
v, err := mappers[j](rows)
if err != nil {
return nil, errors.Wrap(err, "translating extracted table value")
}
data[j] = v
}
colTrans, err := translateCol(col.ColumnID)
if err != nil {
return nil, errors.Wrap(err, "translating column ID in extracted table")
}
cols[i] = featurebase.ExtractedTableColumn{
Column: colTrans,
Rows: data,
}
}
return featurebase.ExtractedTable{
Fields: fields,
Columns: cols,
}, nil
}
return result, nil
}
// validateQueryContext returns a query-appropriate error if the context is done.
func validateQueryContext(ctx context.Context) error {
select {
case <-ctx.Done():
switch err := ctx.Err(); err {
case context.Canceled:
return featurebase.ErrQueryCancelled
case context.DeadlineExceeded:
return featurebase.ErrQueryTimeout
default:
return err
}
default:
return nil
}
}
type reduceFunc func(ctx context.Context, prev, v interface{}) interface{}
type mapResponse struct {
node dax.Address
shards []uint64
result interface{}
err error
}
func callArgString(call *pql.Call, key string) string {
value, ok := call.Args[key]
if !ok {
return ""
}
s, _ := value.(string)
return s
}
type qualifiedOrchestrator struct {
*orchestrator
qual dax.TableQualifier
schemar schemar.Schemar
}
func newQualifiedOrchestrator(orch *orchestrator, qual dax.TableQualifier, schemar schemar.Schemar) *qualifiedOrchestrator {
return &qualifiedOrchestrator{
orchestrator: orch,
qual: qual,
schemar: schemar,
}
}
func (o *qualifiedOrchestrator) Execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) (featurebase.QueryResponse, error) {
resp := featurebase.QueryResponse{}
tkey, err := o.indexToQualifiedTableKey(ctx, index)
if err != nil {
return resp, errors.Wrap(err, "converting index to qualified table key")
}
return o.orchestrator.Execute(ctx, string(tkey), q, shards, opt)
}
func (o *qualifiedOrchestrator) indexToQualifiedTableKey(ctx context.Context, index string) (dax.TableKey, error) {
qtid, err := o.schemar.TableID(ctx, o.qual, dax.TableName(index))
if err != nil {
return "", errors.Wrap(err, "converting index to qualified table id")
}
return qtid.Key(), nil
}