@@ -238,6 +238,10 @@ def display_single_url(
238
238
for _ , row in pandas_df .iterrows ():
239
239
display_single_url (row ["read_url" ], row ["content_type" ])
240
240
241
+ @property
242
+ def session (self ):
243
+ return self ._block .session
244
+
241
245
def _resolve_connection (self , connection : Optional [str ] = None ) -> str :
242
246
"""Resovle the BigQuery connection.
243
247
@@ -286,6 +290,13 @@ def _get_runtime_json_str(
286
290
runtime = self ._get_runtime (mode = mode , with_metadata = with_metadata )
287
291
return runtime ._apply_unary_op (ops .ToJSONString ())
288
292
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
+
289
300
def image_blur (
290
301
self ,
291
302
ksize : tuple [int , int ],
@@ -367,6 +378,8 @@ def image_blur(
367
378
res = df .apply (image_blur_udf , axis = 1 )
368
379
res .cache () # to execute the udf
369
380
381
+ self ._add_to_cleanup_set (image_blur_udf )
382
+
370
383
return dst
371
384
372
385
def image_resize (
@@ -463,6 +476,8 @@ def image_resize(
463
476
res = df .apply (image_resize_udf , axis = 1 )
464
477
res .cache () # to execute the udf
465
478
479
+ self ._add_to_cleanup_set (image_resize_udf )
480
+
466
481
return dst
467
482
468
483
def image_normalize (
@@ -554,6 +569,8 @@ def image_normalize(
554
569
res = df .apply (image_normalize_udf , axis = 1 )
555
570
res .cache () # to execute the udf
556
571
572
+ self ._add_to_cleanup_set (image_normalize_udf )
573
+
557
574
return dst
558
575
559
576
def pdf_extract (
@@ -598,6 +615,9 @@ def pdf_extract(
598
615
599
616
src_rt = self ._get_runtime_json_str (mode = "R" )
600
617
res = src_rt .apply (pdf_extract_udf )
618
+
619
+ self ._add_to_cleanup_set (pdf_extract_udf )
620
+
601
621
return res
602
622
603
623
def pdf_chunk (
@@ -664,4 +684,7 @@ def pdf_chunk(
664
684
res = df .apply (pdf_chunk_udf , axis = 1 )
665
685
666
686
res_array = bbq .json_extract_string_array (res )
687
+
688
+ self ._add_to_cleanup_set (pdf_chunk_udf )
689
+
667
690
return res_array
0 commit comments