diff --git a/executor.go b/executor.go index f4e5bdd45..6536c17f6 100644 --- a/executor.go +++ b/executor.go @@ -2613,6 +2613,93 @@ func (r RowIDs) merge(other RowIDs, limit int) RowIDs { return result } +// 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 []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/Sum + if gci.Sum < gcj.Sum { + return fieldOrder == asc + } else if gci.Sum > gcj.Sum { + 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 +} + func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) ([]GroupCount, error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeGroupBy") defer span.Finish() @@ -2631,6 +2718,25 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c 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) + } + idx := e.Holder.Index(index) if idx == nil { return nil, newNotFoundError(ErrIndexNotFound, index) @@ -2703,43 +2809,21 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c } results, _ := other.([]GroupCount) - // Apply having. - if having, hasHaving, err := c.CallArg("having"); err != nil { - return nil, err - } else if hasHaving { - // parse the condition as PQL - if having.Name != "Condition" { - return nil, errors.New("the only supported having call is Condition()") - } - if len(having.Args) != 1 { - return nil, errors.New("Condition() must contain a single condition") - } - for subj, cond := range having.Args { - switch subj { - case "count", "sum": - results = applyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) - default: - return nil, errors.New("Condition() only supports count or sum") - } + // 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") } } - // 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] - } - } + // 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 { @@ -2759,10 +2843,10 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c countDistinctIntersect := &pql.Call{ Name: "Count", Children: []*pql.Call{ - &pql.Call{ + { Name: "Distinct", Children: []*pql.Call{ - &pql.Call{ + { Name: "Intersect", Children: intersectRows, }, @@ -2781,6 +2865,61 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c results[n].Sum = int64(aggregateCount[0].(uint64)) } } + + // Apply having. + if hasHaving && !opt.Remote { + // parse the condition as PQL + if having.Name != "Condition" { + return nil, errors.New("the only supported having call is Condition()") + } + if len(having.Args) != 1 { + return nil, errors.New("Condition() must contain a single condition") + } + for subj, cond := range having.Args { + switch subj { + case "count", "sum": + results = applyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) + default: + return nil, errors.New("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") + } + + } + + return results, nil +} + +func applyLimitAndOffsetToGroupByResult(c *pql.Call, results []GroupCount) ([]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 } diff --git a/executor_internal_test.go b/executor_internal_test.go index 06e271fde..6de260293 100644 --- a/executor_internal_test.go +++ b/executor_internal_test.go @@ -18,6 +18,7 @@ import ( "context" "encoding/json" "fmt" + "reflect" "strconv" "testing" @@ -384,3 +385,93 @@ func TestToInt64(t *testing.T) { } } } + +func TestGetSorter(t *testing.T) { + tests := []struct { + sortSpec string + expGCS *groupCountSorter + expErr string + }{ + { + sortSpec: "count asc", + expGCS: &groupCountSorter{fields: []int{-1}, order: []order{asc}}, + }, + { + sortSpec: "count asc", + expGCS: &groupCountSorter{fields: []int{-1}, order: []order{asc}}, + }, + { + sortSpec: " count asc", + expGCS: &groupCountSorter{fields: []int{-1}, order: []order{asc}}, + }, + { + sortSpec: " count asc ", + expGCS: &groupCountSorter{fields: []int{-1}, order: []order{asc}}, + }, + { + sortSpec: "count", + expGCS: &groupCountSorter{fields: []int{-1}, order: []order{desc}}, + }, + { + sortSpec: "sum asc", + expGCS: &groupCountSorter{fields: []int{-2}, order: []order{asc}}, + }, + { + sortSpec: "aggregate asc", + expGCS: &groupCountSorter{fields: []int{-2}, order: []order{asc}}, + }, + { + sortSpec: "boondoggle asc", + expErr: "sorting is only supported on count, aggregate, or sum, not 'boondoggle'", + }, + { + sortSpec: "sum asc, count desc", + expGCS: &groupCountSorter{fields: []int{-2, -1}, order: []order{asc, desc}}, + }, + { + sortSpec: "count asc, sum desc", + expGCS: &groupCountSorter{fields: []int{-1, -2}, order: []order{asc, desc}}, + }, + { + sortSpec: " count asc , sum desc ", + expGCS: &groupCountSorter{fields: []int{-1, -2}, order: []order{asc, desc}}, + }, + { + sortSpec: " count asc , sum desc blah", + expErr: "parsing sort directive: 'sum desc blah': too many elements", + }, + { + sortSpec: "count asc, sum fesc", + expErr: "unknown sort direction 'fesc'", + }, + { + sortSpec: " , sum fesc", + expErr: "invalid sorting directive: ''", + }, + { + // weird and useless, but I guess fine? + sortSpec: "count asc,count asc ", + expGCS: &groupCountSorter{fields: []int{-1, -1}, order: []order{asc, asc}}, + }, + } + + for i, tst := range tests { + t.Run(fmt.Sprintf("%s_%d", tst.sortSpec, i), func(t *testing.T) { + gcs, err := getSorter(tst.sortSpec) + if err != nil { + if tst.expErr == "" { + t.Errorf("unexpected error: %v", err) + return + } + if tst.expErr != err.Error() { + t.Errorf("mismatched errors got: '%v', exp: '%s'", err, tst.expErr) + } + return + } + if !reflect.DeepEqual(gcs, tst.expGCS) { + t.Errorf("exp:\n%+v\ngot:\n%v\n", tst.expGCS, gcs) + } + + }) + } +} diff --git a/executor_test.go b/executor_test.go index 5d18bdaac..cfb2d555d 100644 --- a/executor_test.go +++ b/executor_test.go @@ -6789,12 +6789,12 @@ func TestMissingKeyRegression(t *testing.T) { func TestVariousQueries(t *testing.T) { for _, clusterSize := range []int{1, 3, 4, 7} { t.Run(fmt.Sprintf("%d-node", clusterSize), func(t *testing.T) { - testVariousQueries(t, clusterSize) + variousQueries(t, clusterSize) }) } } -func testVariousQueries(t *testing.T, clusterSize int) { +func variousQueries(t *testing.T, clusterSize int) { c := test.MustRunCluster(t, clusterSize) defer c.Close() @@ -6844,6 +6844,18 @@ func testVariousQueries(t *testing.T, clusterSize int) { {Val: 0, Key: "userE"}, }) + // Create and populate "affinity" int field with negative, positive, zero and null values. + + c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "net_worth", pilosa.OptFieldTypeInt(-100000000, 100000000)) + c.ImportIntKey(t, "users", "net_worth", []test.IntKey{ + {Val: 1, Key: "userA"}, + {Val: 10, Key: "userB"}, + {Val: 100, Key: "userC"}, + {Val: 1000, Key: "userD"}, + {Val: 10000, Key: "userE"}, + {Val: 100000, Key: "userF"}, + }) + c.CreateField(t, "users", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "zip_code", pilosa.OptFieldTypeInt(0, 100000)) c.ImportIntKey(t, "users", "zip_code", []test.IntKey{ {Val: 78739, Key: "userA"}, @@ -7008,6 +7020,16 @@ dog,1,0 icecream,6,0 `, }, + { + query: "GroupBy(Rows(field=likes), aggregate=Sum(field=net_worth), limit=2, having=Condition(sum>10))", + csvVerifier: `pangolin,1,100 +zebra,1,1000 +`, + }, + { + query: "GroupBy(Rows(field=likes), having=Condition(count>5))", + csvVerifier: "icecream,6,0\n", + }, { query: "GroupBy(Rows(field=likes), filter=Row(affinity>-7))", csvVerifier: `molecula,1,0 @@ -7028,6 +7050,10 @@ dog,1,0 icecream,6,3 `, }, + { + query: "GroupBy(Rows(field=likes), aggregate=Count(Distinct(field=zip_code)), having=Condition(sum>2))", + csvVerifier: "icecream,6,3\n", + }, { query: "GroupBy(Rows(field=likes), filter=Row(affinity>-11), aggregate=Count(Distinct(field=zip_code)))", csvVerifier: `molecula,1,1 @@ -7046,6 +7072,42 @@ pangolin,1,1 zebra,1,1 toucan,1,1 icecream,5,3 +`, + }, + { + query: "GroupBy(Rows(field=likes), sort=\"count desc\")", + csvVerifier: `icecream,6,0 +molecula,1,0 +pilosa,1,0 +pangolin,1,0 +zebra,1,0 +toucan,1,0 +dog,1,0 +`, + }, + { + query: "GroupBy(Rows(field=likes), aggregate=Sum(field=net_worth), sort=\"aggregate desc, count asc\")", + csvVerifier: `icecream,6,111111 +dog,1,100000 +toucan,1,10000 +zebra,1,1000 +pangolin,1,100 +pilosa,1,10 +molecula,1,1 +`, + }, + { + query: "GroupBy(Rows(field=likes), aggregate=Sum(field=net_worth), sort=\"aggregate desc, count asc\", limit=3)", + csvVerifier: `icecream,6,111111 +dog,1,100000 +toucan,1,10000 +`, + }, + { + query: "GroupBy(Rows(field=likes), aggregate=Sum(field=net_worth),sort=\"aggregate desc, count asc\",limit=3,offset=2)", + csvVerifier: `toucan,1,10000 +zebra,1,1000 +pangolin,1,100 `, }, } @@ -7064,7 +7126,7 @@ icecream,5,3 // verify everything after header got := csvString[strings.Index(csvString, "\n")+1:] if got != tst.csvVerifier { - t.Errorf("expected '%s', got '%s'", tst.csvVerifier, got) + t.Errorf("expected:\n%s\ngot:\n%s", tst.csvVerifier, got) } // TODO: add HTTP and Postgres and ability to convert diff --git a/pql/ast.go b/pql/ast.go index 50b38e5c5..40e2acaba 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -441,9 +441,11 @@ var callInfoByFunc = map[string]callInfo{ prototypes: map[string]interface{}{ "filter": nil, "limit": int64(0), + "offset": int64(0), "previous": nil, "aggregate": nil, "having": nil, + "sort": "", }, }, "Options": {