diff --git a/spanner/spannertest/db_query.go b/spanner/spannertest/db_query.go index ae2657d9ba3..d46462b62f3 100644 --- a/spanner/spannertest/db_query.go +++ b/spanner/spannertest/db_query.go @@ -518,8 +518,7 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter } // Handle aggregation. - // TODO: Support more than one aggregation function; does Spanner support that? - aggI := -1 + var aggI []int for i, e := range sel.List { // Supported aggregate funcs have exactly one arg. f, ok := e.(spansql.Func) @@ -530,12 +529,9 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter if !ok { continue } - if aggI > -1 { - return nil, fmt.Errorf("only one aggregate function is supported") - } - aggI = i + aggI = append(aggI, i) } - if aggI > -1 { + if len(aggI) > 0 { raw, err := toRawIter(ri) if err != nil { return nil, err @@ -545,20 +541,6 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter // This may result in a [0,0) entry for empty inputs. rowGroups = [][2]int{{0, len(raw.rows)}} } - fexpr := sel.List[aggI].(spansql.Func) - fn := aggregateFuncs[fexpr.Name] - starArg := fexpr.Args[0] == spansql.Star - if starArg && !fn.AcceptStar { - return nil, fmt.Errorf("aggregate function %s does not accept * as an argument", fexpr.Name) - } - var argType spansql.Type - if !starArg { - ci, err := ec.colInfo(fexpr.Args[0]) - if err != nil { - return nil, fmt.Errorf("evaluating aggregate function %s arg type: %v", fexpr.Name, err) - } - argType = ci.Type - } // Prepare output. rawOut := &rawIter{ @@ -569,29 +551,8 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter cols: append([]colInfo(nil), raw.cols...), } - var aggType spansql.Type + aggType := make([]*spansql.Type, len(aggI)) for _, rg := range rowGroups { - // Compute aggregate value across this group. - var values []interface{} - for i := rg[0]; i < rg[1]; i++ { - ec.row = raw.rows[i] - if starArg { - // A non-NULL placeholder is sufficient for aggregation. - values = append(values, 1) - } else { - x, err := ec.evalExpr(fexpr.Args[0]) - if err != nil { - return nil, err - } - values = append(values, x) - } - } - x, typ, err := fn.Eval(values, argType) - if err != nil { - return nil, err - } - aggType = typ - var outRow row // Output for the row group is the first row of the group (arbitrary, // but it should be representative), and the aggregate value. @@ -609,27 +570,68 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter outRow = append(outRow, nil) } } - outRow = append(outRow, x) + + for j, aggI := range aggI { + fexpr := sel.List[aggI].(spansql.Func) + fn := aggregateFuncs[fexpr.Name] + starArg := fexpr.Args[0] == spansql.Star + if starArg && !fn.AcceptStar { + return nil, fmt.Errorf("aggregate function %s does not accept * as an argument", fexpr.Name) + } + var argType spansql.Type + if !starArg { + ci, err := ec.colInfo(fexpr.Args[0]) + if err != nil { + return nil, fmt.Errorf("evaluating aggregate function %s arg type: %v", fexpr.Name, err) + } + argType = ci.Type + } + + // Compute aggregate value across this group. + var values []interface{} + for i := rg[0]; i < rg[1]; i++ { + ec.row = raw.rows[i] + if starArg { + // A non-NULL placeholder is sufficient for aggregation. + values = append(values, 1) + } else { + x, err := ec.evalExpr(fexpr.Args[0]) + if err != nil { + return nil, err + } + values = append(values, x) + } + } + x, typ, err := fn.Eval(values, argType) + if err != nil { + return nil, err + } + aggType[j] = &typ + + outRow = append(outRow, x) + } rawOut.rows = append(rawOut.rows, outRow) } - if aggType == (spansql.Type{}) { - // Fallback; there might not be any groups. - // TODO: Should this be in aggregateFunc? - aggType = int64Type + for j, aggI := range aggI { + fexpr := sel.List[aggI].(spansql.Func) + if aggType[j] == nil { + // Fallback; there might not be any groups. + // TODO: Should this be in aggregateFunc? + aggType[j] = &int64Type + } + rawOut.cols = append(rawOut.cols, colInfo{ + Name: spansql.ID(fexpr.SQL()), // TODO: this is a bit hokey, but it is output only + Type: *aggType[j], + AggIndex: aggI + 1, + }) + sel.List[aggI] = aggSentinel{ // Mutate query so evalExpr in selIter picks out the new value. + Type: *aggType[j], + AggIndex: aggI + 1, + } } - rawOut.cols = append(raw.cols, colInfo{ - Name: spansql.ID(fexpr.SQL()), // TODO: this is a bit hokey, but it is output only - Type: aggType, - AggIndex: aggI + 1, - }) - ri = rawOut ec.cols = rawOut.cols - sel.List[aggI] = aggSentinel{ // Mutate query so evalExpr in selIter picks out the new value. - Type: aggType, - AggIndex: aggI + 1, - } } // TODO: Support table sampling. diff --git a/spanner/spannertest/integration_test.go b/spanner/spannertest/integration_test.go index 6c20fed33e0..bdcd4cfe651 100644 --- a/spanner/spannertest/integration_test.go +++ b/spanner/spannertest/integration_test.go @@ -1145,6 +1145,62 @@ func TestIntegration_ReadsAndQueries(t *testing.T) { {"Daniel"}, }, }, + { + `SELECT MIN(Name), MAX(Name) FROM Staff`, + nil, + [][]interface{}{ + {"Daniel", "Teal'c"}, + }, + }, + { + `SELECT Cool, MIN(Name), MAX(Name), COUNT(*) FROM Staff GROUP BY Cool ORDER BY Cool`, + nil, + [][]interface{}{ + {nil, "George", "Jack", int64(2)}, + {false, "Daniel", "Sam", int64(2)}, + {true, "Teal'c", "Teal'c", int64(1)}, + }, + }, + { + `SELECT Tenure/2, Cool, Name FROM Staff WHERE Tenure/2 > 5`, + nil, + [][]interface{}{ + {float64(5.5), false, "Daniel"}, + }, + }, + { + `SELECT Tenure/2, MAX(Cool) FROM Staff WHERE Tenure/2 > 5 GROUP BY Tenure/2`, + nil, + [][]interface{}{ + {float64(5.5), false}, + }, + }, + { + `SELECT Tenure/2, Cool, MIN(Name) FROM Staff WHERE Tenure/2 >= 4 GROUP BY Tenure/2, Cool ORDER BY Cool DESC, Tenure/2`, + nil, + [][]interface{}{ + {float64(4), true, "Teal'c"}, + {float64(4.5), false, "Sam"}, + {float64(5.5), false, "Daniel"}, + {float64(5), nil, "Jack"}, + }, + }, + { + `SELECT MIN(Cool), MAX(Cool), MIN(Tenure), MAX(Tenure), MIN(Height), MAX(Height), MIN(Name), MAX(Name), COUNT(*) FROM Staff`, + nil, + [][]interface{}{ + {false, true, int64(6), int64(11), 1.73, 1.91, "Daniel", "Teal'c", int64(5)}, + }, + }, + { + `SELECT Cool, MIN(Tenure), MAX(Tenure), MIN(Height), MAX(Height), MIN(Name), MAX(Name), COUNT(*) FROM Staff GROUP BY Cool ORDER BY Cool`, + nil, + [][]interface{}{ + {nil, int64(6), int64(10), 1.73, 1.85, "George", "Jack", int64(2)}, + {false, int64(9), int64(11), 1.75, 1.83, "Daniel", "Sam", int64(2)}, + {true, int64(8), int64(8), 1.91, 1.91, "Teal'c", "Teal'c", int64(1)}, + }, + }, } var failures int for _, test := range tests {