From b4cbd45b84285b9563b8b306ff7ed84687f725ba Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Thu, 16 Sep 2021 15:21:25 -0600 Subject: [PATCH] Implement SQL GROUP BY --- planner.go | 191 ++++++++++++++++++++++++++++++++++++++++++++---- planner_test.go | 68 +++++++++++++++++ 2 files changed, 245 insertions(+), 14 deletions(-) diff --git a/planner.go b/planner.go index 7dd39096e..4b080d279 100644 --- a/planner.go +++ b/planner.go @@ -80,27 +80,87 @@ func (p *Planner) planAggregateSelectStatement(ctx context.Context, stmt *sql2.S return nil, err } - // TODO: Support multiple aggregate calls. - if len(stmt.Columns) > 1 { - return nil, fmt.Errorf("only one call allowed in aggregate query") + // Extract calls and grouped expressions from column list. + // TODO: Recursively traverse all expression trees. + var calls []*sql2.Call + var aliases []string + // var groupByCols []*sql2.Ident // TODO: Convert to QualifiedRef + for _, c := range stmt.Columns { + aliases = append(aliases, c.Name()) + + switch c := c.Expr.(type) { + case *sql2.Call: + calls = append(calls, c) + case *sql2.Ident: + // groupByCols = append(groupByCols, c) + default: + return nil, fmt.Errorf("unsupported expression type in aggregate query: %T", c) + } } - // Extract aggregate call. - col := stmt.Columns[0] - var call *sql2.Call - switch expr := col.Expr.(type) { - case *sql2.Call: - call = expr - default: - return nil, fmt.Errorf("unsupported expression in aggregate query: %T", expr) + // TODO: Support multiple calls per query. + if len(calls) > 1 { + return nil, fmt.Errorf("only one aggregate call allowed") } - callName := strings.ToUpper(sql2.IdentName(call.Name)) + // Extract column names in GROUP BY clause. + var groupByColNames []string + for _, expr := range stmt.GroupByExprs { + switch expr := expr.(type) { + case *sql2.Ident: + groupByColNames = append(groupByColNames, expr.Name) + default: + return nil, fmt.Errorf("unsupported expression type in GROUP BY clause: %T", expr) + } + } + + // Extract aggregate call and build execution node. + callName := strings.ToUpper(sql2.IdentName(calls[0].Name)) switch callName { case "COUNT": - return NewCountNode(p.executor, indexName, col.Name(), cond), nil + if len(groupByColNames) == 0 { + return NewCountNode(p.executor, indexName, aliases[0], cond), nil + } + + var aggregate *pql.Call + if calls[0].Distinct.IsValid() { + if len(calls[0].Args) != 1 { + return nil, fmt.Errorf("distinct count must have exactly one field specified") + } + field, ok := calls[0].Args[0].(*sql2.Ident) + if !ok { + return nil, fmt.Errorf("distinct count argument must be a field name") + } + + aggregate = &pql.Call{ + Name: "Count", + Children: []*pql.Call{{ + Name: "Distinct", + Args: map[string]interface{}{"field": field.Name}, + }}, + } + } + + return NewGroupByNode(p.executor, indexName, groupByColNames, aliases, aggregate, cond), nil + + case "SUM": + if len(calls[0].Args) != 1 { + return nil, fmt.Errorf("sum must have exactly one field specified") + } + field, ok := calls[0].Args[0].(*sql2.Ident) + if !ok { + return nil, fmt.Errorf("sum argument must be a field name") + } + + aggregate := &pql.Call{ + Name: "Sum", + Args: map[string]interface{}{"field": field.Name}, + } + + return NewGroupByNode(p.executor, indexName, groupByColNames, aliases, aggregate, cond), nil + default: - return nil, fmt.Errorf("unsupported call in aggregate query: %T", callName) + return nil, fmt.Errorf("unsupported call in aggregate query: %s", callName) } // TODO: Support HAVING @@ -699,3 +759,106 @@ func (n *CountNode) Next(ctx context.Context) error { } func (n *CountNode) Row() []interface{} { return n.row } + +// GroupByNode executes an aggregate with a GROUP BY against a FeatureBase index. +type GroupByNode struct { + executor *executor + indexName string + columns []string + aliases []string + aggregate *pql.Call + cond *pql.Call + + result *GroupCounts + index int + + row []interface{} +} + +func NewGroupByNode(executor *executor, indexName string, columns, aliases []string, aggregate, cond *pql.Call) *GroupByNode { + return &GroupByNode{ + executor: executor, + indexName: indexName, + columns: columns, + aliases: aliases, + aggregate: aggregate, + cond: cond, + row: make([]interface{}, len(columns)+1), + } +} + +func (n *GroupByNode) Columns() []string { + return append([]string{"_aggregate"}, n.columns...) +} + +func (n *GroupByNode) First(ctx context.Context) error { + n.result = nil + return nil +} + +func (n *GroupByNode) Next(ctx context.Context) (err error) { + // Fetch resultset if it doesn't exist yet. + if n.result == nil { + if n.result, err = n.fetch(ctx); err != nil { + return err + } + } + + // Exit if no more rows exist. + if n.index >= len(n.result.groups) { + return sql.ErrNoRows + } + + // Copy results into current row. + group := n.result.groups[n.index] + n.index++ + + if n.aggregate != nil { + n.row[0] = int64(group.Agg) + } else { + n.row[0] = int64(group.Count) + } + + for i, g := range group.Group { + if g.Value != nil { + n.row[i+1] = *g.Value + } else if g.RowKey != "" { + n.row[i+1] = g.RowKey + } else { + n.row[i+1] = int64(g.RowID) + } + } + + return nil +} + +// fetch executes a call to compute the PQL results. +func (n *GroupByNode) fetch(ctx context.Context) (*GroupCounts, error) { + call := &pql.Call{ + Name: "GroupBy", + Args: map[string]interface{}{}, + } + + // Choose fields to group by. + for _, col := range n.columns { + call.Children = append(call.Children, &pql.Call{ + Name: "Rows", Args: map[string]interface{}{"_field": col}, + }) + } + + // Apply filter & aggregate, if set. + if n.aggregate != nil { + call.Args["aggregate"] = n.aggregate + } + if n.cond != nil { + call.Args["filter"] = n.cond + } + + result, err := n.executor.Execute(ctx, n.indexName, &pql.Query{Calls: []*pql.Call{call}}, nil, nil) + if err != nil { + return nil, err + } + return result.Results[0].(*GroupCounts), nil +} + +func (n *GroupByNode) Row() []interface{} { return n.row } diff --git a/planner_test.go b/planner_test.go index bb9ff3bd0..4748046b6 100644 --- a/planner_test.go +++ b/planner_test.go @@ -263,6 +263,74 @@ func TestPlanner_Select(t *testing.T) { }) } +func TestPlanner_GroupBy(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + + i0, err := c.GetHolder(0).CreateIndex("i0", pilosa.IndexOptions{TrackExistence: true}) + if err != nil { + t.Fatal(err) + } + defer i0.Close() + + if _, err := i0.CreateField("x"); err != nil { + t.Fatal(err) + } else if _, err := i0.CreateField("y", pilosa.OptFieldTypeInt(0, 1000)); err != nil { + t.Fatal(err) + } else if _, err := i0.CreateField("z", pilosa.OptFieldTypeInt(0, 1000)); err != nil { + t.Fatal(err) + } + + // Populate with data. + if _, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{ + Index: "i0", + Query: ` + Set(1, x=10) + Set(1, x=20) + Set(1, y=100) + Set(1, z=500) + + Set(2, x=10) + Set(2, y=200) + Set(2, z=500) + + Set(3, x=20) + Set(3, z=600) + `}); err != nil { + t.Fatal(err) + } + + t.Run("Count", func(t *testing.T) { + results := mustQueryRows(t, c.GetNode(0).Server, `SELECT COUNT(*), x FROM i0 GROUP BY x`) + if diff := cmp.Diff([][]interface{}{ + {int64(2), int64(10)}, + {int64(2), int64(20)}, + }, results); diff != "" { + t.Fatal(diff) + } + }) + + t.Run("DistinctCount", func(t *testing.T) { + results := mustQueryRows(t, c.GetNode(0).Server, `SELECT COUNT(DISTINCT z), x FROM i0 GROUP BY x`) + if diff := cmp.Diff([][]interface{}{ + {int64(1), int64(10)}, + {int64(2), int64(20)}, + }, results); diff != "" { + t.Fatal(diff) + } + }) + + t.Run("Sum", func(t *testing.T) { + results := mustQueryRows(t, c.GetNode(0).Server, `SELECT sum(y), x FROM i0 GROUP BY x`) + if diff := cmp.Diff([][]interface{}{ + {int64(300), int64(10)}, + {int64(100), int64(20)}, + }, results); diff != "" { + t.Fatal(diff) + } + }) +} + func mustQueryRows(tb testing.TB, svr *pilosa.Server, q string) [][]interface{} { tb.Helper()