refactor: add compile_readlocal for SQLGlotCompiler (#1663)

This change is adding compile_readlocal for SQLGlotCompiler class and also introducing snapshot as dev dependency for unit tests.

Author: chelsea-lin

bigframes/core/compile/compiled.py

@@ -288,13 +288,7 @@ def _reproject_to_table(self) -> UnorderedIR:
     def from_polars(
         cls, pa_table: pa.Table, schema: Sequence[bigquery.SchemaField]
     ) -> UnorderedIR:
-        # TODO: add offsets
-        """
-        Builds an in-memory only (SQL only) expr from a pandas dataframe.
-
-        Assumed that the dataframe has unique string column names and bigframes-suppported
-        dtypes.
-        """
+        """Builds an in-memory only (SQL only) expr from a pyarrow table."""
         import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes
 
         # derive the ibis schema from the original pandas schema

bigframes/core/compile/compiler.py

@@ -22,7 +22,6 @@
 import bigframes_vendored.ibis.expr.api as ibis_api
 import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
 import bigframes_vendored.ibis.expr.types as ibis_types
-import google.cloud.bigquery
 import pyarrow as pa
 
 from bigframes import dtypes, operations
@@ -169,7 +168,7 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
     pa_table = node.local_data_source.data
     bq_schema = node.schema.to_bigquery()
 
-    pa_table = pa_table.select(list(item.source_id for item in node.scan_list.items))
+    pa_table = pa_table.select([item.source_id for item in node.scan_list.items])
     pa_table = pa_table.rename_columns(
         {item.source_id: item.id.sql for item in node.scan_list.items}
     )
@@ -178,7 +177,6 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
         pa_table = pa_table.append_column(
             offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
         )
-        bq_schema = (*bq_schema, google.cloud.bigquery.SchemaField(offsets, "INT64"))
     return compiled.UnorderedIR.from_polars(pa_table, bq_schema)
 
 

bigframes/core/compile/sqlglot/compiler.py

@@ -15,23 +15,22 @@
 
 import dataclasses
 import functools
+import itertools
 import typing
 
-import google.cloud.bigquery as bigquery
-import sqlglot.expressions as sge
+from google.cloud import bigquery
+import pyarrow as pa
 
-from bigframes.core import expression, nodes, rewrite
+from bigframes.core import expression, identifiers, nodes, rewrite
 from bigframes.core.compile import configs
-from bigframes.core.compile.sqlglot import sql_gen
+import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 import bigframes.core.ordering as bf_ordering
 
 
 @dataclasses.dataclass(frozen=True)
 class SQLGlotCompiler:
     """Compiles BigFrame nodes into SQL using SQLGlot."""
 
-    sql_gen = sql_gen.SQLGen()
-
     def compile(
         self,
         node: nodes.BigFrameNode,
@@ -81,6 +80,7 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
             result_node = typing.cast(
                 nodes.ResultNode, rewrite.column_pruning(result_node)
             )
+            result_node = _remap_variables(result_node)
             sql = self._compile_result_node(result_node)
             return configs.CompileResult(
                 sql, result_node.schema.to_bigquery(), result_node.order_by
@@ -89,6 +89,8 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
         ordering: typing.Optional[bf_ordering.RowOrdering] = result_node.order_by
         result_node = dataclasses.replace(result_node, order_by=None)
         result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
+
+        result_node = _remap_variables(result_node)
         sql = self._compile_result_node(result_node)
         # Return the ordering iff no extra columns are needed to define the row order
         if ordering is not None:
@@ -103,9 +105,9 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
         )
 
     def _compile_result_node(self, root: nodes.ResultNode) -> str:
-        sqlglot_expr = compile_node(root.child)
+        sqlglot_ir = compile_node(root.child)
         # TODO: add order_by, limit, and selections to sqlglot_expr
-        return self.sql_gen.sql(sqlglot_expr)
+        return sqlglot_ir.sql
 
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):
@@ -115,27 +117,52 @@ def _replace_unsupported_ops(node: nodes.BigFrameNode):
     return node
 
 
+def _remap_variables(node: nodes.ResultNode) -> nodes.ResultNode:
+    """Remaps `ColumnId`s in the BFET of a `ResultNode` to produce deterministic UIDs."""
+
+    def anonymous_column_ids() -> typing.Generator[identifiers.ColumnId, None, None]:
+        for i in itertools.count():
+            yield identifiers.ColumnId(name=f"bfcol_{i}")
+
+    result_node, _ = rewrite.remap_variables(node, anonymous_column_ids())
+    return typing.cast(nodes.ResultNode, result_node)
+
+
 @functools.lru_cache(maxsize=5000)
-def compile_node(node: nodes.BigFrameNode) -> sge.Expression:
-    """Compile node into CompileArrayValue. Caches result."""
+def compile_node(node: nodes.BigFrameNode) -> ir.SQLGlotIR:
+    """Compiles node into CompileArrayValue. Caches result."""
     return node.reduce_up(lambda node, children: _compile_node(node, *children))
 
 
 @functools.singledispatch
 def _compile_node(
-    node: nodes.BigFrameNode, *compiled_children: sge.Expression
-) -> sge.Expression:
+    node: nodes.BigFrameNode, *compiled_children: ir.SQLGlotIR
+) -> ir.SQLGlotIR:
     """Defines transformation but isn't cached, always use compile_node instead"""
     raise ValueError(f"Can't compile unrecognized node: {node}")
 
 
 @_compile_node.register
-def compile_readlocal(node: nodes.ReadLocalNode, *args) -> sge.Expression:
-    # TODO: add support for reading from local files
-    return sge.select()
+def compile_readlocal(node: nodes.ReadLocalNode, *args) -> ir.SQLGlotIR:
+    offsets = node.offsets_col.sql if node.offsets_col else None
+    schema_names = node.schema.names
+    schema_dtypes = node.schema.dtypes
+
+    pa_table = node.local_data_source.data
+    pa_table = pa_table.select([item.source_id for item in node.scan_list.items])
+    pa_table = pa_table.rename_columns(
+        {item.source_id: item.id.sql for item in node.scan_list.items}
+    )
+
+    if offsets:
+        pa_table = pa_table.append_column(
+            offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
+        )
+
+    return ir.SQLGlotIR.from_pandas(pa_table.to_pandas(), schema_names, schema_dtypes)
 
 
 @_compile_node.register
-def compile_selection(node: nodes.SelectionNode, child: sge.Expression):
+def compile_selection(node: nodes.SelectionNode, child: ir.SQLGlotIR):
     # TODO: add support for selection
     return child

bigframes/core/compile/sqlglot/sql_gen.py

@@ -0,0 +1,103 @@
+# Copyright 2025 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.
+
+from __future__ import annotations
+
+import dataclasses
+import typing
+
+import pandas as pd
+import sqlglot as sg
+import sqlglot.dialects.bigquery
+import sqlglot.expressions as sge
+
+from bigframes import dtypes
+import bigframes.core.compile.sqlglot.sqlglot_types as sgt
+
+
+@dataclasses.dataclass(frozen=True)
+class SQLGlotIR:
+    """Helper class to build SQLGlot Query and generate SQL string."""
+
+    expr: sge.Expression = sge.Expression()
+    """The SQLGlot expression representing the query."""
+
+    dialect = sqlglot.dialects.bigquery.BigQuery
+    """The SQL dialect used for generation."""
+
+    quoted: bool = True
+    """Whether to quote identifiers in the generated SQL."""
+
+    pretty: bool = True
+    """Whether to pretty-print the generated SQL."""
+
+    @property
+    def sql(self) -> str:
+        """Generate SQL string from the given expression."""
+        return self.expr.sql(dialect=self.dialect, pretty=self.pretty)
+
+    @classmethod
+    def from_pandas(
+        cls,
+        pd_df: pd.DataFrame,
+        schema_names: typing.Sequence[str],
+        schema_dtypes: typing.Sequence[dtypes.Dtype],
+    ) -> SQLGlotIR:
+        """Builds SQLGlot expression from pyarrow table."""
+        dtype_expr = sge.DataType(
+            this=sge.DataType.Type.STRUCT,
+            expressions=[
+                sge.ColumnDef(
+                    this=sge.to_identifier(name, quoted=True),
+                    kind=sgt.SQLGlotType.from_bigframes_dtype(dtype),
+                )
+                for name, dtype in zip(schema_names, schema_dtypes)
+            ],
+            nested=True,
+        )
+        data_expr = [
+            sge.Tuple(
+                expressions=tuple(
+                    _literal(
+                        value=value,
+                        dtype=sgt.SQLGlotType.from_bigframes_dtype(dtype),
+                    )
+                    for value, dtype in zip(row, schema_dtypes)
+                )
+            )
+            for _, row in pd_df.iterrows()
+        ]
+        expr = sge.Unnest(
+            expressions=[
+                sge.DataType(
+                    this=sge.DataType.Type.ARRAY,
+                    expressions=[dtype_expr],
+                    nested=True,
+                    values=data_expr,
+                ),
+            ],
+        )
+        return cls(expr=sg.select(sge.Star()).from_(expr))
+
+
+def _literal(value: typing.Any, dtype: str) -> sge.Expression:
+    if value is None:
+        return _cast(sge.Null(), dtype)
+
+    # TODO: handle other types like visit_DefaultLiteral
+    return sge.convert(value)
+
+
+def _cast(arg, to) -> sge.Cast:
+    return sge.Cast(this=arg, to=to)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -32,7 +32,7 @@ def from_bigframes_dtype(
         bigframes_dtype: typing.Union[
             bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype, np.dtype[typing.Any]
         ],
-    ):
+    ) -> str:
         if bigframes_dtype == bigframes.dtypes.INT_DTYPE:
             return "INT64"
         elif bigframes_dtype == bigframes.dtypes.FLOAT_DTYPE:

bigframes/core/compile/sqlglot/sqlglot_types.py

@@ -1636,11 +1636,11 @@ def remap_vars(
     def remap_refs(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> ResultNode:
-        output_names = tuple(
+        output_cols = tuple(
             (ref.remap_column_refs(mappings), name) for ref, name in self.output_cols
         )
         order_by = self.order_by.remap_column_refs(mappings) if self.order_by else None
-        return dataclasses.replace(self, output_names=output_names, order_by=order_by)  # type: ignore
+        return dataclasses.replace(self, output_cols=output_cols, order_by=order_by)  # type: ignore
 
     @property
     def consumed_ids(self) -> COLUMN_SET:

bigframes/core/nodes.py

@@ -76,7 +76,7 @@
 ]
 UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = []
 UNIT_TEST_DEPENDENCIES: List[str] = []
-UNIT_TEST_EXTRAS: List[str] = []
+UNIT_TEST_EXTRAS: List[str] = ["tests"]
 UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {
     "3.12": ["polars", "scikit-learn"],
 }
@@ -203,7 +203,7 @@ def install_unittest_dependencies(session, install_test_extra, *constraints):
 
     if install_test_extra and UNIT_TEST_EXTRAS_BY_PYTHON:
         extras = UNIT_TEST_EXTRAS_BY_PYTHON.get(session.python, [])
-    elif install_test_extra and UNIT_TEST_EXTRAS:
+    if install_test_extra and UNIT_TEST_EXTRAS:
         extras = UNIT_TEST_EXTRAS
     else:
         extras = []

noxfile.py

@@ -72,7 +72,7 @@
 ]
 extras = {
     # Optional test dependencies packages. If they're missed, may skip some tests.
-    "tests": [],
+    "tests": ["freezegun", "pytest-snapshot"],
     # used for local engine, which is only needed for unit tests at present.
     "polars": ["polars >= 1.7.0"],
     "scikit-learn": ["scikit-learn>=1.2.2"],
@@ -82,7 +82,6 @@
         "pre-commit",
         "nox",
         "google-cloud-testutils",
-        "freezegun",
     ],
 }
 extras["all"] = list(sorted(frozenset(itertools.chain.from_iterable(extras.values()))))

setup.py

@@ -26,4 +26,15 @@ def compiler_session():
 @pytest.fixture(scope="module")
 def all_types_df() -> pd.DataFrame:
     # TODO: all types pandas dataframes
-    return pd.DataFrame({})
+    # TODO: add tests for empty dataframes
+    df = pd.DataFrame(
+        {
+            "int1": pd.Series([1, 2, 3], dtype="Int64"),
+            "int2": pd.Series([-10, 20, 30], dtype="Int64"),
+            "bools": pd.Series([True, None, False], dtype="boolean"),
+            "strings": pd.Series(["b", "aa", "ccc"], dtype="string[pyarrow]"),
+        },
+    )
+    # add more complexity index.
+    df.index = df.index.astype("Int64")
+    return df

tests/unit/core/compile/sqlglot/conftest.py

@@ -0,0 +1,3 @@
+SELECT
+  *
+FROM UNNEST(ARRAY<STRUCT<`bfcol_0` INT64, `bfcol_1` INT64, `bfcol_2` BOOLEAN, `bfcol_3` STRING, `bfcol_4` INT64>>[(1, -10, TRUE, 'b', 0), (2, 20, CAST(NULL AS BOOLEAN), 'aa', 1), (3, 30, FALSE, 'ccc', 2)])