impl/memory: Implement Queries

impl/memory/datastore_index_selection.go

Index: impl/memory/datastore_index_selection.go
diff --git a/impl/memory/datastore_index_selection.go b/impl/memory/datastore_index_selection.go
new file mode 100644
index 0000000000000000000000000000000000000000..07ec2a309fab645fb6d28c65593826bf72dcc55d
--- /dev/null
+++ b/impl/memory/datastore_index_selection.go
@@ -0,0 +1,552 @@
+// 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]map[string]struct{}

[dnj (Google), 2015/08/28 17:54:22]

nit: Consider defining a type for "map[string]struct{}", maybe just "markerMap".
Has "put(k string)" and "has(k string)". I think the code would appear cleaner.

[iannucci, 2015/08/28 19:48:55]

done, though it doesn't really do too much.

+
+	// 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.DESCENDING {

[dnj (Google), 2015/08/28 17:54:21]

Why "false" if we're both descending? Did you mean "ASCENDING" for the latter
condition?

[iannucci, 2015/08/28 19:48:55]

oops. Done.

+			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 map[string]struct{}, 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 indicies 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 {
+		requireSomeEqFilter := len(q.eqFilters) > 1

[dnj (Google), 2015/08/28 17:54:21]

if len(q.eqFilters) > 1 || (len(q.eqFilters) == 1 && q.eqFilters["..."] == nil {
  return false
}

[iannucci, 2015/08/28 19:48:56]

done

+		if !requireSomeEqFilter && len(q.eqFilters) == 1 && q.eqFilters["__ancestor__"] == nil {
+			requireSomeEqFilter = true
+		}
+		if requireSomeEqFilter {
+			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++
+	}
+
+	// See if we actually have this index for the correct namespace
+	coll := s.GetCollection(

[dnj (Google), 2015/08/28 17:54:21]

Need to test if "coll" is nil?

[iannucci, 2015/08/28 19:48:55]

nope. Clarified in comment.

+		fmt.Sprintf("idx:%s:%s", q.ns, serialize.ToBytes(*id.PrepForIdxTable())))
+
+	// ok, we can actually use this
+
+	// 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 {
+		delete(missingTerms, 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
+}
+
+// getRelevantIndicies retrieves the relevant indices which could be used to

[dnj (Google), 2015/08/28 17:54:22]

Indexes?

[iannucci, 2015/08/28 19:48:56]

done

+// service q. It returns nil if it's not possible to service q with the current
+// indicies.
+func getRelevantIndicies(q *reducedQuery, s *memStore) (IndexDefinitionSortableSlice, error) {
+	missingTerms := map[string]struct{}{}
+	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 indices (and has special logic for this in
+			// the addDefinition logic)
+			continue
+		}
+		missingTerms[k] = struct{}{}
+	}
+	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 := map[string]struct{}{}
+	for prop := range q.eqFilters {
+		props[prop] = struct{}{}
+	}
+	for _, col := range q.suffixFormat[:len(q.suffixFormat)-1] {
+		props[col.Property] = struct{}{}
+	}
+	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 indicies 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 indicies. 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 is implied
+			remains.SortBy = remains.SortBy[:len(remains.SortBy)-1]

[dnj (Google), 2015/08/28 17:54:21]

Can just assign to "last".

[iannucci, 2015/08/28 19:48:55]

That would have no effect? Did you mean something else?

[dnj, 2015/08/28 19:57:53]

Oh nope, just overlooked the colon. Nevermind.

+		}
+		if remains.Builtin() {
+			impossible(
+				fmt.Errorf("recommended missing index would be a builtin: %s", remains))
+		}
+		return nil, fmt.Errorf(
+			"Your indicies are insufficient! Try adding:\n  %s", remains)

[dnj (Google), 2015/08/28 17:54:22]

indexes! Also not sure this is the best way to return this information. Consider
a custom "error" implementation that actually includes the list so the caller
can type assert and programmatically pull it apart. The Error() function can
still generate this string :)

[iannucci, 2015/08/28 19:48:55]

yes, I was planning on changing that in a later CL (in fact, the next CL).

+	}
+
+	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 := [][]byte{}

[dnj (Google), 2015/08/28 17:54:22]

Can pre-allocate capacity to len(idx.eqFilts)

[iannucci, 2015/08/28 19:48:56]

Done

+	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))
+		}
+
+		// chop the terminal null byte off the q.ancestor key... we can accept

[dnj (Google), 2015/08/28 17:54:22]

You don't do this here.

[iannucci, 2015/08/28 19:48:55]

oops. moved comment.

+		// anything which is a descendant or an exact match.
+
+		// This silly construction gets the __ancestor__ value, because it's a
+		// map[string]struct{} instead of a [][]byte{} (otherwise we'd just get
+		// the value at the 0th index).
+		anc := ""
+		for k := range q.eqFilters["__ancestor__"] {
+			anc = k
+			break
+		}
+
+		// Intentionally do NOT update prefixLen. This allows multiIterator to
+		// correctly include the entire key in the shared iterator suffix, instead
+		// of just the remainder.
+
+		// 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! start:%v ancestor:%v",

[dnj (Google), 2015/08/28 17:54:22]

end:

[iannucci, 2015/08/28 19:48:56]

done

+					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)]

[dnj (Google), 2015/08/28 17:54:21]

Please document why this is done (pseudorandom padding field). Also maybe add a
PRNG instance to "constraints" and just use.

[iannucci, 2015/08/28 19:48:56]

It's not random. A PRNG would be slower and serve no purpose. The only reason I
round robin (besides being deterministic) is so that we'll 'hopefully' end up
picking a residual value which scans fewer rows. Until we actually execute the
query, all values are equally likely to be the best choice. A PRNG would just
make this slower and more complicated without actually adding any value.  The
reason I said that in the TODO is because I hadn't thought it through completely
(hence the TODO :)).

+	}
+	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 indices, and so its residual value
+			// in ret.original above will be sufficient.
+			continue
+		}
+		ret.constraints[prop] = bvals
+	}
+	return ret
+}
+
+// getIndicies returns a set of iterator definitions. Iterating over these
+// will result in matching suffixes.
+func getIndicies(q *reducedQuery, s *memStore) ([]*iterDefinition, error) {

[dnj (Google), 2015/08/28 17:54:22]

Indexes

[iannucci, 2015/08/28 19:48:55]

You know that both spellings are correct, right?

[dnj, 2015/08/28 19:57:53]

Yes, but you've been generally consistent using "indexes" with a few spot
exceptions.

+	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 = getRelevantIndicies(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)
+
+	// TODO(riannucci): When filling from residual, choose pseudorandom elements

[dnj (Google), 2015/08/28 17:54:22]

You do this with the aforementioned modulus, so no longer TODO?

[iannucci, 2015/08/28 19:48:56]

yes

+	//							    from the original constraints. This will allow the use
+	//							    of large indexes with good hitrates.
+	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
+}