bigframes 2.51.0__tar.gz → 2.52.0__tar.gz

{bigframes-2.51.0/bigframes.egg-info → bigframes-2.52.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bigframes
-Version: 2.51.0
+Version: 2.52.0
 Summary: BigQuery DataFrames -- scalable analytics and machine learning with BigQuery
 Home-page: https://dataframes.bigquery.dev
 Download-URL: https://github.com/google/bigframes/releases
@@ -50,7 +50,6 @@ Requires-Dist: tabulate>=0.9
 Requires-Dist: humanize>=4.6.0
 Requires-Dist: matplotlib>=3.7.1
 Requires-Dist: db-dtypes>=1.4.2
-Requires-Dist: pyiceberg>=0.7.1
 Requires-Dist: atpublic<6,>=2.3
 Requires-Dist: python-dateutil<3,>=2.8.2
 Requires-Dist: pytz>=2022.7

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/bigquery/__init__.py

@@ -134,7 +134,9 @@ from bigframes.operations.googlesql.global_namespace.date import (
     date_from_unix_date,
     date_sub,
     date_trunc,
+    extract,
     format_date,
+    generate_date_array,
     last_day,
     parse_date,
     unix_date,
@@ -177,7 +179,9 @@ _functions = [
     date_from_unix_date,
     date_sub,
     date_trunc,
+    extract,
     format_date,
+    generate_date_array,
     last_day,
     parse_date,
     unix_date,
@@ -273,7 +277,9 @@ __all__ = [
     "date_from_unix_date",
     "date_sub",
     "date_trunc",
+    "extract",
     "format_date",
+    "generate_date_array",
     "last_day",
     "parse_date",
     "unix_date",

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/array_value.py

@@ -17,7 +17,7 @@ import datetime
 import functools
 import typing
 from dataclasses import dataclass
-from typing import Iterable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple
 
 import pandas
 import pyarrow as pa
@@ -90,7 +90,7 @@ class ArrayValue:
     @classmethod
     def from_table(
         cls,
-        table: Union[bq_data.BiglakeIcebergTable, bq_data.GbqNativeTable],
+        table: bq_data.GbqNativeTable,
         session: Session,
         *,
         columns: Optional[Sequence[str]] = None,

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/bq_data.py

@@ -23,7 +23,7 @@ import os
 import queue
 import threading
 import typing
-from typing import Any, Iterator, List, Literal, Optional, Sequence, Tuple, Union
+from typing import Any, Iterator, List, Literal, Optional, Sequence, Tuple
 
 import google.cloud.bigquery as bq
 import google.cloud.bigquery_storage_v1.types as bq_storage_types
@@ -31,7 +31,6 @@ import pyarrow as pa
 from google.cloud import bigquery_storage_v1
 from google.protobuf import timestamp_pb2
 
-import bigframes.constants
 import bigframes.core.schema
 from bigframes.core import pyarrow_utils
 
@@ -39,44 +38,9 @@ if typing.TYPE_CHECKING:
     import bigframes.core.ordering as orderings
 
 
-def _resolve_standard_gcp_region(bq_region: str):
-    """
-    Resolve bq regions to standardized
-    """
-    if bq_region.casefold() == "US":
-        return "us-central1"
-    elif bq_region.casefold() == "EU":
-        return "europe-west4"
-    return bq_region
-
-
-def is_irc_table(table_id: str):
-    """
-    Determines if a table id should be resolved through the iceberg rest catalog.
-    """
-    return len(table_id.split(".")) == 4
-
-
-def is_compatible(
-    data_region: Union[GcsRegion, BigQueryRegion], session_location: str
-) -> bool:
+def is_compatible(data_region: BigQueryRegion, session_location: str) -> bool:
     # based on https://docs.cloud.google.com/bigquery/docs/locations#storage-location-considerations
-    if isinstance(data_region, BigQueryRegion):
-        return data_region.name == session_location
-    else:
-        assert isinstance(data_region, GcsRegion)
-        # TODO(b/463675088): Multi-regions don't yet support rest catalog tables
-        if session_location in bigframes.constants.BIGQUERY_MULTIREGIONS:
-            return False
-        return _resolve_standard_gcp_region(session_location) in data_region.included
-
-
-def get_default_bq_region(data_region: Union[GcsRegion, BigQueryRegion]) -> str:
-    if isinstance(data_region, BigQueryRegion):
-        return data_region.name
-    elif isinstance(data_region, GcsRegion):
-        # should maybe try to track and prefer primary replica?
-        return data_region.included[0]
+    return data_region.name == session_location
 
 
 @dataclasses.dataclass(frozen=True)
@@ -84,19 +48,11 @@ class BigQueryRegion:
     name: str
 
 
-@dataclasses.dataclass(frozen=True)
-class GcsRegion:
-    # this is the name of gcs regions, which may be names for multi-regions, so shouldn't be compared with non-gcs locations
-    storage_regions: tuple[str, ...]
-    # this tracks all the included standard, specific regions (eg us-east1), and should be comparable to bq regions (except non-standard US, EU, omni regions)
-    included: tuple[str, ...]
-
-
 # what is the line between metadata and core fields? Mostly metadata fields are optional or unreliable, but its fuzzy
 @dataclasses.dataclass(frozen=True)
 class TableMetadata:
     # this size metadata might be stale, don't use where strict correctness is needed
-    location: Union[BigQueryRegion, GcsRegion]
+    location: BigQueryRegion
     type: Literal["TABLE", "EXTERNAL", "VIEW", "MATERIALIZE_VIEW", "SNAPSHOT"]
     numBytes: Optional[int] = None
     numRows: Optional[int] = None
@@ -214,50 +170,6 @@ class GbqNativeTable:
         return {col.name: col for col in self.physical_schema}
 
 
-@dataclasses.dataclass(frozen=True)
-class BiglakeIcebergTable:
-    project_id: str = dataclasses.field()
-    catalog_id: str = dataclasses.field()
-    namespace_id: str = dataclasses.field()
-    table_id: str = dataclasses.field()
-    physical_schema: Tuple[bq.SchemaField, ...] = dataclasses.field()
-    cluster_cols: typing.Optional[Tuple[str, ...]]
-    metadata: TableMetadata
-
-    def get_full_id(self, quoted: bool = False) -> str:
-        if quoted:
-            return f"`{self.project_id}`.`{self.catalog_id}`.`{self.namespace_id}`.`{self.table_id}`"
-        return (
-            f"{self.project_id}.{self.catalog_id}.{self.namespace_id}.{self.table_id}"
-        )
-
-    @property
-    @functools.cache
-    def schema_by_id(self):
-        return {col.name: col for col in self.physical_schema}
-
-    @property
-    def partition_col(self) -> Optional[str]:
-        # TODO: Use iceberg partition metadata
-        return None
-
-    @property
-    def dataset_id(self) -> str:
-        """
-        Not a true dataset, but serves as the dataset component of the identifer in sql queries
-        """
-        return f"{self.catalog_id}.{self.namespace_id}"
-
-    @property
-    def primary_key(self) -> Optional[Tuple[str, ...]]:
-        return None
-
-    def get_table_ref(self) -> bq.TableReference:
-        return bq.TableReference(
-            bq.DatasetReference(self.project_id, self.dataset_id), self.table_id
-        )
-
-
 @dataclasses.dataclass(frozen=True)
 class BigqueryDataSource:
     """
@@ -272,7 +184,7 @@ class BigqueryDataSource:
             self.schema.names
         )
 
-    table: Union[GbqNativeTable, BiglakeIcebergTable]
+    table: GbqNativeTable
     schema: bigframes.core.schema.ArraySchema
     at_time: typing.Optional[datetime.datetime] = None
     # Added for backwards compatibility, not validated

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -189,7 +189,7 @@ def _cut_ops_w_int_bins(
         condition: sge.Expression
         if this_bin == bins - 1:
             condition = sge.Is(
-                this=sge.paren(column.expr, copy=False),
+                this=column.expr,
                 expression=sg.not_(sge.Null(), copy=False),
             )
         else:

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/compiler.py

@@ -377,4 +377,5 @@ def _replace_unsupported_ops(node: nodes.BigFrameNode):
     node = nodes.bottom_up(node, rewrite.rewrite_slice)
     node = nodes.bottom_up(node, rewrite.rewrite_range_rolling)
     node = nodes.bottom_up(node, rewrite.lower_udfs)
+    node = nodes.bottom_up(node, rewrite.lower_googlesql)
     return node

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expression_compiler.py

@@ -21,16 +21,23 @@ import bigframes_vendored.sqlglot.expressions as sge
 import bigframes.core.agg_expressions as agg_exprs
 import bigframes.core.expression as ex
 import bigframes.operations as ops
-from bigframes.core.compile.sqlglot import sql
-from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
+
+if typing.TYPE_CHECKING:
+    from bigframes.core.compile.sqlglot.expressions import typed_expr
+
+
+expression_compiler: ExpressionCompiler
 
 
 class ExpressionCompiler:
-    # Mapping of operation name to implemenations
-    _registry: dict[
-        str,
-        typing.Callable[[typing.Sequence[TypedExpr], ops.RowOp], sge.Expression],
-    ] = {}
+    def __init__(self) -> None:
+        # Mapping of operation name to implementations
+        self._registry: dict[
+            str,
+            typing.Callable[
+                [typing.Sequence[typed_expr.TypedExpr], ops.RowOp], sge.Expression
+            ],
+        ] = {}
 
     # A set of SQLGlot classes that may need to be parenthesized
     SQLGLOT_NEEDS_PARENS = {
@@ -41,6 +48,7 @@ class ExpressionCompiler:
         sge.Div,
         sge.Mod,
         sge.Pow,
+        sge.Neg,
         # Comparison operations
         sge.GTE,
         sge.GT,
@@ -56,6 +64,7 @@ class ExpressionCompiler:
         sge.And,
         sge.Or,
         sge.Xor,
+        sge.Not,
         # Bitwise operations
         sge.BitwiseAnd,
         sge.BitwiseOr,
@@ -71,20 +80,30 @@ class ExpressionCompiler:
     def compile_expression(
         self,
         expression: ex.Expression,
+        *,
+        dtype: typing.Any = None,
     ) -> sge.Expression:
         """Compiles BigFrames scalar expression into SQLGlot expression."""
         raise NotImplementedError(f"Unrecognized expression: {expression}")
 
     @compile_expression.register
-    def _(self, expr: ex.DerefOp) -> sge.Expression:
+    def _(self, expr: ex.DerefOp, *, dtype: typing.Any = None) -> sge.Expression:
         return sge.Column(this=sge.to_identifier(expr.id.sql, quoted=True))
 
     @compile_expression.register
-    def _(self, expr: ex.ScalarConstantExpression) -> sge.Expression:
-        return sql.literal(expr.value, expr.dtype)
+    def _(
+        self, expr: ex.ScalarConstantExpression, *, dtype: typing.Any = None
+    ) -> sge.Expression:
+        # Avoid circular import since expression_compiler is imported as part of
+        # the bigframes.core.compile.sqlglot import.
+        from bigframes.core.compile.sqlglot import sql
+
+        return sql.literal(expr.value, dtype=dtype or expr.dtype)
 
     @compile_expression.register
-    def _(self, expr: agg_exprs.WindowExpression) -> sge.Expression:
+    def _(
+        self, expr: agg_exprs.WindowExpression, *, dtype: typing.Any = None
+    ) -> sge.Expression:
         import bigframes.core.compile.sqlglot.aggregate_compiler as agg_compile
 
         return agg_compile.compile_analytic(
@@ -93,17 +112,34 @@ class ExpressionCompiler:
         )
 
     @compile_expression.register
-    def _(self, expr: ex.OpExpression) -> sge.Expression:
-        inputs = tuple(
-            TypedExpr(self.compile_expression(sub_expr), sub_expr.output_type)
-            if not isinstance(sub_expr, ex.OmittedArg)
-            else TypedExpr(sge.Null(), None, is_omitted=True)
-            for sub_expr in expr.inputs
-        )
-        return self.compile_row_op(expr.op, inputs)
+    def _(self, expr: ex.OpExpression, *, dtype: typing.Any = None) -> sge.Expression:
+        # Avoid circular imports since this module is also a part of
+        # bigframes.core.compile.sqlglot.
+        from bigframes.core.compile.sqlglot.expressions import typed_expr
+
+        arg_specs = getattr(expr.op, "args", None)
+        num_args = len(arg_specs) if arg_specs is not None else 0
+
+        inputs = []
+        for i, sub_expr in enumerate(expr.inputs):
+            if isinstance(sub_expr, ex.OmittedArg):
+                inputs.append(typed_expr.TypedExpr(sge.Null(), None, is_omitted=True))
+                continue
+
+            sub_dtype = None
+            if arg_specs is not None and num_args > 0:
+                if i < num_args:
+                    sub_dtype = getattr(arg_specs[i], "dtype", None)
+                elif arg_specs[-1].is_vararg:
+                    sub_dtype = getattr(arg_specs[-1], "dtype", None)
+
+            compiled_expr = self.compile_expression(sub_expr, dtype=sub_dtype)
+            inputs.append(typed_expr.TypedExpr(compiled_expr, sub_expr.output_type))
+
+        return self.compile_row_op(expr.op, tuple(inputs))
 
     def compile_row_op(
-        self, op: ops.RowOp, inputs: typing.Sequence[TypedExpr]
+        self, op: ops.RowOp, inputs: typing.Sequence[typed_expr.TypedExpr]
     ) -> sge.Expression:
         impl = self._registry[op.name]
         return impl(inputs, op)
@@ -126,11 +162,14 @@ class ExpressionCompiler:
         key = typing.cast(str, op_ref.name)
 
         def decorator(impl: typing.Callable[..., sge.Expression]):
-            def normalized_impl(args: typing.Sequence[TypedExpr], op: ops.RowOp):
+            def normalized_impl(
+                args: typing.Sequence[typed_expr.TypedExpr], op: ops.RowOp
+            ):
+                arg = self._add_parentheses(args[0])
                 if pass_op:
-                    return impl(args[0], op)
+                    return impl(arg, op)
                 else:
-                    return impl(args[0])
+                    return impl(arg)
 
             self._register(key, normalized_impl)
             return impl
@@ -155,7 +194,9 @@ class ExpressionCompiler:
         key = typing.cast(str, op_ref.name)
 
         def decorator(impl: typing.Callable[..., sge.Expression]):
-            def normalized_impl(args: typing.Sequence[TypedExpr], op: ops.RowOp):
+            def normalized_impl(
+                args: typing.Sequence[typed_expr.TypedExpr], op: ops.RowOp
+            ):
                 left = self._add_parentheses(args[0])
                 right = self._add_parentheses(args[1])
                 if pass_op:
@@ -181,8 +222,13 @@ class ExpressionCompiler:
         key = typing.cast(str, op_ref.name)
 
         def decorator(impl: typing.Callable[..., sge.Expression]):
-            def normalized_impl(args: typing.Sequence[TypedExpr], op: ops.RowOp):
-                return impl(args[0], args[1], args[2])
+            def normalized_impl(
+                args: typing.Sequence[typed_expr.TypedExpr], op: ops.RowOp
+            ):
+                arg0 = self._add_parentheses(args[0])
+                arg1 = self._add_parentheses(args[1])
+                arg2 = self._add_parentheses(args[2])
+                return impl(arg0, arg1, arg2)
 
             self._register(key, normalized_impl)
             return impl
@@ -205,7 +251,10 @@ class ExpressionCompiler:
         key = typing.cast(str, op_ref.name)
 
         def decorator(impl: typing.Callable[..., sge.Expression]):
-            def normalized_impl(args: typing.Sequence[TypedExpr], op: ops.RowOp):
+            def normalized_impl(
+                args: typing.Sequence[typed_expr.TypedExpr], op: ops.RowOp
+            ):
+                args = [self._add_parentheses(arg) for arg in args]
                 if pass_op:
                     return impl(*args, op=op)
                 else:
@@ -219,16 +268,22 @@ class ExpressionCompiler:
     def _register(
         self,
         op_name: str,
-        impl: typing.Callable[[typing.Sequence[TypedExpr], ops.RowOp], sge.Expression],
+        impl: typing.Callable[
+            [typing.Sequence[typed_expr.TypedExpr], ops.RowOp], sge.Expression
+        ],
     ):
         if op_name in self._registry:
             raise ValueError(f"Operation name {op_name} already registered")
         self._registry[op_name] = impl
 
     @classmethod
-    def _add_parentheses(cls, expr: TypedExpr) -> TypedExpr:
+    def _add_parentheses(cls, expr: typed_expr.TypedExpr) -> typed_expr.TypedExpr:
+        # Avoid circular imports since this module is also a part of
+        # bigframes.core.compile.sqlglot.
+        from bigframes.core.compile.sqlglot.expressions import typed_expr
+
         if type(expr.expr) in cls.SQLGLOT_NEEDS_PARENS:
-            return TypedExpr(sge.paren(expr.expr, copy=False), expr.dtype)
+            return typed_expr.TypedExpr(sge.paren(expr.expr, copy=False), expr.dtype)
         return expr
 
 

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expressions/common.py

@@ -54,6 +54,6 @@ def _to_nullable_bool(expr: typed_expr.TypedExpr) -> sge.Expression:
         return sge.GT(this=sge.func("ARRAY_LENGTH", sg_expr), expression=sge.convert(0))
 
     return sge.Is(
-        this=sge.paren(sg_expr, copy=False),
+        this=sg_expr,
         expression=sg.not_(sge.Null(), copy=False),
     )

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expressions/comparison_ops.py

@@ -154,12 +154,12 @@ def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     if sql.is_null_literal(left.expr):
         return sge.Is(
-            this=sge.paren(right.expr, copy=False),
+            this=right.expr,
             expression=sg.not_(sge.Null(), copy=False),
         )
     if sql.is_null_literal(right.expr):
         return sge.Is(
-            this=sge.paren(left.expr, copy=False),
+            this=left.expr,
             expression=sg.not_(sge.Null(), copy=False),
         )
 

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expressions/datetime_ops.py

@@ -25,6 +25,23 @@ from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
 register_unary_op = expression_compiler.expression_compiler.register_unary_op
 register_binary_op = expression_compiler.expression_compiler.register_binary_op
+register_nary_op = expression_compiler.expression_compiler.register_nary_op
+
+
+@register_nary_op(ops.ExtractOp, pass_op=True)
+def _(*operands: TypedExpr, op: ops.ExtractOp) -> sge.Expression:
+    part: sge.Expression
+    if op.part is not None:
+        part = sge.var(op.part.upper())
+        date_expr = operands[0].expr
+        if len(operands) > 1 and not operands[1].is_omitted:
+            date_expr = sge.AtTimeZone(this=date_expr, zone=operands[1].expr)
+    else:
+        date_expr = operands[0].expr
+        part = operands[1].expr
+        if len(operands) > 2 and not operands[2].is_omitted:
+            date_expr = sge.AtTimeZone(this=date_expr, zone=operands[2].expr)
+    return sge.Extract(this=part, expression=date_expr)
 
 
 @register_binary_op(ops.DatetimeToIntegerLabelOp, pass_op=True)

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expressions/generic_ops.py

@@ -73,8 +73,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.invert_op)
 def _(expr: TypedExpr) -> sge.Expression:
     if expr.dtype == dtypes.BOOL_DTYPE:
-        return sge.Not(this=sge.paren(expr.expr))
-    return sge.BitwiseNot(this=sge.paren(expr.expr))
+        return sge.Not(this=expr.expr)
+    return sge.BitwiseNot(this=expr.expr)
 
 
 @register_nary_op(ops.GoogleSqlScalarOp, pass_op=True)
@@ -95,7 +95,7 @@ def _(*operands: TypedExpr, op: ops.GoogleSqlScalarOp) -> sge.Expression:
             args.append(sge.Kwarg(this=arg_spec.arg_name, expression=operand.expr))
         else:
             args.append(operand.expr)
-    return sg.func(op.sql_name, *args)
+    return sge.Anonymous(this=op.sql_name, expressions=args)
 
 
 @register_nary_op(ops.SqlScalarOp, pass_op=True)
@@ -110,7 +110,7 @@ def _(*operands: TypedExpr, op: ops.SqlScalarOp) -> sge.Expression:
 
 @register_unary_op(ops.isnull_op)
 def _(expr: TypedExpr) -> sge.Expression:
-    return sge.Is(this=sge.paren(expr.expr), expression=sge.Null())
+    return sge.Is(this=expr.expr, expression=sge.Null())
 
 
 @register_unary_op(ops.MapOp, pass_op=True)
@@ -144,7 +144,7 @@ def _(expr: TypedExpr, op: ops.MapOp) -> sge.Expression:
 @register_unary_op(ops.notnull_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.Is(
-        this=sge.paren(expr.expr, copy=False),
+        this=expr.expr,
         expression=sg.not_(sge.Null(), copy=False),
     )
 

{bigframes-2.51.0 → bigframes-2.52.0}/bigframes/core/compile/sqlglot/expressions/numeric_ops.py

@@ -240,7 +240,7 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.neg_op)
 def _(expr: TypedExpr) -> sge.Expression:
-    return sge.Neg(this=sge.paren(expr.expr))
+    return sge.Neg(this=expr.expr)
 
 
 @register_unary_op(ops.pos_op)