improve data source handling, compile predicates again

Author: TrevorBergeron

bigframes/core/compile/sqlglot/compiler.py

@@ -141,7 +141,7 @@ def _compile_node(
 
 
 @_compile_node.register
-def compile_sql_select(node: sql_nodes.SelectNode, child: ir.SQLGlotIR):
+def compile_sql_select(node: sql_nodes.SqlSelectNode, child: ir.SQLGlotIR):
     sqlglot_ir = child
     if node.sorting is not None:
         ordering_cols = tuple(
@@ -165,6 +165,12 @@ def compile_sql_select(node: sql_nodes.SelectNode, child: ir.SQLGlotIR):
     )
     sqlglot_ir = sqlglot_ir.select(projected_cols)
 
+    if len(node.predicates) > 0:
+        sge_predicates = tuple(
+            scalar_compiler.scalar_op_compiler.compile_expression(expression)
+            for expression in node.predicates
+        )
+        sqlglot_ir = sqlglot_ir.filter(sge_predicates)
     if node.limit is not None:
         sqlglot_ir = sqlglot_ir.limit(node.limit)
 
@@ -185,14 +191,12 @@ def compile_readlocal(node: nodes.ReadLocalNode, child: ir.SQLGlotIR) -> ir.SQLG
 
 
 @_compile_node.register
-def compile_readtable(node: nodes.ReadTableNode, child: ir.SQLGlotIR):
+def compile_readtable(node: sql_nodes.SqlDataSource, child: ir.SQLGlotIR):
     table = node.source.table
     return ir.SQLGlotIR.from_table(
         table.project_id,
         table.dataset_id,
         table.table_id,
-        col_names=[col.source_id for col in node.scan_list.items],
-        alias_names=[col.id.sql for col in node.scan_list.items],
         uid_gen=child.uid_gen,
         sql_predicate=node.source.sql_predicate,
         system_time=node.source.at_time,

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -116,8 +116,6 @@ def from_table(
         project_id: str,
         dataset_id: str,
         table_id: str,
-        col_names: typing.Sequence[str],
-        alias_names: typing.Sequence[str],
         uid_gen: guid.SequentialUIDGenerator,
         sql_predicate: typing.Optional[str] = None,
         system_time: typing.Optional[datetime.datetime] = None,
@@ -134,15 +132,6 @@ def from_table(
             sql_predicate (typing.Optional[str]): An optional SQL predicate for filtering.
             system_time (typing.Optional[str]): An optional system time for time-travel queries.
         """
-        selections = [
-            sge.Alias(
-                this=sge.to_identifier(col_name, quoted=cls.quoted),
-                alias=sge.to_identifier(alias_name, quoted=cls.quoted),
-            )
-            if col_name != alias_name
-            else sge.to_identifier(col_name, quoted=cls.quoted)
-            for col_name, alias_name in zip(col_names, alias_names)
-        ]
         version = (
             sge.Version(
                 this="TIMESTAMP",
@@ -158,12 +147,14 @@ def from_table(
             catalog=sg.to_identifier(project_id, quoted=cls.quoted),
             version=version,
         )
-        select_expr = sge.Select().select(*selections).from_(table_expr)
         if sql_predicate:
+            select_expr = sge.Select().select(sge.Star()).from_(table_expr)
             select_expr = select_expr.where(
                 sg.parse_one(sql_predicate, dialect="bigquery"), append=False
             )
-        return cls(expr=select_expr, uid_gen=uid_gen)
+            return cls(expr=select_expr, uid_gen=uid_gen)
+
+        return cls(expr=table_expr, uid_gen=uid_gen)
 
     @classmethod
     def from_query_string(

bigframes/core/rewrite/as_sql.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Google LLC
+# Copyright 2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,14 +27,14 @@
 import bigframes.core.rewrite
 
 
-def _limit(select: sql_nodes.SelectNode, limit: int) -> sql_nodes.SelectNode:
+def _limit(select: sql_nodes.SqlSelectNode, limit: int) -> sql_nodes.SqlSelectNode:
     new_limit = limit if select.limit is None else min([select.limit, limit])
     return dataclasses.replace(select, limit=new_limit)
 
 
 def _try_sort(
-    select: sql_nodes.SelectNode, sort_by: Sequence[ordering.OrderingExpression]
-) -> Optional[sql_nodes.SelectNode]:
+    select: sql_nodes.SqlSelectNode, sort_by: Sequence[ordering.OrderingExpression]
+) -> Optional[sql_nodes.SqlSelectNode]:
     new_order_exprs = []
     for sort_expr in sort_by:
         new_expr = _try_bind(
@@ -50,8 +50,8 @@ def _try_sort(
 
 def _sort(
     node: nodes.BigFrameNode, sort_by: Sequence[ordering.OrderingExpression]
-) -> sql_nodes.SelectNode:
-    if isinstance(node, sql_nodes.SelectNode):
+) -> sql_nodes.SqlSelectNode:
+    if isinstance(node, sql_nodes.SqlSelectNode):
         merged = _try_sort(node, sort_by)
         if merged:
             return merged
@@ -73,8 +73,8 @@ def _try_bind(
 
 
 def _try_add_cdefs(
-    select: sql_nodes.SelectNode, cdefs: Sequence[nodes.ColumnDef]
-) -> Optional[sql_nodes.SelectNode]:
+    select: sql_nodes.SqlSelectNode, cdefs: Sequence[nodes.ColumnDef]
+) -> Optional[sql_nodes.SqlSelectNode]:
     # TODO: add up complexity measure while inlining refs
     new_defs = []
     for cdef in cdefs:
@@ -91,8 +91,8 @@ def _try_add_cdefs(
 
 def _add_cdefs(
     node: nodes.BigFrameNode, cdefs: Sequence[nodes.ColumnDef]
-) -> sql_nodes.SelectNode:
-    if isinstance(node, sql_nodes.SelectNode):
+) -> sql_nodes.SqlSelectNode:
+    if isinstance(node, sql_nodes.SqlSelectNode):
         merged = _try_add_cdefs(node, cdefs)
         if merged:
             return merged
@@ -103,8 +103,8 @@ def _add_cdefs(
 
 
 def _try_add_filter(
-    select: sql_nodes.SelectNode, predicates: Sequence[expression.Expression]
-) -> Optional[sql_nodes.SelectNode]:
+    select: sql_nodes.SqlSelectNode, predicates: Sequence[expression.Expression]
+) -> Optional[sql_nodes.SqlSelectNode]:
     # Constraint: filters can only be merged if they are scalar expression after binding
     new_predicates = []
     # bind variables, merge predicates
@@ -118,8 +118,8 @@ def _try_add_filter(
 
 def _add_filter(
     node: nodes.BigFrameNode, predicates: Sequence[expression.Expression]
-) -> sql_nodes.SelectNode:
-    if isinstance(node, sql_nodes.SelectNode):
+) -> sql_nodes.SqlSelectNode:
+    if isinstance(node, sql_nodes.SqlSelectNode):
         result = _try_add_filter(node, predicates)
         if result:
             return result
@@ -128,8 +128,8 @@ def _add_filter(
     return new_node
 
 
-def _create_noop_select(node: nodes.BigFrameNode) -> sql_nodes.SelectNode:
-    return sql_nodes.SelectNode(
+def _create_noop_select(node: nodes.BigFrameNode) -> sql_nodes.SqlSelectNode:
+    return sql_nodes.SqlSelectNode(
         node,
         selections=tuple(
             nodes.ColumnDef(expression.ResolvedDerefOp.from_field(field), field.id)
@@ -139,7 +139,7 @@ def _create_noop_select(node: nodes.BigFrameNode) -> sql_nodes.SelectNode:
 
 
 def _try_remap_select_cols(
-    select: sql_nodes.SelectNode, cols: Sequence[nodes.AliasedRef]
+    select: sql_nodes.SqlSelectNode, cols: Sequence[nodes.AliasedRef]
 ):
     new_defs = []
     for aliased_ref in cols:
@@ -151,7 +151,7 @@ def _try_remap_select_cols(
 
 
 def _remap_select_cols(node: nodes.BigFrameNode, cols: Sequence[nodes.AliasedRef]):
-    if isinstance(node, sql_nodes.SelectNode):
+    if isinstance(node, sql_nodes.SqlSelectNode):
         result = _try_remap_select_cols(node, cols)
         if result:
             return result
@@ -183,7 +183,14 @@ def _get_added_cdefs(node: Union[nodes.ProjectionNode, nodes.WindowOpNode]):
 
 def _as_sql_node(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
     # case one, can be converted to select
-    if isinstance(node, (nodes.ProjectionNode, nodes.WindowOpNode)):
+    if isinstance(node, nodes.ReadTableNode):
+        leaf = sql_nodes.SqlDataSource(source=node.source)
+        mappings = [
+            nodes.AliasedRef(expression.deref(scan_item.source_id), scan_item.id)
+            for scan_item in node.scan_list.items
+        ]
+        return _remap_select_cols(leaf, mappings)
+    elif isinstance(node, (nodes.ProjectionNode, nodes.WindowOpNode)):
         cdefs = _get_added_cdefs(node)
         return _add_cdefs(node.child, cdefs)
     elif isinstance(node, (nodes.SelectionNode)):

bigframes/core/sql_nodes.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,26 +18,74 @@
 import functools
 from typing import Mapping, Optional, Sequence, Tuple
 
-from bigframes.core import identifiers, nodes
+from bigframes.core import bq_data, identifiers, nodes
 import bigframes.core.expression as ex
 from bigframes.core.ordering import OrderingExpression
 import bigframes.dtypes
 
-# A fixed number of variable to assume for overhead on some operations
-OVERHEAD_VARIABLES = 5
 
+# TODO: Join node, union node
+@dataclasses.dataclass(frozen=True)
+class SqlDataSource(nodes.LeafNode):
+    source: bq_data.BigqueryDataSource
+
+    @functools.cached_property
+    def fields(self) -> Sequence[nodes.Field]:
+        return tuple(
+            nodes.Field(
+                identifiers.ColumnId(source_id),
+                self.source.schema.get_type(source_id),
+                self.source.table.schema_by_id[source_id].is_nullable,
+            )
+            for source_id in self.source.schema.names
+        )
+
+    @property
+    def variables_introduced(self) -> int:
+        # This operation only renames variables, doesn't actually create new ones
+        return 0
+
+    @property
+    def defines_namespace(self) -> bool:
+        return True
+
+    @property
+    def explicitly_ordered(self) -> bool:
+        return False
+
+    @property
+    def order_ambiguous(self) -> bool:
+        return True
+
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.source.n_rows
+
+    @property
+    def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
+        return tuple(self.ids)
+
+    @property
+    def consumed_ids(self):
+        return ()
 
-@dataclasses.dataclass(frozen=True, eq=True)
-class ColumnDef:
-    expression: ex.Expression
-    id: identifiers.ColumnId
+    @property
+    def _node_expressions(self):
+        return ()
 
+    def remap_vars(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlSelectNode:
+        raise NotImplementedError()
 
-# TODO: Raw data source node, join node, union node
+    def remap_refs(
+        self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
+    ) -> SqlSelectNode:
+        raise NotImplementedError()  # type: ignore
 
 
 @dataclasses.dataclass(frozen=True)
-class SelectNode(nodes.UnaryNode):
+class SqlSelectNode(nodes.UnaryNode):
     selections: tuple[nodes.ColumnDef, ...] = ()
     predicates: tuple[ex.Expression, ...] = ()
     sorting: tuple[OrderingExpression, ...] = ()
@@ -106,10 +154,10 @@ def get_id_mapping(self) -> dict[identifiers.ColumnId, ex.Expression]:
 
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
-    ) -> SelectNode:
+    ) -> SqlSelectNode:
         raise NotImplementedError()
 
     def remap_refs(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
-    ) -> SelectNode:
+    ) -> SqlSelectNode:
         raise NotImplementedError()  # type: ignore