PIVOT ON support

proto/gen/rill/runtime/v1/runtime.swagger.yaml

@@ -1053,6 +1053,10 @@ paths:
                 title: Deprecated in favor of time_range
               filter:
                 $ref: '#/definitions/v1MetricsViewFilter'
+              pivotOn:
+                type: array
+                items:
+                  type: string
               limit:
                 type: string
                 format: int64
@@ -3383,6 +3387,10 @@ definitions:
         title: Deprecated in favor of time_range
       filter:
         $ref: '#/definitions/v1MetricsViewFilter'
+      pivotOn:
+        type: array
+        items:
+          type: string
       limit:
         type: string
         format: int64

proto/rill/runtime/v1/queries.proto

@@ -280,6 +280,7 @@ message MetricsViewAggregationRequest {
   google.protobuf.Timestamp time_start = 6; // Deprecated in favor of time_range
   google.protobuf.Timestamp time_end = 7; // Deprecated in favor of time_range
   MetricsViewFilter filter = 8;
+  repeated string pivot_on = 13;
   int64 limit = 9 [(validate.rules).int64.gte = 0];
   int64 offset = 10 [(validate.rules).int64.gte = 0];
   int32 priority = 11;

runtime/drivers/duckdb/transporter_sqlstore_to_duckDB.go

@@ -122,7 +122,7 @@ func (s *sqlStoreToDuckDB) transferFromRowIterator(ctx context.Context, iter dri
 		p.Target(int64(total), drivers.ProgressUnitRecord)
 	}
 	// create table
-	qry, err := createTableQuery(schema, table)
+	qry, err := CreateTableQuery(schema, table)
 	if err != nil {
 		return err
 	}
@@ -132,7 +132,7 @@ func (s *sqlStoreToDuckDB) transferFromRowIterator(ctx context.Context, iter dri
 	}
 
 	return s.to.WithConnection(ctx, 1, true, false, func(ctx, ensuredCtx context.Context, conn *sql.Conn) error {
-		return rawConn(conn, func(conn driver.Conn) error {
+		return RawConn(conn, func(conn driver.Conn) error {
 			a, err := duckdb.NewAppenderFromConn(conn, "", table)
 			if err != nil {
 				return err
@@ -175,7 +175,7 @@ func (s *sqlStoreToDuckDB) transferFromRowIterator(ctx context.Context, iter dri
 	})
 }
 
-func createTableQuery(schema *runtimev1.StructType, name string) (string, error) {
+func CreateTableQuery(schema *runtimev1.StructType, name string) (string, error) {
 	query := fmt.Sprintf("CREATE OR REPLACE TABLE %s(", safeName(name))
 	for i, s := range schema.Fields {
 		i++

runtime/drivers/duckdb/utils.go

@@ -13,7 +13,7 @@ import (
 )
 
 // rawConn is similar to *sql.Conn.Raw, but additionally unwraps otelsql (which we use for instrumentation).
-func rawConn(conn *sql.Conn, f func(driver.Conn) error) error {
+func RawConn(conn *sql.Conn, f func(driver.Conn) error) error {
 	return conn.Raw(func(raw any) error {
 		// For details, see: https://github.com/XSAM/otelsql/issues/98
 		if c, ok := raw.(interface{ Raw() driver.Conn }); ok {

runtime/queries/metricsview_aggregation.go

@@ -2,15 +2,28 @@ package queries
 
 import (
 	"context"
+	databasesql "database/sql"
+	"database/sql/driver"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"strings"
+	"sync"
 
+	"github.com/marcboeker/go-duckdb"
 	runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
 	"github.com/rilldata/rill/runtime"
 	"github.com/rilldata/rill/runtime/drivers"
+	duckdbolap "github.com/rilldata/rill/runtime/drivers/duckdb"
+	"github.com/rilldata/rill/runtime/pkg/activity"
+	"go.uber.org/zap"
+)
+
+var (
+	dbOnce    sync.Once
+	handle    drivers.Handle
+	errHandle error
 )
 
 type MetricsViewAggregation struct {
@@ -25,6 +38,7 @@ type MetricsViewAggregation struct {
 	Offset             int64                                        `json:"offset,omitempty"`
 	MetricsView        *runtimev1.MetricsViewSpec                   `json:"-"`
 	ResolvedMVSecurity *runtime.ResolvedMetricsViewSecurity         `json:"security"`
+	PivotOn            []string                                     `json:"pivot_on,omitempty"`
 
 	Result *runtimev1.MetricsViewAggregationResponse `json:"-"`
 }
@@ -76,14 +90,159 @@ func (q *MetricsViewAggregation) Resolve(ctx context.Context, rt *runtime.Runtim
 		return fmt.Errorf("metrics view '%s' does not have a time dimension", q.MetricsView)
 	}
 
-	// Build query
-	sql, args, err := q.buildMetricsAggregationSQL(q.MetricsView, olap.Dialect(), q.ResolvedMVSecurity)
+	// execute druid query
+	sqlString, args, err := q.buildMetricsAggregationSQL(q.MetricsView, olap.Dialect(), q.ResolvedMVSecurity)
 	if err != nil {
 		return fmt.Errorf("error building query: %w", err)
 	}
 
-	// Execute
-	schema, data, err := olapQuery(ctx, olap, priority, sql, args)
+	if len(q.PivotOn) == 0 {
+		schema, data, err := olapQuery(ctx, olap, priority, sqlString, args)
+		if err != nil {
+			return err
+		}
+
+		q.Result = &runtimev1.MetricsViewAggregationResponse{
+			Schema: schema,
+			Data:   data,
+		}
+		return nil
+	}
+
+	if olap.Dialect() == drivers.DialectDuckDB {
+		return olap.WithConnection(ctx, priority, false, false, func(ctx context.Context, ensuredCtx context.Context, conn *databasesql.Conn) error {
+			temporaryTableName := tempName("_for_pivot_")
+
+			err = olap.Exec(ctx, &drivers.Statement{
+				Query:    fmt.Sprintf("CREATE TEMPORARY TABLE %s AS %s", temporaryTableName, sqlString),
+				Args:     args,
+				Priority: priority,
+			})
+			if err != nil {
+				return err
+			}
+			defer func() {
+				_ = olap.Exec(ensuredCtx, &drivers.Statement{
+					Query: `DROP TABLE "` + temporaryTableName + `"`,
+				})
+			}()
+
+			return q.pivotOn(ctx, olap, temporaryTableName)
+		})
+	}
+	dbOnce.Do(func() {
+		handle, errHandle = duckdbolap.Driver{}.Open(map[string]any{"pool_size": 10}, false, activity.NewNoopClient(), zap.NewNop())
+	})
+	if errHandle != nil {
+		return errHandle
+	}
+
+	schema, data, err := olapQuery(ctx, olap, priority, sqlString, args)
+	if err != nil {
+		return err
+	}
+
+	duckDBOLAP, _ := handle.AsOLAP("")
+	err = duckDBOLAP.WithConnection(ctx, priority, false, false, func(ctx context.Context, ensuredCtx context.Context, conn *databasesql.Conn) error {
+		temporaryTableName := tempName("_for_pivot_")
+		createTableSQL, err := duckdbolap.CreateTableQuery(schema, temporaryTableName)
+		if err != nil {
+			return err
+		}
+
+		err = duckDBOLAP.Exec(ctx, &drivers.Statement{
+			Query: createTableSQL,
+		})
+		if err != nil {
+			return err
+		}
+		defer func() {
+			_ = duckDBOLAP.Exec(ensuredCtx, &drivers.Statement{
+				Query: `DROP TABLE "` + temporaryTableName + `"`,
+			})
+		}()
+
+		err = duckdbolap.RawConn(conn, func(conn driver.Conn) error {
+			appender, err := duckdb.NewAppenderFromConn(conn, "", temporaryTableName)
+			if err != nil {
+				return err
+			}
+			defer appender.Close()
+
+			batchSize := 10000
+			arr := make([]driver.Value, 0, len(schema.Fields))
+			count := 0
+			for _, row := range data {
+				for _, key := range schema.Fields {
+					arr = append(arr, row.Fields[key.Name].AsInterface())
+				}
+				err = appender.AppendRowArray(arr)
+				if err != nil {
+					return err
+				}
+				arr = arr[:0]
+				count++
+				if count >= batchSize {
+					appender.Flush()
+					count = 0
+				}
+			}
+			appender.Flush()
+
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+
+		return q.pivotOn(ctx, duckDBOLAP, temporaryTableName)
+	})
+	return err
+}
+
+func (q *MetricsViewAggregation) pivotOn(ctx context.Context, olap drivers.OLAPStore, temporaryTableName string) error {
+	measureCols := make([]string, 0, len(q.Measures))
+	for _, m := range q.Measures {
+		sn := safeName(m.Name)
+		measureCols = append(measureCols, fmt.Sprintf("LAST(%s) as %s", sn, sn))
+	}
+
+	sortingCriteria := make([]string, 0, len(q.Sort))
+	for _, s := range q.Sort {
+		sortCriterion := safeName(s.Name)
+		if s.Desc {
+			sortCriterion += " DESC"
+		}
+		if olap.Dialect() == drivers.DialectDuckDB {
+			sortCriterion += " NULLS LAST"
+		}
+		sortingCriteria = append(sortingCriteria, sortCriterion)
+	}
+
+	orderClause := ""
+	if len(sortingCriteria) > 0 {
+		orderClause = "ORDER BY " + strings.Join(sortingCriteria, ", ")
+	}
+
+	var limitClause string
+	if q.Limit != nil {
+		if *q.Limit == 0 {
+			*q.Limit = 100
+		}
+		limitClause = fmt.Sprintf("LIMIT %d", *q.Limit)
+	}
+
+	// execute duckdb pivot
+	//	PIVOT t ON year USING LAST(ap) ap;
+	pivotSQL := fmt.Sprintf("PIVOT %[1]s ON %[2]s USING %[3]s %[4]s %[5]s OFFSET %[6]d",
+		temporaryTableName,              // 1
+		strings.Join(q.PivotOn, ", "),   // 2
+		strings.Join(measureCols, ", "), // 3
+		orderClause,                     // 4
+		limitClause,                     // 5
+		q.Offset,                        // 6
+	)
+	schema, data, err := olapQuery(ctx, olap, int(q.Priority), pivotSQL, nil)
 	if err != nil {
 		return err
 	}
@@ -135,11 +294,12 @@ func (q *MetricsViewAggregation) buildMetricsAggregationSQL(mv *runtimev1.Metric
 		return "", nil, errors.New("no dimensions or measures specified")
 	}
 
-	selectCols := make([]string, 0, len(q.Dimensions)+len(q.Measures))
+	cols := len(q.Dimensions) + len(q.Measures)
+	selectCols := make([]string, 0, cols)
+
 	groupCols := make([]string, 0, len(q.Dimensions))
 	unnestClauses := make([]string, 0)
 	args := []any{}
-
 	for _, d := range q.Dimensions {
 		// Handle regular dimensions
 		if d.TimeGrain == runtimev1.TimeGrain_TIME_GRAIN_UNSPECIFIED {
@@ -173,13 +333,14 @@ func (q *MetricsViewAggregation) buildMetricsAggregationSQL(mv *runtimev1.Metric
 	}
 
 	for _, m := range q.Measures {
+		sn := safeName(m.Name)
 		switch m.BuiltinMeasure {
 		case runtimev1.BuiltinMeasure_BUILTIN_MEASURE_UNSPECIFIED:
 			expr, err := metricsViewMeasureExpression(mv, m.Name)
 			if err != nil {
 				return "", nil, err
 			}
-			selectCols = append(selectCols, fmt.Sprintf("%s as %s", expr, safeName(m.Name)))
+			selectCols = append(selectCols, fmt.Sprintf("%s as %s", expr, sn))
 		case runtimev1.BuiltinMeasure_BUILTIN_MEASURE_COUNT:
 			selectCols = append(selectCols, fmt.Sprintf("COUNT(*) as %s", safeName(m.Name)))
 		case runtimev1.BuiltinMeasure_BUILTIN_MEASURE_COUNT_DISTINCT:
@@ -246,16 +407,30 @@ func (q *MetricsViewAggregation) buildMetricsAggregationSQL(mv *runtimev1.Metric
 		limitClause = fmt.Sprintf("LIMIT %d", *q.Limit)
 	}
 
-	sql := fmt.Sprintf("SELECT %s FROM %s %s %s %s %s %s OFFSET %d",
-		strings.Join(selectCols, ", "),
-		safeName(mv.Table),
-		strings.Join(unnestClauses, ""),
-		whereClause,
-		groupClause,
-		orderClause,
-		limitClause,
-		q.Offset,
-	)
+	var sql string
+	if q.PivotOn != nil {
+		// select m1, m2, d1, d2 from t, lateral unnest(t.d1) tbl(unnested_d1_) where d1 = 'a' group by d1, d2
+		sql = fmt.Sprintf("SELECT %[1]s FROM %[2]s %[3]s %[4]s %[5]s %[6]s %[7]s",
+			strings.Join(selectCols, ", "),  // 1
+			safeName(mv.Table),              // 2
+			strings.Join(unnestClauses, ""), // 3
+			whereClause,                     // 4
+			groupClause,                     // 5
+			orderClause,                     // 6
+			limitClause,                     // 7
+		)
+	} else {
+		sql = fmt.Sprintf("SELECT %s FROM %s %s %s %s %s %s OFFSET %d",
+			strings.Join(selectCols, ", "),
+			safeName(mv.Table),
+			strings.Join(unnestClauses, ""),
+			whereClause,
+			groupClause,
+			orderClause,
+			limitClause,
+			q.Offset,
+		)
+	}
 
 	return sql, args, nil
 }