bigframes 2.41.0__tar.gz → 2.43.0__tar.gz

{bigframes-2.41.0/bigframes.egg-info → bigframes-2.43.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bigframes
-Version: 2.41.0
+Version: 2.43.0
 Summary: BigQuery DataFrames -- scalable analytics and machine learning with BigQuery
 Home-page: https://dataframes.bigquery.dev
 Download-URL: https://github.com/googleapis/google-cloud-python/tree/main/packages/bigframes/releases
@@ -30,7 +30,7 @@ Requires-Dist: cloudpickle>=2.0.0
 Requires-Dist: fsspec>=2023.3.0
 Requires-Dist: gcsfs!=2025.5.0,!=2026.2.0,!=2026.3.0,>=2023.3.0
 Requires-Dist: geopandas>=0.12.2
-Requires-Dist: google-auth<3.0,>=2.15.0
+Requires-Dist: google-auth[pyopenssl]<3.0,>=2.15.0
 Requires-Dist: google-cloud-bigquery[bqstorage,pandas]>=3.36.0
 Requires-Dist: google-cloud-bigquery-storage<3.0.0,>=2.30.0
 Requires-Dist: google-cloud-functions>=1.12.0
@@ -42,7 +42,7 @@ Requires-Dist: grpc-google-iam-v1>=0.14.2
 Requires-Dist: numpy>=1.24.0
 Requires-Dist: pandas>=1.5.3
 Requires-Dist: pandas-gbq>=0.26.1
-Requires-Dist: pyarrow>=15.0.2
+Requires-Dist: pyarrow>=23.0.1
 Requires-Dist: pydata-google-auth>=1.8.2
 Requires-Dist: requests>=2.27.1
 Requires-Dist: shapely>=1.8.5
@@ -62,6 +62,7 @@ Requires-Dist: freezegun; extra == "tests"
 Requires-Dist: pytest-snapshot; extra == "tests"
 Requires-Dist: google-cloud-bigtable>=2.24.0; extra == "tests"
 Requires-Dist: google-cloud-pubsub>=2.21.4; extra == "tests"
+Requires-Dist: tzdata; extra == "tests"
 Provides-Extra: polars
 Requires-Dist: polars>=1.21.0; extra == "polars"
 Provides-Extra: scikit-learn
@@ -87,6 +88,7 @@ Requires-Dist: pytest; extra == "all"
 Requires-Dist: pytest-snapshot; extra == "all"
 Requires-Dist: scikit-learn>=1.2.2; extra == "all"
 Requires-Dist: traitlets>=5.0.0; extra == "all"
+Requires-Dist: tzdata; extra == "all"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/__init__.py

@@ -42,6 +42,7 @@ import bigframes.operations  # noqa: E402 # isort: skip
 
 # Register pandas extensions
 import bigframes.extensions.pandas.dataframe_accessor  # noqa: F401, E402
+import bigframes.extensions.pandas.series_accessor  # noqa: F401, E402
 from bigframes._config.bigquery_options import BigQueryOptions  # noqa: E402
 from bigframes.core.global_session import (  # noqa: E402
     close_session,

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/bigquery/__init__.py

@@ -114,6 +114,18 @@ from bigframes.operations.googlesql.global_namespace.array import (
     flatten,
     generate_array,
 )
+from bigframes.operations.googlesql.global_namespace.bit import (
+    bit_count,
+)
+from bigframes.operations.googlesql.global_namespace.conversion import (
+    bool_,
+    double,
+    float64,
+    int64,
+    parse_bignumeric,
+    parse_numeric,
+    string,
+)
 
 _functions = [
     # approximate aggregate ops
@@ -134,6 +146,16 @@ _functions = [
     array_to_string,
     flatten,
     generate_array,
+    # bit ops
+    bit_count,
+    # conversion ops
+    bool_,
+    double,
+    float64,
+    int64,
+    parse_bignumeric,
+    parse_numeric,
+    string,
     # datetime ops
     unix_micros,
     unix_millis,
@@ -208,6 +230,16 @@ __all__ = [
     "array_to_string",
     "flatten",
     "generate_array",
+    # bit ops
+    "bit_count",
+    # conversion ops
+    "bool_",
+    "double",
+    "float64",
+    "int64",
+    "parse_bignumeric",
+    "parse_numeric",
+    "string",
     # datetime ops
     "unix_micros",
     "unix_millis",

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/bigquery/_operations/ai.py

@@ -61,7 +61,7 @@ def generate(
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
         >>> country = bpd.Series(["Japan", "Canada"])
-        >>> bbq.ai.generate(("What's the capital city of ", country, " one word only"))
+        >>> bbq.ai.generate(("What's the capital city of ", country, " one word only")) # doctest: +ELLIPSIS
         0    {'result': 'Tokyo', 'full_response': '{"cand...
         1    {'result': 'Ottawa', 'full_response': '{"can...
         dtype: struct<result: string, full_response: extension<dbjson<JSONArrowType>>, status: string>[pyarrow]
@@ -231,8 +231,8 @@ def generate_int(
 
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
-        >>> animal = bpd.Series(["Kangaroo", "Rabbit", "Spider"])
-        >>> bbq.ai.generate_int(("How many legs does a ", animal, " have?"))
+        >>> animal = bpd.Series(["Ostrich", "Rabbit", "Spider"])
+        >>> bbq.ai.generate_int(("How many legs does a ", animal, " have?")) # doctest: +ELLIPSIS
         0    {'result': 2, 'full_response': '{"candidates":...
         1    {'result': 4, 'full_response': '{"candidates":...
         2    {'result': 8, 'full_response': '{"candidates":...
@@ -305,8 +305,8 @@ def generate_double(
 
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
-        >>> animal = bpd.Series(["Kangaroo", "Rabbit", "Spider"])
-        >>> bbq.ai.generate_double(("How many legs does a ", animal, " have?"))
+        >>> animal = bpd.Series(["Ostrich", "Rabbit", "Spider"])
+        >>> bbq.ai.generate_double(("How many legs does a ", animal, " have?")) # doctest: +ELLIPSIS
         0    {'result': 2.0, 'full_response': '{"candidates...
         1    {'result': 4.0, 'full_response': '{"candidates...
         2    {'result': 8.0, 'full_response': '{"candidates...
@@ -383,7 +383,7 @@ def generate_embedding(
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
         >>> df = bpd.DataFrame({"content": ["apple", "bear", "pear"]})
-        >>> bbq.ai.generate_embedding(
+        >>> bbq.ai.generate_embedding( # doctest: +SKIP
         ...     "project.dataset.model_name",
         ...     df
         ... )
@@ -486,7 +486,7 @@ def generate_text(
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
         >>> df = bpd.DataFrame({"prompt": ["write a poem about apples"]})
-        >>> bbq.ai.generate_text(
+        >>> bbq.ai.generate_text( # doctest: +SKIP
         ...     "project.dataset.model_name",
         ...     df
         ... )
@@ -601,7 +601,7 @@ def generate_table(
         >>> # the necessary columns for the model's prompt. For example, a
         >>> # DataFrame with a 'prompt' column for text classification.
         >>> df = bpd.DataFrame({'prompt': ["some text to classify"]})
-        >>> result = bbq.ai.generate_table(
+        >>> result = bbq.ai.generate_table( # doctest: +SKIP
         ...     "project.dataset.model_name",
         ...     data=df,
         ...     output_schema="category STRING"
@@ -708,12 +708,14 @@ def embed(
 
         >>> import bigframes.pandas as bpd
         >>> import bigframes.bigquery as bbq
-        >>> bbq.ai.embed("dog", endpoint="text-embedding-005")
+        >>> bbq.ai.embed("dog", endpoint="text-embedding-005") # doctest: +ELLIPSIS
         0    {'result': array([ 1.78243860e-03, -1.10658340...
+        dtype: struct<result: list<item: double>, status: string>[pyarrow]
 
         >>> s = bpd.Series(['dog'])
-        >>> bbq.ai.embed(s, endpoint='text-embedding-005')
+        >>> bbq.ai.embed(s, endpoint='text-embedding-005') # doctest: +ELLIPSIS
         0    {'result': array([ 1.78243860e-03, -1.10658340...
+        dtype: struct<result: list<item: double>, status: string>[pyarrow]
 
     Args:
         content (str | Series):
@@ -1004,6 +1006,7 @@ def similarity(
         >>> bbq.ai.similarity(df['word'], 'glad', endpoint='text-embedding-005')
         0    0.916601
         1    0.660579
+        Name: word, dtype: Float64
 
     Args:
         content1 (str | Series):
@@ -1082,8 +1085,8 @@ def forecast(
         >>> df = pd.DataFrame({"value": [1, 2, 3], "time": pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03"])})
         >>> bpd.options.display.progress_bar = None
         >>> forecasted_pandas_df = df.bigquery.ai.forecast(data_col="value", timestamp_col="time", horizon=2)
-        >>> type(forecasted_pandas_df)
-        <class 'pandas.core.frame.DataFrame'>
+        >>> type(forecasted_pandas_df) # doctest: +ELLIPSIS
+        <class 'pandas...DataFrame'>
 
         Forecast using a BigFrames DataFrame:
 

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/bigquery/_operations/struct.py

@@ -57,5 +57,5 @@ def struct(value: dataframe.DataFrame) -> series.Series:
     block, result_id = block.apply_nary_op(
         block.value_columns, ops.StructOp(column_names=tuple(block.column_labels))
     )
-    block = block.select_column(result_id)
+    block = block.select_column(result_id).with_column_labels([None])
     return series.Series(block)

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/blocks.py

@@ -696,6 +696,7 @@ class Block:
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         allow_large_results: Optional[bool] = None,
+        cell_execution_count: Optional[int] = None,
     ) -> PandasBatches:
         """Download results one message at a time.
 
@@ -713,6 +714,7 @@ class Block:
             execution_spec.ExecutionSpec(
                 promise_under_10gb=under_10gb,
                 ordered=True,
+                cell_execution_count=cell_execution_count,
             ),
         )
         result_batches = execution_result.batches()
@@ -1989,6 +1991,10 @@ class Block:
                 )
             level = level or 0
             col_id = self.index.resolve_level(level)[0]
+            if isinstance(level, int):
+                resample_label = self.index.names[level]
+            else:
+                resample_label = level
             # Reset index to make the resampling level a column, then drop all other index columns.
             # This simplifies processing by focusing solely on the column required for resampling.
             block = self.reset_index(drop=False)
@@ -2007,6 +2013,7 @@ class Block:
                 raise KeyError(f"The grouper name {on} is not found")
 
             col_id = matches[0]
+            resample_label = on
             block = self
         if level is None:
             dtype = self._column_type(col_id)
@@ -2099,6 +2106,7 @@ class Block:
             block.value_columns[0],
             block.value_columns[1],
             op=ops.IntegerLabelToDatetimeOp(freq=freq, label=label, origin=origin),
+            result_label=resample_label,
         )
 
         # After multiple merges, the columns:

bigframes-2.43.0/bigframes/core/bytecode.py

@@ -0,0 +1,263 @@
+# 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.
+# 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 dis
+import operator
+import sys
+from types import ModuleType
+from typing import Callable
+
+import bigframes.core.py_expressions as py_exprs
+from bigframes.core import expression
+
+_BINARY_OP_MAP = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+    "//": operator.floordiv,
+    "%": operator.mod,
+    "**": operator.pow,
+}
+
+_COMPARE_OP_MAP = {
+    "==": operator.eq,
+    "!=": operator.ne,
+    "<": operator.lt,
+    "<=": operator.le,
+    ">": operator.gt,
+    ">=": operator.ge,
+}
+
+_OLD_BINARY_OP_MAP = {
+    "BINARY_ADD": operator.add,
+    "INPLACE_ADD": operator.add,
+    "BINARY_SUBTRACT": operator.sub,
+    "INPLACE_SUBTRACT": operator.sub,
+    "BINARY_MULTIPLY": operator.mul,
+    "INPLACE_MULTIPLY": operator.mul,
+    "BINARY_TRUE_DIVIDE": operator.truediv,
+    "INPLACE_TRUE_DIVIDE": operator.truediv,
+    "BINARY_FLOOR_DIVIDE": operator.floordiv,
+    "INPLACE_FLOOR_DIVIDE": operator.floordiv,
+    "BINARY_MODULO": operator.mod,
+    "INPLACE_MODULO": operator.mod,
+    "BINARY_POWER": operator.pow,
+    "INPLACE_POWER": operator.pow,
+}
+
+
+_NULL = py_exprs.PyObject(None)
+
+
+def _compile_bytecode_to_py_expr(func: Callable) -> expression.Expression:
+    instructions = list(dis.get_instructions(func))
+
+    stack: list[expression.Expression] = []
+    globals_dict = func.__globals__
+    import builtins
+
+    builtins_dict = builtins.__dict__
+
+    closure_dict = {}
+    if func.__closure__:
+        free_vars = func.__code__.co_freevars
+        for var, cell in zip(free_vars, func.__closure__):
+            try:
+                closure_dict[var] = cell.cell_contents
+            except ValueError:
+                pass
+
+    for inst in instructions:
+        opname = inst.opname
+
+        if opname in ("RESUME", "PRECALL"):
+            continue
+
+        elif opname in ("LOAD_FAST_LOAD_FAST", "LOAD_FAST_BORROW_LOAD_FAST_BORROW"):
+            var1, var2 = inst.argval
+            stack.append(expression.UnboundVariableExpression(var1))
+            stack.append(expression.UnboundVariableExpression(var2))
+
+        elif opname.startswith("LOAD_FAST"):
+            stack.append(expression.UnboundVariableExpression(inst.argval))
+
+        elif opname in ("LOAD_CONST", "LOAD_SMALL_INT"):
+            stack.append(py_exprs.PyObject(inst.argval))
+
+        elif opname == "LOAD_GLOBAL":
+            # In Python 3.11+, the lowest bit of inst.arg indicates that a NULL
+            # should be pushed before the global variable.
+            if sys.version_info >= (3, 11) and inst.arg is not None and (inst.arg & 1):
+                stack.append(_NULL)
+            name = inst.argval
+            found = False
+            val = None
+            if name in closure_dict:
+                val = closure_dict[name]
+                found = True
+            elif name in globals_dict:
+                val = globals_dict[name]
+                found = True
+            elif name in builtins_dict:
+                val = builtins_dict[name]
+                found = True
+
+            if found:
+                if isinstance(val, ModuleType):
+                    stack.append(py_exprs.Module(val))
+                else:
+                    stack.append(py_exprs.PyObject(val))
+            else:
+                stack.append(expression.UnboundVariableExpression(name))
+
+        elif opname in ("LOAD_ATTR", "LOAD_METHOD"):
+            if not stack:
+                raise ValueError("Stack is empty")
+            target = stack.pop()
+            stack.append(py_exprs.GetAttr(target, inst.argval))
+            if opname == "LOAD_METHOD":
+                if isinstance(target, py_exprs.Module):
+                    stack.append(_NULL)
+                else:
+                    stack.append(target)
+
+        elif opname == "PUSH_NULL":
+            stack.append(_NULL)
+
+        elif opname == "BINARY_OP":
+            if len(stack) < 2:
+                raise ValueError("Stack is empty")
+            right = stack.pop()
+            left = stack.pop()
+            op_symbol = inst.argrepr
+            if not op_symbol and isinstance(inst.argval, str):
+                op_symbol = inst.argval
+            if op_symbol and op_symbol.endswith("="):
+                op_symbol = op_symbol[:-1]
+
+            if op_symbol not in _BINARY_OP_MAP:
+                raise ValueError(f"Unsupported binary operator: {op_symbol}")
+            stack.append(
+                py_exprs.Call(
+                    py_exprs.PyObject(_BINARY_OP_MAP[op_symbol]), (left, right)
+                )
+            )
+
+        # Support older Python versions compatibility
+        elif opname in _OLD_BINARY_OP_MAP:
+            if len(stack) < 2:
+                raise ValueError("Stack has < 2 elements")
+            right = stack.pop()
+            left = stack.pop()
+            stack.append(
+                py_exprs.Call(
+                    py_exprs.PyObject(_OLD_BINARY_OP_MAP[opname]), (left, right)
+                )
+            )
+
+        elif opname == "COMPARE_OP":
+            if len(stack) < 2:
+                raise ValueError("Stack has < 2 elements")
+            right = stack.pop()
+            left = stack.pop()
+            op_symbol = inst.argval
+            if op_symbol not in _COMPARE_OP_MAP:
+                raise ValueError(f"Unsupported compare operator: {op_symbol}")
+            stack.append(
+                py_exprs.Call(
+                    py_exprs.PyObject(_COMPARE_OP_MAP[op_symbol]), (left, right)
+                )
+            )
+
+        elif opname in ("UNARY_NEGATIVE", "UNARY_INVERT"):
+            if not stack:
+                raise ValueError("Stack is empty")
+            target = stack.pop()
+            stack.append(
+                py_exprs.Call(
+                    py_exprs.PyObject(
+                        operator.neg if opname == "UNARY_NEGATIVE" else operator.invert
+                    ),
+                    (target,),
+                )
+            )
+
+        elif opname == "UNARY_POSITIVE":
+            if not stack:
+                raise ValueError("Stack is empty")
+            target = stack.pop()
+            stack.append(py_exprs.Call(py_exprs.PyObject(operator.pos), (target,)))
+
+        elif opname == "CALL_INTRINSIC_1":
+            if inst.argrepr == "INTRINSIC_UNARY_POSITIVE":
+                if not stack:
+                    raise ValueError("Stack is empty")
+                target = stack.pop()
+                stack.append(py_exprs.Call(py_exprs.PyObject(operator.pos), (target,)))
+            else:
+                raise ValueError(f"Unsupported intrinsic: {inst.argrepr}")
+
+        elif opname in ("CALL", "CALL_FUNCTION", "CALL_METHOD"):
+            num_args = inst.arg
+            assert num_args is not None
+            if len(stack) < num_args:
+                raise ValueError("Stack has < 2 elements")
+            args = [stack.pop() for _ in range(num_args)][::-1]
+            # In Python 3.11, LOAD_GLOBAL with NULL push puts NULL below the global.
+            # If NULL is below the callable on the stack, swap them to match
+            # the expected layout [callable, NULL].
+            if len(stack) >= 2 and stack[-2] == _NULL:
+                stack[-1], stack[-2] = stack[-2], stack[-1]
+            if stack and stack[-1] == _NULL:
+                stack.pop()
+            elif (
+                stack
+                and stack[-1] != _NULL
+                and isinstance(stack[-1], expression.Expression)
+            ):
+                self_arg = stack.pop()
+                args = [self_arg] + args
+            if not stack:
+                raise ValueError("Stack is empty")
+            callable_expr = stack.pop()
+            stack.append(py_exprs.Call(callable_expr, tuple(args)))
+
+        elif opname == "RETURN_VALUE":
+            if not stack:
+                raise ValueError("Stack is empty")
+            return stack[-1]
+
+        elif opname in ("STORE_FAST", "POP_TOP"):
+            if stack:
+                stack.pop()
+
+        else:
+            raise ValueError(f"Unsupported opcode: {opname}")
+
+    raise ValueError("No return value found")
+
+
+def dis_to_expr(func: Callable, unpack_mode: bool = False) -> expression.Expression:
+    """
+    Try to convert a python function to a BigQuery expression.
+
+    Unpack mode is whether SQL columns are addressed as attributes of a single
+    python argument (e.g. row.col1), or as separate arguments (e.g. col1).
+
+    This is "best effort" - if the function contains operations that cannot
+    be converted to BigQuery expressions, it will raise an Exception.
+    """
+    py_expr = _compile_bytecode_to_py_expr(func)
+    return py_exprs.resolve_py_exprs(py_expr, unpack_mode=unpack_mode)

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/compile/polars/compiler.py

@@ -178,8 +178,59 @@ if polars_installed:
             self,
             expression: ex.OpExpression,
         ) -> pl.Expr:
-            # TODO: Complete the implementation
+            import datetime
+
+            import pyarrow as pa
+
             op = expression.op
+
+            # Polars panics on nulls from pandas objects in timezone-aware
+            # datetimes for certain ops. Convert to timezone-naive temporarily
+            # to avoid this issue.
+            # TODO(tswast): Remove workaround when
+            # https://github.com/pola-rs/polars/issues/27862 has been fixed.
+            is_problematic_op = type(op) in (
+                date_ops.YearOp,
+                date_ops.QuarterOp,
+                date_ops.MonthOp,
+                date_ops.DayOp,
+                date_ops.IsoWeekOp,
+            )
+
+            if is_problematic_op and len(expression.inputs) == 1:
+                input_expr = expression.inputs[0]
+                if (
+                    input_expr.is_resolved
+                    and isinstance(input_expr.output_type, pd.ArrowDtype)
+                    and isinstance(
+                        input_expr.output_type.pyarrow_dtype, pa.TimestampType
+                    )
+                    and input_expr.output_type.pyarrow_dtype.tz is not None
+                ):
+                    tz_str = input_expr.output_type.pyarrow_dtype.tz
+                    if tz_str == "UTC":
+                        dummy_tz = datetime.timezone.utc
+                    else:
+                        try:
+                            from zoneinfo import ZoneInfo
+
+                            dummy_tz = ZoneInfo(tz_str)  # type: ignore
+                        except Exception:
+                            dummy_tz = datetime.timezone.utc
+
+                    dummy_val = datetime.datetime(1970, 1, 1, tzinfo=dummy_tz)
+
+                    compiled_input = self.compile_expression(input_expr)
+                    filled_input = compiled_input.fill_null(dummy_val)
+                    compiled_op_with_fill = self.compile_op(op, filled_input)
+
+                    return (
+                        pl.when(compiled_input.is_null())
+                        .then(None)
+                        .otherwise(compiled_op_with_fill)
+                    )
+
+            # TODO: Complete the implementation
             args = tuple(map(self.compile_expression, expression.inputs))
             return self.compile_op(op, *args)
 

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/compile/sqlglot/expressions/ai_ops.py

@@ -111,7 +111,8 @@ def _construct_prompt(
         else:
             prompt.append(sge.Literal.string(elem))
 
-    return sge.Kwarg(this=param_name, expression=sge.Tuple(expressions=prompt))
+    # Need Struct rather than tuple syntax, as tuple syntax is ambiguous for single arg
+    return sge.Kwarg(this=param_name, expression=sge.Struct(expressions=prompt))
 
 
 def _construct_named_args(op: ops.ScalarOp) -> list[sge.Kwarg]:

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/compile/sqlglot/expressions/comparison_ops.py

@@ -46,7 +46,7 @@ def _(expr: TypedExpr, op: ops.IsInOp) -> sge.Expression:
         if dtypes.can_compare(expr.dtype, dtype):
             if must_upcast_bools and dtype == dtypes.BOOL_DTYPE:
                 value = int(value)
-            values.append(sge.convert(value))
+            values.append(sql.literal(value))
 
     sg_lexpr: sge.Expression = expr.expr
     if expr.dtype == dtypes.BOOL_DTYPE and must_upcast_bools:

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -249,12 +249,13 @@ class SQLGlotIR:
         # TODO: Explicitly insert CTEs into plan
         if len(selections) > 0:
             to_select = [
-                sge.Alias(
-                    this=expr,
+                expr
+                if (isinstance(expr, sge.Alias) and expr.alias == id)
+                or (isinstance(expr, sge.Column) and expr.name == id)
+                else sge.Alias(
+                    this=expr.this if isinstance(expr, sge.Alias) else expr,
                     alias=sql.identifier(id),
                 )
-                if expr.alias_or_name != id
-                else expr
                 for id, expr in selections
             ]
             new_expr = self.expr.select(*to_select)

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/events.py

@@ -20,7 +20,7 @@ import dataclasses
 import datetime
 import threading
 import uuid
-from typing import Any, Callable, Literal, Set
+from typing import Any, Callable, Literal, Optional, Set
 
 import google.cloud.bigquery._job_helpers
 import google.cloud.bigquery.job.query
@@ -127,8 +127,22 @@ class Event:
 
 @dataclasses.dataclass(frozen=True)
 class EventEnvelope:
+    """An envelope that wraps an execution event with metadata and display options.
+
+    Attributes:
+        event:
+            The actual execution event details (e.g., ExecutionStarted, BigQuerySentEvent).
+        progress_bar:
+            Specifies the style of progress bar to display during execution.
+        cell_execution_count:
+            The 1-indexed IPython/Jupyter notebook cell execution number (e.g. the 'x' in 'In [x]').
+            This is NOT a job count, but rather the sequential number of the cell execution in the
+            current notebook session, used to group and filter execution history on a per-cell basis.
+    """
+
     event: Event
     progress_bar: ProgressBarType = _DEFAULT
+    cell_execution_count: Optional[int] = None
 
 
 @dataclasses.dataclass(frozen=True)

{bigframes-2.41.0 → bigframes-2.43.0}/bigframes/core/global_session.py

@@ -19,7 +19,7 @@ from __future__ import annotations
 import threading
 import traceback
 import warnings
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Callable, Iterable, Optional, TypeVar
 
 import google.auth.exceptions
 
@@ -124,12 +124,20 @@ def with_default_session(func_: Callable[..., _T], *args, **kwargs) -> _T:
     return func_(get_global_session(), *args, **kwargs)
 
 
-def execution_history() -> "bigframes.session._ExecutionHistory":
-    import pandas  # noqa: F401
-
+def execution_history(
+    *,
+    events: Optional[Iterable[bigframes.core.events.Event]] = None,
+    job_ids: Optional[Iterable[str]] = None,
+    all_cells: bool = True,
+) -> "bigframes.session._ExecutionHistory":
     import bigframes.session
 
-    return with_default_session(bigframes.session.Session.execution_history)
+    return with_default_session(
+        bigframes.session.Session.execution_history,
+        events=events,
+        job_ids=job_ids,
+        all_cells=all_cells,
+    )
 
 
 class _GlobalSessionContext: