address comments

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -21,20 +21,22 @@
 import pyarrow as pa
 import sqlglot.expressions as sge
 
-from bigframes.core import expression, guid, nodes, rewrite
+from bigframes.core import expression, guid, identifiers, nodes, rewrite
 from bigframes.core.compile import configs
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 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."""
 
     uid_gen: guid.SequentialUIDGenerator
     """Generator for unique identifiers."""
 
+    def __init__(self):
+        self.uid_gen = guid.SequentialUIDGenerator()
+
     def compile(
         self,
         node: nodes.BigFrameNode,
@@ -84,8 +86,8 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
             result_node = typing.cast(
                 nodes.ResultNode, rewrite.column_pruning(result_node)
             )
-            remap_node, _ = rewrite.remap_variables(result_node, self.uid_gen)
-            sql = self._compile_result_node(typing.cast(nodes.ResultNode, remap_node))
+            result_node = self._remap_variables(result_node)
+            sql = self._compile_result_node(result_node)
             return configs.CompileResult(
                 sql, result_node.schema.to_bigquery(), result_node.order_by
             )
@@ -94,8 +96,8 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
         result_node = dataclasses.replace(result_node, order_by=None)
         result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node))
 
-        remap_node, _ = rewrite.remap_variables(result_node, self.uid_gen)
-        sql = self._compile_result_node(typing.cast(nodes.ResultNode, remap_node))
+        result_node = self._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:
             output_order = (
@@ -108,6 +110,14 @@ def _compile_sql(self, request: configs.CompileRequest) -> configs.CompileResult
             sql, result_node.schema.to_bigquery(), output_order
         )
 
+    def _remap_variables(self, node: nodes.ResultNode) -> nodes.ResultNode:
+        """Remaps `ColumnId`s in the BFET of a `ResultNode` to produce deterministic UIDs."""
+
+        result_node, _ = rewrite.remap_variables(
+            node, map(identifiers.ColumnId, self.uid_gen.get_uid_stream("bfcol_"))
+        )
+        return typing.cast(nodes.ResultNode, result_node)
+
     def _compile_result_node(self, root: nodes.ResultNode) -> str:
         sqlglot_ir = self.compile_node(root.child)
         # TODO: add order_by, limit, and selections to sqlglot_expr

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -128,7 +128,7 @@ def _encapsulate_as_cte(
 
         existing_ctes = select_expr.args.pop("with", [])
         new_cte_name = sge.to_identifier(
-            self.uid_gen.generate_sequential_uid("bfcte_"), quoted=self.quoted
+            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
         )
         new_cte = sge.CTE(
             this=select_expr,

bigframes/core/guid.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import typing
+
 _GUID_COUNTER = 0
 
 
@@ -22,18 +24,17 @@ def generate_guid(prefix="col_"):
 
 
 class SequentialUIDGenerator:
-    """
-    Generates sequential-like UIDs with multiple prefixes, e.g., "t0", "t1", "c0", "t2", etc.
-    """
+    """Yields a continuous stream of raw UID strings for the given prefix."""
 
     def __init__(self):
-        self.prefix_counters = {}
+        self.prefix_counters: typing.Dict[str, int] = {}
 
-    def generate_sequential_uid(self, prefix: str) -> str:
-        """Generates a sequential UID with specified prefix."""
+    def get_uid_stream(self, prefix: str) -> typing.Generator[str, None, None]:
+        """Returns next sequential UID with specified prefix."""
         if prefix not in self.prefix_counters:
             self.prefix_counters[prefix] = 0
 
-        uid = f"{prefix}{self.prefix_counters[prefix]}"
-        self.prefix_counters[prefix] += 1
-        return uid
+        while True:
+            uid = f"{prefix}{self.prefix_counters[prefix]}"
+            self.prefix_counters[prefix] += 1
+            yield uid

bigframes/core/rewrite/identifiers.py

@@ -13,16 +13,19 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Tuple
+import typing
 
-from bigframes.core import guid, identifiers, nodes
+from bigframes.core import identifiers, nodes
 
 
 # TODO: May as well just outright remove selection nodes in this process.
 def remap_variables(
     root: nodes.BigFrameNode,
-    uid_gen: guid.SequentialUIDGenerator,
-) -> Tuple[nodes.BigFrameNode, dict[identifiers.ColumnId, identifiers.ColumnId],]:
+    id_generator: typing.Iterator[identifiers.ColumnId],
+) -> typing.Tuple[
+    nodes.BigFrameNode,
+    dict[identifiers.ColumnId, identifiers.ColumnId],
+]:
     """Remaps `ColumnId`s in the BFET to produce deterministic and sequential UIDs.
 
     Note: this will convert a DAG to a tree.
@@ -31,7 +34,7 @@ def remap_variables(
     ref_mapping = dict()
     # Sequential ids are assigned bottom-up left-to-right
     for child in root.child_nodes:
-        new_child, child_var_mapping = remap_variables(child, uid_gen=uid_gen)
+        new_child, child_var_mapping = remap_variables(child, id_generator=id_generator)
         child_replacement_map[child] = new_child
         ref_mapping.update(child_var_mapping)
 
@@ -42,10 +45,7 @@ def remap_variables(
 
     with_new_refs = with_new_children.remap_refs(ref_mapping)
 
-    node_var_mapping = {
-        old_id: identifiers.ColumnId(name=uid_gen.generate_sequential_uid("bfcol_"))
-        for old_id in root.node_defined_ids
-    }
+    node_var_mapping = {old_id: next(id_generator) for old_id in root.node_defined_ids}
     with_new_vars = with_new_refs.remap_vars(node_var_mapping)
     with_new_vars._validate()
 

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

@@ -18,7 +18,6 @@
 
 import bigframes.core
 import bigframes.core.compile.sqlglot as sqlglot
-import bigframes.core.guid
 import bigframes.dataframe
 import bigframes.session.executor
 import bigframes.session.metrics
@@ -42,9 +41,9 @@ def to_sql(
 
         # Compared with BigQueryCachingExecutor, SQLCompilerExecutor skips
         # caching the subtree.
-        return self.compiler.SQLGlotCompiler(
-            uid_gen=bigframes.core.guid.SequentialUIDGenerator()
-        ).compile(array_value.node, ordered=ordered)
+        return self.compiler.SQLGlotCompiler().compile(
+            array_value.node, ordered=ordered
+        )
 
 
 class SQLCompilerSession(bigframes.session.Session):