diff --git a/bigframes/functions/_function_client.py b/bigframes/functions/_function_client.py index d98633e1cc..a8c9f9c301 100644 --- a/bigframes/functions/_function_client.py +++ b/bigframes/functions/_function_client.py @@ -245,7 +245,7 @@ def provision_bq_managed_function( # Augment user package requirements with any internal package # requirements. - packages = _utils._get_updated_package_requirements( + packages = _utils.get_updated_package_requirements( packages, is_row_processor, capture_references, ignore_package_version=True ) if packages: @@ -258,7 +258,7 @@ def provision_bq_managed_function( bq_function_name = name if not bq_function_name: # Compute a unique hash representing the user code. - function_hash = _utils._get_hash(func, packages) + function_hash = _utils.get_hash(func, packages) bq_function_name = _utils.get_bigframes_function_name( function_hash, session_id, @@ -539,12 +539,12 @@ def provision_bq_remote_function( """Provision a BigQuery remote function.""" # Augment user package requirements with any internal package # requirements - package_requirements = _utils._get_updated_package_requirements( + package_requirements = _utils.get_updated_package_requirements( package_requirements, is_row_processor ) # Compute a unique hash representing the user code - function_hash = _utils._get_hash(def_, package_requirements) + function_hash = _utils.get_hash(def_, package_requirements) # If reuse of any existing function with the same name (indicated by the # same hash of its source code) is not intended, then attach a unique diff --git a/bigframes/functions/_function_session.py b/bigframes/functions/_function_session.py index 64d2f0edfd..29e175d02f 100644 --- a/bigframes/functions/_function_session.py +++ b/bigframes/functions/_function_session.py @@ -597,7 +597,7 @@ def wrapper(func): bqrf_metadata = _utils.get_bigframes_metadata( python_output_type=py_sig.return_annotation ) - post_process_routine = _utils._build_unnest_post_routine( + post_process_routine = _utils.build_unnest_post_routine( py_sig.return_annotation ) py_sig = py_sig.replace(return_annotation=str) diff --git a/bigframes/functions/_utils.py b/bigframes/functions/_utils.py index 8d0ea57c92..37ced54d49 100644 --- a/bigframes/functions/_utils.py +++ b/bigframes/functions/_utils.py @@ -63,7 +63,7 @@ def get_remote_function_locations(bq_location): return bq_location, cloud_function_region -def _get_updated_package_requirements( +def get_updated_package_requirements( package_requirements=None, is_row_processor=False, capture_references=True, @@ -105,7 +105,7 @@ def _get_updated_package_requirements( return requirements -def _clean_up_by_session_id( +def clean_up_by_session_id( bqclient: bigquery.Client, gcfclient: functions_v2.FunctionServiceClient, dataset: bigquery.DatasetReference, @@ -169,7 +169,7 @@ def _clean_up_by_session_id( pass -def _get_hash(def_, package_requirements=None): +def get_hash(def_, package_requirements=None): "Get hash (32 digits alphanumeric) of a function." # There is a known cell-id sensitivity of the cloudpickle serialization in # notebooks https://github.com/cloudpipe/cloudpickle/issues/538. Because of @@ -279,7 +279,7 @@ def get_python_version(is_compat: bool = False) -> str: return f"python{major}{minor}" if is_compat else f"python-{major}.{minor}" -def _build_unnest_post_routine(py_list_type: type[list]): +def build_unnest_post_routine(py_list_type: type[list]): sdk_type = function_typing.sdk_array_output_type_from_python_type(py_list_type) assert sdk_type.array_element_type is not None inner_sdk_type = sdk_type.array_element_type diff --git a/bigframes/functions/function.py b/bigframes/functions/function.py index b695bcd250..a62da57075 100644 --- a/bigframes/functions/function.py +++ b/bigframes/functions/function.py @@ -90,7 +90,7 @@ def _try_import_routine( return BigqueryCallableRoutine( udf_def, session, - post_routine=_utils._build_unnest_post_routine(override_type), + post_routine=_utils.build_unnest_post_routine(override_type), ) return BigqueryCallableRoutine(udf_def, session, is_managed=not is_remote) @@ -107,7 +107,7 @@ def _try_import_row_routine( return BigqueryCallableRowRoutine( udf_def, session, - post_routine=_utils._build_unnest_post_routine(override_type), + post_routine=_utils.build_unnest_post_routine(override_type), ) return BigqueryCallableRowRoutine(udf_def, session, is_managed=not is_remote) diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 76e0f8719b..6ffed5b53f 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -293,7 +293,7 @@ def clean_up_by_session_id( session.bqclient, dataset, session_id ) - bff_utils._clean_up_by_session_id( + bff_utils.clean_up_by_session_id( session.bqclient, session.cloudfunctionsclient, dataset, session_id ) diff --git a/bigframes/testing/utils.py b/bigframes/testing/utils.py index c3a8008465..5da24c5b9b 100644 --- a/bigframes/testing/utils.py +++ b/bigframes/testing/utils.py @@ -440,11 +440,11 @@ def get_function_name(func, package_requirements=None, is_row_processor=False): """Get a bigframes function name for testing given a udf.""" # Augment user package requirements with any internal package # requirements. - package_requirements = bff_utils._get_updated_package_requirements( + package_requirements = bff_utils.get_updated_package_requirements( package_requirements, is_row_processor ) # Compute a unique hash representing the user code. - function_hash = bff_utils._get_hash(func, package_requirements) + function_hash = bff_utils.get_hash(func, package_requirements) return f"bigframes_{function_hash}" diff --git a/tests/system/large/functions/test_remote_function.py b/tests/system/large/functions/test_remote_function.py index 0d6029bd2c..a93435d11a 100644 --- a/tests/system/large/functions/test_remote_function.py +++ b/tests/system/large/functions/test_remote_function.py @@ -527,8 +527,8 @@ def add_one(x): add_one_uniq, add_one_uniq_dir = make_uniq_udf(add_one) # Expected cloud function name for the unique udf - package_requirements = bff_utils._get_updated_package_requirements() - add_one_uniq_hash = bff_utils._get_hash(add_one_uniq, package_requirements) + package_requirements = bff_utils.get_updated_package_requirements() + add_one_uniq_hash = bff_utils.get_hash(add_one_uniq, package_requirements) add_one_uniq_cf_name = bff_utils.get_cloud_function_name( add_one_uniq_hash, session.session_id )