diff --git a/internal/formatter.go b/internal/formatter.go index b077ace..cb6f323 100644 --- 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 index f92065e..62a0f19 100644 --- 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`,