impl/memory: Implement Queries

impl/memory/datastore_index_selection.go

+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package memory
+
+import (
+        "bytes"
+        "fmt"
+        "sort"
+        "strings"
+
+        ds "github.com/luci/gae/service/datastore"
+        "github.com/luci/gae/service/datastore/serialize"
+)
+
+// reducedQuery contains only the pieces of the query necessary to iterate for
+// results.
+//   deduplication is applied externally
+//   projection / keysonly / entity retrieval is done externally
+type reducedQuery struct {
+        ns   string
+        kind string
+
+        // eqFilters indicate the set of all prefix constraints which need to be
+        // fulfilled in the composite query. All of these will translate into prefix
+        // bytes for SOME index.
+        eqFilters map[string]stringSet
+
+        // suffixFormat is the PRECISE listing of the suffix columns that ALL indexes
+        //   in the multi query will have.
+        //
+        // suffixFormat ALWAYS includes the inequality filter (if any) as the 0th
+        //   element
+        // suffixFormat ALWAYS includes any additional projections (in ascending
+        //   order) after all user defined sort orders
+        // suffixFormat ALWAYS has __key__ as the last column
+        suffixFormat []ds.IndexColumn
+
+        // limits of the inequality and/or full sort order. This is ONLY a suffix,
+        // and it will be appended to the prefix during iteration.
+        start []byte
+        end   []byte
+
+        // metadata describing the total number of columns that this query requires to
+        // execute perfectly.
+        numCols int
+}
+
+type IndexDefinitionSortable struct {
+        // eqFilts is the list of ACTUAL prefix columns. Note that it may contain
+        // redundant columns! (e.g. (tag, tag) is a perfectly valid prefix, becuase
+        // (tag=1, tag=2) is a perfectly valid query).
+        eqFilts []ds.IndexColumn
+        coll    *memCollection
+}
+
+func (i *IndexDefinitionSortable) hasAncestor() bool {
+        return len(i.eqFilts) > 0 && i.eqFilts[0].Property == "__ancestor__"
+}
+
+func (i *IndexDefinitionSortable) numEqHits(c *constraints) int {
+        ret := 0
+        for _, filt := range i.eqFilts {
+                if _, ok := c.constraints[filt.Property]; ok {
+                        ret++
+                }
+        }
+        return ret
+}
+
+type IndexDefinitionSortableSlice []IndexDefinitionSortable
+
+func (s IndexDefinitionSortableSlice) Len() int      { return len(s) }
+func (s IndexDefinitionSortableSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
+func (s IndexDefinitionSortableSlice) Less(i, j int) bool {
+        a, b := s[i], s[j]
+        if a.coll == nil && b.coll != nil {
+                return true
+        } else if a.coll != nil && b.coll == nil {
+                return false
+        }
+
+        cmp := len(a.eqFilts) - len(b.eqFilts)
+        if cmp < 0 {
+                return true
+        } else if cmp > 0 {
+                return false
+        }
+        for k, col := range a.eqFilts {
+                ocol := b.eqFilts[k]
+                if col.Direction == ds.ASCENDING && ocol.Direction == ds.DESCENDING {
+                        return true
+                } else if col.Direction == ds.DESCENDING && ocol.Direction == ds.ASCENDING {
+                        return false
+                }
+                if col.Property < ocol.Property {
+                        return true
+                } else if col.Property > ocol.Property {
+                        return false
+                }
+        }
+        return false
+}
+
+// maybeAddDefinition possibly adds a new IndexDefinitionSortable to this slice.
+// It's only added if it could be useful in servicing q, otherwise this function
+// is a noop.
+//
+// This returns true iff the proposed index is OK and depletes missingTerms to
+// empty.
+//
+// If the proposed index is PERFECT (e.g. contains enough columns to cover all
+// equality filters, and also has the correct suffix), idxs will be replaced
+// with JUST that index, and this will return true.
+func (idxs *IndexDefinitionSortableSlice) maybeAddDefinition(q *reducedQuery, s *memStore, missingTerms stringSet, id *ds.IndexDefinition) bool {
+        // Kindless queries are handled elsewhere.
+        if id.Kind != q.kind {
+                impossible(
+                        fmt.Errorf("maybeAddDefinition given index with wrong kind %q v %q", id.Kind, q.kind))
+        }
+
+        // If we're an ancestor query, and the index is compound, but doesn't include
+        // an Ancestor field, it doesn't work. Builtin indexes can be used for
+        // ancestor queries (and have !Ancestor), assuming that it's only equality
+        // filters (plus inequality on __key__), or a single inequality.
+        if q.eqFilters["__ancestor__"] != nil && !id.Ancestor && !id.Builtin() {
+                impossible(
+                        fmt.Errorf("maybeAddDefinition given compound index with wrong ancestor info: %s %#v", id, q))
+        }
+
+        // add __ancestor__ if necessary
+        sortBy := id.GetFullSortOrder()
+
+        // If the index has fewer fields than we need for the suffix, it can't
+        // possibly help.
+        if len(sortBy) < len(q.suffixFormat) {
+                return false
+        }
+
+        numEqFilts := len(sortBy) - len(q.suffixFormat)
+        // make sure the orders are precisely the same
+        for i, sb := range sortBy[numEqFilts:] {
+                if q.suffixFormat[i] != sb {
+                        return false
+                }
+        }
+
+        if id.Builtin() && numEqFilts == 0 {
+                if len(q.eqFilters) > 1 || (len(q.eqFilters) == 1 && q.eqFilters["__ancestor__"] == nil) {
+                        return false
+                }
+        }
+
+        // Make sure the equalities section doesn't contain any properties we don't
+        // want in our query.
+        //
+        // numByProp && totalEqFilts will be used to see if this is a perfect match
+        // later.
+        numByProp := make(map[string]int, len(q.eqFilters))
+        totalEqFilts := 0
+
+        eqFilts := sortBy[:numEqFilts]
+        for _, p := range eqFilts {
+                if _, ok := q.eqFilters[p.Property]; !ok {
+                        return false
+                }
+                numByProp[p.Property]++
+                totalEqFilts++
+        }
+
+        // ok, we can actually use this
+
+        // Grab the collection for convenience later. We don't want to invalidate this
+        // index's potential just because the collection doesn't exist. If it's
+        // a builtin and it doesn't exist, it still needs to be one of the 'possible'
+        // indexes... it just means that the user's query will end up with no results.
+        coll := s.GetCollection(
+                fmt.Sprintf("idx:%s:%s", q.ns, serialize.ToBytes(*id.PrepForIdxTable())))
+
+        // First, see if it's a perfect match. If it is, then our search is over.
+        //
+        // A perfect match contains ALL the equality filter columns (or more, since
+        // we can use residuals to fill in the extras).
+        toAdd := IndexDefinitionSortable{coll: coll}
+        toAdd.eqFilts = eqFilts
+        for _, sb := range toAdd.eqFilts {
+                missingTerms.rm(sb.Property)
+        }
+
+        perfect := false
+        if len(sortBy) == q.numCols {
+                perfect = true
+                for k, num := range numByProp {
+                        if num < len(q.eqFilters[k]) {
+                                perfect = false
+                                break
+                        }
+                }
+        }
+        if perfect {
+                *idxs = IndexDefinitionSortableSlice{toAdd}
+        } else {
+                *idxs = append(*idxs, toAdd)
+        }
+        return len(missingTerms) == 0
+}
+
+// getRelevantIndexes retrieves the relevant indexes which could be used to
+// service q. It returns nil if it's not possible to service q with the current
+// indexes.
+func getRelevantIndexes(q *reducedQuery, s *memStore) (IndexDefinitionSortableSlice, error) {
+        missingTerms := stringSet{}
+        for k := range q.eqFilters {
+                if k == "__ancestor__" {
+                        // ancestor is not a prefix which can be satisfied by a single index. It
+                        // must be satisfied by ALL indexes (and has special logic for this in
+                        // the addDefinition logic)
+                        continue
+                }
+                missingTerms.add(k)
+        }
+        idxs := IndexDefinitionSortableSlice{}
+
+        // First we add builtins
+        // add
+        //   idx:KIND
+        if idxs.maybeAddDefinition(q, s, missingTerms, &ds.IndexDefinition{
+                Kind: q.kind,
+        }) {
+                return idxs, nil
+        }
+
+        // add
+        //   idx:KIND:prop
+        //   idx:KIND:-prop
+        props := stringSet{}
+        for prop := range q.eqFilters {
+                props.add(prop)
+        }
+        for _, col := range q.suffixFormat[:len(q.suffixFormat)-1] {
+                props.add(col.Property)
+        }
+        for prop := range props {
+                if strings.HasPrefix(prop, "__") && strings.HasSuffix(prop, "__") {
+                        continue
+                }
+                if idxs.maybeAddDefinition(q, s, missingTerms, &ds.IndexDefinition{
+                        Kind: q.kind,
+                        SortBy: []ds.IndexColumn{
+                                {Property: prop},
+                        },
+                }) {
+                        return idxs, nil
+                }
+                if idxs.maybeAddDefinition(q, s, missingTerms, &ds.IndexDefinition{
+                        Kind: q.kind,
+                        SortBy: []ds.IndexColumn{
+                                {Property: prop, Direction: ds.DESCENDING},
+                        },
+                }) {
+                        return idxs, nil
+                }
+        }
+
+        // Try adding all compound indexes whose suffix matches.
+        suffix := &ds.IndexDefinition{
+                Kind:     q.kind,
+                Ancestor: q.eqFilters["__ancestor__"] != nil,
+                SortBy:   q.suffixFormat,
+        }
+        walkCompIdxs(s, suffix, func(def *ds.IndexDefinition) bool {
+                // keep walking until we find a perfect index.
+                return !idxs.maybeAddDefinition(q, s, missingTerms, def)
+        })
+
+        // this query is impossible to fulfil with the current indexes. Not all the
+        // terms (equality + projection) are satisfied.
+        if len(missingTerms) < 0 || len(idxs) == 0 {
+                remains := &ds.IndexDefinition{
+                        Kind:     q.kind,
+                        Ancestor: q.eqFilters["__ancestor__"] != nil,
+                }
+                terms := make([]string, 0, len(missingTerms))
+                for mt := range missingTerms {
+                        terms = append(terms, mt)
+                }
+                if serializationDeterministic {
+                        sort.Strings(terms)
+                }
+                for _, term := range terms {
+                        remains.SortBy = append(remains.SortBy, ds.IndexColumn{Property: term})
+                }
+                remains.SortBy = append(remains.SortBy, q.suffixFormat...)
+                last := remains.SortBy[len(remains.SortBy)-1]
+                if last.Direction == ds.ASCENDING {
+                        // this removes the __key__ column, since it's implicit.
+                        remains.SortBy = remains.SortBy[:len(remains.SortBy)-1]
+                }
+                if remains.Builtin() {
+                        impossible(
+                                fmt.Errorf("recommended missing index would be a builtin: %s", remains))
+                }
+                return nil, fmt.Errorf(
+                        "Your indexes are insufficient! Try adding:\n  %s", remains)
+        }
+
+        return idxs, nil
+}
+
+// generate generates a single iterDefinition for the given index.
+func generate(q *reducedQuery, idx *IndexDefinitionSortable, c *constraints) *iterDefinition {
+        def := &iterDefinition{
+                c:     idx.coll,
+                start: q.start,
+                end:   q.end,
+        }
+        toJoin := make([][]byte, len(idx.eqFilts))
+        for _, sb := range idx.eqFilts {
+                val := c.peel(sb.Property)
+                if sb.Direction == ds.DESCENDING {
+                        val = invert(val)
+                }
+                toJoin = append(toJoin, val)
+        }
+        def.prefix = bjoin(toJoin...)
+        def.prefixLen = len(def.prefix)
+
+        if q.eqFilters["__ancestor__"] != nil && !idx.hasAncestor() {
+                // The query requires an ancestor, but the index doesn't explicitly have it
+                // as part of the prefix (otherwise it would have been the first eqFilt
+                // above). This happens when it's a builtin index, or if it's the primary
+                // index (for a kindless query), or if it's the Kind index (for a filterless
+                // query).
+                //
+                // builtin indexes are:
+                //   Kind/__key__
+                //   Kind/Prop/__key__
+                //   Kind/Prop/-__key__
+                if len(q.suffixFormat) > 2 || q.suffixFormat[len(q.suffixFormat)-1].Property != "__key__" {
+                        // This should never happen. One of the previous validators would have
+                        // selected a different index. But just in case.
+                        impossible(fmt.Errorf("cannot supply an implicit ancestor for %#v", idx))
+                }
+
+                // get the only value out of __ancestor__
+                anc := q.eqFilters["__ancestor__"].getOne()
+
+                // Intentionally do NOT update prefixLen. This allows multiIterator to
+                // correctly include the entire key in the shared iterator suffix, instead
+                // of just the remainder.
+
+                // chop the terminal null byte off the q.ancestor key... we can accept
+                // anything which is a descendant or an exact match.  Removing the last byte
+                // from the key (the terminating null) allows this trick to work. Otherwise
+                // it would be a closed range of EXACTLY this key.
+                chopped := []byte(anc[:len(anc)-1])
+                if q.suffixFormat[0].Direction == ds.DESCENDING {
+                        chopped = invert(chopped)
+                }
+                def.prefix = bjoin(def.prefix, chopped)
+
+                // Update start and end, since we know that if they contain anything, they
+                // contain values for the __key__ field.
+                if def.start != nil {
+                        offset := 0
+                        if len(q.suffixFormat) > 1 {
+                                chunks, _ := parseSuffix(q.ns, q.suffixFormat, def.start, 1)
+                                offset = len(chunks[0])
+                        }
+                        if !bytes.HasPrefix(def.start[offset:], chopped) {
+                                // again, shouldn't happen, but if it does, we want to know about it.
+                                impossible(fmt.Errorf(
+                                        "start suffix for implied ancestor doesn't start with ancestor! start:%v ancestor:%v",
+                                        def.start, chopped))
+                        }
+                        def.start = def.start[:offset+len(chopped)]
+                }
+                if def.end != nil {
+                        offset := 0
+                        if len(q.suffixFormat) > 1 {
+                                chunks, _ := parseSuffix(q.ns, q.suffixFormat, def.end, 1)
+                                offset = len(chunks[0])
+                        }
+                        if !bytes.HasPrefix(def.end[offset:], chopped) {
+                                impossible(fmt.Errorf(
+                                        "end suffix for implied ancestor doesn't start with ancestor! end:%v ancestor:%v",
+                                        def.end, chopped))
+                        }
+                        def.end = def.end[:offset+len(chopped)]
+                }
+        }
+
+        return def
+}
+
+type constraints struct {
+        constraints     map[string][][]byte
+        original        map[string][][]byte
+        residualMapping map[string]int
+}
+
+// peel picks a constraint value for the property. It then removes this value
+// from constraints (possibly removing the entire row from constraints if it
+// was the last value). If the value wasn't available in constraints, it picks
+// the value from residuals.
+func (c *constraints) peel(prop string) []byte {
+        ret := []byte(nil)
+        if vals, ok := c.constraints[prop]; ok {
+                ret = vals[0]
+                if len(vals) == 1 {
+                        delete(c.constraints, prop)
+                } else {
+                        c.constraints[prop] = vals[1:]
+                }
+        } else {
+                row := c.original[prop]
+                idx := c.residualMapping[prop]
+                c.residualMapping[prop]++
+                ret = row[idx%len(row)]
+        }
+        return ret
+}
+
+func (c *constraints) empty() bool {
+        return len(c.constraints) == 0
+}
+
+// calculateConstraints produces a mapping of all equality filters to the values
+// that they're constrained to. It also calculates residuals, which are an
+// arbitrary value for filling index prefixes which have more equality fields
+// than are necessary. The value doesn't matter, as long as its an equality
+// constraint in the original query.
+func calculateConstraints(q *reducedQuery) *constraints {
+        ret := &constraints{
+                original:        make(map[string][][]byte, len(q.eqFilters)),
+                constraints:     make(map[string][][]byte, len(q.eqFilters)),
+                residualMapping: make(map[string]int),
+        }
+        for prop, vals := range q.eqFilters {
+                bvals := make([][]byte, 0, len(vals))
+                for val := range vals {
+                        bvals = append(bvals, []byte(val))
+                }
+                ret.original[prop] = bvals
+                if prop == "__ancestor__" {
+                        // exclude __ancestor__ from the constraints.
+                        //
+                        // This is because it's handled specially during index proposal and
+                        // generation. Ancestor is used by ALL indexes, and so its residual value
+                        // in ret.original above will be sufficient.
+                        continue
+                }
+                ret.constraints[prop] = bvals
+        }
+        return ret
+}
+
+// getIndexes returns a set of iterator definitions. Iterating over these
+// will result in matching suffixes.
+func getIndexes(q *reducedQuery, s *memStore) ([]*iterDefinition, error) {
+        relevantIdxs := IndexDefinitionSortableSlice(nil)
+        if q.kind == "" {
+                if coll := s.GetCollection("ents:" + q.ns); coll != nil {
+                        relevantIdxs = IndexDefinitionSortableSlice{{coll: coll}}
+                }
+        } else {
+                err := error(nil)
+                relevantIdxs, err = getRelevantIndexes(q, s)
+                if err != nil {
+                        return nil, err
+                }
+        }
+        if len(relevantIdxs) == 0 {
+                return nil, errQueryDone
+        }
+
+        // This sorts it so that relevantIdxs goes less filters -> more filters. We
+        // traverse this list backwards, however, so we traverse it in more filters ->
+        // less filters order.
+        sort.Sort(relevantIdxs)
+
+        constraints := calculateConstraints(q)
+
+        ret := []*iterDefinition{}
+        for !constraints.empty() || len(ret) == 0 {
+                bestIdx := (*IndexDefinitionSortable)(nil)
+                if len(ret) == 0 {
+                        // if ret is empty, take the biggest relevantIdx. It's guaranteed to have
+                        // the greatest number of equality filters of any index in the list, and
+                        // we know that every equality filter will be pulled from constraints and
+                        // not residual.
+                        //
+                        // This also takes care of the case when the query has no equality filters,
+                        // in which case relevantIdxs will actually only contain one index anyway
+                        // :)
+                        bestIdx = &relevantIdxs[len(relevantIdxs)-1]
+                        if bestIdx.coll == nil {
+                                return nil, errQueryDone
+                        }
+                } else {
+                        // If ret's not empty, then we need to find the best index we can. The
+                        // best index will be the one with the most matching equality columns.
+                        // Since relevantIdxs is sorted primarially by the number of equality
+                        // columns, we walk down the list until the number of possible columns is
+                        // worse than our best-so-far.
+                        //
+                        // Traversing the list backwards goes from more filters -> less filters,
+                        // but also allows us to remove items from the list as we iterate over it.
+                        bestNumEqHits := 0
+                        for i := len(relevantIdxs) - 1; i >= 0; i-- {
+                                idx := &relevantIdxs[i]
+                                if len(idx.eqFilts) < bestNumEqHits {
+                                        // if the number of filters drops below our best hit, it's never going
+                                        // to get better than that. This index might be helpful on a later
+                                        // loop though, so don't remove it.
+                                        break
+                                }
+                                numHits := 0
+                                if idx.coll != nil {
+                                        numHits = idx.numEqHits(constraints)
+                                }
+                                if numHits > bestNumEqHits {
+                                        bestNumEqHits = numHits
+                                        bestIdx = idx
+                                } else if numHits == 0 {
+                                        // This index will never become useful again, so remove it.
+                                        relevantIdxs = append(relevantIdxs[:i], relevantIdxs[i+1:]...)
+                                }
+                        }
+                }
+                if bestIdx == nil {
+                        // something is really wrong here... if relevantIdxs is !nil, then we
+                        // should always be able to make progress in this loop.
+                        impossible(fmt.Errorf("deadlock: cannot fulfil query?"))
+                }
+                ret = append(ret, generate(q, bestIdx, constraints))
+        }
+
+        return ret, nil
+}