Merge pull request #527 from jaddr2line/like

Add `Rows(like=...)` and `UnionRows` queries

docs/query-language.md

@@ -52,6 +52,7 @@ curl localhost:10101/index/repository/query \
 * `CALL` Any query.
 * `ROW_CALL` Any query which returns a row, such as `Row`, `Union`, `Difference`, `Xor`, `Intersect`, `Not`.
 * `ROWS_CALL` A query that returns a `Rows` result (i.e. a list of row IDs). Currently only the `Rows` query.
+* `ROWSET_CALL` A query that returns a set of rows. Currently only the `Rows` and `TopN` queries.
 * `[]ATTR_VALUE` Denotes an array of `ATTR_VALUE`s. (e.g. `["a", "b", "c"]`).
 
 ### Write Operations
@@ -796,7 +797,7 @@ Options(Row(f1=10), shards=[0, 2])
 **Spec:**
 
 ```
-Rows(<FIELD>, previous=<UINT|STRING>, limit=<UINT>, column=<UINT|STRING>, from=<TIMESTAMP>, to=<TIMESTAMP>)
+Rows(<FIELD>, previous=<UINT|STRING>, limit=<UINT>, column=<UINT|STRING>, from=<TIMESTAMP>, to=<TIMESTAMP>, like=<STRING>)
 ```
 
 **Description:**
@@ -817,6 +818,10 @@ If the field is of type `time`, the `from` and `to` arguments can be provided
 to restrict the result to a specific time span. If `from` and `to` are
 not provided, the full range of existing data will be queried.
 
+If `like` is given, only keys matching a pattern will be selected.
+A `like` pattern may use `_` as a placeholder to match a single UTF-8 codepoint, and `%` to match 0 or more codepoints.
+All other characters will be matched exactly.
+
 **Result Type:** Object with `"rows" or "keys" and an array of integers or strings respectively.`
 
 **Examples:**
@@ -834,7 +839,15 @@ With keys:
 Rows(job)
 ```
 ```response
-{"rows":null,"keys":["engineer","management","student""]}
+{"rows":null,"keys":["engineer","management","student"]}
+```
+
+With `like`:
+```request
+Rows(job, like="%t")
+```
+```response
+{"rows":null,"keys":["management","student"]}
 ```
 
 #### Group By
@@ -923,3 +936,31 @@ GroupBy(Rows(age), Rows(job), limit=7, filter=Row(country=USA))
  {"group":[{"field":"age","rowID":22},{"field":"job","rowKey":"student"}],"count":3},
  {"group":[{"field":"age","rowID":29},{"field":"job","rowKey":"management"}],"count":7}]
 ```
+
+#### UnionRows
+
+**Spec:**
+
+```
+UnionRows([ROWSET_CALL ...])
+```
+
+**Description:**
+
+UnionRows performs a logical OR on the rows matched by the results of all `ROWSET_CALL` queries passed to it.
+
+**Result Type:** object with attrs and bits
+
+attrs will always be empty
+
+**Examples:**
+
+Query columns with a bit set in any row (repositories that are starred by any user):
+```request
+UnionRows(Rows(stargazer))
+```
+```response
+{"attrs":{},"columns":[10, 20, 30]}
+```
+
+* columns are repositories that were starred by any user

executor.go

@@ -495,6 +495,111 @@ func (e *executor) execute(ctx context.Context, tx Tx, index string, q *pql.Quer
 	return results, nil
 }
 
+// preprocessQuery expands any calls that need preprocessing.
+// So far, this only needs to process UnionRows.
+func (e *executor) preprocessQuery(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*pql.Call, error) {
+	switch c.Name {
+	case "UnionRows":
+		// 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 := e.executeCall(ctx, tx, index, child, shards, opt)
+			if err != nil {
+				return nil, err
+			}
+
+			// Turn the results into rows calls.
+			var resultRows []*pql.Call
+			switch rowsResult := rowsResult.(type) {
+			case *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 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.
+		return &pql.Call{
+			Name:     "Union",
+			Children: rows,
+		}, nil
+
+	default:
+		// Recurse through child calls.
+		out := make([]*pql.Call, len(c.Children))
+		var changed bool
+		for i, child := range c.Children {
+			res, err := e.preprocessQuery(ctx, tx, index, child, shards, opt)
+			if err != nil {
+				return nil, err
+			}
+			if res != child {
+				changed = true
+			}
+			out[i] = res
+		}
+		if changed {
+			c = c.Clone()
+			c.Children = out
+		}
+		return c, nil
+	}
+}
+
 // executeCall executes a call.
 func (e *executor) executeCall(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (interface{}, error) {
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCall")
@@ -534,6 +639,12 @@ func (e *executor) executeCall(ctx context.Context, tx Tx, index string, c *pql.
 		}
 	}
 
+	// Preprocess the query.
+	c, err := e.preprocessQuery(ctx, tx, index, c, shards, opt)
+	if err != nil {
+		return nil, err
+	}
+
 	// Special handling for mutation and top-n calls.
 	if op, ok := e.additionalCountOps[c.Name]; ok {
 		statFn()
@@ -2431,6 +2542,14 @@ func (e *executor) executeRowsShard(ctx context.Context, tx Tx, index string, fi
 		limit = int(lim)
 	}
 
+	var likeErr chan error
+	if like, hasLike, err := c.StringArg("like"); err != nil {
+		return nil, errors.Wrap(err, "getting like pattern")
+	} else if hasLike {
+		likeErr = make(chan error, 1)
+		filters = append(filters, filterLike(like, f.TranslateStore(), likeErr))
+	}
+
 	for _, view := range views {
 		if err := ctx.Err(); err != nil {
 			return nil, err
@@ -2444,6 +2563,11 @@ func (e *executor) executeRowsShard(ctx context.Context, tx Tx, index string, fi
 		if err != nil {
 			return nil, err
 		}
+		select {
+		case err = <-likeErr:
+			return nil, err
+		default:
+		}
 		rowIDs = rowIDs.merge(viewRows, limit)
 	}
 

executor_test.go

@@ -4742,6 +4742,10 @@ func TestExecutor_Execute_Rows_Keys(t *testing.T) {
 			q:   `Rows(f, previous="1", limit=0, column="0")`,
 			exp: []string{},
 		},
+		{
+			q: `Rows(f, like="__")`,
+			exp: []string{"10", "11", "12", "13", "14", "15", "16", "17", "18"},
+		},
 	}
 
 	for i, test := range tests {
@@ -5843,6 +5847,31 @@ func TestExecutor_Execute_TopNDistinct(t *testing.T) {
 	})
 }
 
+func Test_Executor_Execute_UnionRows(t *testing.T) {
+	c := test.MustRunCluster(t, 2)
+	defer c.Close()
+
+	c.CreateField(t, "i", pilosa.IndexOptions{}, "s",
+		pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 50000),
+	)
+
+	// Populate data.
+	c.Query(t, "i", `
+			Set(0, s=1)
+			Set(1, s=2)
+			Set(2, s=3)
+			Set(3, s=1)
+			Set(3, s=5)
+		`)
+
+	if res := c.Query(t, "i", `Count(UnionRows(TopN(s, n=1)))`); res.Results[0] != uint64(2) {
+		t.Errorf("expected 2 columns, got %v", res.Results[0])
+	}
+	if res := c.Query(t, "i", `Count(UnionRows(Rows(s)))`); res.Results[0] != uint64(4) {
+		t.Errorf("expected 4 columns, got %v", res.Results[0])
+	}
+}
+
 func TestTimelessClearRegression(t *testing.T) {
 	data, err := ioutil.ReadFile("testdata/timeRegressionSchema.json")
 	if err != nil {

fragment.go

@@ -2803,6 +2803,22 @@ func filterColumn(col uint64) rowFilter {
 	}
 }
 
+func filterLike(like string, t TranslateStore, e chan error) rowFilter {
+	plan := planLike(like)
+
+	return func(rowID, key uint64, c *roaring.Container) (include, done bool) {
+		keyStr, err := t.TranslateID(rowID)
+		if err != nil {
+			select {
+			case e <- err:
+			default:
+			}
+			return false, true
+		}
+		return matchLike(keyStr, plan...), false
+	}
+}
+
 // TODO: this works, but it would be more performant if the fragment could seek
 // to the next row in the rows list rather than asking the filter for each
 // container serially. The container iterator would need to expose a seek

like.go

@@ -0,0 +1,239 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pilosa
+
+import (
+	"strings"
+	"unicode/utf8"
+)
+
+// tokenizeLike turns a "like" pattern into a list of tokens.
+// Every token is either a string to exactly match or a combination of % and _ placeholders.
+func tokenizeLike(like string) []string {
+	var tokens []string
+	for like != "" {
+		var token string
+		i := strings.IndexAny(like, "%_")
+		switch i {
+		case 0:
+			// Generate a token of placeholders.
+
+			// Iterate bytewise over the string to find the end of the token.
+			// The % and _ characters are ASCII, so we do not have to worry about Unicode right here.
+			j := 1
+			for j < len(like) && (like[j] == '%' || like[j] == '_') {
+				j++
+			}
+			token, like = like[:j], like[j:]
+		case -1:
+			// There are no more placeholders - generate the last token.
+			token, like = like, ""
+		default:
+			// Generate an exact match token.
+			token, like = like[:i], like[i:]
+		}
+		tokens = append(tokens, token)
+	}
+	return tokens
+}
+
+// filterStepKind is a kind of step in a like filter.
+type filterStepKind uint8
+
+const (
+	filterStepPrefix      filterStepKind = iota // x...
+	filterStepSkipN                             // __...
+	filterStepSkipThrough                       // %x...
+	filterStepSuffix                            // %x
+	filterStepMinLength                         // _%
+)
+
+// filterStep is a step in a like filter.
+type filterStep struct {
+	// kind is the step kind.
+	kind filterStepKind
+
+	// str is the substring for a prefix/skipthrough/suffix step.
+	str string
+
+	// n is the number of underscores in the step (if relevant).
+	n int
+}
+
+// planLike generates a filtering plan for a like pattern.
+func planLike(like string) []filterStep {
+	// Tokenize the like pattern.
+	tokens := tokenizeLike(like)
+
+	steps := make([]filterStep, 0, len(tokens))
+	var merged bool
+	for i, t := range tokens {
+		if merged {
+			// The token was already merged into the previous step.
+			merged = false
+			continue
+		}
+
+		// Convert the token to a step.
+		var step filterStep
+		hasPercent := strings.ContainsRune(t, '%')
+		underscores := strings.Count(t, "_")
+		switch {
+		case hasPercent && i+1 < len(tokens):
+			// Generate a step to skip through the next token.
+			step = filterStep{
+				kind: filterStepSkipThrough,
+				str:  tokens[i+1],
+				n:    underscores,
+			}
+			merged = true
+		case hasPercent:
+			// Generate a terminating step to absorb the remainder of the string.
+			step = filterStep{
+				kind: filterStepMinLength,
+				n:    underscores,
+			}
+		case underscores > 0:
+			// Generate a step to absorb _ placeholders.
+			step = filterStep{
+				kind: filterStepSkipN,
+				n:    underscores,
+			}
+		default:
+			// Generate a step to process an exact match of the beginning of a string.
+			step = filterStep{
+				kind: filterStepPrefix,
+				str:  t,
+			}
+		}
+		steps = append(steps, step)
+	}
+
+	// Optimize suffix matching.
+	if len(steps) > 0 && steps[len(steps)-1].kind == filterStepSkipThrough {
+		steps[len(steps)-1].kind = filterStepSuffix
+	}
+
+	return steps
+}
+
+// matchLike matches a string using a like plan.
+func matchLike(key string, like ...filterStep) bool {
+	for i, step := range like {
+		switch step.kind {
+		case filterStepPrefix:
+			// Match a prefix.
+			if !strings.HasPrefix(key, step.str) {
+				return false
+			}
+			key = key[len(step.str):]
+		case filterStepSkipN:
+			// Skip some placeholders.
+			n := step.n
+			for j := 0; j < n; j++ {
+				_, len := utf8.DecodeRuneInString(key)
+				if len == 0 {
+					return false
+				}
+				key = key[len:]
+			}
+		case filterStepSkipThrough:
+			// Skip through a string.
+
+			// Skip through placeholders.
+			var skipped int
+			for skipped < step.n {
+				j := strings.Index(key, step.str)
+				switch j {
+				case -1:
+					// There are no more matches.
+					return false
+				case 0:
+					// Skip a single rune to ensure forward progress.
+					// This is somewhat inefficient since we have to search the string again next time.
+					// This will hopefully not have to be used very frequently.
+					_, len := utf8.DecodeRuneInString(key)
+					key = key[len:]
+					skipped += len
+				default:
+					// Skip until the substring and count the skipped runes.
+					k := -1
+					for k = range key[:j] {
+					}
+					skipped += k + 1
+
+					key = key[j:]
+				}
+			}
+
+			// Iterate through the substring matches until the rest of the pattern matches.
+			remaining := like[i+1:]
+			for {
+				// Find the next substring match.
+				j := strings.Index(key, step.str)
+				switch {
+				case j == -1:
+					// There are no more matches.
+					return false
+				case j > 0:
+					// Skip the data before the substring.
+					key = key[j:]
+				}
+
+				// Apply the rest of the filter.
+				if matchLike(key[len(step.str):], remaining...) {
+					// This instance matches, no need to search any more.
+					return true
+				}
+
+				// Skip the first rune of the substring so we do not rescan this substring match.
+				_, len := utf8.DecodeRuneInString(key)
+				key = key[len:]
+			}
+		case filterStepSuffix:
+			// Match a suffix.
+			if !strings.HasSuffix(key, step.str) {
+				// Suffix not present.
+				return false
+			}
+			if step.n <= 0 {
+				// No skip length check necessary.
+				return true
+			}
+
+			// Check length of the substring before the suffix.
+			key = key[:len(key)-len(step.str)]
+			fallthrough
+		case filterStepMinLength:
+			if len(key) < step.n {
+				// The string is definitely too short.
+				return false
+			}
+
+			// Count the runes.
+			j := -1
+			for j = range key {
+			}
+
+			// Check if the string is long enough.
+			return j+1 >= step.n
+		default:
+			panic("invalid step")
+		}
+	}
+
+	// If there is any unmatched data left, this is not a match.
+	return key == ""
+}

like_test.go

@@ -0,0 +1,248 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pilosa
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestPlanLike(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name            string
+		like            string
+		plan            []filterStep
+		match, nonmatch []string
+	}{
+		{
+			name:     "Empty",
+			like:     "",
+			plan:     []filterStep{},
+			match:    []string{""},
+			nonmatch: []string{"a", " "},
+		},
+		{
+			name: "Exact",
+			like: "x",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "x",
+				},
+			},
+			match:    []string{"x"},
+			nonmatch: []string{"", "y", "z", "xy", "yx"},
+		},
+		{
+			name: "Anything",
+			like: "%",
+			plan: []filterStep{
+				{
+					kind: filterStepMinLength,
+					n:    0,
+				},
+			},
+			match: []string{"", "a", "b", "ab"},
+		},
+		{
+			name: "Prefix",
+			like: "x%",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "x",
+				},
+				{
+					kind: filterStepMinLength,
+					n:    0,
+				},
+			},
+			match:    []string{"xy", "xyz", "xyzzy"},
+			nonmatch: []string{"plugh", "yx", ""},
+		},
+		{
+			name: "Suffix",
+			like: "%x",
+			plan: []filterStep{
+				{
+					kind: filterStepSuffix,
+					str:  "x",
+				},
+			},
+			match:    []string{"x", "xx", "ax"},
+			nonmatch: []string{"", "a", "x^"},
+		},
+		{
+			name: "Sandwich",
+			like: "x%y",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "x",
+				},
+				{
+					kind: filterStepSuffix,
+					str:  "y",
+				},
+			},
+			match:    []string{"xy", "xzy", "xyzzy"},
+			nonmatch: []string{"plugh", ".xy.", ".x.y", "x.y."},
+		},
+		{
+			name: "DoubleDeckerSandwich",
+			like: "x%y%z",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "x",
+				},
+				{
+					kind: filterStepSkipThrough,
+					str:  "y",
+				},
+				{
+					kind: filterStepSuffix,
+					str:  "z",
+				},
+			},
+			match:    []string{"xyz", "xzyzz", "x.y.z", "x.y.y..z"},
+			nonmatch: []string{"plugh", ".xyz.", ".x.y.z", "x.y.z."},
+		},
+		{
+			name: "Skips",
+			like: "a_b_%_c_%_%_d",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "a",
+				},
+				{
+					kind: filterStepSkipN,
+					n:    1,
+				},
+				{
+					kind: filterStepPrefix,
+					str:  "b",
+				},
+				{
+					kind: filterStepSkipThrough,
+					str:  "c",
+					n:    2,
+				},
+				{
+					kind: filterStepSuffix,
+					str:  "d",
+					n:    3,
+				},
+			},
+			match:    []string{"a1b234c5678d"},
+			nonmatch: []string{"abcd", "a1b2345678d", "a1b2c5678d"},
+		},
+		{
+			name: "SingleRune",
+			like: "_",
+			plan: []filterStep{
+				{
+					kind: filterStepSkipN,
+					n:    1,
+				},
+			},
+			match:    []string{"a", "á", "☺"},
+			nonmatch: []string{"ab", "á", "h̷", ""},
+		},
+		{
+			name: "DoubleRune",
+			like: "__",
+			plan: []filterStep{
+				{
+					kind: filterStepSkipN,
+					n:    2,
+				},
+			},
+			match:    []string{"ab", "á", "h̷"},
+			nonmatch: []string{"a", "á", "☺", "abc"},
+		},
+		{
+			name: "MiddleBlank",
+			like: "x_y",
+			plan: []filterStep{
+				{
+					kind: filterStepPrefix,
+					str:  "x",
+				},
+				{
+					kind: filterStepSkipN,
+					n:    1,
+				},
+				{
+					kind: filterStepPrefix,
+					str:  "y",
+				},
+			},
+			match:    []string{"x.y", "xay", "x y", "x⊕y"},
+			nonmatch: []string{"x++y", "", "a"},
+		},
+		{
+			name: "MinLength",
+			like: "_%_",
+			plan: []filterStep{
+				{
+					kind: filterStepMinLength,
+					n:    2,
+				},
+			},
+			match:    []string{"ab", "á", "abc", "pilosa"},
+			nonmatch: []string{"h", "á", ".", "☺"},
+		},
+	}
+	t.Run("Plan", func(t *testing.T) {
+		t.Parallel()
+
+		for _, c := range cases {
+			c := c
+			t.Run(c.name, func(t *testing.T) {
+				t.Parallel()
+
+				plan := planLike(c.like)
+				if !reflect.DeepEqual(plan, c.plan) {
+					t.Errorf("incorrect plan: %v", plan)
+				}
+			})
+		}
+	})
+	t.Run("Match", func(t *testing.T) {
+		t.Parallel()
+
+		for _, c := range cases {
+			c := c
+			t.Run(c.name, func(t *testing.T) {
+				t.Parallel()
+
+				for _, m := range c.match {
+					if !matchLike(m, c.plan...) {
+						t.Errorf("key %q was not matched", m)
+					}
+				}
+				for _, nm := range c.nonmatch {
+					if matchLike(nm, c.plan...) {
+						t.Errorf("key %q was matched", nm)
+					}
+				}
+			})
+		}
+	})
+}

pql/ast.go

@@ -384,6 +384,7 @@ var callInfoByFunc = map[string]callInfo{
 			"previous": nil,
 			"from":     nil,
 			"to":       nil,
+			"like":     "",
 		},
 	},
 	"Shift": {allowUnknown: false,
@@ -391,8 +392,9 @@ var callInfoByFunc = map[string]callInfo{
 			"n": int64(0),
 		},
 	},
-	"Union": {allowUnknown: false},
-	"Xor":   {allowUnknown: false},
+	"Union":     {allowUnknown: false},
+	"UnionRows": {allowUnknown: false},
+	"Xor":       {allowUnknown: false},
 
 	// things that take _field
 	"TopN": allowUnderField,
@@ -662,6 +664,19 @@ func (c *Call) UintSliceArg(key string) ([]uint64, bool, error) {
 	}
 }
 
+func (c *Call) StringArg(key string) (string, bool, error) {
+	val, ok := c.Args[key]
+	if !ok {
+		return "", false, nil
+	}
+	switch tval := val.(type) {
+	case string:
+		return tval, true, nil
+	default:
+		return "", true, fmt.Errorf("unexpected type %T in StringArg, val %v", tval, tval)
+	}
+}
+
 // CallArg is for reading the value at key from call.Args as a Call. If the
 // key is not in Call.Args, the value of the returned value will be nil, and
 // the error will be nil. An error is returned if the value is not a Call.