Support UNNEST WITH OFFSET

Recidiviz · Mar 28, 2024 · 0723438 · 0723438
1 parent ed775cb
commit 0723438
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/internal/formatter.go b/internal/formatter.go
@@ -727,6 +727,12 @@ func (n *ArrayScanNode) FormatSQL(ctx context.Context) (string, error) {
 		return "", err
 	}
 	colName := uniqueColumnName(ctx, n.node.ElementColumn())
+	columns := []string{fmt.Sprintf("json_each.value AS `%s`", colName)}
+
+	if offsetColumn := n.node.ArrayOffsetColumn(); offsetColumn != nil {
+		offsetColName := uniqueColumnName(ctx, offsetColumn.Column())
+		columns = append(columns, fmt.Sprintf("json_each.key AS `%s`", offsetColName))
+	}
 	if n.node.InputScan() != nil {
 		input, err := newNode(n.node.InputScan()).FormatSQL(ctx)
 		if err != nil {
@@ -762,15 +768,15 @@ func (n *ArrayScanNode) FormatSQL(ctx context.Context) (string, error) {
 		}
 
 		return fmt.Sprintf(
-			"SELECT *, json_each.value AS `%s` %s %s",
-			colName,
+			"SELECT *, %s %s %s",
+			strings.Join(columns, ","),
 			formattedInput,
 			arrayJoinExpr,
 		), nil
 	}
 	return fmt.Sprintf(
-		"SELECT json_each.value AS `%s` FROM json_each(zetasqlite_decode_array(%s))",
-		colName,
+		"SELECT %s FROM json_each(zetasqlite_decode_array(%s))",
+		strings.Join(columns, ","),
 		arrayExpr,
 	), nil
 }

diff --git a/query_test.go b/query_test.go
@@ -2392,6 +2392,15 @@ FROM UNNEST([
 			query:        `SELECT a, b FROM UNNEST([STRUCT(DATE(2022, 1, 1) AS a, 1 AS b)])`,
 			expectedRows: [][]interface{}{{"2022-01-01", int64(1)}},
 		},
+		{
+			name: "unnest with offset",
+			query: `SELECT *
+FROM UNNEST(['foo', 'bar', 'baz'])
+  AS element
+WITH OFFSET AS offset
+ORDER BY offset DESC;`,
+			expectedRows: [][]interface{}{{"baz", int64(2)}, {"bar", int64(1)}, {"foo", int64(0)}},
+		},
 		{
 			name:  "array function",
 			query: `SELECT ARRAY (SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3) AS new_array`,