Skip to content

Commit 06d343c

Browse files
authored
chore: cleanup udf resources in blob transform functions (#1530)
* chore: cleanup udf resources in blob transform functions * cleanup on session close * use remote function cleanup
1 parent 6ab4ffd commit 06d343c

File tree

1 file changed

+23
-0
lines changed

1 file changed

+23
-0
lines changed

bigframes/operations/blob.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,10 @@ def display_single_url(
238238
for _, row in pandas_df.iterrows():
239239
display_single_url(row["read_url"], row["content_type"])
240240

241+
@property
242+
def session(self):
243+
return self._block.session
244+
241245
def _resolve_connection(self, connection: Optional[str] = None) -> str:
242246
"""Resovle the BigQuery connection.
243247
@@ -286,6 +290,13 @@ def _get_runtime_json_str(
286290
runtime = self._get_runtime(mode=mode, with_metadata=with_metadata)
287291
return runtime._apply_unary_op(ops.ToJSONString())
288292

293+
# TODO(b/404605969): remove cleanups when UDF fixes dataset deletion.
294+
def _add_to_cleanup_set(self, udf):
295+
"""Add udf name to session cleanup set. Won't need this after UDF fixes dataset deletion."""
296+
self.session._function_session._update_temp_artifacts(
297+
udf.bigframes_bigquery_function, ""
298+
)
299+
289300
def image_blur(
290301
self,
291302
ksize: tuple[int, int],
@@ -367,6 +378,8 @@ def image_blur(
367378
res = df.apply(image_blur_udf, axis=1)
368379
res.cache() # to execute the udf
369380

381+
self._add_to_cleanup_set(image_blur_udf)
382+
370383
return dst
371384

372385
def image_resize(
@@ -463,6 +476,8 @@ def image_resize(
463476
res = df.apply(image_resize_udf, axis=1)
464477
res.cache() # to execute the udf
465478

479+
self._add_to_cleanup_set(image_resize_udf)
480+
466481
return dst
467482

468483
def image_normalize(
@@ -554,6 +569,8 @@ def image_normalize(
554569
res = df.apply(image_normalize_udf, axis=1)
555570
res.cache() # to execute the udf
556571

572+
self._add_to_cleanup_set(image_normalize_udf)
573+
557574
return dst
558575

559576
def pdf_extract(
@@ -598,6 +615,9 @@ def pdf_extract(
598615

599616
src_rt = self._get_runtime_json_str(mode="R")
600617
res = src_rt.apply(pdf_extract_udf)
618+
619+
self._add_to_cleanup_set(pdf_extract_udf)
620+
601621
return res
602622

603623
def pdf_chunk(
@@ -664,4 +684,7 @@ def pdf_chunk(
664684
res = df.apply(pdf_chunk_udf, axis=1)
665685

666686
res_array = bbq.json_extract_string_array(res)
687+
688+
self._add_to_cleanup_set(pdf_chunk_udf)
689+
667690
return res_array

0 commit comments

Comments
 (0)