diff --git a/executor.go b/executor.go index f4e5bdd45..32a1b2f77 100644 --- a/executor.go +++ b/executor.go @@ -2613,6 +2613,135 @@ func (r RowIDs) merge(other RowIDs, limit int) RowIDs { return result } +type order bool + +const ( + asc order = true + desc order = false +) + +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 { + continue + } + if fieldOrder == asc { + return gci.Count < gcj.Count + } + return gcj.Count < gci.Count + case -2: // aggregate/Sum + if gci.Sum == gcj.Sum { + continue + } + if fieldOrder == asc { + return gci.Sum < gcj.Sum + } + return gcj.Sum < gci.Sum + default: + switch compareFieldRows(gci.Group[fieldIndex], gcj.Group[fieldIndex]) { + case 0: + continue + case -1: + return fieldOrder == asc + case 1: + return fieldOrder == desc + } + + } + } + return true +} + +// compareFieldRows returns -1 if a < b, +1 if a > b, and 0 if they +// are equal. It checks Value, RowKey and RowID, but assumes that +// Field is equal. +func compareFieldRows(a, b FieldRow) int { + if a.Value != nil { + if a.Value == b.Value { + return 0 + } + if *a.Value < *b.Value { + return -1 + } + return +1 + } + if a.RowKey != "" { + if a.RowKey == b.RowKey { + return 0 + } + if a.RowKey < b.RowKey { + return -1 + } + return +1 + } + if a.RowID != 0 { + if a.RowID == b.RowID { + return 0 + } + if a.RowID < b.RowID { + return -1 + } + return +1 + } + // OK, everything in a is zero... so b must be greater unless it's + // also zero (because if a.Value is nil, b.Value must also be + // nil... or something fishy is happening) + if b.RowKey == "" && b.RowID == 0 { + return 0 + } + return -1 +} + +// getSorter hackily parses the sortSpec and figures out how to sort +// the GroupBy results. TODO Probably has about a billion edge case +// bugs. Also needs to take in the Call so it can figure out if fields +// are valid and what order they're in. +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.Split(sortField, " ") + defaultOrder := asc + if fieldDir[0] == "count" { + gcs.fields = append(gcs.fields, -1) + defaultOrder = desc + } else if fieldDir[0] == "aggregate" { + gcs.fields = append(gcs.fields, -2) + defaultOrder = desc + } else { + gcs.fields = append(gcs.fields, 0) // TODO actually figure out which field. probably need the call + } + + if len(fieldDir) == 0 { + gcs.order = append(gcs.order, defaultOrder) + } + 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 +2760,16 @@ 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") + } + } + idx := e.Holder.Index(index) if idx == nil { return nil, newNotFoundError(ErrIndexNotFound, index) @@ -2781,6 +2920,12 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c results[n].Sum = int64(aggregateCount[0].(uint64)) } } + + if sorter != nil { + sorter.data = results + sort.Sort(sorter) + } + return results, nil } diff --git a/executor_test.go b/executor_test.go index 5d18bdaac..f6d711971 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() @@ -7046,6 +7046,30 @@ pangolin,1,1 zebra,1,1 toucan,1,1 icecream,5,3 +`, + }, + { + query: "GroupBy(Rows(field=likes), sort=\"count desc, likes asc\")", + // note, sort is in order of rowID rather than rowKey + 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), sort=\"count desc, likes desc\")", + // note, sort is in order of rowID rather than rowKey + csvVerifier: `icecream,6,0 +dog,1,0 +toucan,1,0 +zebra,1,0 +pangolin,1,0 +pilosa,1,0 +molecula,1,0 `, }, } diff --git a/pql/ast.go b/pql/ast.go index 50b38e5c5..056a8fc85 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -444,6 +444,7 @@ var callInfoByFunc = map[string]callInfo{ "previous": nil, "aggregate": nil, "having": nil, + "sort": "", }, }, "Options": {