bigframes 2.2.0__tar.gz → 2.3.0__tar.gz

{bigframes-2.2.0/bigframes.egg-info → bigframes-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bigframes
-Version: 2.2.0
+Version: 2.3.0
 Summary: BigQuery DataFrames -- scalable analytics and machine learning with BigQuery
 Home-page: https://github.com/googleapis/python-bigquery-dataframes
 Author: Google LLC

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/core/blocks.py

@@ -22,7 +22,6 @@ circular dependencies.
 from __future__ import annotations
 
 import ast
-import copy
 import dataclasses
 import datetime
 import functools
@@ -30,17 +29,7 @@ import itertools
 import random
 import textwrap
 import typing
-from typing import (
-    Any,
-    Iterable,
-    List,
-    Literal,
-    Mapping,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, Union
 import warnings
 
 import bigframes_vendored.constants as constants
@@ -69,6 +58,8 @@ import bigframes.dtypes
 import bigframes.exceptions as bfe
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
+from bigframes.session import dry_runs
+from bigframes.session import executor as executors
 
 # Type constraint for wherever column labels are used
 Label = typing.Hashable
@@ -821,59 +812,18 @@ class Block:
         if sampling.enable_downsampling:
             raise NotImplementedError("Dry run with sampling is not supported")
 
-        index: List[Any] = []
-        values: List[Any] = []
-
-        index.append("columnCount")
-        values.append(len(self.value_columns))
-        index.append("columnDtypes")
-        values.append(
-            {
-                col: self.expr.get_column_type(self.resolve_label_exact_or_error(col))
-                for col in self.column_labels
-            }
-        )
-
-        index.append("indexLevel")
-        values.append(self.index.nlevels)
-        index.append("indexDtypes")
-        values.append(self.index.dtypes)
-
         expr = self._apply_value_keys_to_expr(value_keys=value_keys)
         query_job = self.session._executor.dry_run(expr, ordered)
-        job_api_repr = copy.deepcopy(query_job._properties)
-
-        job_ref = job_api_repr["jobReference"]
-        for key, val in job_ref.items():
-            index.append(key)
-            values.append(val)
-
-        index.append("jobType")
-        values.append(job_api_repr["configuration"]["jobType"])
-
-        query_config = job_api_repr["configuration"]["query"]
-        for key in ("destinationTable", "useLegacySql"):
-            index.append(key)
-            values.append(query_config.get(key))
-
-        query_stats = job_api_repr["statistics"]["query"]
-        for key in (
-            "referencedTables",
-            "totalBytesProcessed",
-            "cacheHit",
-            "statementType",
-        ):
-            index.append(key)
-            values.append(query_stats.get(key))
 
-        index.append("creationTime")
-        values.append(
-            pd.Timestamp(
-                job_api_repr["statistics"]["creationTime"], unit="ms", tz="UTC"
-            )
-        )
+        column_dtypes = {
+            col: self.expr.get_column_type(self.resolve_label_exact_or_error(col))
+            for col in self.column_labels
+        }
 
-        return pd.Series(values, index=index), query_job
+        dry_run_stats = dry_runs.get_query_stats_with_dtypes(
+            query_job, column_dtypes, self.index.dtypes
+        )
+        return dry_run_stats, query_job
 
     def _apply_value_keys_to_expr(self, value_keys: Optional[Iterable[str]] = None):
         expr = self._expr
@@ -1560,12 +1510,19 @@ class Block:
         """
 
         # head caches full underlying expression, so row_count will be free after
-        head_result = self.session._executor.head(self.expr, max_results)
+        executor = self.session._executor
+        executor.cached(
+            array_value=self.expr,
+            config=executors.CacheConfig(optimize_for="head", if_cached="reuse-strict"),
+        )
+        head_result = self.session._executor.execute(
+            self.expr.slice(start=None, stop=max_results, step=None)
+        )
         row_count = self.session._executor.execute(self.expr.row_count()).to_py_scalar()
 
-        df = head_result.to_pandas()
-        self._copy_index_to_pandas(df)
-        return df, row_count, head_result.query_job
+        head_df = head_result.to_pandas()
+        self._copy_index_to_pandas(head_df)
+        return head_df, row_count, head_result.query_job
 
     def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]:
         expr, result_id = self._expr.promote_offsets()
@@ -2535,9 +2492,12 @@ class Block:
         # use a heuristic for whether something needs to be cached
         self.session._executor.cached(
             self.expr,
-            force=force,
-            use_session=session_aware,
-            cluster_cols=self.index_columns,
+            config=executors.CacheConfig(
+                optimize_for="auto"
+                if session_aware
+                else executors.HierarchicalKey(tuple(self.index_columns)),
+                if_cached="replace" if force else "reuse-any",
+            ),
         )
 
     def _is_monotonic(

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/core/compile/sqlglot/compiler.py

@@ -15,24 +15,28 @@ from __future__ import annotations
 
 import dataclasses
 import functools
-import itertools
 import typing
 
 from google.cloud import bigquery
 import pyarrow as pa
 import sqlglot.expressions as sge
 
-from bigframes.core import expression, identifiers, 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,
@@ -82,7 +86,7 @@ class SQLGlotCompiler:
             result_node = typing.cast(
                 nodes.ResultNode, rewrite.column_pruning(result_node)
             )
-            result_node = _remap_variables(result_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
@@ -92,7 +96,7 @@ class SQLGlotCompiler:
         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)
+        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:
@@ -106,63 +110,62 @@ class SQLGlotCompiler:
             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 = compile_node(root.child)
+        sqlglot_ir = self.compile_node(root.child)
         # TODO: add order_by, limit, and selections to sqlglot_expr
         return sqlglot_ir.sql
 
+    @functools.lru_cache(maxsize=5000)
+    def compile_node(self, node: nodes.BigFrameNode) -> ir.SQLGlotIR:
+        """Compiles node into CompileArrayValue. Caches result."""
+        return node.reduce_up(
+            lambda node, children: self._compile_node(node, *children)
+        )
 
-def _replace_unsupported_ops(node: nodes.BigFrameNode):
-    node = nodes.bottom_up(node, rewrite.rewrite_slice)
-    node = nodes.bottom_up(node, rewrite.rewrite_timedelta_expressions)
-    node = nodes.bottom_up(node, rewrite.rewrite_range_rolling)
-    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) -> 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: ir.SQLGlotIR
-) -> ir.SQLGlotIR:
-    """Defines transformation but isn't cached, always use compile_node instead"""
-    raise ValueError(f"Can't compile unrecognized node: {node}")
+    @functools.singledispatchmethod
+    def _compile_node(
+        self, 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(self, node: nodes.ReadLocalNode, *args) -> ir.SQLGlotIR:
+        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.id.sql for item in node.scan_list.items]
+        )
 
+        offsets = node.offsets_col.sql if node.offsets_col else None
+        if offsets:
+            pa_table = pa_table.append_column(
+                offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
+            )
 
-@_compile_node.register
-def compile_readlocal(node: nodes.ReadLocalNode, *args) -> ir.SQLGlotIR:
-    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.id.sql for item in node.scan_list.items])
+        return ir.SQLGlotIR.from_pyarrow(pa_table, node.schema, uid_gen=self.uid_gen)
 
-    offsets = node.offsets_col.sql if node.offsets_col else None
-    if offsets:
-        pa_table = pa_table.append_column(
-            offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
+    @_compile_node.register
+    def compile_selection(
+        self, node: nodes.SelectionNode, child: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        selected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
+            (id.sql, scalar_compiler.compile_scalar_expression(expr))
+            for expr, id in node.input_output_pairs
         )
+        return child.select(selected_cols)
 
-    return ir.SQLGlotIR.from_pyarrow(pa_table, node.schema)
 
-
-@_compile_node.register
-def compile_selection(node: nodes.SelectionNode, child: ir.SQLGlotIR) -> ir.SQLGlotIR:
-    select_cols: typing.Dict[str, sge.Expression] = {
-        id.name: scalar_compiler.compile_scalar_expression(expr)
-        for expr, id in node.input_output_pairs
-    }
-    return child.select(select_cols)
+def _replace_unsupported_ops(node: nodes.BigFrameNode):
+    node = nodes.bottom_up(node, rewrite.rewrite_slice)
+    node = nodes.bottom_up(node, rewrite.rewrite_timedelta_expressions)
+    node = nodes.bottom_up(node, rewrite.rewrite_range_rolling)
+    return node

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -23,6 +23,7 @@ import sqlglot.dialects.bigquery
 import sqlglot.expressions as sge
 
 from bigframes import dtypes
+from bigframes.core import guid
 import bigframes.core.compile.sqlglot.sqlglot_types as sgt
 import bigframes.core.local_data as local_data
 import bigframes.core.schema as schemata
@@ -52,6 +53,9 @@ class SQLGlotIR:
     pretty: bool = True
     """Whether to pretty-print the generated SQL."""
 
+    uid_gen: guid.SequentialUIDGenerator = guid.SequentialUIDGenerator()
+    """Generator for unique identifiers."""
+
     @property
     def sql(self) -> str:
         """Generate SQL string from the given expression."""
@@ -59,7 +63,10 @@ class SQLGlotIR:
 
     @classmethod
     def from_pyarrow(
-        cls, pa_table: pa.Table, schema: schemata.ArraySchema
+        cls,
+        pa_table: pa.Table,
+        schema: schemata.ArraySchema,
+        uid_gen: guid.SequentialUIDGenerator,
     ) -> SQLGlotIR:
         """Builds SQLGlot expression from pyarrow table."""
         dtype_expr = sge.DataType(
@@ -95,21 +102,44 @@ class SQLGlotIR:
                 ),
             ],
         )
-        return cls(expr=sg.select(sge.Star()).from_(expr))
+        return cls(expr=sg.select(sge.Star()).from_(expr), uid_gen=uid_gen)
 
     def select(
         self,
-        select_cols: typing.Dict[str, sge.Expression],
+        selected_cols: tuple[tuple[str, sge.Expression], ...],
     ) -> SQLGlotIR:
-        selected_cols = [
+        cols_expr = [
             sge.Alias(
                 this=expr,
                 alias=sge.to_identifier(id, quoted=self.quoted),
             )
-            for id, expr in select_cols.items()
+            for id, expr in selected_cols
         ]
-        expr = self.expr.select(*selected_cols, append=False)
-        return SQLGlotIR(expr=expr)
+        new_expr = self._encapsulate_as_cte().select(*cols_expr, append=False)
+        return SQLGlotIR(expr=new_expr)
+
+    def _encapsulate_as_cte(
+        self,
+    ) -> sge.Select:
+        """Transforms a given sge.Select query by pushing its main SELECT statement
+        into a new CTE and then generates a 'SELECT * FROM new_cte_name'
+        for the new query."""
+        select_expr = self.expr.copy()
+
+        existing_ctes = select_expr.args.pop("with", [])
+        new_cte_name = sge.to_identifier(
+            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
+        )
+        new_cte = sge.CTE(
+            this=select_expr,
+            alias=new_cte_name,
+        )
+        new_with_clause = sge.With(expressions=existing_ctes + [new_cte])
+        new_select_expr = (
+            sge.Select().select(sge.Star()).from_(sge.Table(this=new_cte_name))
+        )
+        new_select_expr.set("with", new_with_clause)
+        return new_select_expr
 
 
 def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:

bigframes-2.3.0/bigframes/core/guid.py

@@ -0,0 +1,46 @@
+# Copyright 2023 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 threading
+import typing
+
+_GUID_LOCK = threading.Lock()
+_GUID_COUNTER = 0
+
+
+def generate_guid(prefix="col_"):
+    global _GUID_LOCK
+    with _GUID_LOCK:
+        global _GUID_COUNTER
+        _GUID_COUNTER += 1
+        return f"bfuid_{prefix}{_GUID_COUNTER}"
+
+
+class SequentialUIDGenerator:
+    """Produces a sequence of UIDs, such as {"t0", "t1", "c0", "t2", ...}, by
+    cycling through provided prefixes (e.g., "t" and "c").
+    Note: this function is not thread-safe.
+    """
+
+    def __init__(self):
+        self.prefix_counters: typing.Dict[str, int] = {}
+
+    def get_uid_stream(self, prefix: str) -> typing.Generator[str, None, None]:
+        """Yields a continuous stream of raw UID strings for the given prefix."""
+        if prefix not in self.prefix_counters:
+            self.prefix_counters[prefix] = 0
+
+        while True:
+            uid = f"{prefix}{self.prefix_counters[prefix]}"
+            self.prefix_counters[prefix] += 1
+            yield uid

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/core/rewrite/identifiers.py

@@ -13,22 +13,20 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Generator, Tuple
+import typing
 
-import bigframes.core.identifiers
-import bigframes.core.nodes
+from bigframes.core import identifiers, nodes
 
 
 # TODO: May as well just outright remove selection nodes in this process.
 def remap_variables(
-    root: bigframes.core.nodes.BigFrameNode,
-    id_generator: Generator[bigframes.core.identifiers.ColumnId, None, None],
-) -> Tuple[
-    bigframes.core.nodes.BigFrameNode,
-    dict[bigframes.core.identifiers.ColumnId, bigframes.core.identifiers.ColumnId],
+    root: nodes.BigFrameNode,
+    id_generator: typing.Iterator[identifiers.ColumnId],
+) -> typing.Tuple[
+    nodes.BigFrameNode,
+    dict[identifiers.ColumnId, identifiers.ColumnId],
 ]:
-    """
-    Remap all variables in the BFET using the id_generator.
+    """Remaps `ColumnId`s in the BFET to produce deterministic and sequential UIDs.
 
     Note: this will convert a DAG to a tree.
     """

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/core/tools/datetimes.py

@@ -52,7 +52,7 @@ def to_datetime(
             f"to datetime is not implemented. {constants.FEEDBACK_LINK}"
         )
 
-    arg = bigframes.series.Series(arg)._cached()
+    arg = bigframes.series.Series(arg)
 
     if format and unit and arg.dtype in (bigframes.dtypes.INT_DTYPE, bigframes.dtypes.FLOAT_DTYPE):  # type: ignore
         raise ValueError("cannot specify both format and unit")
@@ -74,6 +74,11 @@ def to_datetime(
             )
 
         assert unit is None
+
+        # The following operations evaluate individual values to infer a format,
+        # so cache if needed.
+        arg = arg._cached(force=False)
+
         as_datetime = arg._apply_unary_op(  # type: ignore
             ops.ToDatetimeOp(
                 format=format,

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/ml/llm.py

@@ -617,7 +617,7 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
                 It creates a struct column of the items of the iterable, and use the concatenated result as the input prompt. No-op if set to None.
             output_schema (Mapping[str, str] or None, default None):
                 The schema used to generate structured output as a bigframes DataFrame. The schema is a string key-value pair of <column_name>:<type>.
-                Supported types are int64, float64, bool and string. If None, output text result.
+                Supported types are int64, float64, bool, string, array<type> and struct<column type>. If None, output text result.
         Returns:
             bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values.
         """

{bigframes-2.2.0 → bigframes-2.3.0}/bigframes/operations/datetime_ops.py

@@ -84,7 +84,7 @@ class UnixSeconds(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_seconds"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -94,7 +94,7 @@ class UnixMillis(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_millis"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -104,7 +104,7 @@ class UnixMicros(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "unix_micros"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not dtypes.TIMESTAMP_DTYPE:
+        if input_types[0] != dtypes.TIMESTAMP_DTYPE:
             raise TypeError("expected timestamp input")
         return dtypes.INT_DTYPE
 
@@ -114,7 +114,7 @@ class TimestampDiff(base_ops.BinaryOp):
     name: typing.ClassVar[str] = "timestamp_diff"
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
-        if input_types[0] is not input_types[1]:
+        if input_types[0] != input_types[1]:
             raise TypeError(
                 f"two inputs have different types. left: {input_types[0]}, right: {input_types[1]}"
             )