refactor: replace a few functions in sql.py with AST. (#784)

* refactor: update Select, update sql.py

* update test

* update test

* update test

Author: Genesis929

bigframes/core/blocks.py

@@ -40,6 +40,7 @@
 import bigframes.constants
 import bigframes.constants as constants
 import bigframes.core as core
+import bigframes.core.compile.googlesql as googlesql
 import bigframes.core.expression as ex
 import bigframes.core.expression as scalars
 import bigframes.core.guid as guid
@@ -2417,7 +2418,9 @@ def _get_rows_as_json_values(self) -> Block:
         select_columns = (
             [ordering_column_name] + list(self.index_columns) + [row_json_column_name]
         )
-        select_columns_csv = sql.csv([sql.identifier(col) for col in select_columns])
+        select_columns_csv = sql.csv(
+            [googlesql.identifier(col) for col in select_columns]
+        )
         json_sql = f"""\
 With T0 AS (
 {textwrap.indent(expr_sql, "    ")}
@@ -2430,7 +2433,7 @@ def _get_rows_as_json_values(self) -> Block:
                "values", [{column_references_csv}],
                "indexlength", {index_columns_count},
                "dtype", {pandas_row_dtype}
-           ) AS {sql.identifier(row_json_column_name)} FROM T0
+           ) AS {googlesql.identifier(row_json_column_name)} FROM T0
 )
 SELECT {select_columns_csv} FROM T1
 """

bigframes/core/compile/compiled.py

@@ -28,6 +28,7 @@
 import pandas
 
 import bigframes.core.compile.aggregate_compiler as agg_compiler
+import bigframes.core.compile.googlesql
 import bigframes.core.compile.ibis_types
 import bigframes.core.compile.scalar_op_compiler as op_compilers
 import bigframes.core.expression as ex
@@ -905,7 +906,12 @@ def to_sql(
             output_columns = [
                 col_id_overrides.get(col, col) for col in baked_ir.column_ids
             ]
-            sql = bigframes.core.sql.select_from_subquery(output_columns, sql)
+            sql = (
+                bigframes.core.compile.googlesql.Select()
+                .from_(sql)
+                .select(output_columns)
+                .sql()
+            )
 
             # Single row frames may not have any ordering columns
             if len(baked_ir._ordering.all_ordering_columns) > 0:

bigframes/core/compile/googlesql/__init__.py

@@ -22,6 +22,7 @@
     AliasExpression,
     ColumnExpression,
     CTEExpression,
+    identifier,
     StarExpression,
     TableExpression,
 )
@@ -38,6 +39,7 @@
 
 __all__ = [
     "_escape_chars",
+    "identifier",
     "AliasExpression",
     "AsAlias",
     "ColumnExpression",

bigframes/core/compile/googlesql/expression.py

@@ -45,8 +45,8 @@ class ColumnExpression(Expression):
 
     def sql(self) -> str:
         if self.parent is not None:
-            return f"{self.parent.sql()}.`{self.name}`"
-        return f"`{self.name}`"
+            return f"{self.parent.sql()}.{identifier(self.name)}"
+        return identifier(self.name)
 
 
 @dataclasses.dataclass
@@ -72,10 +72,10 @@ def __post_init__(self):
     def sql(self) -> str:
         text = []
         if self.project_id is not None:
-            text.append(f"`{_escape_chars(self.project_id)}`")
+            text.append(identifier(self.project_id))
         if self.dataset_id is not None:
-            text.append(f"`{_escape_chars(self.dataset_id)}`")
-        text.append(f"`{_escape_chars(self.table_id)}`")
+            text.append(identifier(self.dataset_id))
+        text.append(identifier(self.table_id))
         return ".".join(text)
 
 
@@ -84,15 +84,22 @@ class AliasExpression(Expression):
     alias: str
 
     def sql(self) -> str:
-        return f"`{_escape_chars(self.alias)}`"
+        return identifier(self.alias)
 
 
 @dataclasses.dataclass
 class CTEExpression(Expression):
     name: str
 
     def sql(self) -> str:
-        return f"`{_escape_chars(self.name)}`"
+        return identifier(self.name)
+
+
+def identifier(id: str) -> str:
+    """Return a string representing column reference in a SQL."""
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#identifiers
+    # Just always escape, otherwise need to check against every reserved sql keyword
+    return f"`{_escape_chars(id)}`"
 
 
 def _escape_chars(value: str):

bigframes/core/compile/googlesql/query.py

@@ -16,19 +16,19 @@
 
 import dataclasses
 import typing
-from typing import TYPE_CHECKING
+
+import google.cloud.bigquery as bigquery
 
 import bigframes.core.compile.googlesql.abc as abc
 import bigframes.core.compile.googlesql.expression as expr
 
-if TYPE_CHECKING:
-    import google.cloud.bigquery as bigquery
-
 """This module provides a structured representation of GoogleSQL syntax using nodes.
 Each node's name and child nodes are designed to strictly follow the official GoogleSQL
 syntax rules outlined in the documentation:
 https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax"""
 
+TABLE_SOURCE_TYPE = typing.Union[str, bigquery.TableReference]
+
 
 @dataclasses.dataclass
 class QueryExpr(abc.SQLSyntax):
@@ -53,11 +53,47 @@ def sql(self) -> str:
 class Select(abc.SQLSyntax):
     """This class represents GoogleSQL `select` syntax."""
 
-    select_list: typing.Sequence[typing.Union[SelectExpression, SelectAll]]
-    from_clause_list: typing.Sequence[FromClause] = ()
+    select_list: typing.Sequence[
+        typing.Union[SelectExpression, SelectAll]
+    ] = dataclasses.field(default_factory=list)
+    from_clause_list: typing.Sequence[FromClause] = dataclasses.field(
+        default_factory=list
+    )
     distinct: bool = False
 
+    def select(
+        self,
+        columns: typing.Union[typing.Iterable[str], str, None] = None,
+        distinct: bool = False,
+    ) -> Select:
+        if isinstance(columns, str):
+            columns = [columns]
+        self.select_list: typing.List[typing.Union[SelectExpression, SelectAll]] = (
+            [
+                SelectExpression(expression=expr.ColumnExpression(name=column))
+                for column in columns
+            ]
+            if columns
+            else [SelectAll(expression=expr.StarExpression())]
+        )
+        self.distinct = distinct
+        return self
+
+    def from_(
+        self,
+        sources: typing.Union[TABLE_SOURCE_TYPE, typing.Iterable[TABLE_SOURCE_TYPE]],
+    ) -> Select:
+        if (not isinstance(sources, typing.Iterable)) or isinstance(sources, str):
+            sources = [sources]
+        self.from_clause_list = [
+            FromClause(FromItem.from_source(source)) for source in sources
+        ]
+        return self
+
     def sql(self) -> str:
+        if (self.select_list is not None) and (not self.select_list):
+            raise ValueError("Select clause has not been properly initialized.")
+
         text = ["SELECT"]
 
         if self.distinct:
@@ -66,7 +102,7 @@ def sql(self) -> str:
         select_list_sql = ",\n".join([select.sql() for select in self.select_list])
         text.append(select_list_sql)
 
-        if self.from_clause_list is not None:
+        if self.from_clause_list:
             from_clauses_sql = ",\n".join(
                 [clause.sql() for clause in self.from_clause_list]
             )
@@ -118,19 +154,27 @@ class FromItem(abc.SQLSyntax):
     as_alias: typing.Optional[AsAlias] = None
 
     @classmethod
-    def from_table_ref(
+    def from_source(
         cls,
-        table_ref: bigquery.TableReference,
+        subquery_or_tableref: typing.Union[bigquery.TableReference, str],
         as_alias: typing.Optional[AsAlias] = None,
     ):
-        return cls(
-            expression=expr.TableExpression(
-                table_id=table_ref.table_id,
-                dataset_id=table_ref.dataset_id,
-                project_id=table_ref.project,
-            ),
-            as_alias=as_alias,
-        )
+        if isinstance(subquery_or_tableref, bigquery.TableReference):
+            return cls(
+                expression=expr.TableExpression(
+                    table_id=subquery_or_tableref.table_id,
+                    dataset_id=subquery_or_tableref.dataset_id,
+                    project_id=subquery_or_tableref.project,
+                ),
+                as_alias=as_alias,
+            )
+        elif isinstance(subquery_or_tableref, str):
+            return cls(
+                expression=subquery_or_tableref,
+                as_alias=as_alias,
+            )
+        else:
+            raise ValueError("The source must be bigquery.TableReference or str.")
 
     def sql(self) -> str:
         if isinstance(self.expression, (expr.TableExpression, expr.CTEExpression)):

bigframes/core/sql.py

@@ -23,10 +23,6 @@
 
 import bigframes.core.compile.googlesql as googlesql
 
-# Literals and identifiers matching this pattern can be unquoted
-unquoted = r"^[A-Za-z_][A-Za-z_0-9]*$"
-
-
 if TYPE_CHECKING:
     import google.cloud.bigquery as bigquery
 
@@ -62,23 +58,16 @@ def multi_literal(*values: str):
     return "(" + ", ".join(literal_strings) + ")"
 
 
-def identifier(id: str) -> str:
-    """Return a string representing column reference in a SQL."""
-    # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#identifiers
-    # Just always escape, otherwise need to check against every reserved sql keyword
-    return f"`{googlesql._escape_chars(id)}`"
-
-
 def cast_as_string(column_name: str) -> str:
     """Return a string representing string casting of a column."""
 
-    return f"CAST({identifier(column_name)} AS STRING)"
+    return f"CAST({googlesql.identifier(column_name)} AS STRING)"
 
 
 def to_json_string(column_name: str) -> str:
     """Return a string representing JSON version of a column."""
 
-    return f"TO_JSON_STRING({identifier(column_name)})"
+    return f"TO_JSON_STRING({googlesql.identifier(column_name)})"
 
 
 def csv(values: Iterable[str]) -> str:
@@ -91,55 +80,12 @@ def infix_op(opname: str, left_arg: str, right_arg: str):
     return f"{left_arg} {opname} {right_arg}"
 
 
-### Writing SELECT expressions
-def select_from_subquery(columns: Iterable[str], subquery: str, distinct: bool = False):
-    select_list = [
-        googlesql.SelectExpression(expression=googlesql.ColumnExpression(name=column))
-        for column in columns
-    ]
-    from_clause_list = [googlesql.FromClause(googlesql.FromItem(expression=subquery))]
-
-    select_expr = googlesql.Select(
-        select_list=select_list, from_clause_list=from_clause_list, distinct=distinct
-    )
-    return select_expr.sql()
-
-
-def select_from_table_ref(
-    columns: Iterable[str], table_ref: bigquery.TableReference, distinct: bool = False
-):
-    select_list = [
-        googlesql.SelectExpression(expression=googlesql.ColumnExpression(name=column))
-        for column in columns
-    ]
-    from_clause_list = [
-        googlesql.FromClause(googlesql.FromItem.from_table_ref(table_ref))
-    ]
-
-    select_expr = googlesql.Select(
-        select_list=select_list, from_clause_list=from_clause_list, distinct=distinct
-    )
-    return select_expr.sql()
-
-
-def select_table(table_ref: bigquery.TableReference):
-    select_list = [googlesql.SelectAll(expression=googlesql.StarExpression())]
-    from_clause_list = [
-        googlesql.FromClause(googlesql.FromItem.from_table_ref(table_ref))
-    ]
-
-    select_expr = googlesql.Select(
-        select_list=select_list, from_clause_list=from_clause_list
-    )
-    return select_expr.sql()
-
-
 def is_distinct_sql(columns: Iterable[str], table_ref: bigquery.TableReference) -> str:
     is_unique_sql = f"""WITH full_table AS (
-        {select_from_table_ref(columns, table_ref)}
+        {googlesql.Select().from_(table_ref).select(columns).sql()}
     ),
     distinct_table AS (
-        {select_from_table_ref(columns, table_ref, distinct=True)}
+        {googlesql.Select().from_(table_ref).select(columns, distinct=True).sql()}
     )
 
     SELECT (SELECT COUNT(*) FROM full_table) AS `total_count`,

bigframes/session/_io/bigquery/__init__.py

@@ -32,6 +32,7 @@
 
 import bigframes
 from bigframes.core import log_adapter
+import bigframes.core.compile.googlesql as googlesql
 import bigframes.core.sql
 import bigframes.formatting_helpers as formatting_helpers
 
@@ -480,7 +481,7 @@ def compile_filters(filters: third_party_pandas_gbq.FiltersType) -> str:
 
             operator_str = valid_operators[operator]
 
-            column_ref = bigframes.core.sql.identifier(column)
+            column_ref = googlesql.identifier(column)
             if operator_str in ["IN", "NOT IN"]:
                 value_literal = bigframes.core.sql.multi_literal(*value)
             else: