feat(spanner/spannertest): implement SELECT ... FROM UNNEST(...)

spanner/spannertest/db_query.go

@@ -412,6 +412,9 @@ func (d *database) queryContext(q spansql.Query, params queryParams) (*queryCont
 				return err
 			}
 			return findTables(sf.RHS)
+		case spansql.SelectFromUnnest:
+			// TODO: if array paths get supported, this will need more work.
+			return nil
 		}
 	}
 	for _, sf := range q.Select.From {
@@ -552,7 +555,7 @@ func (d *database) evalSelect(sel spansql.Select, qc *queryContext) (si *selIter
 		if !starArg {
 			ci, err := ec.colInfo(fexpr.Args[0])
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("evaluating aggregate function %s arg type: %v", fexpr.Name, err)
 			}
 			argType = ci.Type
 		}
@@ -708,6 +711,39 @@ func (d *database) evalSelectFrom(qc *queryContext, ec evalContext, sf spansql.S
 			return ec, nil, err
 		}
 		return ec, ji, nil
+	case spansql.SelectFromUnnest:
+		// TODO: Do all relevant types flow through here? Path expressions might be tricky here.
+		col, err := ec.colInfo(sf.Expr)
+		if err != nil {
+			return ec, nil, fmt.Errorf("evaluating type of UNNEST arg: %v", err)
+		}
+		if !col.Type.Array {
+			return ec, nil, fmt.Errorf("type of UNNEST arg is non-array %s", col.Type.SQL())
+		}
+		// The output of this UNNEST is the non-array version.
+		col.Name = sf.Alias // may be empty
+		col.Type.Array = false
+
+		// Evaluate the expression, and yield a virtual table with one column.
+		e, err := ec.evalExpr(sf.Expr)
+		if err != nil {
+			return ec, nil, fmt.Errorf("evaluating UNNEST arg: %v", err)
+		}
+		arr, ok := e.([]interface{})
+		if !ok {
+			return ec, nil, fmt.Errorf("evaluating UNNEST arg gave %t, want array", e)
+		}
+		var rows []row
+		for _, v := range arr {
+			rows = append(rows, row{v})
+		}
+
+		ri := &rawIter{
+			cols: []colInfo{col},
+			rows: rows,
+		}
+		ec.cols = ri.cols
+		return ec, ri, nil
 	}
 }
 

spanner/spannertest/integration_test.go

@@ -890,10 +890,12 @@ func TestIntegration_ReadsAndQueries(t *testing.T) {
 			},
 		},
 		{
-			`SELECT AVG(Height) FROM Staff WHERE ID <= 2`,
-			nil,
+			// From https://cloud.google.com/spanner/docs/aggregate_functions#avg
+			// but using a param for the array since array literals aren't supported yet.
+			`SELECT AVG(x) AS avg FROM UNNEST(@p) AS x`,
+			map[string]interface{}{"p": []int64{0, 2, 4, 4, 5}},
 			[][]interface{}{
-				{float64(1.84)},
+				{float64(3)},
 			},
 		},
 		{

spanner/spansql/parser.go

@@ -1936,8 +1936,31 @@ func (p *parser) parseSelectFrom() (SelectFrom, *parseError) {
 			{ INNER | CROSS | FULL [OUTER] | LEFT [OUTER] | RIGHT [OUTER] }
 	*/
 
+	if p.eat("UNNEST") {
+		if err := p.expect("("); err != nil {
+			return nil, err
+		}
+		e, err := p.parseExpr()
+		if err != nil {
+			return nil, err
+		}
+		if err := p.expect(")"); err != nil {
+			return nil, err
+		}
+		sfu := SelectFromUnnest{Expr: e}
+		if p.eat("AS") { // TODO: The "AS" keyword is optional.
+			alias, err := p.parseAlias()
+			if err != nil {
+				return nil, err
+			}
+			sfu.Alias = alias
+		}
+		// TODO: hint, offset
+		return sfu, nil
+	}
+
 	// A join starts with a from_item, so that can't be detected in advance.
-	// TODO: Support more than table name or join.
+	// TODO: Support subquery, field_path, array_path, WITH.
 	// TODO: Verify associativity of multile joins.
 
 	tname, err := p.parseTableOrIndexOrColumnName()

spanner/spansql/parser_test.go

@@ -190,6 +190,17 @@ func TestParseQuery(t *testing.T) {
 				},
 			},
 		},
+		{`SELECT * FROM UNNEST (@p) AS data`, // array literals aren't yet supported
+			Query{
+				Select: Select{
+					List: []Expr{Star},
+					From: []SelectFrom{SelectFromUnnest{
+						Expr:  Param("p"),
+						Alias: ID("data"),
+					}},
+				},
+			},
+		},
 	}
 	for _, test := range tests {
 		got, err := ParseQuery(test.in)

spanner/spansql/sql.go

@@ -345,6 +345,14 @@ var joinTypes = map[JoinType]string{
 	RightJoin: "RIGHT",
 }
 
+func (sfu SelectFromUnnest) SQL() string {
+	str := "UNNEST(" + sfu.Expr.SQL() + ")"
+	if sfu.Alias != "" {
+		str += " AS " + sfu.Alias.SQL()
+	}
+	return str
+}
+
 func (o Order) SQL() string { return buildSQL(o) }
 func (o Order) addSQL(sb *strings.Builder) {
 	o.Expr.addSQL(sb)

spanner/spansql/types.go

@@ -392,6 +392,17 @@ const (
 	RightJoin
 )
 
+// SelectFromUnnest is a SelectFrom that yields a virtual table from an array.
+// https://cloud.google.com/spanner/docs/query-syntax#unnest
+type SelectFromUnnest struct {
+	Expr  Expr
+	Alias ID // empty if not aliased
+
+	// TODO: Implicit
+}
+
+func (SelectFromUnnest) isSelectFrom() {}
+
 // TODO: SelectFromSubquery, etc.
 
 type Order struct {