From 9930e5ec00b23c3687b642674bb3dc53655e45d2 Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Thu, 28 Mar 2024 11:58:21 -0700
Subject: [PATCH 1/2] Support `UNNEST WITH OFFSET`

---
 internal/formatter.go | 14 ++++++++++----
 query_test.go         |  9 +++++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/internal/formatter.go b/internal/formatter.go
index fbf9c09..56c60f0 100644
--- a/internal/formatter.go
+++ b/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
 }
diff --git a/query_test.go b/query_test.go
index 0d9102b..7e50e23 100644
--- a/query_test.go
+++ b/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`,

From b565721336de192e56551d7185e0af98d6198737 Mon Sep 17 00:00:00 2001
From: Dan Hansen <ohaibbq@gmail.com>
Date: Thu, 28 Mar 2024 14:13:41 -0700
Subject: [PATCH 2/2] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 0483dde..0037e61 100644
--- a/README.md
+++ b/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