Skip to content

Commit 80ff0b2

Browse files
committed
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.
1 parent 173b83d commit 80ff0b2

File tree

9 files changed

+291
-5
lines changed

9 files changed

+291
-5
lines changed

bigframes/bigquery/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
import sys
2020

21-
from bigframes.bigquery import ai, ml
21+
from bigframes.bigquery import ai, ml, obj
2222
from bigframes.bigquery._operations.approx_agg import approx_top_count
2323
from bigframes.bigquery._operations.array import (
2424
array_agg,
@@ -158,4 +158,5 @@
158158
# Modules / SQL namespaces
159159
"ai",
160160
"ml",
161+
"obj",
161162
]
Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
"""
17+
ObjectRef functions defined from
18+
https://cloud.google.com/bigquery/docs/reference/standard-sql/object-ref-functions
19+
"""
20+
21+
22+
from __future__ import annotations
23+
24+
import datetime
25+
from typing import Optional, Union
26+
27+
import numpy as np
28+
import pandas as pd
29+
30+
from bigframes.core import log_adapter
31+
import bigframes.core.utils as utils
32+
import bigframes.operations as ops
33+
import bigframes.series as series
34+
35+
36+
@log_adapter.method_logger(custom_base_name="bigquery_obj")
37+
def fetch_metadata(
38+
objectref: series.Series,
39+
) -> series.Series:
40+
"""The OBJ.FETCH_METADATA function returns Cloud Storage metadata for a partially populated ObjectRef value.
41+
42+
Args:
43+
objectref (bigframes.series.Series):
44+
A partially populated ObjectRef value, in which the uri and authorizer fields are populated and the details field isn't.
45+
46+
Returns:
47+
bigframes.series.Series: A fully populated ObjectRef value. The metadata is provided in the details field of the returned ObjectRef value.
48+
"""
49+
return objectref._apply_unary_op(ops.obj_fetch_metadata_op)
50+
51+
52+
@log_adapter.method_logger(custom_base_name="bigquery_obj")
53+
def get_access_url(
54+
objectref: series.Series,
55+
mode: str,
56+
duration: Optional[
57+
Union[datetime.timedelta, pd.Timedelta, np.timedelta64]
58+
] = None,
59+
) -> series.Series:
60+
"""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.
61+
62+
Args:
63+
objectref (bigframes.series.Series):
64+
An ObjectRef value that represents a Cloud Storage object.
65+
mode (str):
66+
A STRING value that identifies the type of URL that you want to be returned. The following values are supported:
67+
'r': Returns a URL that lets you read the object.
68+
'rw': Returns two URLs, one that lets you read the object, and one that lets you modify the object.
69+
duration (Union[datetime.timedelta, pandas.Timedelta, numpy.timedelta64], optional):
70+
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.
71+
72+
Returns:
73+
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.
74+
"""
75+
duration_micros = None
76+
if duration is not None:
77+
duration_micros = utils.timedelta_to_micros(duration)
78+
79+
return objectref._apply_unary_op(
80+
ops.ObjGetAccessUrl(mode=mode, duration=duration_micros)
81+
)
82+
83+
84+
@log_adapter.method_logger(custom_base_name="bigquery_obj")
85+
def make_ref(
86+
uri_or_json: series.Series,
87+
authorizer: Optional[series.Series] = None,
88+
) -> series.Series:
89+
"""Use the OBJ.MAKE_REF function to create an ObjectRef value that contains reference information for a Cloud Storage object.
90+
91+
Args:
92+
uri_or_json (bigframes.series.Series):
93+
A STRING value that contains the URI for the Cloud Storage object, for example, gs://mybucket/flowers/12345.jpg.
94+
OR
95+
A JSON value that represents a Cloud Storage object.
96+
authorizer (bigframes.series.Series, optional):
97+
A STRING value that contains the Cloud Resource connection used to access the Cloud Storage object.
98+
Required if uri_or_json is a URI string.
99+
100+
Returns:
101+
bigframes.series.Series: An ObjectRef value.
102+
"""
103+
if authorizer is not None:
104+
return uri_or_json._apply_binary_op(authorizer, ops.obj_make_ref_op)
105+
106+
# If authorizer is not provided, we assume uri_or_json is a JSON objectref
107+
return uri_or_json._apply_unary_op(ops.obj_make_ref_json_op)

bigframes/bigquery/obj.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""This module exposes `BigQuery ObjectRef
16+
<https://cloud.google.com/bigquery/docs/object-table-object-ref-functions>`_ functions.
17+
"""
18+
19+
from bigframes.bigquery._operations.obj import (
20+
fetch_metadata,
21+
get_access_url,
22+
make_ref,
23+
)
24+
25+
__all__ = [
26+
"fetch_metadata",
27+
"get_access_url",
28+
"make_ref",
29+
]

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1247,6 +1247,11 @@ def obj_fetch_metadata_op_impl(obj_ref: ibis_types.Value):
12471247

12481248
@scalar_op_compiler.register_unary_op(ops.ObjGetAccessUrl, pass_op=True)
12491249
def obj_get_access_url_op_impl(obj_ref: ibis_types.Value, op: ops.ObjGetAccessUrl):
1250+
if op.duration is not None:
1251+
duration_value = ibis_types.literal(op.duration).to_interval("us")
1252+
return obj_get_access_url_with_duration(
1253+
obj_ref=obj_ref, mode=op.mode, duration=duration_value
1254+
)
12501255
return obj_get_access_url(obj_ref=obj_ref, mode=op.mode)
12511256

12521257

@@ -1807,6 +1812,11 @@ def obj_make_ref_op(x: ibis_types.Value, y: ibis_types.Value):
18071812
return obj_make_ref(uri=x, authorizer=y)
18081813

18091814

1815+
@scalar_op_compiler.register_unary_op(ops.obj_make_ref_json_op)
1816+
def obj_make_ref_json_op(x: ibis_types.Value):
1817+
return obj_make_ref_json(objectref_json=x)
1818+
1819+
18101820
# Ternary Operations
18111821
@scalar_op_compiler.register_ternary_op(ops.where_op)
18121822
def where_op(
@@ -2141,11 +2151,23 @@ def obj_make_ref(uri: str, authorizer: str) -> _OBJ_REF_IBIS_DTYPE: # type: ign
21412151
"""Make ObjectRef Struct from uri and connection."""
21422152

21432153

2154+
@ibis_udf.scalar.builtin(name="OBJ.MAKE_REF")
2155+
def obj_make_ref_json(objectref_json: ibis_dtypes.JSON) -> _OBJ_REF_IBIS_DTYPE: # type: ignore
2156+
"""Make ObjectRef Struct from json."""
2157+
2158+
21442159
@ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
21452160
def obj_get_access_url(obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String) -> ibis_dtypes.JSON: # type: ignore
21462161
"""Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
21472162

21482163

2164+
@ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
2165+
def obj_get_access_url_with_duration(
2166+
obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String, duration: ibis_dtypes.Interval(unit="us") # type: ignore
2167+
) -> ibis_dtypes.JSON: # type: ignore
2168+
"""Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
2169+
2170+
21492171
@ibis_udf.scalar.builtin(name="ltrim")
21502172
def str_lstrip_op( # type: ignore[empty-body]
21512173
x: ibis_dtypes.String, to_strip: ibis_dtypes.String

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,24 @@ def _(expr: TypedExpr) -> sge.Expression:
2929
return sge.func("OBJ.FETCH_METADATA", expr.expr)
3030

3131

32-
@register_unary_op(ops.ObjGetAccessUrl)
33-
def _(expr: TypedExpr) -> sge.Expression:
34-
return sge.func("OBJ.GET_ACCESS_URL", expr.expr)
32+
@register_unary_op(ops.ObjGetAccessUrl, pass_op=True)
33+
def _(expr: TypedExpr, op: ops.ObjGetAccessUrl) -> sge.Expression:
34+
args = [expr.expr, sge.Literal.string(op.mode)]
35+
if op.duration is not None:
36+
args.append(
37+
sge.Interval(
38+
this=sge.Literal.number(op.duration),
39+
unit=sge.Var(this="MICROSECOND"),
40+
)
41+
)
42+
return sge.func("OBJ.GET_ACCESS_URL", *args)
3543

3644

3745
@register_binary_op(ops.obj_make_ref_op)
3846
def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
3947
return sge.func("OBJ.MAKE_REF", left.expr, right.expr)
48+
49+
50+
@register_unary_op(ops.obj_make_ref_json_op)
51+
def _(expr: TypedExpr) -> sge.Expression:
52+
return sge.func("OBJ.MAKE_REF", expr.expr)

bigframes/operations/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
)
4141
from bigframes.operations.blob_ops import (
4242
obj_fetch_metadata_op,
43+
obj_make_ref_json_op,
4344
obj_make_ref_op,
4445
ObjGetAccessUrl,
4546
)
@@ -365,6 +366,7 @@
365366
"ArrayToStringOp",
366367
# Blob ops
367368
"ObjGetAccessUrl",
369+
"obj_make_ref_json_op",
368370
"obj_make_ref_op",
369371
"obj_fetch_metadata_op",
370372
# Struct ops

bigframes/operations/blob_ops.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
class ObjGetAccessUrl(base_ops.UnaryOp):
3030
name: typing.ClassVar[str] = "obj_get_access_url"
3131
mode: str # access mode, e.g. R read, W write, RW read & write
32+
duration: typing.Optional[int] = None # duration in microseconds
3233

3334
def output_type(self, *input_types):
3435
return dtypes.JSON_DTYPE
@@ -46,3 +47,14 @@ def output_type(self, *input_types):
4647

4748

4849
obj_make_ref_op = ObjMakeRef()
50+
51+
52+
@dataclasses.dataclass(frozen=True)
53+
class ObjMakeRefJson(base_ops.UnaryOp):
54+
name: typing.ClassVar[str] = "obj_make_ref_json"
55+
56+
def output_type(self, *input_types):
57+
return dtypes.OBJ_REF_DTYPE
58+
59+
60+
obj_make_ref_json_op = ObjMakeRefJson()

tests/unit/bigquery/test_obj.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import datetime
16+
from unittest.mock import MagicMock
17+
18+
import pytest
19+
20+
import bigframes.bigquery.obj as obj
21+
import bigframes.operations as ops
22+
import bigframes.series as series
23+
24+
25+
def test_fetch_metadata_op_structure():
26+
op = ops.obj_fetch_metadata_op
27+
assert op.name == "obj_fetch_metadata"
28+
29+
def test_get_access_url_op_structure():
30+
op = ops.ObjGetAccessUrl(mode="r")
31+
assert op.name == "obj_get_access_url"
32+
assert op.mode == "r"
33+
assert op.duration is None
34+
35+
def test_get_access_url_with_duration_op_structure():
36+
op = ops.ObjGetAccessUrl(mode="rw", duration=3600000000)
37+
assert op.name == "obj_get_access_url"
38+
assert op.mode == "rw"
39+
assert op.duration == 3600000000
40+
41+
def test_make_ref_op_structure():
42+
op = ops.obj_make_ref_op
43+
assert op.name == "obj_make_ref"
44+
45+
def test_make_ref_json_op_structure():
46+
op = ops.obj_make_ref_json_op
47+
assert op.name == "obj_make_ref_json"
48+
49+
def test_fetch_metadata_calls_apply_unary_op():
50+
s = MagicMock(spec=series.Series)
51+
52+
obj.fetch_metadata(s)
53+
54+
s._apply_unary_op.assert_called_once()
55+
args, _ = s._apply_unary_op.call_args
56+
assert args[0] == ops.obj_fetch_metadata_op
57+
58+
def test_get_access_url_calls_apply_unary_op_without_duration():
59+
s = MagicMock(spec=series.Series)
60+
61+
obj.get_access_url(s, mode="r")
62+
63+
s._apply_unary_op.assert_called_once()
64+
args, _ = s._apply_unary_op.call_args
65+
assert isinstance(args[0], ops.ObjGetAccessUrl)
66+
assert args[0].mode == "r"
67+
assert args[0].duration is None
68+
69+
def test_get_access_url_calls_apply_unary_op_with_duration():
70+
s = MagicMock(spec=series.Series)
71+
duration = datetime.timedelta(hours=1)
72+
73+
obj.get_access_url(s, mode="rw", duration=duration)
74+
75+
s._apply_unary_op.assert_called_once()
76+
args, kwargs = s._apply_unary_op.call_args
77+
assert isinstance(args[0], ops.ObjGetAccessUrl)
78+
assert args[0].mode == "rw"
79+
# 1 hour = 3600 seconds = 3600 * 1000 * 1000 microseconds
80+
assert args[0].duration == 3600000000
81+
82+
def test_make_ref_calls_apply_binary_op_with_authorizer():
83+
uri = MagicMock(spec=series.Series)
84+
auth = MagicMock(spec=series.Series)
85+
86+
obj.make_ref(uri, authorizer=auth)
87+
88+
uri._apply_binary_op.assert_called_once()
89+
args, _ = uri._apply_binary_op.call_args
90+
assert args[0] == auth
91+
assert args[1] == ops.obj_make_ref_op
92+
93+
def test_make_ref_calls_apply_unary_op_without_authorizer():
94+
json_val = MagicMock(spec=series.Series)
95+
96+
obj.make_ref(json_val)
97+
98+
json_val._apply_unary_op.assert_called_once()
99+
args, _ = json_val._apply_unary_op.call_args
100+
assert args[0] == ops.obj_make_ref_json_op

tests/unit/core/compile/sqlglot/expressions/snapshots/test_blob_ops/test_obj_get_access_url/out.sql

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ WITH `bfcte_0` AS (
1111
), `bfcte_2` AS (
1212
SELECT
1313
*,
14-
OBJ.GET_ACCESS_URL(`bfcol_4`) AS `bfcol_7`
14+
OBJ.GET_ACCESS_URL(`bfcol_4`, 'R') AS `bfcol_7`
1515
FROM `bfcte_1`
1616
), `bfcte_3` AS (
1717
SELECT

0 commit comments

Comments
 (0)