16
16
import pandas
17
17
import pyarrow
18
18
import pytest
19
+ import test_utils .prefixer
19
20
20
21
import bigframes
21
22
import bigframes .exceptions as bfe
22
23
import bigframes .pandas as bpd
23
24
from tests .system .utils import cleanup_function_assets
24
25
26
+ prefixer = test_utils .prefixer .Prefixer ("bigframes" , "" )
27
+
25
28
26
29
def test_managed_function_multiply_with_ibis (
27
30
session ,
@@ -37,6 +40,7 @@ def test_managed_function_multiply_with_ibis(
37
40
input_types = [int , int ],
38
41
output_type = int ,
39
42
dataset = dataset_id ,
43
+ name = prefixer .create_prefix (),
40
44
)
41
45
def multiply (x , y ):
42
46
return x * y
@@ -87,6 +91,7 @@ def test_managed_function_stringify_with_ibis(
87
91
input_types = [int ],
88
92
output_type = str ,
89
93
dataset = dataset_id ,
94
+ name = prefixer .create_prefix (),
90
95
)
91
96
def stringify (x ):
92
97
return f"I got { x } "
@@ -123,7 +128,10 @@ def stringify(x):
123
128
def test_managed_function_array_output (session , scalars_dfs , dataset_id ):
124
129
try :
125
130
126
- @session .udf (dataset = dataset_id )
131
+ @session .udf (
132
+ dataset = dataset_id ,
133
+ name = prefixer .create_prefix (),
134
+ )
127
135
def featurize (x : int ) -> list [float ]:
128
136
return [float (i ) for i in [x , x + 1 , x + 2 ]]
129
137
@@ -160,10 +168,10 @@ def featurize(x: int) -> list[float]:
160
168
cleanup_function_assets (featurize , session .bqclient , ignore_failures = False )
161
169
162
170
163
- def test_managed_function_series_apply (session , scalars_dfs ):
171
+ def test_managed_function_series_apply (session , dataset_id , scalars_dfs ):
164
172
try :
165
173
166
- @session .udf ()
174
+ @session .udf (dataset = dataset_id , name = prefixer . create_prefix () )
167
175
def foo (x : int ) -> bytes :
168
176
return bytes (abs (x ))
169
177
@@ -214,13 +222,14 @@ def foo(x: int) -> bytes:
214
222
215
223
def test_managed_function_series_apply_array_output (
216
224
session ,
225
+ dataset_id ,
217
226
scalars_dfs ,
218
227
):
219
228
try :
220
229
221
230
with pytest .warns (bfe .PreviewWarning , match = "udf is in preview." ):
222
231
223
- @session .udf ()
232
+ @session .udf (dataset = dataset_id , name = prefixer . create_prefix () )
224
233
def foo_list (x : int ) -> list [float ]:
225
234
return [float (abs (x )), float (abs (x ) + 1 )]
226
235
@@ -243,7 +252,7 @@ def foo_list(x: int) -> list[float]:
243
252
cleanup_function_assets (foo_list , session .bqclient , ignore_failures = False )
244
253
245
254
246
- def test_managed_function_series_combine (session , scalars_dfs ):
255
+ def test_managed_function_series_combine (session , dataset_id , scalars_dfs ):
247
256
try :
248
257
# This function is deliberately written to not work with NA input.
249
258
def add (x : int , y : int ) -> int :
@@ -258,7 +267,9 @@ def add(x: int, y: int) -> int:
258
267
# make sure there are NA values in the test column.
259
268
assert any ([pandas .isna (val ) for val in bf_df [int_col_name_with_nulls ]])
260
269
261
- add_managed_func = session .udf ()(add )
270
+ add_managed_func = session .udf (
271
+ dataset = dataset_id , name = prefixer .create_prefix ()
272
+ )(add )
262
273
263
274
# with nulls in the series the managed function application would fail.
264
275
with pytest .raises (
@@ -301,7 +312,7 @@ def add(x: int, y: int) -> int:
301
312
)
302
313
303
314
304
- def test_managed_function_series_combine_array_output (session , scalars_dfs ):
315
+ def test_managed_function_series_combine_array_output (session , dataset_id , scalars_dfs ):
305
316
try :
306
317
307
318
def add_list (x : int , y : int ) -> list [int ]:
@@ -316,7 +327,9 @@ def add_list(x: int, y: int) -> list[int]:
316
327
# Make sure there are NA values in the test column.
317
328
assert any ([pandas .isna (val ) for val in bf_df [int_col_name_with_nulls ]])
318
329
319
- add_list_managed_func = session .udf ()(add_list )
330
+ add_list_managed_func = session .udf (
331
+ dataset = dataset_id , name = prefixer .create_prefix ()
332
+ )(add_list )
320
333
321
334
# After filtering out nulls the managed function application should work
322
335
# similar to pandas.
@@ -364,7 +377,7 @@ def add_list(x: int, y: int) -> list[int]:
364
377
)
365
378
366
379
367
- def test_managed_function_dataframe_map (session , scalars_dfs ):
380
+ def test_managed_function_dataframe_map (session , dataset_id , scalars_dfs ):
368
381
try :
369
382
370
383
def add_one (x ):
@@ -373,6 +386,8 @@ def add_one(x):
373
386
mf_add_one = session .udf (
374
387
input_types = [int ],
375
388
output_type = int ,
389
+ dataset = dataset_id ,
390
+ name = prefixer .create_prefix (),
376
391
)(add_one )
377
392
378
393
scalars_df , scalars_pandas_df = scalars_dfs
@@ -398,9 +413,7 @@ def add_one(x):
398
413
cleanup_function_assets (mf_add_one , session .bqclient , ignore_failures = False )
399
414
400
415
401
- def test_managed_function_dataframe_map_array_output (
402
- session , scalars_dfs , dataset_id_permanent
403
- ):
416
+ def test_managed_function_dataframe_map_array_output (session , scalars_dfs , dataset_id ):
404
417
try :
405
418
406
419
def add_one_list (x ):
@@ -409,6 +422,8 @@ def add_one_list(x):
409
422
mf_add_one_list = session .udf (
410
423
input_types = [int ],
411
424
output_type = list [int ],
425
+ dataset = dataset_id ,
426
+ name = prefixer .create_prefix (),
412
427
)(add_one_list )
413
428
414
429
scalars_df , scalars_pandas_df = scalars_dfs
@@ -439,7 +454,7 @@ def add_one_list(x):
439
454
)
440
455
441
456
442
- def test_managed_function_dataframe_apply_axis_1 (session , scalars_dfs ):
457
+ def test_managed_function_dataframe_apply_axis_1 (session , dataset_id , scalars_dfs ):
443
458
try :
444
459
scalars_df , scalars_pandas_df = scalars_dfs
445
460
series = scalars_df ["int64_too" ]
@@ -451,6 +466,8 @@ def add_ints(x, y):
451
466
add_ints_mf = session .udf (
452
467
input_types = [int , int ],
453
468
output_type = int ,
469
+ dataset = dataset_id ,
470
+ name = prefixer .create_prefix (),
454
471
)(add_ints )
455
472
assert add_ints_mf .bigframes_bigquery_function # type: ignore
456
473
@@ -475,7 +492,7 @@ def add_ints(x, y):
475
492
cleanup_function_assets (add_ints_mf , session .bqclient , ignore_failures = False )
476
493
477
494
478
- def test_managed_function_dataframe_apply_axis_1_array_output (session ):
495
+ def test_managed_function_dataframe_apply_axis_1_array_output (session , dataset_id ):
479
496
bf_df = bigframes .dataframe .DataFrame (
480
497
{
481
498
"Id" : [1 , 2 , 3 ],
@@ -498,6 +515,8 @@ def test_managed_function_dataframe_apply_axis_1_array_output(session):
498
515
@session .udf (
499
516
input_types = [int , float , str ],
500
517
output_type = list [str ],
518
+ dataset = dataset_id ,
519
+ name = prefixer .create_prefix (),
501
520
)
502
521
def foo (x , y , z ):
503
522
return [str (x ), str (y ), z ]
@@ -591,12 +610,16 @@ def foo(x, y, z):
591
610
],
592
611
)
593
612
def test_managed_function_with_connection (
594
- session , scalars_dfs , request , connection_fixture
613
+ session , scalars_dfs , dataset_id , request , connection_fixture
595
614
):
596
615
try :
597
616
bigquery_connection = request .getfixturevalue (connection_fixture )
598
617
599
- @session .udf (bigquery_connection = bigquery_connection )
618
+ @session .udf (
619
+ bigquery_connection = bigquery_connection ,
620
+ dataset = dataset_id ,
621
+ name = prefixer .create_prefix (),
622
+ )
600
623
def foo (x : int ) -> int :
601
624
return x + 10
602
625
0 commit comments