feat: Add BigQuery ObjectRef functions to bigframes.bigquery.obj

This change introduces support for BigQuery ObjectRef functions:
- `OBJ.FETCH_METADATA`
- `OBJ.GET_ACCESS_URL`
- `OBJ.MAKE_REF`

These are exposed via a new `bigframes.bigquery.obj` module.

Changes:
- Added `ObjMakeRefJson` operation in `bigframes/operations/blob_ops.py`.
- Updated `ObjGetAccessUrl` operation to support optional duration.
- Updated `bigframes/operations/__init__.py` to export new operations.
- Updated `bigframes/core/compile/ibis_compiler/scalar_op_registry.py` and `bigframes/core/compile/sqlglot/expressions/blob_ops.py` to support new operations.
- Created `bigframes/bigquery/_operations/obj.py` with the implementation of `fetch_metadata`, `get_access_url`, and `make_ref`.
- Created `bigframes/bigquery/obj.py` to expose the functions.
- Exposed `obj` module in `bigframes/bigquery/__init__.py`.
- Added unit tests in `tests/unit/bigquery/test_obj.py`.

Author: google-labs-jules[bot]

bigframes/bigquery/_operations/obj.py

@@ -21,8 +21,12 @@
 
 from __future__ import annotations
 
+import datetime
 from typing import Optional, Union
 
+import numpy as np
+import pandas as pd
+
 import bigframes.core.utils as utils
 import bigframes.operations as ops
 import bigframes.series as series
@@ -48,7 +52,9 @@ def fetch_metadata(
 def get_access_url(
     objectref: series.Series,
     mode: str,
-    duration: Optional[series.Series] = None,
+    duration: Optional[
+        Union[datetime.timedelta, pd.Timedelta, np.timedelta64]
+    ] = None,
 ) -> series.Series:
     """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.
 
@@ -59,17 +65,19 @@ def get_access_url(
             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 (bigframes.series.Series, optional):
+        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.series.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.
     """
+    duration_micros = None
     if duration is not None:
-        return objectref._apply_binary_op(
-            duration, ops.ObjGetAccessUrlWithDuration(mode=mode)
-        )
-    return objectref._apply_unary_op(ops.ObjGetAccessUrl(mode=mode))
+        duration_micros = utils.timedelta_to_micros(duration)
+
+    return objectref._apply_unary_op(
+        ops.ObjGetAccessUrl(mode=mode, duration=duration_micros)
+    )
 
 
 @utils.preview(name="The ObjectRef API `make_ref`")

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1247,18 +1247,14 @@ 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 = 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)
 
 
-@scalar_op_compiler.register_binary_op(ops.ObjGetAccessUrlWithDuration, pass_op=True)
-def obj_get_access_url_with_duration_op_impl(
-    obj_ref: ibis_types.Value, duration: ibis_types.Value, op: ops.ObjGetAccessUrlWithDuration
-):
-    return obj_get_access_url_with_duration(
-        obj_ref=obj_ref, mode=op.mode, duration=duration
-    )
-
-
 ### Binary Ops
 def short_circuit_nulls(type_override: typing.Optional[ibis_dtypes.DataType] = None):
     """Wraps a binary operator to generate nulls of the expected type if either input is a null scalar."""

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

@@ -31,16 +31,15 @@ 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.Literal.string(op.mode))
-
-
-@register_binary_op(ops.ObjGetAccessUrlWithDuration, pass_op=True)
-def _(
-    left: TypedExpr, right: TypedExpr, op: ops.ObjGetAccessUrlWithDuration
-) -> sge.Expression:
-    return sge.func(
-        "OBJ.GET_ACCESS_URL", left.expr, sge.Literal.string(op.mode), right.expr
-    )
+    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)

bigframes/operations/__init__.py

@@ -43,7 +43,6 @@
     obj_make_ref_json_op,
     obj_make_ref_op,
     ObjGetAccessUrl,
-    ObjGetAccessUrlWithDuration,
 )
 from bigframes.operations.bool_ops import and_op, or_op, xor_op
 from bigframes.operations.comparison_ops import (
@@ -367,7 +366,6 @@
     "ArrayToStringOp",
     # Blob ops
     "ObjGetAccessUrl",
-    "ObjGetAccessUrlWithDuration",
     "obj_make_ref_json_op",
     "obj_make_ref_op",
     "obj_fetch_metadata_op",

bigframes/operations/blob_ops.py

@@ -29,15 +29,7 @@
 class ObjGetAccessUrl(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "obj_get_access_url"
     mode: str  # access mode, e.g. R read, W write, RW read & write
-
-    def output_type(self, *input_types):
-        return dtypes.JSON_DTYPE
-
-
-@dataclasses.dataclass(frozen=True)
-class ObjGetAccessUrlWithDuration(base_ops.BinaryOp):
-    name: typing.ClassVar[str] = "obj_get_access_url_with_duration"
-    mode: str  # access mode, e.g. R read, W write, RW read & write
+    duration: typing.Optional[int] = None  # duration in microseconds
 
     def output_type(self, *input_types):
         return dtypes.JSON_DTYPE

tests/unit/bigquery/test_obj.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import pytest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import bigframes.bigquery.obj as obj
 import bigframes.operations as ops
@@ -27,11 +28,13 @@ def test_get_access_url_op_structure():
     op = ops.ObjGetAccessUrl(mode="r")
     assert op.name == "obj_get_access_url"
     assert op.mode == "r"
+    assert op.duration is None
 
 def test_get_access_url_with_duration_op_structure():
-    op = ops.ObjGetAccessUrlWithDuration(mode="rw")
-    assert op.name == "obj_get_access_url_with_duration"
+    op = ops.ObjGetAccessUrl(mode="rw", duration=3600000000)
+    assert op.name == "obj_get_access_url"
     assert op.mode == "rw"
+    assert op.duration == 3600000000
 
 def test_make_ref_op_structure():
     op = ops.obj_make_ref_op
@@ -59,18 +62,20 @@ def test_get_access_url_calls_apply_unary_op_without_duration():
     args, _ = s._apply_unary_op.call_args
     assert isinstance(args[0], ops.ObjGetAccessUrl)
     assert args[0].mode == "r"
+    assert args[0].duration is None
 
-def test_get_access_url_calls_apply_binary_op_with_duration():
+def test_get_access_url_calls_apply_unary_op_with_duration():
     s = MagicMock(spec=series.Series)
-    duration = MagicMock(spec=series.Series)
+    duration = datetime.timedelta(hours=1)
 
     obj.get_access_url(s, mode="rw", duration=duration)
 
-    s._apply_binary_op.assert_called_once()
-    args, kwargs = s._apply_binary_op.call_args
-    assert args[0] == duration
-    assert isinstance(args[1], ops.ObjGetAccessUrlWithDuration)
-    assert args[1].mode == "rw"
+    s._apply_unary_op.assert_called_once()
+    args, kwargs = s._apply_unary_op.call_args
+    assert isinstance(args[0], ops.ObjGetAccessUrl)
+    assert args[0].mode == "rw"
+    # 1 hour = 3600 seconds = 3600 * 1000 * 1000 microseconds
+    assert args[0].duration == 3600000000
 
 def test_make_ref_calls_apply_binary_op_with_authorizer():
     uri = MagicMock(spec=series.Series)