Fix(bigquery): allow SPLIT call with 1 argument (#1770)

* Fix(bigquery): allow SPLIT call with 1 argument * Formatting * Reorder
tobymao · Jun 13, 2023 · 697c8b1 · 697c8b1
1 parent 88ab724
commit 697c8b1
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 28 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -166,13 +166,21 @@ class Parser(parser.Parser):
 
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
+            "DATE_ADD": parse_date_delta_with_interval(exp.DateAdd),
+            "DATE_SUB": parse_date_delta_with_interval(exp.DateSub),
             "DATE_TRUNC": lambda args: exp.DateTrunc(
                 unit=exp.Literal.string(str(seq_get(args, 1))),
                 this=seq_get(args, 0),
             ),
-            "DATE_ADD": parse_date_delta_with_interval(exp.DateAdd),
             "DATETIME_ADD": parse_date_delta_with_interval(exp.DatetimeAdd),
+            "DATETIME_SUB": parse_date_delta_with_interval(exp.DatetimeSub),
             "DIV": lambda args: exp.IntDiv(this=seq_get(args, 0), expression=seq_get(args, 1)),
+            "PARSE_DATE": lambda args: format_time_lambda(exp.StrToDate, "bigquery")(
+                [seq_get(args, 1), seq_get(args, 0)]
+            ),
+            "PARSE_TIMESTAMP": lambda args: format_time_lambda(exp.StrToTime, "bigquery")(
+                [seq_get(args, 1), seq_get(args, 0)]
+            ),
             "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list,
             "REGEXP_EXTRACT": lambda args: exp.RegexpExtract(
                 this=seq_get(args, 0),
@@ -183,24 +191,15 @@ class Parser(parser.Parser):
                 if re.compile(str(seq_get(args, 1))).groups == 1
                 else None,
             ),
+            "SPLIT": lambda args: exp.Split(
+                # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#split
+                this=seq_get(args, 0),
+                expression=seq_get(args, 1) or exp.Literal.string(","),
+            ),
             "TIME_ADD": parse_date_delta_with_interval(exp.TimeAdd),
-            "TIMESTAMP_ADD": parse_date_delta_with_interval(exp.TimestampAdd),
-            "DATE_SUB": parse_date_delta_with_interval(exp.DateSub),
-            "DATETIME_SUB": parse_date_delta_with_interval(exp.DatetimeSub),
             "TIME_SUB": parse_date_delta_with_interval(exp.TimeSub),
+            "TIMESTAMP_ADD": parse_date_delta_with_interval(exp.TimestampAdd),
             "TIMESTAMP_SUB": parse_date_delta_with_interval(exp.TimestampSub),
-            "PARSE_DATE": lambda args: format_time_lambda(exp.StrToDate, "bigquery")(
-                [
-                    seq_get(args, 1),
-                    seq_get(args, 0),
-                ]
-            ),
-            "PARSE_TIMESTAMP": lambda args: format_time_lambda(exp.StrToTime, "bigquery")(
-                [
-                    seq_get(args, 1),
-                    seq_get(args, 0),
-                ]
-            ),
         }
 
         FUNCTION_PARSERS = {

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -6,18 +6,6 @@ class TestBigQuery(Validator):
     dialect = "bigquery"
 
     def test_bigquery(self):
-        self.validate_all(
-            "cast(x as date format 'MM/DD/YYYY')",
-            write={
-                "bigquery": "PARSE_DATE('%m/%d/%Y', x)",
-            },
-        )
-        self.validate_all(
-            "cast(x as time format 'YYYY.MM.DD HH:MI:SSTZH')",
-            write={
-                "bigquery": "PARSE_TIMESTAMP('%Y.%m.%d %I:%M:%S%z', x)",
-            },
-        )
         self.validate_identity("SELECT * FROM x-0.a")
         self.validate_identity("SELECT * FROM pivot CROSS JOIN foo")
         self.validate_identity("SAFE_CAST(x AS STRING)")
@@ -49,6 +37,19 @@ def test_bigquery(self):
             "CREATE TABLE IF NOT EXISTS foo AS SELECT * FROM bla EXCEPT DISTINCT (SELECT * FROM bar) LIMIT 0"
         )
 
+        self.validate_all("SELECT SPLIT(foo)", write={"bigquery": "SELECT SPLIT(foo, ',')"})
+        self.validate_all(
+            "cast(x as date format 'MM/DD/YYYY')",
+            write={
+                "bigquery": "PARSE_DATE('%m/%d/%Y', x)",
+            },
+        )
+        self.validate_all(
+            "cast(x as time format 'YYYY.MM.DD HH:MI:SSTZH')",
+            write={
+                "bigquery": "PARSE_TIMESTAMP('%Y.%m.%d %I:%M:%S%z', x)",
+            },
+        )
         self.validate_all("SELECT 1 AS hash", write={"bigquery": "SELECT 1 AS `hash`"})
         self.validate_all('x <> ""', write={"bigquery": "x <> ''"})
         self.validate_all('x <> """"""', write={"bigquery": "x <> ''"})