// Copyright 2021 Molecula Corp. All rights reserved. package queryer import ( "context" "fmt" "math" "sort" "strings" "time" featurebase "github.com/featurebasedb/featurebase/v3" "github.com/featurebasedb/featurebase/v3/dax" "github.com/featurebasedb/featurebase/v3/errors" "github.com/featurebasedb/featurebase/v3/logger" "github.com/featurebasedb/featurebase/v3/pql" "github.com/featurebasedb/featurebase/v3/tracing" "github.com/prometheus/client_golang/prometheus" "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) ([]dax.ComputeNode, error) } type ServerlessTopology struct { controller dax.Controller } func (m *ServerlessTopology) ComputeNodes(ctx context.Context, index string, shards []uint64) ([]dax.ComputeNode, error) { var daxShards = make(dax.ShardNums, len(shards)) for i, s := range shards { daxShards[i] = dax.ShardNum(s) } // TODO(tlt): this needs review; ServerlessTopology is converting from // string/uint64 to qtid/shardNum?? Perhaps we can get rid of the Topologer // interface altogether and replace it with dax.Noder. qtid := dax.TableKey(index).QualifiedTableID() return m.controller.ComputeNodes(ctx, qtid, daxShards...) } // Translator serves the translation portion of a query request. // 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 Controller 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, tableKeyer dax.TableKeyer, 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.SchemaAPI topology Topologer trans Translator // Client used for remote requests. client *featurebase.InternalClient 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, tableKeyer dax.TableKeyer, 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{} qtbl, ok := tableKeyer.(*dax.QualifiedTable) if !ok { return resp, errors.New(errors.ErrUncoded, "orchestrator.Execute expects a dax.QualifiedTable") } // Check for query cancellation. if err := validateQueryContext(ctx); err != nil { return resp, err } // Default options. if opt == nil { opt = &featurebase.ExecOptions{} } results, err := o.execute(ctx, tableKeyer, 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, qtbl, 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, tableKeyer dax.TableKeyer, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) ([]interface{}, error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.execute") defer span.Finish() index := string(tableKeyer.Key()) // 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, tableKeyer, 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, tableKeyer, child, shards, opt) if err != nil { return nil, err } } } else { err := o.handlePreCallChildren(ctx, tableKeyer, 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() newTableKeyer := dax.StringTableKeyer(newIndex) if newIndex != "" && newIndex != index { v, err = o.executeCall(ctx, newTableKeyer, call, nil, opt) } else { v, err = o.executeCall(ctx, tableKeyer, 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, tableKeyer dax.TableKeyer, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) error { index := string(tableKeyer.Key()) 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 tableKeyer = dax.StringTableKeyer(index) // we need to recompute shards, then shards = nil } if err := o.handlePreCallChildren(ctx, tableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer, call, shards, opt); err != nil { return err } } } return nil } // preprocessQuery expands any calls that need preprocessing. func (o *orchestrator) preprocessQuery(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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") } labels := prometheus.Labels{"index": string(tableKeyer.Key())} statFn := func(ctr *prometheus.CounterVec) { if !opt.Remote { ctr.With(labels).Inc() } } // Preprocess the query. c, err := o.preprocessQuery(ctx, tableKeyer, c, shards, opt) if err != nil { return nil, err } switch c.Name { case "Sum": statFn(featurebase.CounterQuerySumTotal) res, err := o.executeSum(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeSum") case "Min": statFn(featurebase.CounterQueryMinTotal) res, err := o.executeMin(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeMin") case "Max": statFn(featurebase.CounterQueryMaxTotal) res, err := o.executeMax(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeMax") case "MinRow": statFn(featurebase.CounterQueryMinRowTotal) res, err := o.executeMinRow(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeMinRow") case "MaxRow": statFn(featurebase.CounterQueryMaxRowTotal) res, err := o.executeMaxRow(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeMaxRow") // case "Clear": // statFn(featurebase.CounterQueryClearTotal) // res, err := o.executeClearBit(ctx, index, c, opt) // return res, errors.Wrap(err, "executeClearBit") // case "ClearRow": // statFn(featurebase.CounterQueryClearRowTotal) // res, err := o.executeClearRow(ctx, index, c, shards, opt) // return res, errors.Wrap(err, "executeClearRow") case "Distinct": statFn(featurebase.CounterQueryDistinctTotal) res, err := o.executeDistinct(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeDistinct") // case "Store": // statFn(featurebase.CounterQueryStoreTotal) // res, err := o.executeSetRow(ctx, index, c, shards, opt) // return res, errors.Wrap(err, "executeSetRow") case "Count": statFn(featurebase.CounterQueryCountTotal) res, err := o.executeCount(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeCount") // case "Set": // statFn(featurebase.CounterQuerySetTotal) // res, err := o.executeSet(ctx, index, c, opt) // return res, errors.Wrap(err, "executeSet") case "TopK": statFn(featurebase.CounterQueryTopKTotal) res, err := o.executeTopK(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeTopK") case "TopN": statFn(featurebase.CounterQueryTopNTotal) res, err := o.executeTopN(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeTopN") case "Rows": statFn(featurebase.CounterQueryRowsTotal) res, err := o.executeRows(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeRows") case "Extract": statFn(featurebase.CounterQueryExtractTotal) res, err := o.executeExtract(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeExtract") case "GroupBy": statFn(featurebase.CounterQueryGroupByTotal) res, err := o.executeGroupBy(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeGroupBy") case "Options": statFn(featurebase.CounterQueryOptionsTotal) res, err := o.executeOptionsCall(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeOptionsCall") case "IncludesColumn": statFn(featurebase.CounterQueryIncludesColumnTotal) res, err := o.executeIncludesColumnCall(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeIncludesColumnCall") case "FieldValue": statFn(featurebase.CounterQueryFieldValueTotal) res, err := o.executeFieldValueCall(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeFieldValueCall") case "Precomputed": statFn(featurebase.CounterQueryPrecomputedTotal) res, err := o.executePrecomputedCall(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executePrecomputedCall") case "UnionRows": statFn(featurebase.CounterQueryUnionRowsTotal) res, err := o.executeUnionRows(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeUnionRows") case "ConstRow": statFn(featurebase.CounterQueryConstRowTotal) res, err := o.executeConstRow(ctx, tableKeyer, c) return res, errors.Wrap(err, "executeConstRow") case "Limit": statFn(featurebase.CounterQueryLimitTotal) res, err := o.executeLimitCall(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executeLimitCall") case "Percentile": statFn(featurebase.CounterQueryPercentileTotal) res, err := o.executePercentile(ctx, tableKeyer, c, shards, opt) return res, errors.Wrap(err, "executePercentile") // case "Delete": // statFn(featurebase.CounterQueryDeleteTotal) // 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. res, err := o.executeBitmapCall(ctx, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, c.Children[0], shards, optCopy) } // executeIncludesColumnCall executes an IncludesColumn() call. func (o *orchestrator) executeIncludesColumnCall(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, []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, tableKeyer dax.TableKeyer, 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, tableKeyer, []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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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.schemaFieldInfo(ctx, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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 } // executePercentile executes a Percentile() call. This logic is mirrored from // featurebase executor, but we should probably replace it with a smarter algorithm. func (o *orchestrator) executePercentile(ctx context.Context, tableKeyer dax.TableKeyer, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ interface{}, err error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executePercentile") defer span.Finish() // get nth var nthFloat float64 nthArg := c.Args["nth"] switch nthArg := nthArg.(type) { case pql.Decimal: nthFloat = nthArg.Float64() case int64: nthFloat = float64(nthArg) case nil: return nil, errors.New(errors.ErrUncoded, "Percentile(): nth required") default: return nil, 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 nil, 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 nil, errors.New(errors.ErrUncoded, "Percentile(): field required") } field, err := o.schemaFieldInfo(ctx, tableKeyer, fieldName) if err != nil { return nil, ErrFieldNotFound } // filter call for min & max var filterCall *pql.Call // We want to know the total number of values, so that when we check // for values X, we are also able to infer the number of values // equal to X. var totalCountCall *pql.Call // check if filter provided if filterArg, ok := c.Args["filter"].(*pql.Call); ok && filterArg != nil { // You could supply a filter like `Not(x=3)` which would yield values // which exist in the database but are null in this field, we don't // want that. filterCall = filterArg totalCountCall = &pql.Call{ Name: "Count", Children: []*pql.Call{ { Name: "Intersect", Children: []*pql.Call{ filterCall, { Name: "Row", Args: map[string]interface{}{ fieldName: &pql.Condition{ Op: pql.NEQ, Value: nil, }, }, }, }, }, }, } } else { // request a count of IS NOT NULL, aka Row(field!=null). We care about // the actual number of results that should exist. totalCountCall = &pql.Call{ Name: "Count", Children: []*pql.Call{ { Name: "Row", Args: map[string]interface{}{ fieldName: &pql.Condition{ Op: pql.NEQ, Value: nil, }, }, }, }, } } // total values matched by the filter (if it exists) or that aren't null totalCountInterface, err := o.executeCall(ctx, tableKeyer, totalCountCall, shards, opt) totalCount, ok := totalCountInterface.(uint64) if !ok || totalCount == 0 { // it's not an error, but the median of nothing is NULL. return nil, nil } // We have totalCount values. If nth is 50, we want half the values to be // above us, and half below us. So for instance, if we have 6 values, we want // 3 above us, and 3 below us. For odd numbers, we can round these *both* // down -- for 7 values, we'd want 3 higher, and 3 lower. desiredLess := uint64((float64(totalCount) * nthFloat) / 100.0) desiredGreater := uint64((float64(totalCount) * (100 - nthFloat)) / 100.0) // get min var minVal featurebase.ValCount if desiredGreater != 0 { q, err := pql.ParseString(fmt.Sprintf(`Min(field="%s")`, fieldName)) if err != nil { return nil, errors.Wrap(err, "parsing max call for Percentile") } minCall := q.Calls[0] if filterCall != nil { minCall.Children = append(minCall.Children, filterCall) } minVal, err = o.executeMin(ctx, tableKeyer, minCall, shards, opt) if err != nil { return nil, errors.Wrap(err, "executing Min call for Percentile") } if desiredLess == 0 { if minVal.DecimalVal != nil { minVal.FloatVal = minVal.DecimalVal.Float64() } return minVal, nil } } // get max q, err := pql.ParseString(fmt.Sprintf(`Max(field="%s")`, fieldName)) if err != nil { return nil, errors.Wrap(err, "parsing max call for Percentile") } maxCall := q.Calls[0] if filterCall != nil { maxCall.Children = append(maxCall.Children, filterCall) } maxVal, err := o.executeMax(ctx, tableKeyer, maxCall, shards, opt) if err != nil { return nil, errors.Wrap(err, "executing Max call for Percentile") } if desiredGreater == 0 { if maxVal.DecimalVal != nil { maxVal.FloatVal = maxVal.DecimalVal.Float64() } return maxVal, nil } // o.executeCount(ctx, tableKeyer, countCall, shards, opt) // cookValCount(possibleNthVal, 1, field), nil // the logic here is basically identical whether we're doing a decimal field // or an integer field, but the actual code used to compare maximum and minimum // values, or extract values from valCount objects, differs. // So we set up generic functions which will produce the right values. var averageMinMax func() interface{} var minLessthanMax func() bool var maxValueUnder func(interface{}) var minValueOver func(interface{}) if field.Options.Type == FieldTypeDecimal { minPtr := minVal.DecimalVal maxPtr := maxVal.DecimalVal if minPtr == nil { return nil, fmt.Errorf("unexpectedly nil min value in percentile") } if maxPtr == nil { return nil, fmt.Errorf("unexpectedly nil max value in percentile") } min := *minPtr max := *maxPtr two := pql.NewDecimal(2, 0) one := pql.NewDecimal(1, field.Options.Scale) averageMinMax = func() interface{} { return pql.DivideDecimal(pql.AddDecimal(min, max), two) } minLessthanMax = func() bool { return min.LessThan(max) } maxValueUnder = func(v interface{}) { max = pql.SubtractDecimal(v.(pql.Decimal), one) } minValueOver = func(v interface{}) { min = pql.AddDecimal(v.(pql.Decimal), one) } } else { // plain BSI field min := minVal.Val max := maxVal.Val averageMinMax = func() interface{} { // min+max could overflow, in theory, but if they're both odd, we want one // higher than min/2 + max/2. return (min / 2) + (max / 2) + (((min % 2) + (max % 2)) / 2) } minLessthanMax = func() bool { return min < max } maxValueUnder = func(v interface{}) { max = v.(int64) - 1 } minValueOver = func(v interface{}) { min = v.(int64) + 1 } } // set up reusable pql.Call objects representing a count (or intersectioncount, // if we have a filter) with a condition we can alter. var countCall, rangeCall *pql.Call rangeCondition := pql.Condition{ Op: pql.LT, Value: nil, } rangeCall = &pql.Call{ Name: "Row", Args: map[string]interface{}{ fieldName: &rangeCondition, }, } if filterCall == nil { countCall = &pql.Call{ Name: "Count", Children: []*pql.Call{rangeCall}, } } else { countCall = &pql.Call{ Name: "Count", Children: []*pql.Call{ { Name: "Intersect", Children: []*pql.Call{rangeCall, filterCall}, }, }, } } // estimate nth val, eg median when nth=0.5 // we start with a blind guess of minVal, so if min and max are equal, // we just fall out of the loop. If they're not, we compute the middle value // of whatever range we're looking at, and compare it to our expectations of // how many var possibleNthVal interface{} if minVal.DecimalVal != nil { possibleNthVal = minVal.DecimalVal } else { possibleNthVal = minVal.Val } for minLessthanMax() { // compute average without integer overflow, then correct for division of // odd numbers by 2 possibleNthVal = averageMinMax() rangeCondition.Value = possibleNthVal rangeCondition.Op = pql.LT leftCount, err := o.executeCount(ctx, tableKeyer, countCall, shards, opt) if err != nil { return nil, errors.Wrap(err, "executing Count call L for Percentile") } // If there's more things less than possibleNthVal than our desired number // of things less, we need to look at the left side of this. if leftCount > desiredLess { maxValueUnder(possibleNthVal) continue } rangeCondition.Op = pql.GT rightCount, err := o.executeCount(ctx, tableKeyer, countCall, shards, opt) if err != nil { return nil, errors.Wrap(err, "executing Count call R for Percentile") } // If there's more things greater than the desired number, we need to look to the right. if rightCount > desiredGreater { minValueOver(possibleNthVal) continue } // min and max may be different, but the number of values above and below this // value are both reasonable. For instance, with 7 items and looking for median, // we'd have 3 less and 3 greater, and we can't really do better than that. break } switch v := possibleNthVal.(type) { case int64: return featurebase.ValCount{ Val: v, Count: 1, }, nil case pql.Decimal: return featurebase.ValCount{ DecimalVal: &v, FloatVal: v.Float64(), Count: 1, }, nil default: return nil, fmt.Errorf("unexpected percentile Nth value type %T", possibleNthVal) } } // executeMinRow executes a MinRow() call. func (o *orchestrator) executeMinRow(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, shards, c, opt, reduceFn) } // executeMaxRow executes a MaxRow() call. func (o *orchestrator) executeMaxRow(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, shards, c, opt, reduceFn) } // executePrecomputedCall pretends to execute a call that we have a precomputed value for. func (o *orchestrator) executePrecomputedCall(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer dax.TableKeyer, 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() labels := prometheus.Labels{"index": string(tableKeyer.Key())} statFn := func(ctr *prometheus.CounterVec) { if !opt.Remote { ctr.With(labels).Inc() } } if !opt.Remote { switch c.Name { case "Row": if c.HasConditionArg() { statFn(featurebase.CounterQueryRowBSITotal) } else { statFn(featurebase.CounterQueryRowTotal) } case "Range": statFn(featurebase.CounterQueryRangeTotal) case "Difference": statFn(featurebase.CounterQueryBitmapTotal) case "Intersect": statFn(featurebase.CounterQueryIntersectTotal) case "Union": statFn(featurebase.CounterQueryUnionTotal) case "InnerUnionRows": statFn(featurebase.CounterQueryInnerUnionRowsTotal) case "Xor": statFn(featurebase.CounterQueryXorTotal) case "Not": statFn(featurebase.CounterQueryNotTotal) case "Shift": statFn(featurebase.CounterQueryShiftTotal) case "All": statFn(featurebase.CounterQueryAllTotal) default: statFn(featurebase.CounterQueryBitmapTotal) } } // 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, tableKeyer, 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 ErrViewNotFound = Error("view not found") const ErrFragmentNotFound = Error("fragment not found") func (o *orchestrator) executeTopK(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer, 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, tableKeyer, &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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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 } // TODO(tlt): is `fields` used? // 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, tableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer, c, shards, opt) } // executeCount executes a count() call. func (o *orchestrator) executeCount(ctx context.Context, tableKeyer dax.TableKeyer, 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, tableKeyer, 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, tableKeyer, 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, } resp, err := o.client.QueryNode(ctx, node, 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, tableKeyer dax.TableKeyer, shards []uint64, c *pql.Call, opt *featurebase.ExecOptions, reduceFn reduceFunc) (result interface{}, err error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapReduce") defer span.Finish() index := string(tableKeyer.Key()) 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 []dax.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 Controller 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, tableKeyer dax.TableKeyer, columnKeys map[string]map[string]uint64, rowKeys map[string]map[string]map[string]uint64) (*pql.Call, error) { index := string(tableKeyer.Key()) // Check for an overriding 'index' argument. // This also applies to all child calls. if callIndex := c.CallIndex(); callIndex != "" { index = callIndex // TODO(tlt): checking for prefix like this is bad form. Ideally, the // argument stored in the Call.Args map would be of type TableKey // (currently they are restricted to type: string). In that case we // could just pass it through without doing this conversion. (This would // require changing the logic in Queryer.convertIndex() to set "index" // to a TableKeyer). if strings.HasPrefix(index, dax.PrefixTable+dax.TableKeyDelimiter) { qtid, err := dax.QualifiedTableIDFromKey(index) if err != nil { return nil, errors.Wrapf(err, "getting qtid from key: %s", index) } tableKeyer = qtid } else { tableKeyer = dax.StringTableKeyer(index) } } idx, err := o.schemaIndexInfo(ctx, tableKeyer) 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.schemaFieldInfo(ctx, tableKeyer, 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.schemaFieldInfo(ctx, tableKeyer, 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.schemaFieldInfo(ctx, tableKeyer, 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.schemaFieldInfo(ctx, tableKeyer, 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.schemaFieldInfo(ctx, tableKeyer, 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, tableKeyer, 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, tableKeyer, 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, qtbl *dax.QualifiedTable, calls []*pql.Call, results []interface{}, memoryAvailable int64) (err error) { span, _ := tracing.StartSpanFromContext(ctx, "Executor.translateResults") defer span.Finish() idx := featurebase.TableToIndexInfo(&qtbl.Table) 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, string(qtbl.Key()), idSet); err != nil { return err } } for i := range results { results[i], err = o.translateResult(ctx, qtbl, 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.schemaIndexInfo(ctx, dax.TableKey(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.schemaFieldInfo(ctx, dax.TableKey(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.schemaIndexInfo(ctx, dax.StringTableKeyer(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, tableKeyer dax.TableKeyer, 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, tableKeyer, field, ids) } func (o *orchestrator) translateResult(ctx context.Context, qtbl *dax.QualifiedTable, call *pql.Call, result interface{}, idSet map[uint64]string) (_ interface{}, err error) { idx := featurebase.TableToIndexInfo(&qtbl.Table) 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, result.Index, rowField.Name, result.Columns()) if err != nil { return nil, errors.Wrap(err, "translating Row to field keys") } result.Keys = keys case byRowFieldForeignIndex: if _, err := o.schemaIndexInfo(ctx, dax.StringTableKeyer(rowField.Options.ForeignIndex)); 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.schemaFieldInfo(ctx, qtbl, 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.schemaFieldInfo(ctx, qtbl, 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.schemaFieldInfo(ctx, qtbl, 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.schemaFieldInfo(ctx, qtbl, 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, qtbl, 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.schemaFieldInfo(ctx, qtbl, 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.schemaFieldInfo(ctx, qtbl, 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), qtbl, field.Name) if err != nil { return nil, errors.Wrapf(err, "orch: translating IDs of field %q", v) } mapper = func(ids []uint64) (interface{}, error) { if ids == nil { return []string(nil), nil } 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) { return ids, nil } } case FieldTypeMutex: if field.Options.Keys { datatype = "string" translations, err := o.preTranslateMatrixSet(ctx, result, uint(i), qtbl, field.Name) if err != nil { return nil, errors.Wrapf(err, "orch: 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 qdbid dax.QualifiedDatabaseID } func newQualifiedOrchestrator(orch *orchestrator, qdbid dax.QualifiedDatabaseID) *qualifiedOrchestrator { return &qualifiedOrchestrator{ orchestrator: orch, qdbid: qdbid, } } func (o *qualifiedOrchestrator) Execute(ctx context.Context, tableKeyer dax.TableKeyer, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) (featurebase.QueryResponse, error) { resp := featurebase.QueryResponse{} var qtbl *dax.QualifiedTable switch keyer := tableKeyer.(type) { case *dax.Table: qtbl = dax.NewQualifiedTable(o.qdbid, keyer) case *dax.QualifiedTable: qtbl = keyer default: return resp, errors.Errorf("qualifiedOrchestrator.Execute expects a *dax.Table or *dax.QualifiedTable, but got: %T", tableKeyer) } return o.orchestrator.Execute(ctx, qtbl, q, shards, opt) } // schemaFieldInfo is a function introduced when we replaced // `schema.FieldInfo()` calls, where schema was a `featurebase.SchemaInfoAPI` to // `schema.Table().Field()` calls, where schema is a `pilosa.SchemaAPI`. In the // future, when we're no longer dealing with IndexInfo and FieldInfo, and // instead use dax.Table and dax.Field, this helper function can be factored // out. func (o *orchestrator) schemaFieldInfo(ctx context.Context, tableKeyer dax.TableKeyer, fieldName string) (*featurebase.FieldInfo, error) { var tbl *dax.Table var err error switch v := tableKeyer.(type) { case *dax.QualifiedTable: tbl = &v.Table case *dax.Table: tbl = v case dax.QualifiedTableID: tbl, err = o.schema.TableByID(ctx, v.ID) if err != nil { return nil, errors.Wrapf(err, "getting table by id: %s", v.ID) } case dax.StringTableKeyer: tbl, err = o.schema.TableByName(ctx, dax.TableName(v)) if err != nil { return nil, errors.Wrapf(err, "getting table by name: %s", v) } case dax.TableKey: qtid := v.QualifiedTableID() tbl, err = o.schema.TableByID(ctx, qtid.ID) if err != nil { return nil, errors.Wrapf(err, "getting table by ID from TableKey: %s", v) } default: return nil, errors.Errorf("unsupport table keyer type in schemaFieldInfo: %T", tableKeyer) } fld, ok := tbl.Field(dax.FieldName(fieldName)) if !ok { return nil, errors.Errorf("field not found: %s", fieldName) } return featurebase.FieldToFieldInfo(fld), nil } // schemaIndexInfo - see comment on schemaFieldInfo. func (o *orchestrator) schemaIndexInfo(ctx context.Context, tableKeyer dax.TableKeyer) (*featurebase.IndexInfo, error) { var tbl *dax.Table var err error switch v := tableKeyer.(type) { case *dax.QualifiedTable: tbl = &v.Table case *dax.Table: tbl = v case dax.QualifiedTableID: tbl, err = o.schema.TableByID(ctx, v.ID) if err != nil { return nil, errors.Wrapf(err, "getting table by id: %s", v.ID) } case dax.TableKey: qtid := v.QualifiedTableID() tbl, err = o.schema.TableByID(ctx, qtid.ID) if err != nil { return nil, errors.Wrapf(err, "getting table by ID from TableKey: %s", v) } case dax.StringTableKeyer: tbl, err = o.schema.TableByName(ctx, dax.TableName(v)) if err != nil { return nil, errors.Wrapf(err, "getting table by name: %s", v) } default: return nil, errors.Errorf("unsupport table keyer type in schemaIndexInfo: %T", tableKeyer) } return featurebase.TableToIndexInfo(tbl), nil }