Merge remote-tracking branch 'origin/main' into b401261155-autosummary

Author: tswast

bigframes/core/compile/sqlglot/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from bigframes.core.compile.sqlglot.compiler import SQLGlotCompiler
+from bigframes.core.compile.sqlglot.compiler import compile_sql
 import bigframes.core.compile.sqlglot.expressions.ai_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.array_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.blob_ops  # noqa: F401
@@ -29,4 +29,4 @@
 import bigframes.core.compile.sqlglot.expressions.struct_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.timedelta_ops  # noqa: F401
 
-__all__ = ["SQLGlotCompiler"]
+__all__ = ["compile_sql"]

bigframes/core/compile/sqlglot/compiler.py

@@ -134,6 +134,11 @@ def _(
     )
 
 
+@register_binary_op(ops.fillna_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.Coalesce(this=left.expr, expressions=[right.expr])
+
+
 @register_nary_op(ops.case_when_op)
 def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     # Need to upcast BOOL to INT if any output is numeric

bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -305,6 +305,18 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
         return result
 
 
+@register_binary_op(ops.euclidean_distance_op)
+def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
+    return sge.Anonymous(
+        this="ML.DISTANCE",
+        expressions=[
+            left.expr,
+            right.expr,
+            sge.Literal.string("EUCLIDEAN"),
+        ],
+    )
+
+
 @register_binary_op(ops.floordiv_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     left_expr = _coerce_bool_to_int(left)

bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import dataclasses
+import datetime
 import functools
 import typing
 
@@ -118,6 +119,7 @@ def from_table(
         col_names: typing.Sequence[str],
         alias_names: typing.Sequence[str],
         uid_gen: guid.SequentialUIDGenerator,
+        system_time: typing.Optional[datetime.datetime] = None,
     ) -> SQLGlotIR:
         """Builds a SQLGlotIR expression from a BigQuery table.
 
@@ -128,18 +130,31 @@ def from_table(
             col_names (typing.Sequence[str]): The names of the columns to select.
             alias_names (typing.Sequence[str]): The aliases for the selected columns.
             uid_gen (guid.SequentialUIDGenerator): A generator for unique identifiers.
+            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",
+                expression=sge.Literal(this=system_time.isoformat(), is_string=True),
+                kind="AS OF",
+            )
+            if system_time
+            else None
+        )
         table_expr = sge.Table(
             this=sg.to_identifier(table_id, quoted=cls.quoted),
             db=sg.to_identifier(dataset_id, quoted=cls.quoted),
             catalog=sg.to_identifier(project_id, quoted=cls.quoted),
+            version=version,
         )
         select_expr = sge.Select().select(*selections).from_(table_expr)
         return cls(expr=select_expr, uid_gen=uid_gen)
@@ -227,6 +242,8 @@ def select(
                 this=expr,
                 alias=sge.to_identifier(id, quoted=self.quoted),
             )
+            if expr.alias_or_name != id
+            else expr
             for id, expr in selected_cols
         ]
 

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -593,6 +593,7 @@ def _agg_func(self, func) -> df.DataFrame:
     def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
         aggregations: typing.List[agg_expressions.Aggregation] = []
         column_labels = []
+        function_labels = []
 
         want_aggfunc_level = any(utils.is_list_like(aggs) for aggs in func.values())
 
@@ -602,8 +603,10 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
                 funcs_for_id if utils.is_list_like(funcs_for_id) else [funcs_for_id]
             )
             for f in func_list:
-                aggregations.append(aggs.agg(col_id, agg_ops.lookup_agg_func(f)[0]))
+                f_op, f_label = agg_ops.lookup_agg_func(f)
+                aggregations.append(aggs.agg(col_id, f_op))
                 column_labels.append(label)
+                function_labels.append(f_label)
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
             aggregations=aggregations,
@@ -613,10 +616,7 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
             agg_block = agg_block.with_column_labels(
                 utils.combine_indices(
                     pd.Index(column_labels),
-                    pd.Index(
-                        typing.cast(agg_ops.AggregateOp, agg.op).name
-                        for agg in aggregations
-                    ),
+                    pd.Index(function_labels),
                 )
             )
         else:

bigframes/core/groupby/dataframe_group_by.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import dataclasses
+import functools
 from typing import cast
 
-from bigframes.core import expression, nodes
+from bigframes.core import expression, identifiers, nodes
 
 
 def defer_selection(
@@ -26,12 +27,19 @@ def defer_selection(
 
     In many cases, these nodes will be merged or eliminated entirely, simplifying the overall tree.
     """
-    return nodes.bottom_up(root, pull_up_select)
+    return nodes.bottom_up(
+        root, functools.partial(pull_up_select, prefer_source_names=True)
+    )
 
 
-def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+def pull_up_select(
+    node: nodes.BigFrameNode, prefer_source_names: bool
+) -> nodes.BigFrameNode:
     if isinstance(node, nodes.LeafNode):
-        return node
+        if prefer_source_names and isinstance(node, nodes.ReadTableNode):
+            return pull_up_source_ids(node)
+        else:
+            return node
     if isinstance(node, nodes.JoinNode):
         return pull_up_selects_under_join(node)
     if isinstance(node, nodes.ConcatNode):
@@ -42,6 +50,32 @@ def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
     return node
 
 
+def pull_up_source_ids(node: nodes.ReadTableNode) -> nodes.BigFrameNode:
+    if all(id.sql == source_id for id, source_id in node.scan_list.items):
+        return node
+    else:
+        source_ids = sorted(
+            set(scan_item.source_id for scan_item in node.scan_list.items)
+        )
+        new_scan_list = nodes.ScanList.from_items(
+            [
+                nodes.ScanItem(identifiers.ColumnId(source_id), source_id)
+                for source_id in source_ids
+            ]
+        )
+        new_source = dataclasses.replace(node, scan_list=new_scan_list)
+        new_selection = nodes.SelectionNode(
+            new_source,
+            tuple(
+                nodes.AliasedRef(
+                    expression.DerefOp(identifiers.ColumnId(source_id)), id
+                )
+                for id, source_id in node.scan_list.items
+            ),
+        )
+        return new_selection
+
+
 def pull_up_select_unary(node: nodes.UnaryNode) -> nodes.BigFrameNode:
     child = node.child
     if not isinstance(child, nodes.SelectionNode):

bigframes/core/rewrite/select_pullup.py

@@ -717,9 +717,15 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
     np.all: all_op,
     np.any: any_op,
     np.unique: nunique_op,
-    # TODO(b/443252872): Solve
-    # list: ArrayAggOp(),
     np.size: size_op,
+    # TODO(b/443252872): Solve
+    list: ArrayAggOp(),
+    len: size_op,
+    sum: sum_op,
+    min: min_op,
+    max: max_op,
+    any: any_op,
+    all: all_op,
 }
 
 

bigframes/operations/aggregations.py

@@ -402,6 +402,7 @@ def get_index_cols(
     | bigframes.enums.DefaultIndexKind,
     *,
     rename_to_schema: Optional[Dict[str, str]] = None,
+    default_index_type: bigframes.enums.DefaultIndexKind = bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64,
 ) -> List[str]:
     """
     If we can get a total ordering from the table, such as via primary key
@@ -471,7 +472,11 @@ def get_index_cols(
         # find index_cols to use. This is to avoid unexpected performance and
         # resource utilization because of the default sequential index. See
         # internal issue 335727141.
-        if _is_table_clustered_or_partitioned(table) and not primary_keys:
+        if (
+            _is_table_clustered_or_partitioned(table)
+            and not primary_keys
+            and default_index_type == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
+        ):
             msg = bfe.format_message(
                 f"Table '{str(table.reference)}' is clustered and/or "
                 "partitioned, but BigQuery DataFrames was not able to find a "

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -40,9 +40,7 @@ def __init__(
     ):
         self.bqclient = bqclient
         self._compile_fn = (
-            compile.compile_sql
-            if compiler == "ibis"
-            else sqlglot.SQLGlotCompiler()._compile_sql
+            compile.compile_sql if compiler == "ibis" else sqlglot.compile_sql
         )
         self._publisher = publisher