refactor: add googlesql cast function (#794)

Author: chelsea-lin

bigframes/core/compile/googlesql/__init__.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+from bigframes.core.compile.googlesql.datatype import DataType
 from bigframes.core.compile.googlesql.expression import (
     _escape_chars,
     AliasExpression,
@@ -26,6 +27,7 @@
     StarExpression,
     TableExpression,
 )
+from bigframes.core.compile.googlesql.function import Cast
 from bigframes.core.compile.googlesql.query import (
     AsAlias,
     FromClause,
@@ -42,8 +44,10 @@
     "identifier",
     "AliasExpression",
     "AsAlias",
+    "Cast",
     "ColumnExpression",
     "CTEExpression",
+    "DataType",
     "FromClause",
     "FromItem",
     "NonRecursiveCTE",
@@ -52,5 +56,6 @@
     "SelectAll",
     "SelectExpression",
     "StarExpression",
+    "StringType",
     "TableExpression",
 ]

bigframes/core/compile/googlesql/datatype.py

@@ -0,0 +1,23 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+
+"""This module represents all GoogleSQL for BigQuery data types:
+https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types"""
+
+
+class DataType(enum.Enum):
+    STRING = 1
+    FLOAT64 = 2

bigframes/core/compile/googlesql/function.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dataclasses
+
+import bigframes.core.compile.googlesql.datatype as datatype
+import bigframes.core.compile.googlesql.expression as expr
+
+# Conversion functions:
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/conversion_functions
+
+
+@dataclasses.dataclass
+class Cast(expr.Expression):
+    """This class represents the `cast` function."""
+
+    expression: expr.ColumnExpression
+    type: datatype.DataType
+
+    def sql(self) -> str:
+        return f"CAST ({self.expression.sql()} AS {self.type.name})"

bigframes/core/sql.py

@@ -61,7 +61,9 @@ def multi_literal(*values: str):
 def cast_as_string(column_name: str) -> str:
     """Return a string representing string casting of a column."""
 
-    return f"CAST({googlesql.identifier(column_name)} AS STRING)"
+    return googlesql.Cast(
+        googlesql.ColumnExpression(column_name), googlesql.DataType.STRING
+    ).sql()
 
 
 def to_json_string(column_name: str) -> str:

tests/unit/core/compiler/__init__.py renamed to tests/unit/core/compile/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unit/core/compile/googlesql/__init__.py

@@ -0,0 +1,37 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import bigframes.core.compile.googlesql as sql
+
+
+@pytest.mark.parametrize(
+    ("table_id", "dataset_id", "project_id", "expected"),
+    [
+        pytest.param("a", None, None, "`a`"),
+        pytest.param("a", "b", None, "`b`.`a`"),
+        pytest.param("a", "b", "c", "`c`.`b`.`a`"),
+        pytest.param("a", None, "c", None, marks=pytest.mark.xfail(raises=ValueError)),
+    ],
+)
+def test_table_expression(table_id, dataset_id, project_id, expected):
+    expr = sql.TableExpression(
+        table_id=table_id, dataset_id=dataset_id, project_id=project_id
+    )
+    assert expr.sql() == expected
+
+
+def test_escape_chars():
+    assert sql._escape_chars("\a\b\f\n\r\t\v\\?'\"`") == r"\a\b\f\n\r\t\v\\\?\'\"\`"

tests/unit/core/compile/googlesql/test_expression.py

@@ -0,0 +1,21 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import bigframes.core.compile.googlesql as sql
+
+
+def test_cast():
+    col = sql.ColumnExpression("col")
+    assert sql.Cast(col, sql.DataType.STRING).sql() == "CAST (`col` AS STRING)"
+    assert sql.Cast(col, sql.DataType.FLOAT64).sql() == "CAST (`col` AS FLOAT64)"