sql/parser: Introduce IndirectionExpr and support for array subscripting

pkg/sql/parser/eval.go

@@ -32,6 +32,7 @@ import (
 	"gopkg.in/inf.v0"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/decimal"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -1939,6 +1940,48 @@ func (expr *CastExpr) Eval(ctx *EvalContext) (Datum, error) {
 	return nil, fmt.Errorf("invalid cast: %s -> %s", d.ResolvedType(), expr.Type)
 }
 
+// Eval implements the TypedExpr interface.
+func (expr *IndirectionExpr) Eval(ctx *EvalContext) (Datum, error) {
+	var subscriptIdx int
+	for i, part := range expr.Indirection {
+		switch t := part.(type) {
+		case *ArraySubscript:
+			if t.Slice {
+				return nil, util.UnimplementedWithIssueErrorf(2115, "ARRAY slicing in %s", expr)
+			}
+			if i > 0 {
+				return nil, util.UnimplementedWithIssueErrorf(2115, "multidimensional ARRAY %s", expr)
+			}
+
+			d, err := t.Begin.(TypedExpr).Eval(ctx)
+			if err != nil {
+				return nil, err
+			}
+			if d == DNull {
+				return d, nil
+			}
+			subscriptIdx = int(*d.(*DInt))
+		default:
+			return nil, errors.Errorf("syntax not yet supported: %s", expr.Indirection)
+		}
+	}
+
+	d, err := expr.Expr.(TypedExpr).Eval(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if d == DNull {
+		return d, nil
+	}
+
+	// Index into the DArray, using 1-indexing.
+	arr := d.(*DArray)
+	if subscriptIdx < 1 || subscriptIdx > len(*arr) {
+		return DNull, nil
+	}
+	return (*arr)[subscriptIdx-1], nil
+}
+
 // Eval implements the TypedExpr interface.
 func (expr *AnnotateTypeExpr) Eval(ctx *EvalContext) (Datum, error) {
 	return expr.Expr.(TypedExpr).Eval(ctx)

pkg/sql/parser/expr.go

@@ -1002,6 +1002,20 @@ func (node *CastExpr) castTypeAndValidArgTypes() (Type, []Type) {
 	return colTypeToTypeAndValidArgTypes(node.Type)
 }
 
+// IndirectionExpr represents a subscript expression.
+type IndirectionExpr struct {
+	Expr        Expr
+	Indirection UnresolvedName
+
+	typeAnnotation
+}
+
+// Format implements the NodeFormatter interface.
+func (node *IndirectionExpr) Format(buf *bytes.Buffer, f FmtFlags) {
+	FormatNode(buf, f, node.Expr)
+	FormatNode(buf, f, node.Indirection)
+}
+
 type annotateSyntaxMode int
 
 const (
@@ -1088,6 +1102,7 @@ func (node Exprs) String() string             { return AsString(node) }
 func (node *FuncExpr) String() string         { return AsString(node) }
 func (node *IfExpr) String() string           { return AsString(node) }
 func (node *IndexedVar) String() string       { return AsString(node) }
+func (node *IndirectionExpr) String() string  { return AsString(node) }
 func (node *IsOfTypeExpr) String() string     { return AsString(node) }
 func (node Name) String() string              { return AsString(node) }
 func (node *NotExpr) String() string          { return AsString(node) }

pkg/sql/parser/name_part.go

@@ -127,15 +127,20 @@ func (u UnqualifiedStar) String() string                     { return AsString(u
 type ArraySubscript struct {
 	Begin Expr
 	End   Expr
+	Slice bool
 }
 
 // Format implements the NodeFormatter interface.
 func (a *ArraySubscript) Format(buf *bytes.Buffer, f FmtFlags) {
 	buf.WriteByte('[')
-	FormatNode(buf, f, a.Begin)
-	if a.End != nil {
+	if a.Begin != nil {
+		FormatNode(buf, f, a.Begin)
+	}
+	if a.Slice {
 		buf.WriteByte(':')
-		FormatNode(buf, f, a.End)
+		if a.End != nil {
+			FormatNode(buf, f, a.End)
+		}
 	}
 	buf.WriteByte(']')
 }

pkg/sql/parser/parse_test.go

@@ -298,8 +298,11 @@ func TestParse(t *testing.T) {
 		{`SELECT true AND NULL`},
 		{`SELECT true = false`},
 		{`SELECT (true = false)`},
+		{`SELECT (ARRAY['a', 'b'])[2]`},
 		{`SELECT (SELECT 1)`},
 		{`SELECT ((SELECT 1))`},
+		{`SELECT (SELECT ARRAY['a', 'b'])[2]`},
+		{`SELECT ((SELECT ARRAY['a', 'b']))[2]`},
 		{`SELECT ((((VALUES (1)))))`},
 		{`SELECT EXISTS (SELECT 1)`},
 		{`SELECT (VALUES (1))`},
@@ -326,6 +329,9 @@ func TestParse(t *testing.T) {
 		{`SELECT a.b.* FROM t`},
 		{`SELECT a.b[1] FROM t`},
 		{`SELECT a.b[1 + 1:4][3] FROM t`},
+		{`SELECT a.b[:4][3] FROM t`},
+		{`SELECT a.b[1 + 1:][3] FROM t`},
+		{`SELECT a.b[:][3] FROM t`},
 		{`SELECT 'a' FROM t`},
 		{`SELECT 'a' FROM t@bar`},
 		{`SELECT 'a' FROM t@{NO_INDEX_JOIN}`},