Support UNNEST WITH OFFSET

README.md

@@ -159,7 +159,7 @@ A list of ZetaSQL ( Google Standard SQL ) specifications and features supported
   - [x] UNNEST and STRUCTs
   - [ ] Explicit and implicit UNNEST
   - [ ] UNNEST and NULLs
-  - [ ] UNNEST and WITH OFFSET
+  - [x] UNNEST and WITH OFFSET
 - [ ] PIVOT operator
 - [ ] UNPIVOT operator
 - [ ] TABLESAMPLE operator

internal/formatter.go

@@ -660,6 +660,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 {
@@ -695,15 +701,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
 }

query_test.go

@@ -2264,6 +2264,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`,