chore: add experimental blob read and write URL functions (#1308)

* chore: add experimental blob read and write URL functions * fix * fix

Author: GarrettWu

bigframes/dataframe.py

@@ -748,9 +748,8 @@ def _repr_html_(self) -> str:
  if df[col].dtype == bigframes.dtypes.OBJ_REF_DTYPE
  ]
  for col in blob_cols:
- df[col] = df[col]._apply_unary_op(ops.obj_fetch_metadata_op)
  # TODO(garrettwu): Not necessary to get access urls for all the rows. Update when having a to get URLs from local data.
- df[col] = df[col]._apply_unary_op(ops.ObjGetAccessUrl(mode="R"))
+ df[col] = df[col].blob._get_runtime(mode="R", with_metadata=True)
 
  # TODO(swast): pass max_columns and get the true column count back. Maybe
  # get 1 more column than we have requested so that pandas can add the

bigframes/operations/blob.py

@@ -66,11 +66,11 @@ def version(self) -> bigframes.series.Series:
 
  Returns:
  BigFrames Series: Version as string."""
- # version must be retrived after fetching metadata
+ # version must be retrieved after fetching metadata
  return self._apply_unary_op(ops.obj_fetch_metadata_op).struct.field("version")
 
  def metadata(self) -> bigframes.series.Series:
- """Retrive the metadata of the Blob.
+ """Retrieve the metadata of the Blob.
 
  .. note::
  BigFrames Blob is still under experiments. It may not work and subject to change in the future.
@@ -85,7 +85,7 @@ def metadata(self) -> bigframes.series.Series:
  return bbq.json_extract(details_json, "$.gcs_metadata").rename("metadata")
 
  def content_type(self) -> bigframes.series.Series:
- """Retrive the content type of the Blob.
+ """Retrieve the content type of the Blob.
 
  .. note::
  BigFrames Blob is still under experiments. It may not work and subject to change in the future.
@@ -99,7 +99,7 @@ def content_type(self) -> bigframes.series.Series:
  )
 
  def md5_hash(self) -> bigframes.series.Series:
- """Retrive the md5 hash of the Blob.
+ """Retrieve the md5 hash of the Blob.
 
  .. note::
  BigFrames Blob is still under experiments. It may not work and subject to change in the future.
@@ -113,7 +113,7 @@ def md5_hash(self) -> bigframes.series.Series:
  )
 
  def size(self) -> bigframes.series.Series:
- """Retrive the file size of the Blob.
+ """Retrieve the file size of the Blob.
 
  .. note::
  BigFrames Blob is still under experiments. It may not work and subject to change in the future.
@@ -128,7 +128,7 @@ def size(self) -> bigframes.series.Series:
  )
 
  def updated(self) -> bigframes.series.Series:
- """Retrive the updated time of the Blob.
+ """Retrieve the updated time of the Blob.
 
  .. note::
  BigFrames Blob is still under experiments. It may not work and subject to change in the future.
@@ -146,6 +146,46 @@ def updated(self) -> bigframes.series.Series:
 
  return bpd.to_datetime(updated, unit="us", utc=True)
 
+ def _get_runtime(
+ self, mode: str, with_metadata: bool = False
+ ) -> bigframes.series.Series:
+ """Retrieve the ObjectRefRuntime as JSON.
+
+ Args:
+ mode (str): mode for the URLs, "R" for read, "RW" for read & write.
+ metadata (bool, default False): whether to fetch the metadata in the ObjectRefRuntime.
+
+ Returns:
+ bigframes Series: ObjectRefRuntime JSON.
+ """
+ s = self._apply_unary_op(ops.obj_fetch_metadata_op) if with_metadata else self
+
+ return s._apply_unary_op(ops.ObjGetAccessUrl(mode=mode))
+
+ def read_url(self) -> bigframes.series.Series:
+ """Retrieve the read URL of the Blob.
+
+ .. note::
+ BigFrames Blob is still under experiments. It may not work and subject to change in the future.
+
+ Returns:
+ BigFrames Series: Read only URLs."""
+ return self._get_runtime(mode="R")._apply_unary_op(
+ ops.JSONValue(json_path="$.access_urls.read_url")
+ )
+
+ def write_url(self) -> bigframes.series.Series:
+ """Retrieve the write URL of the Blob.
+
+ .. note::
+ BigFrames Blob is still under experiments. It may not work and subject to change in the future.
+
+ Returns:
+ BigFrames Series: Writable URLs."""
+ return self._get_runtime(mode="RW")._apply_unary_op(
+ ops.JSONValue(json_path="$.access_urls.write_url")
+ )
+
  def display(self, n: int = 3, *, content_type: str = ""):
  """Display the blob content in the IPython Notebook environment. Only works for image type now.
 
@@ -159,10 +199,7 @@ def display(self, n: int = 3, *, content_type: str = ""):
  # col name doesn't matter here. Rename to avoid column name conflicts
  df = bigframes.series.Series(self._block).rename("blob_col").head(n).to_frame()
 
- obj_ref_runtime = df["blob_col"]._apply_unary_op(ops.ObjGetAccessUrl(mode="R"))
- df["read_url"] = obj_ref_runtime._apply_unary_op(
- ops.JSONValue(json_path="$.access_urls.read_url")
- )
+ df["read_url"] = df["blob_col"].blob.read_url()
 
  if content_type:
  df["content_type"] = content_type
@@ -231,10 +268,8 @@ def image_blur(
  connection=connection,
  ).udf()
 
- src_rt = bigframes.series.Series(self._block)._apply_unary_op(
- ops.ObjGetAccessUrl(mode="R")
- )
- dst_rt = dst._apply_unary_op(ops.ObjGetAccessUrl(mode="RW"))
+ src_rt = self._get_runtime(mode="R")
+ dst_rt = dst.blob._get_runtime(mode="RW")
 
  src_rt = src_rt._apply_unary_op(ops.ToJSONString())
  dst_rt = dst_rt._apply_unary_op(ops.ToJSONString())