Merge remote-tracking branch 'github/main' into unordered_mode

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -507,11 +507,11 @@ def try_align_as_projection(
         join_type: join_def.JoinType,
         mappings: typing.Tuple[join_def.JoinColumnMapping, ...],
     ) -> typing.Optional[ArrayValue]:
-        left_side = bigframes.core.rewrite.SquashedSelect.from_node(self.node)
-        right_side = bigframes.core.rewrite.SquashedSelect.from_node(other.node)
-        result = left_side.maybe_merge(right_side, join_type, mappings)
+        result = bigframes.core.rewrite.join_as_projection(
+            self.node, other.node, mappings, join_type
+        )
         if result is not None:
-            return ArrayValue(result.expand())
+            return ArrayValue(result)
         return None
 
     def explode(self, column_ids: typing.Sequence[str]) -> ArrayValue:
@@ -530,7 +530,3 @@ def _uniform_sampling(self, fraction: float) -> ArrayValue:
             The row numbers of result is non-deterministic, avoid to use.
         """
         return ArrayValue(nodes.RandomSampleNode(self.node, fraction))
-
-    def merge_projections(self) -> ArrayValue:
-        new_node = bigframes.core.rewrite.maybe_squash_projection(self.node)
-        return ArrayValue(new_node)

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

@@ -17,14 +17,17 @@
 
 from __future__ import annotations
 
+from bigframes.core.compile.googlesql.datatype import DataType
 from bigframes.core.compile.googlesql.expression import (
     _escape_chars,
     AliasExpression,
     ColumnExpression,
     CTEExpression,
+    identifier,
     StarExpression,
     TableExpression,
 )
+from bigframes.core.compile.googlesql.function import Cast
 from bigframes.core.compile.googlesql.query import (
     AsAlias,
     FromClause,
@@ -38,10 +41,13 @@
 
 __all__ = [
     "_escape_chars",
+    "identifier",
     "AliasExpression",
     "AsAlias",
+    "Cast",
     "ColumnExpression",
     "CTEExpression",
+    "DataType",
     "FromClause",
     "FromItem",
     "NonRecursiveCTE",
@@ -50,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/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/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/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)):