25
25
import bigframes
26
26
import bigframes .dtypes
27
27
import bigframes .exceptions
28
- from bigframes .functions import _utils as rf_utils
29
- from bigframes .functions import remote_function as rf
28
+ from bigframes .functions import _utils as bff_utils
29
+ from bigframes .functions import function as bff
30
30
from tests .system .utils import assert_pandas_df_equal
31
31
32
32
_prefixer = test_utils .prefixer .Prefixer ("bigframes" , "" )
@@ -94,12 +94,12 @@ def get_rf_name(func, package_requirements=None, is_row_processor=False):
94
94
"""Get a remote function name for testing given a udf."""
95
95
# Augment user package requirements with any internal package
96
96
# requirements
97
- package_requirements = rf_utils ._get_updated_package_requirements (
97
+ package_requirements = bff_utils ._get_updated_package_requirements (
98
98
package_requirements , is_row_processor
99
99
)
100
100
101
101
# Compute a unique hash representing the user code
102
- function_hash = rf_utils ._get_hash (func , package_requirements )
102
+ function_hash = bff_utils ._get_hash (func , package_requirements )
103
103
104
104
return f"bigframes_{ function_hash } "
105
105
@@ -117,7 +117,7 @@ def test_remote_function_direct_no_session_param(
117
117
def square (x ):
118
118
return x * x
119
119
120
- square = rf .remote_function (
120
+ square = bff .remote_function (
121
121
int ,
122
122
int ,
123
123
bigquery_client = bigquery_client ,
@@ -176,7 +176,7 @@ def test_remote_function_direct_no_session_param_location_specified(
176
176
def square (x ):
177
177
return x * x
178
178
179
- square = rf .remote_function (
179
+ square = bff .remote_function (
180
180
int ,
181
181
int ,
182
182
bigquery_client = bigquery_client ,
@@ -235,7 +235,7 @@ def square(x):
235
235
ValueError ,
236
236
match = re .escape ("The location does not match BigQuery connection location:" ),
237
237
):
238
- rf .remote_function (
238
+ bff .remote_function (
239
239
int ,
240
240
int ,
241
241
bigquery_client = bigquery_client ,
@@ -263,7 +263,7 @@ def test_remote_function_direct_no_session_param_location_project_specified(
263
263
def square (x ):
264
264
return x * x
265
265
266
- square = rf .remote_function (
266
+ square = bff .remote_function (
267
267
int ,
268
268
int ,
269
269
bigquery_client = bigquery_client ,
@@ -324,7 +324,7 @@ def square(x):
324
324
"The project_id does not match BigQuery connection gcp_project_id:"
325
325
),
326
326
):
327
- rf .remote_function (
327
+ bff .remote_function (
328
328
int ,
329
329
int ,
330
330
bigquery_client = bigquery_client ,
@@ -346,7 +346,7 @@ def test_remote_function_direct_session_param(
346
346
def square (x ):
347
347
return x * x
348
348
349
- square = rf .remote_function (
349
+ square = bff .remote_function (
350
350
int ,
351
351
int ,
352
352
session = session_with_bq_connection ,
@@ -636,7 +636,7 @@ def add_one(x):
636
636
def test_read_gbq_function_detects_invalid_function (session , dataset_id ):
637
637
dataset_ref = bigquery .DatasetReference .from_string (dataset_id )
638
638
with pytest .raises (ValueError ) as e :
639
- rf .read_gbq_function (
639
+ bff .read_gbq_function (
640
640
str (dataset_ref .routine ("not_a_function" )),
641
641
session = session ,
642
642
)
@@ -658,7 +658,7 @@ def test_read_gbq_function_like_original(
658
658
def square1 (x ):
659
659
return x * x
660
660
661
- square1 = rf .remote_function (
661
+ square1 = bff .remote_function (
662
662
[int ],
663
663
int ,
664
664
bigquery_client = bigquery_client ,
@@ -674,7 +674,7 @@ def square1(x):
674
674
# Function should still work normally.
675
675
assert square1 (2 ) == 4
676
676
677
- square2 = rf .read_gbq_function (
677
+ square2 = bff .read_gbq_function (
678
678
function_name = square1 .bigframes_remote_function , # type: ignore
679
679
session = session ,
680
680
)
@@ -745,7 +745,7 @@ def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id):
745
745
for routine in (sql_routine , js_routine ):
746
746
# Create the routine in BigQuery and read it back using read_gbq_function.
747
747
bigquery_client .create_routine (routine , exists_ok = True )
748
- square = rf .read_gbq_function (
748
+ square = bff .read_gbq_function (
749
749
str (routine .reference ),
750
750
session = session ,
751
751
)
@@ -757,7 +757,7 @@ def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id):
757
757
758
758
src = {"x" : [- 5 , - 4 , - 3 , - 2 , - 1 , 0 , 1 , 2 , 3 , 4 , 5 ]}
759
759
760
- routine_ref_str = rf_utils .routine_ref_to_string_for_query (routine .reference )
760
+ routine_ref_str = bff_utils .routine_ref_to_string_for_query (routine .reference )
761
761
direct_sql = " UNION ALL " .join (
762
762
[f"SELECT { x } AS x, { routine_ref_str } ({ x } ) AS y" for x in src ["x" ]]
763
763
)
@@ -818,25 +818,25 @@ def test_read_gbq_function_requires_explicit_types(
818
818
bigquery_client .create_routine (only_arg_type_specified , exists_ok = True )
819
819
bigquery_client .create_routine (neither_type_specified , exists_ok = True )
820
820
821
- rf .read_gbq_function (
821
+ bff .read_gbq_function (
822
822
str (both_types_specified .reference ),
823
823
session = session ,
824
824
)
825
825
with pytest .warns (
826
826
bigframes .exceptions .UnknownDataTypeWarning ,
827
827
match = "missing input data types.*assume default data type" ,
828
828
):
829
- rf .read_gbq_function (
829
+ bff .read_gbq_function (
830
830
str (only_return_type_specified .reference ),
831
831
session = session ,
832
832
)
833
833
with pytest .raises (ValueError ):
834
- rf .read_gbq_function (
834
+ bff .read_gbq_function (
835
835
str (only_arg_type_specified .reference ),
836
836
session = session ,
837
837
)
838
838
with pytest .raises (ValueError ):
839
- rf .read_gbq_function (
839
+ bff .read_gbq_function (
840
840
str (neither_type_specified .reference ),
841
841
session = session ,
842
842
)
@@ -878,13 +878,13 @@ def test_read_gbq_function_respects_python_output_type(
878
878
body = "TO_JSON_STRING([x, x+1, x+2])" ,
879
879
arguments = [arg ],
880
880
return_type = bigquery .StandardSqlDataType (bigquery .StandardSqlTypeNames .STRING ),
881
- description = rf_utils .get_bigframes_metadata (python_output_type = array_type ),
881
+ description = bff_utils .get_bigframes_metadata (python_output_type = array_type ),
882
882
type_ = bigquery .RoutineType .SCALAR_FUNCTION ,
883
883
)
884
884
885
885
# Create the routine in BigQuery and read it back using read_gbq_function.
886
886
bigquery_client .create_routine (sql_routine , exists_ok = True )
887
- func = rf .read_gbq_function (str (sql_routine .reference ), session = session )
887
+ func = bff .read_gbq_function (str (sql_routine .reference ), session = session )
888
888
889
889
# test that the function works as expected
890
890
s = bigframes .series .Series ([1 , 10 , 100 ])
@@ -920,7 +920,7 @@ def test_read_gbq_function_supports_python_output_type_only_for_string_outputs(
920
920
body = "x+1" ,
921
921
arguments = [arg ],
922
922
return_type = bigquery .StandardSqlDataType (bigquery .StandardSqlTypeNames .INT64 ),
923
- description = rf_utils .get_bigframes_metadata (python_output_type = array_type ),
923
+ description = bff_utils .get_bigframes_metadata (python_output_type = array_type ),
924
924
type_ = bigquery .RoutineType .SCALAR_FUNCTION ,
925
925
)
926
926
@@ -933,7 +933,7 @@ def test_read_gbq_function_supports_python_output_type_only_for_string_outputs(
933
933
TypeError ,
934
934
match = "An explicit output_type should be provided only for a BigQuery function with STRING output." ,
935
935
):
936
- rf .read_gbq_function (str (sql_routine .reference ), session = session )
936
+ bff .read_gbq_function (str (sql_routine .reference ), session = session )
937
937
938
938
939
939
@pytest .mark .parametrize (
@@ -959,13 +959,13 @@ def test_read_gbq_function_supported_python_output_type(
959
959
body = "CAST(x AS STRING)" ,
960
960
arguments = [arg ],
961
961
return_type = bigquery .StandardSqlDataType (bigquery .StandardSqlTypeNames .STRING ),
962
- description = rf_utils .get_bigframes_metadata (python_output_type = array_type ),
962
+ description = bff_utils .get_bigframes_metadata (python_output_type = array_type ),
963
963
type_ = bigquery .RoutineType .SCALAR_FUNCTION ,
964
964
)
965
965
966
966
# Create the routine in BigQuery and read it back using read_gbq_function.
967
967
bigquery_client .create_routine (sql_routine , exists_ok = True )
968
- rf .read_gbq_function (str (sql_routine .reference ), session = session )
968
+ bff .read_gbq_function (str (sql_routine .reference ), session = session )
969
969
970
970
971
971
@pytest .mark .flaky (retries = 2 , delay = 120 )
0 commit comments