From d44cbf9898fd8a714ccc3e0b8ac2dfd74980eb8a Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Thu, 28 Mar 2024 11:35:31 -0700
Subject: [PATCH] Support `PIVOT` / `UNPIVOT`

---
 README.md            |  4 ++--
 internal/analyzer.go |  2 ++
 query_test.go        | 45 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 0483dde..957bbca 100644
--- a/README.md
+++ b/README.md
@@ -160,8 +160,8 @@ A list of ZetaSQL ( Google Standard SQL ) specifications and features supported
   - [ ] Explicit and implicit UNNEST
   - [ ] UNNEST and NULLs
   - [ ] UNNEST and WITH OFFSET
-- [ ] PIVOT operator
-- [ ] UNPIVOT operator
+- [x] PIVOT operator
+- [x] UNPIVOT operator
 - [ ] TABLESAMPLE operator
 - [x] JOIN operation
   - [x] INNER JOIN
diff --git a/internal/analyzer.go b/internal/analyzer.go
index c778cf0..950f9bd 100644
--- a/internal/analyzer.go
+++ b/internal/analyzer.go
@@ -69,6 +69,8 @@ func newAnalyzerOptions() *zetasql.AnalyzerOptions {
 		zetasql.FeatureV13ExtendedGeographyParsers,
 		zetasql.FeatureTemplateFunctions,
 		zetasql.FeatureV11WithOnSubquery,
+		zetasql.FeatureV13Pivot,
+		zetasql.FeatureV13Unpivot,
 	})
 	langOpt.SetSupportedStatementKinds([]ast.Kind{
 		ast.BeginStmt,
diff --git a/query_test.go b/query_test.go
index 0d9102b..2c44291 100644
--- a/query_test.go
+++ b/query_test.go
@@ -3961,6 +3961,51 @@ WITH example AS (
 			query:        `SELECT DATE_ADD('2023-01-29', INTERVAL 1 MONTH)`,
 			expectedRows: [][]interface{}{{"2023-02-28"}},
 		},
+		{
+			name: "PIVOT",
+			query: `
+WITH produce AS (
+	SELECT 'Kale' AS product, 51 AS sales, 'Q1' AS quarter, 2020 AS year UNION ALL
+	SELECT 'Kale', 23, 'Q2', 2020 UNION ALL
+	SELECT 'Kale', 45, 'Q3', 2020 UNION ALL
+	SELECT 'Kale', 3, 'Q4', 2020 UNION ALL
+	SELECT 'Kale', 70, 'Q1', 2021 UNION ALL
+	SELECT 'Kale', 85, 'Q2', 2021 UNION ALL
+	SELECT 'Apple', 77, 'Q1', 2020 UNION ALL
+	SELECT 'Apple', 0, 'Q2', 2020 UNION ALL
+	SELECT 'Apple', 1, 'Q1', 2021
+)
+SELECT * FROM
+  Produce
+  PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2', 'Q3', 'Q4'))
+`,
+			expectedRows: [][]interface{}{
+				{"Apple", int64(2020), int64(77), int64(0), nil, nil},
+				{"Apple", int64(2021), int64(1), nil, nil, nil},
+				{"Kale", int64(2020), int64(51), int64(23), int64(45), int64(3)},
+				{"Kale", int64(2021), int64(70), int64(85), nil, nil},
+			},
+		},
+		{
+			name: "UNPIVOT",
+			query: `
+WITH Produce AS (
+  SELECT 'Kale' as product, 51 as Q1, 23 as Q2, 45 as Q3, 3 as Q4 UNION ALL
+  SELECT 'Apple', 77, 0, 25, 2)
+SELECT * FROM Produce
+UNPIVOT(sales FOR quarter IN (Q1, Q2, Q3, Q4))
+`,
+			expectedRows: [][]interface{}{
+				{"Kale", int64(51), "Q1"},
+				{"Kale", int64(23), "Q2"},
+				{"Kale", int64(45), "Q3"},
+				{"Kale", int64(3), "Q4"},
+				{"Apple", int64(77), "Q1"},
+				{"Apple", int64(0), "Q2"},
+				{"Apple", int64(25), "Q3"},
+				{"Apple", int64(2), "Q4"},
+			},
+		},
 		{
 			name:         "date_sub",
 			query:        `SELECT DATE_SUB('2023-03-31', INTERVAL 1 MONTH)`,