feat(bigquery): Support JSON_VALUE_ARRAY(...) (#4356)

tobymao · Nov 6, 2024 · 42da638 · 42da638
1 parent efcbfdb
commit 42da638
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 7 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -231,11 +231,14 @@ def _build_regexp_extract(args: t.List) -> exp.RegexpExtract:
     )
 
 
-def _build_json_extract_scalar(args: t.List, dialect: Dialect) -> exp.JSONExtractScalar:
-    if len(args) == 1:
-        # The default value for the JSONPath is '$' i.e all of the data
-        args.append(exp.Literal.string("$"))
-    return parser.build_extract_json_with_path(exp.JSONExtractScalar)(args, dialect)
+def _build_extract_json_with_default_path(expr_type: t.Type[E]) -> t.Callable[[t.List, Dialect], E]:
+    def _builder(args: t.List, dialect: Dialect) -> E:
+        if len(args) == 1:
+            # The default value for the JSONPath is '$' i.e all of the data
+            args.append(exp.Literal.string("$"))
+        return parser.build_extract_json_with_path(expr_type)(args, dialect)
+
+    return _builder
 
 
 def _str_to_datetime_sql(
@@ -440,9 +443,10 @@ class Parser(parser.Parser):
                 this=exp.TsOrDsToDate(this=seq_get(args, 1)), format=seq_get(args, 0)
             ),
             "GENERATE_ARRAY": exp.GenerateSeries.from_arg_list,
-            "JSON_EXTRACT_SCALAR": _build_json_extract_scalar,
+            "JSON_EXTRACT_SCALAR": _build_extract_json_with_default_path(exp.JSONExtractScalar),
             "JSON_QUERY": parser.build_extract_json_with_path(exp.JSONExtract),
-            "JSON_VALUE": _build_json_extract_scalar,
+            "JSON_VALUE": _build_extract_json_with_default_path(exp.JSONExtractScalar),
+            "JSON_VALUE_ARRAY": _build_extract_json_with_default_path(exp.JSONValueArray),
             "LENGTH": lambda args: exp.Length(this=seq_get(args, 0), binary=True),
             "MD5": exp.MD5Digest.from_arg_list,
             "TO_HEX": _build_to_hex,

diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
@@ -970,3 +970,8 @@ def regexpextract_sql(self, expression: exp.RegexpExtract) -> str:
             return self.func(
                 "REGEXP_EXTRACT", expression.this, expression.expression, group, params
             )
+
+        def jsonvaluearray_sql(self, expression: exp.JSONValueArray) -> str:
+            json_extract = exp.JSONExtract(this=expression.this, expression=expression.expression)
+
+            return self.sql(exp.cast(json_extract, to=exp.DataType.build("ARRAY<STRING>")))
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
@@ -1110,3 +1110,12 @@ def jsonextract_sql(self, expression: exp.JSONExtract):
                 exp.ParseJSON(this=this) if this.is_string else this,
                 expression.expression,
             )
+
+        def jsonvaluearray_sql(self, expression: exp.JSONValueArray):
+            json_extract = exp.JSONExtract(this=expression.this, expression=expression.expression)
+            ident = exp.to_identifier("x")
+            transform_lambda = exp.Lambda(
+                expressions=[ident], this=exp.cast(ident, to=exp.DataType.Type.VARCHAR)
+            )
+
+            return self.func("TRANSFORM", json_extract, transform_lambda)
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -5939,6 +5939,10 @@ class JSONValue(Expression):
     }
 
 
+class JSONValueArray(Func):
+    arg_types = {"this": True, "expression": False}
+
+
 # # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/JSON_TABLE.html
 class JSONTable(Func):
     arg_types = {

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -1603,6 +1603,14 @@ def test_bigquery(self):
                 "snowflake": """SELECT GET_PATH(PARSE_JSON('{"class": {"students": []}}'), 'class')""",
             },
         )
+        self.validate_all(
+            """SELECT JSON_VALUE_ARRAY('{"arr": [1, "a"]}', '$.arr')""",
+            write={
+                "bigquery": """SELECT JSON_VALUE_ARRAY('{"arr": [1, "a"]}', '$.arr')""",
+                "duckdb": """SELECT CAST('{"arr": [1, "a"]}' -> '$.arr' AS TEXT[])""",
+                "snowflake": """SELECT TRANSFORM(GET_PATH(PARSE_JSON('{"arr": [1, "a"]}'), 'arr'), x -> CAST(x AS VARCHAR))""",
+            },
+        )
 
     def test_errors(self):
         with self.assertRaises(TokenError):