Runtime: Simplify code for unnest dimensions (#5154)

* Runtime: Simplify code for unnest dimensions

* Nits

runtime/metricsview/ast.go

@@ -20,6 +20,7 @@ type AST struct {
 	underlyingTable *string
 	underlyingWhere *ExprNode
 	dimFields       []FieldNode
+	unnests         []string
 	nextIdentifier  int
 
 	// Contextual info for building the AST
@@ -44,6 +45,7 @@ type SelectNode struct {
 	LeftJoinSelects      []*SelectNode    // Sub-selects to left join onto FromSelect, to enable "per-dimension" measures
 	JoinComparisonSelect *SelectNode      // Sub-select to join onto FromSelect for comparison measures
 	JoinComparisonType   string           // Type of join to use for JoinComparisonSelect
+	Unnests              []string         // Unnest expressions to add in the FROM clause
 	Group                bool             // Whether the SELECT is grouped. If yes, it will group by all DimFields.
 	Where                *ExprNode        // Expression for the WHERE clause
 	TimeWhere            *ExprNode        // Expression for the time range to add to the WHERE clause
@@ -57,11 +59,9 @@ type SelectNode struct {
 // The Name must always match a the name of a dimension/measure in the metrics view or a computed field specified in the request.
 // This means that if two columns in different places in the AST have the same Name, they're guaranteed to resolve to the same value.
 type FieldNode struct {
-	Name        string
-	Label       string
-	Expr        string
-	Unnest      bool
-	UnnestAlias string
+	Name  string
+	Label string
+	Expr  string
 }
 
 // ExprNode represents an expression for a WHERE clause.
@@ -120,27 +120,35 @@ func NewAST(mv *runtimev1.MetricsViewSpec, sec *runtime.ResolvedMetricsViewSecur
 
 	// Build dimensions to apply against the underlying SELECT.
 	// We cache these in the AST type because when resolving expressions and adding new JOINs, we need the ability to reference these.
-	dimFields := make([]FieldNode, 0, len(ast.query.Dimensions))
+	ast.dimFields = make([]FieldNode, 0, len(ast.query.Dimensions))
 	for _, qd := range ast.query.Dimensions {
 		dim, err := ast.resolveDimension(qd, true)
 		if err != nil {
 			return nil, fmt.Errorf("invalid dimension %q: %w", qd.Name, err)
 		}
 
-		var unnestAlias string
+		f := FieldNode{
+			Name:  dim.Name,
+			Label: dim.Label,
+			Expr:  ast.dialect.MetricsViewDimensionExpression(dim),
+		}
+
 		if dim.Unnest {
-			unnestAlias = ast.generateIdentifier()
+			unnestAlias := ast.generateIdentifier()
+
+			tblWithAlias, auto, err := ast.dialect.LateralUnnest(f.Expr, unnestAlias, f.Name)
+			if err != nil {
+				return nil, fmt.Errorf("failed to unnest field %q: %w", f.Name, err)
+			}
+
+			if !auto {
+				ast.unnests = append(ast.unnests, tblWithAlias)
+				f.Expr = ast.sqlForMember(unnestAlias, f.Name)
+			}
 		}
 
-		dimFields = append(dimFields, FieldNode{
-			Name:        dim.Name,
-			Label:       dim.Label,
-			Expr:        ast.dialect.MetricsViewDimensionExpression(dim),
-			Unnest:      dim.Unnest,
-			UnnestAlias: unnestAlias,
-		})
+		ast.dimFields = append(ast.dimFields, f)
 	}
-	ast.dimFields = dimFields
 
 	// Build underlying SELECT
 	tbl := ast.dialect.EscapeTable(mv.Database, mv.DatabaseSchema, mv.Table)
@@ -523,6 +531,7 @@ func (a *AST) buildBaseSelect(alias string, tr *TimeRange) (*SelectNode, error)
 	n := &SelectNode{
 		Alias:     alias,
 		DimFields: a.dimFields,
+		Unnests:   a.unnests,
 		Group:     true,
 		FromTable: a.underlyingTable,
 		Where:     a.underlyingWhere,
@@ -567,6 +576,7 @@ func (a *AST) buildSpineSelect(alias string, spine *Spine, tr *TimeRange) (*Sele
 		n := &SelectNode{
 			Alias:     alias,
 			DimFields: a.dimFields,
+			Unnests:   a.unnests,
 			Group:     true,
 			FromTable: a.underlyingTable,
 		}
@@ -867,7 +877,6 @@ func (a *AST) wrapSelect(s *SelectNode, innerAlias string) {
 			Name:  f.Name,
 			Label: f.Label,
 			Expr:  a.sqlForMember(cpy.Alias, f.Name),
-			// Not copying Unnest because we should only unnest once (at the innermost level).
 		})
 	}
 
@@ -886,6 +895,7 @@ func (a *AST) wrapSelect(s *SelectNode, innerAlias string) {
 	s.LeftJoinSelects = nil
 	s.JoinComparisonSelect = nil
 	s.JoinComparisonType = ""
+	s.Unnests = nil
 	s.Group = false
 	s.Where = nil
 	s.TimeWhere = nil
@@ -1043,12 +1053,7 @@ func (a *AST) sqlForMeasure(m *runtimev1.MetricsViewSpec_MeasureV2, n *SelectNod
 				b.WriteString(", ")
 			}
 
-			expr := f.Expr
-			if f.Unnest {
-				expr = a.sqlForMember(f.UnnestAlias, f.Name)
-			}
-
-			b.WriteString(expr)
+			b.WriteString(f.Expr)
 		}
 	}
 	if len(orderFields) > 0 {
@@ -1061,12 +1066,7 @@ func (a *AST) sqlForMeasure(m *runtimev1.MetricsViewSpec_MeasureV2, n *SelectNod
 				b.WriteString(", ")
 			}
 
-			expr := f.Expr
-			if f.Unnest {
-				expr = a.sqlForMember(f.UnnestAlias, f.Name)
-			}
-
-			b.WriteString(expr)
+			b.WriteString(f.Expr)
 			if orderDesc[i] {
 				b.WriteString(" DESC")
 			}

runtime/metricsview/astexpr.go

@@ -534,12 +534,8 @@ func (b *sqlExprBuilder) sqlForName(name string) (expr string, unnest bool, err
 		// First, search for the dimension in the ASTs dimension fields (this also covers any computed dimension)
 		for _, f := range b.ast.dimFields {
 			if f.Name == name {
-				if f.Unnest {
-					// Since it's unnested, we need to reference the unnested alias.
-					// Note that we return "false" for "unnest" because it will already have been unnested since it's one of the dimensions included in the query,
-					// so we can filter against it as if it's a normal dimension.
-					return b.ast.sqlForMember(f.UnnestAlias, f.Name), false, nil
-				}
+				// Note that we return "false" even though it may be an unnest dimension because it will already have been unnested since it's one of the dimensions included in the query.
+				// So we can filter against it as if it's a normal dimension.
 				return f.Expr, false, nil
 			}
 		}

runtime/metricsview/astsql.go

@@ -1,7 +1,6 @@
 package metricsview
 
 import (
-	"fmt"
 	"strconv"
 	"strings"
 )
@@ -82,13 +81,8 @@ func (b *sqlBuilder) writeSelect(n *SelectNode) error {
 			b.out.WriteString(", ")
 		}
 
-		expr := f.Expr
-		if f.Unnest {
-			expr = b.ast.sqlForMember(f.UnnestAlias, f.Name)
-		}
-
 		b.out.WriteByte('(')
-		b.out.WriteString(expr)
+		b.out.WriteString(f.Expr)
 		b.out.WriteString(") AS ")
 		b.out.WriteString(b.ast.dialect.EscapeIdentifier(f.Name))
 	}
@@ -109,22 +103,9 @@ func (b *sqlBuilder) writeSelect(n *SelectNode) error {
 		b.out.WriteString(*n.FromTable)
 
 		// Add unnest joins. We only and always apply these against FromPlain (ensuring they are already unnested when referenced in outer SELECTs).
-		for _, f := range n.DimFields {
-			if !f.Unnest {
-				continue
-			}
-
-			tblWithAlias, auto, err := b.ast.dialect.LateralUnnest(f.Expr, f.UnnestAlias, f.Name)
-			if err != nil {
-				return fmt.Errorf("failed to unnest field %q: %w", f.Name, err)
-			}
-
-			if auto {
-				continue
-			}
-
+		for _, u := range n.Unnests {
 			b.out.WriteString(", ")
-			b.out.WriteString(tblWithAlias)
+			b.out.WriteString(u)
 		}
 	} else if n.FromSelect != nil {
 		b.out.WriteByte('(')