Merge branch 'main' into better_uniform_sample

Author: TrevorBergeron

.gitignore

@@ -64,3 +64,4 @@ tests/js/node_modules/
 pylintrc
 pylintrc.test
 dummy.pkl
+.mypy_cache/

bigframes/bigquery/__init__.py

@@ -18,7 +18,7 @@
 
 import sys
 
-from bigframes.bigquery import ai, ml
+from bigframes.bigquery import ai, ml, obj
 from bigframes.bigquery._operations.approx_agg import approx_top_count
 from bigframes.bigquery._operations.array import (
     array_agg,
@@ -158,4 +158,5 @@
     # Modules / SQL namespaces
     "ai",
     "ml",
+    "obj",
 ]

bigframes/bigquery/_operations/obj.py

@@ -0,0 +1,115 @@
+# Copyright 2025 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.
+
+
+"""This module exposes BigQuery ObjectRef functions.
+
+See bigframes.bigquery.obj for public docs.
+"""
+
+
+from __future__ import annotations
+
+import datetime
+from typing import Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+
+from bigframes.core import convert
+from bigframes.core.logging import log_adapter
+import bigframes.core.utils as utils
+import bigframes.operations as ops
+import bigframes.series as series
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_obj")
+def fetch_metadata(
+    objectref: series.Series,
+) -> series.Series:
+    """[Preview] The OBJ.FETCH_METADATA function returns Cloud Storage metadata for a partially populated ObjectRef value.
+
+    Args:
+        objectref (bigframes.pandas.Series):
+            A partially populated ObjectRef value, in which the uri and authorizer fields are populated and the details field isn't.
+
+    Returns:
+        bigframes.pandas.Series: A fully populated ObjectRef value. The metadata is provided in the details field of the returned ObjectRef value.
+    """
+    objectref = convert.to_bf_series(objectref, default_index=None)
+    return objectref._apply_unary_op(ops.obj_fetch_metadata_op)
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_obj")
+def get_access_url(
+    objectref: series.Series,
+    mode: str,
+    duration: Optional[Union[datetime.timedelta, pd.Timedelta, np.timedelta64]] = None,
+) -> series.Series:
+    """[Preview] The OBJ.GET_ACCESS_URL function returns JSON that contains reference information for the input ObjectRef value, and also access URLs that you can use to read or modify the Cloud Storage object.
+
+    Args:
+        objectref (bigframes.pandas.Series):
+            An ObjectRef value that represents a Cloud Storage object.
+        mode (str):
+            A STRING value that identifies the type of URL that you want to be returned. The following values are supported:
+            'r': Returns a URL that lets you read the object.
+            'rw': Returns two URLs, one that lets you read the object, and one that lets you modify the object.
+        duration (Union[datetime.timedelta, pandas.Timedelta, numpy.timedelta64], optional):
+            An optional INTERVAL value that specifies how long the generated access URLs remain valid. You can specify a value between 30 minutes and 6 hours. For example, you could specify INTERVAL 2 HOUR to generate URLs that expire after 2 hours. The default value is 6 hours.
+
+    Returns:
+        bigframes.pandas.Series: A JSON value that contains the Cloud Storage object reference information from the input ObjectRef value, and also one or more URLs that you can use to access the Cloud Storage object.
+    """
+    objectref = convert.to_bf_series(objectref, default_index=None)
+
+    duration_micros = None
+    if duration is not None:
+        duration_micros = utils.timedelta_to_micros(duration)
+
+    return objectref._apply_unary_op(
+        ops.ObjGetAccessUrl(mode=mode, duration=duration_micros)
+    )
+
+
+@log_adapter.method_logger(custom_base_name="bigquery_obj")
+def make_ref(
+    uri_or_json: Union[series.Series, Sequence[str]],
+    authorizer: Union[series.Series, str, None] = None,
+) -> series.Series:
+    """[Preview] Use the OBJ.MAKE_REF function to create an ObjectRef value that contains reference information for a Cloud Storage object.
+
+    Args:
+        uri_or_json (bigframes.pandas.Series or str):
+            A series of STRING values that contains the URI for the Cloud Storage object, for example, gs://mybucket/flowers/12345.jpg.
+            OR
+            A series of JSON value that represents a Cloud Storage object.
+        authorizer (bigframes.pandas.Series or str, optional):
+            A STRING value that contains the Cloud Resource connection used to access the Cloud Storage object.
+            Required if ``uri_or_json`` is a URI string.
+
+    Returns:
+        bigframes.pandas.Series: An ObjectRef value.
+    """
+    uri_or_json = convert.to_bf_series(uri_or_json, default_index=None)
+
+    if authorizer is not None:
+        # Avoid join problems encountered if we try to convert a literal into Series.
+        if not isinstance(authorizer, str):
+            authorizer = convert.to_bf_series(authorizer, default_index=None)
+
+        return uri_or_json._apply_binary_op(authorizer, ops.obj_make_ref_op)
+
+    # If authorizer is not provided, we assume uri_or_json is a JSON objectref
+    return uri_or_json._apply_unary_op(ops.obj_make_ref_json_op)

bigframes/bigquery/obj.py

@@ -0,0 +1,41 @@
+# Copyright 2025 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.
+
+"""This module integrates BigQuery built-in 'ObjectRef' functions for use with Series/DataFrame objects,
+such as OBJ.FETCH_METADATA:
+https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/objectref_functions
+
+
+.. warning::
+
+    This product or feature is subject to the "Pre-GA Offerings Terms" in the
+    General Service Terms section of the `Service Specific Terms
+    <https://cloud.google.com/terms/service-terms>`_. Pre-GA products and
+    features are available "as is" and might have limited support. For more
+    information, see the `launch stage descriptions
+    <https://cloud.google.com/products?hl=en#product-launch-stages>`_.
+
+.. note::
+
+    To provide feedback or request support for this feature, send an email to
+    bq-objectref-feedback@google.com.
+"""
+
+from bigframes.bigquery._operations.obj import fetch_metadata, get_access_url, make_ref
+
+__all__ = [
+    "fetch_metadata",
+    "get_access_url",
+    "make_ref",
+]

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -16,6 +16,7 @@
 
 import functools
 import typing
+from typing import cast
 
 from bigframes_vendored import ibis
 import bigframes_vendored.ibis.expr.api as ibis_api
@@ -1247,6 +1248,13 @@ def obj_fetch_metadata_op_impl(obj_ref: ibis_types.Value):
 
 @scalar_op_compiler.register_unary_op(ops.ObjGetAccessUrl, pass_op=True)
 def obj_get_access_url_op_impl(obj_ref: ibis_types.Value, op: ops.ObjGetAccessUrl):
+    if op.duration is not None:
+        duration_value = cast(
+            ibis_types.IntegerValue, ibis_types.literal(op.duration)
+        ).to_interval("us")
+        return obj_get_access_url_with_duration(
+            obj_ref=obj_ref, mode=op.mode, duration=duration_value
+        )
     return obj_get_access_url(obj_ref=obj_ref, mode=op.mode)
 
 
@@ -1807,6 +1815,11 @@ def obj_make_ref_op(x: ibis_types.Value, y: ibis_types.Value):
     return obj_make_ref(uri=x, authorizer=y)
 
 
+@scalar_op_compiler.register_unary_op(ops.obj_make_ref_json_op)
+def obj_make_ref_json_op(x: ibis_types.Value):
+    return obj_make_ref_json(objectref_json=x)
+
+
 # Ternary Operations
 @scalar_op_compiler.register_ternary_op(ops.where_op)
 def where_op(
@@ -2141,11 +2154,21 @@ def obj_make_ref(uri: str, authorizer: str) -> _OBJ_REF_IBIS_DTYPE:  # type: ign
     """Make ObjectRef Struct from uri and connection."""
 
 
+@ibis_udf.scalar.builtin(name="OBJ.MAKE_REF")
+def obj_make_ref_json(objectref_json: ibis_dtypes.JSON) -> _OBJ_REF_IBIS_DTYPE:  # type: ignore
+    """Make ObjectRef Struct from json."""
+
+
 @ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
 def obj_get_access_url(obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String) -> ibis_dtypes.JSON:  # type: ignore
     """Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
 
 
+@ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
+def obj_get_access_url_with_duration(obj_ref, mode, duration) -> ibis_dtypes.JSON:  # type: ignore
+    """Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
+
+
 @ibis_udf.scalar.builtin(name="ltrim")
 def str_lstrip_op(  # type: ignore[empty-body]
     x: ibis_dtypes.String, to_strip: ibis_dtypes.String

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -23,6 +23,7 @@
 from bigframes.core import window_spec
 import bigframes.core.compile.sqlglot.aggregations.op_registration as reg
 from bigframes.core.compile.sqlglot.aggregations.windows import apply_window_if_present
+from bigframes.core.compile.sqlglot.expressions import constants
 import bigframes.core.compile.sqlglot.expressions.typed_expr as typed_expr
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
 from bigframes.operations import aggregations as agg_ops
@@ -44,9 +45,13 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    # BQ will return null for empty column, result would be false in pandas.
-    result = apply_window_if_present(sge.func("LOGICAL_AND", column.expr), window)
-    return sge.func("IFNULL", result, sge.true())
+    expr = column.expr
+    if column.dtype != dtypes.BOOL_DTYPE:
+        expr = sge.NEQ(this=expr, expression=sge.convert(0))
+    expr = apply_window_if_present(sge.func("LOGICAL_AND", expr), window)
+
+    # BQ will return null for empty column, result would be true in pandas.
+    return sge.func("COALESCE", expr, sge.convert(True))
 
 
 @UNARY_OP_REGISTRATION.register(agg_ops.AnyOp)
@@ -56,6 +61,8 @@ def _(
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
     expr = column.expr
+    if column.dtype != dtypes.BOOL_DTYPE:
+        expr = sge.NEQ(this=expr, expression=sge.convert(0))
     expr = apply_window_if_present(sge.func("LOGICAL_OR", expr), window)
 
     # BQ will return null for empty column, result would be false in pandas.
@@ -326,6 +333,15 @@ def _(
             unit=sge.Identifier(this="MICROSECOND"),
         )
 
+    if column.dtype == dtypes.DATE_DTYPE:
+        date_diff = sge.DateDiff(
+            this=column.expr, expression=shifted, unit=sge.Identifier(this="DAY")
+        )
+        return sge.Cast(
+            this=sge.Floor(this=date_diff * constants._DAY_TO_MICROSECONDS),
+            to="INT64",
+        )
+
     raise TypeError(f"Cannot perform diff on type {column.dtype}")
 
 
@@ -410,24 +426,28 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
     # Need to short-circuit as log with zeroes is illegal sql
-    is_zero = sge.EQ(this=column.expr, expression=sge.convert(0))
+    is_zero = sge.EQ(this=expr, expression=sge.convert(0))
 
     # There is no product sql aggregate function, so must implement as a sum of logs, and then
     # apply power after. Note, log and power base must be equal! This impl uses natural log.
-    logs = (
-        sge.Case()
-        .when(is_zero, sge.convert(0))
-        .else_(sge.func("LN", sge.func("ABS", column.expr)))
+    logs = sge.If(
+        this=is_zero,
+        true=sge.convert(0),
+        false=sge.func("LOG", sge.convert(2), sge.func("ABS", expr)),
     )
     logs_sum = apply_window_if_present(sge.func("SUM", logs), window)
-    magnitude = sge.func("EXP", logs_sum)
+    magnitude = sge.func("POWER", sge.convert(2), logs_sum)
 
     # Can't determine sign from logs, so have to determine parity of count of negative inputs
     is_negative = (
         sge.Case()
         .when(
-            sge.LT(this=sge.func("SIGN", column.expr), expression=sge.convert(0)),
+            sge.EQ(this=sge.func("SIGN", expr), expression=sge.convert(-1)),
             sge.convert(1),
         )
         .else_(sge.convert(0))
@@ -445,11 +465,7 @@ def _(
         .else_(
             sge.Mul(
                 this=magnitude,
-                expression=sge.If(
-                    this=sge.EQ(this=negative_count_parity, expression=sge.convert(1)),
-                    true=sge.convert(-1),
-                    false=sge.convert(1),
-                ),
+                expression=sge.func("POWER", sge.convert(-1), negative_count_parity),
             )
         )
     )
@@ -499,14 +515,18 @@ def _(
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
 ) -> sge.Expression:
-    # TODO: Support interpolation argument
-    # TODO: Support percentile_disc
-    result: sge.Expression = sge.func("PERCENTILE_CONT", column.expr, sge.convert(op.q))
+    expr = column.expr
+    if column.dtype == dtypes.BOOL_DTYPE:
+        expr = sge.Cast(this=expr, to="INT64")
+
+    result: sge.Expression = sge.func("PERCENTILE_CONT", expr, sge.convert(op.q))
     if window is None:
-        # PERCENTILE_CONT is a navigation function, not an aggregate function, so it always needs an OVER clause.
+        # PERCENTILE_CONT is a navigation function, not an aggregate function,
+        # so it always needs an OVER clause.
         result = sge.Window(this=result)
     else:
         result = apply_window_if_present(result, window)
+
     if op.should_floor_result:
         result = sge.Cast(this=sge.func("FLOOR", result), to="INT64")
     return result

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

@@ -31,9 +31,22 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.ObjGetAccessUrl, pass_op=True)
 def _(expr: TypedExpr, op: ops.ObjGetAccessUrl) -> sge.Expression:
-    return sge.func("OBJ.GET_ACCESS_URL", expr.expr, sge.convert(op.mode))
+    args = [expr.expr, sge.Literal.string(op.mode)]
+    if op.duration is not None:
+        args.append(
+            sge.Interval(
+                this=sge.Literal.number(op.duration),
+                unit=sge.Var(this="MICROSECOND"),
+            )
+        )
+    return sge.func("OBJ.GET_ACCESS_URL", *args)
 
 
 @register_binary_op(ops.obj_make_ref_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.func("OBJ.MAKE_REF", left.expr, right.expr)
+
+
+@register_unary_op(ops.obj_make_ref_json_op)
+def _(expr: TypedExpr) -> sge.Expression:
+    return sge.func("OBJ.MAKE_REF", expr.expr)