diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index ce57661919..822f1cd64f 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -3951,7 +3951,7 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs): # Early check whether the dataframe dtypes are currently supported # in the remote function # NOTE: Keep in sync with the value converters used in the gcf code - # generated in remote_function_template.py + # generated in function_template.py remote_function_supported_dtypes = ( bigframes.dtypes.INT_DTYPE, bigframes.dtypes.FLOAT_DTYPE, diff --git a/bigframes/functions/_remote_function_client.py b/bigframes/functions/_function_client.py similarity index 99% rename from bigframes/functions/_remote_function_client.py rename to bigframes/functions/_function_client.py index 0d0cc08128..104119a510 100644 --- a/bigframes/functions/_remote_function_client.py +++ b/bigframes/functions/_function_client.py @@ -29,7 +29,7 @@ from bigframes_vendored import constants import requests -import bigframes.functions.remote_function_template +import bigframes.functions.function_template as bff_template if TYPE_CHECKING: from bigframes.session import Session @@ -215,7 +215,7 @@ def generate_cloud_function_code( f.write("\n".join(package_requirements)) # main.py - entry_point = bigframes.functions.remote_function_template.generate_cloud_function_main_code( + entry_point = bff_template.generate_cloud_function_main_code( def_, directory, input_types=input_types, diff --git a/bigframes/functions/_remote_function_session.py b/bigframes/functions/_function_session.py similarity index 99% rename from bigframes/functions/_remote_function_session.py rename to bigframes/functions/_function_session.py index d6b729bf6e..00626a252f 100644 --- a/bigframes/functions/_remote_function_session.py +++ b/bigframes/functions/_function_session.py @@ -52,8 +52,7 @@ import pandas -from . import _remote_function_client as rf_client -from . import _utils +from . import _function_client, _utils class RemoteFunctionSession: @@ -468,7 +467,7 @@ def wrapper(func): signature, input_types, output_type # type: ignore ) - remote_function_client = rf_client.RemoteFunctionClient( + remote_function_client = _function_client.RemoteFunctionClient( dataset_ref.project, cloud_function_region, cloud_functions_client, diff --git a/bigframes/functions/remote_function.py b/bigframes/functions/function.py similarity index 96% rename from bigframes/functions/remote_function.py rename to bigframes/functions/function.py index 533c93e7cb..57df8f9407 100644 --- a/bigframes/functions/remote_function.py +++ b/bigframes/functions/function.py @@ -35,9 +35,9 @@ import bigframes.core.compile.ibis_types import bigframes.dtypes import bigframes.exceptions as bfe -import bigframes.functions.remote_function_template +import bigframes.functions.function_template -from . import _remote_function_session as rf_session +from . import _function_session as bff_session from . import _utils logger = logging.getLogger(__name__) @@ -120,11 +120,11 @@ def get_routine_reference( def remote_function(*args, **kwargs): - remote_function_session = rf_session.RemoteFunctionSession() + remote_function_session = bff_session.RemoteFunctionSession() return remote_function_session.remote_function(*args, **kwargs) -remote_function.__doc__ = rf_session.RemoteFunctionSession.remote_function.__doc__ +remote_function.__doc__ = bff_session.RemoteFunctionSession.remote_function.__doc__ def read_gbq_function( diff --git a/bigframes/functions/remote_function_template.py b/bigframes/functions/function_template.py similarity index 100% rename from bigframes/functions/remote_function_template.py rename to bigframes/functions/function_template.py diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 395b573916..c744d3b945 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -34,7 +34,7 @@ import bigframes.core.tools import bigframes.dataframe import bigframes.enums -import bigframes.functions._utils as functions_utils +import bigframes.functions._utils as bff_utils from bigframes.pandas.io.api import ( from_glob_path, read_csv, @@ -222,7 +222,7 @@ def clean_up_by_session_id( session.bqclient, dataset, session_id ) - functions_utils._clean_up_by_session_id( + bff_utils._clean_up_by_session_id( session.bqclient, session.cloudfunctionsclient, dataset, session_id ) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 1d85967729..02f79a7d99 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -67,8 +67,8 @@ import bigframes.dtypes import bigframes.exceptions import bigframes.exceptions as bfe -import bigframes.functions._remote_function_session as bigframes_rf_session -import bigframes.functions.remote_function as bigframes_rf +import bigframes.functions._function_session as bff_session +import bigframes.functions.function as bff import bigframes.session._io.bigquery as bf_io_bigquery import bigframes.session.clients import bigframes.session.executor @@ -245,7 +245,7 @@ def __init__( ) self._metrics = bigframes.session.metrics.ExecutionMetrics() - self._remote_function_session = bigframes_rf_session.RemoteFunctionSession() + self._function_session = bff_session.RemoteFunctionSession() self._temp_storage_manager = ( bigframes.session.temp_storage.TemporaryGbqStorageManager( self._clients_provider.bqclient, @@ -377,9 +377,9 @@ def close(self): if temp_storage_manager: self._temp_storage_manager.clean_up_tables() - remote_function_session = getattr(self, "_remote_function_session", None) + remote_function_session = getattr(self, "_function_session", None) if remote_function_session: - self._remote_function_session.clean_up( + self._function_session.clean_up( self.bqclient, self.cloudfunctionsclient, self.session_id ) @@ -1380,7 +1380,7 @@ def remote_function( `bigframes_remote_function` - The bigquery remote function capable of calling into `bigframes_cloud_function`. """ - return self._remote_function_session.remote_function( + return self._function_session.remote_function( input_types, output_type, session=self, @@ -1556,7 +1556,7 @@ def read_gbq_function( not including the `bigframes_cloud_function` property. """ - return bigframes_rf.read_gbq_function( + return bff.read_gbq_function( function_name=function_name, session=self, is_row_processor=is_row_processor, diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index d0eb6c1904..f226143b50 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -32,7 +32,7 @@ import bigframes.dataframe import bigframes.dtypes import bigframes.exceptions -import bigframes.functions._utils as functions_utils +import bigframes.functions._utils as bff_utils import bigframes.pandas as bpd import bigframes.series from tests.system.utils import ( @@ -633,11 +633,9 @@ 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 = functions_utils._get_updated_package_requirements() - add_one_uniq_hash = functions_utils._get_hash( - add_one_uniq, package_requirements - ) - add_one_uniq_cf_name = functions_utils.get_cloud_function_name( + 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 ) diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index c3f3890459..0dc8960f62 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -25,8 +25,8 @@ import bigframes import bigframes.dtypes import bigframes.exceptions -from bigframes.functions import _utils as rf_utils -from bigframes.functions import remote_function as rf +from bigframes.functions import _utils as bff_utils +from bigframes.functions import function as bff from tests.system.utils import assert_pandas_df_equal _prefixer = test_utils.prefixer.Prefixer("bigframes", "") @@ -94,12 +94,12 @@ def get_rf_name(func, package_requirements=None, is_row_processor=False): """Get a remote function name for testing given a udf.""" # Augment user package requirements with any internal package # requirements - package_requirements = rf_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 = rf_utils._get_hash(func, package_requirements) + function_hash = bff_utils._get_hash(func, package_requirements) return f"bigframes_{function_hash}" @@ -117,7 +117,7 @@ def test_remote_function_direct_no_session_param( def square(x): return x * x - square = rf.remote_function( + square = bff.remote_function( int, int, bigquery_client=bigquery_client, @@ -176,7 +176,7 @@ def test_remote_function_direct_no_session_param_location_specified( def square(x): return x * x - square = rf.remote_function( + square = bff.remote_function( int, int, bigquery_client=bigquery_client, @@ -235,7 +235,7 @@ def square(x): ValueError, match=re.escape("The location does not match BigQuery connection location:"), ): - rf.remote_function( + bff.remote_function( int, int, bigquery_client=bigquery_client, @@ -263,7 +263,7 @@ def test_remote_function_direct_no_session_param_location_project_specified( def square(x): return x * x - square = rf.remote_function( + square = bff.remote_function( int, int, bigquery_client=bigquery_client, @@ -324,7 +324,7 @@ def square(x): "The project_id does not match BigQuery connection gcp_project_id:" ), ): - rf.remote_function( + bff.remote_function( int, int, bigquery_client=bigquery_client, @@ -346,7 +346,7 @@ def test_remote_function_direct_session_param( def square(x): return x * x - square = rf.remote_function( + square = bff.remote_function( int, int, session=session_with_bq_connection, @@ -636,7 +636,7 @@ def add_one(x): def test_read_gbq_function_detects_invalid_function(session, dataset_id): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) with pytest.raises(ValueError) as e: - rf.read_gbq_function( + bff.read_gbq_function( str(dataset_ref.routine("not_a_function")), session=session, ) @@ -658,7 +658,7 @@ def test_read_gbq_function_like_original( def square1(x): return x * x - square1 = rf.remote_function( + square1 = bff.remote_function( [int], int, bigquery_client=bigquery_client, @@ -674,7 +674,7 @@ def square1(x): # Function should still work normally. assert square1(2) == 4 - square2 = rf.read_gbq_function( + square2 = bff.read_gbq_function( function_name=square1.bigframes_remote_function, # type: ignore session=session, ) @@ -745,7 +745,7 @@ def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id): for routine in (sql_routine, js_routine): # Create the routine in BigQuery and read it back using read_gbq_function. bigquery_client.create_routine(routine, exists_ok=True) - square = rf.read_gbq_function( + square = bff.read_gbq_function( str(routine.reference), session=session, ) @@ -757,7 +757,7 @@ def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id): src = {"x": [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5]} - routine_ref_str = rf_utils.routine_ref_to_string_for_query(routine.reference) + routine_ref_str = bff_utils.routine_ref_to_string_for_query(routine.reference) direct_sql = " UNION ALL ".join( [f"SELECT {x} AS x, {routine_ref_str}({x}) AS y" for x in src["x"]] ) @@ -818,7 +818,7 @@ def test_read_gbq_function_requires_explicit_types( bigquery_client.create_routine(only_arg_type_specified, exists_ok=True) bigquery_client.create_routine(neither_type_specified, exists_ok=True) - rf.read_gbq_function( + bff.read_gbq_function( str(both_types_specified.reference), session=session, ) @@ -826,17 +826,17 @@ def test_read_gbq_function_requires_explicit_types( bigframes.exceptions.UnknownDataTypeWarning, match="missing input data types.*assume default data type", ): - rf.read_gbq_function( + bff.read_gbq_function( str(only_return_type_specified.reference), session=session, ) with pytest.raises(ValueError): - rf.read_gbq_function( + bff.read_gbq_function( str(only_arg_type_specified.reference), session=session, ) with pytest.raises(ValueError): - rf.read_gbq_function( + bff.read_gbq_function( str(neither_type_specified.reference), session=session, ) @@ -878,13 +878,13 @@ def test_read_gbq_function_respects_python_output_type( body="TO_JSON_STRING([x, x+1, x+2])", arguments=[arg], return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.STRING), - description=rf_utils.get_bigframes_metadata(python_output_type=array_type), + description=bff_utils.get_bigframes_metadata(python_output_type=array_type), type_=bigquery.RoutineType.SCALAR_FUNCTION, ) # Create the routine in BigQuery and read it back using read_gbq_function. bigquery_client.create_routine(sql_routine, exists_ok=True) - func = rf.read_gbq_function(str(sql_routine.reference), session=session) + func = bff.read_gbq_function(str(sql_routine.reference), session=session) # test that the function works as expected s = bigframes.series.Series([1, 10, 100]) @@ -920,7 +920,7 @@ def test_read_gbq_function_supports_python_output_type_only_for_string_outputs( body="x+1", arguments=[arg], return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.INT64), - description=rf_utils.get_bigframes_metadata(python_output_type=array_type), + description=bff_utils.get_bigframes_metadata(python_output_type=array_type), type_=bigquery.RoutineType.SCALAR_FUNCTION, ) @@ -933,7 +933,7 @@ def test_read_gbq_function_supports_python_output_type_only_for_string_outputs( TypeError, match="An explicit output_type should be provided only for a BigQuery function with STRING output.", ): - rf.read_gbq_function(str(sql_routine.reference), session=session) + bff.read_gbq_function(str(sql_routine.reference), session=session) @pytest.mark.parametrize( @@ -959,13 +959,13 @@ def test_read_gbq_function_supported_python_output_type( body="CAST(x AS STRING)", arguments=[arg], return_type=bigquery.StandardSqlDataType(bigquery.StandardSqlTypeNames.STRING), - description=rf_utils.get_bigframes_metadata(python_output_type=array_type), + description=bff_utils.get_bigframes_metadata(python_output_type=array_type), type_=bigquery.RoutineType.SCALAR_FUNCTION, ) # Create the routine in BigQuery and read it back using read_gbq_function. bigquery_client.create_routine(sql_routine, exists_ok=True) - rf.read_gbq_function(str(sql_routine.reference), session=session) + bff.read_gbq_function(str(sql_routine.reference), session=session) @pytest.mark.flaky(retries=2, delay=120) diff --git a/tests/system/utils.py b/tests/system/utils.py index 83d0e683bc..7c12c8033a 100644 --- a/tests/system/utils.py +++ b/tests/system/utils.py @@ -26,7 +26,7 @@ import pyarrow as pa # type: ignore import pytest -import bigframes.functions._utils as functions_utils +import bigframes.functions._utils as bff_utils import bigframes.pandas ML_REGRESSION_METRICS = [ @@ -351,7 +351,7 @@ def get_cloud_functions( not name or not name_prefix ), "Either 'name' or 'name_prefix' can be passed but not both." - _, location = functions_utils.get_remote_function_locations(location) + _, location = bff_utils.get_remote_function_locations(location) parent = f"projects/{project}/locations/{location}" request = functions_v2.ListFunctionsRequest(parent=parent) page_result = functions_client.list_functions(request=request) diff --git a/tests/unit/functions/test_remote_function_template.py b/tests/unit/functions/test_function_template.py similarity index 92% rename from tests/unit/functions/test_remote_function_template.py rename to tests/unit/functions/test_function_template.py index 70b033d938..11db01ed9e 100644 --- a/tests/unit/functions/test_remote_function_template.py +++ b/tests/unit/functions/test_function_template.py @@ -20,7 +20,7 @@ import pytest import bigframes.dtypes -import bigframes.functions.remote_function_template as remote_function_template +import bigframes.functions.function_template as bff_template HELLO_WORLD_BASE64_BYTES = b"SGVsbG8sIFdvcmxkIQ==" HELLO_WORLD_BASE64_STR = "SGVsbG8sIFdvcmxkIQ==" @@ -59,7 +59,7 @@ ), ) def test_convert_from_bq_json(type_, json_value, expected): - got = remote_function_template.convert_from_bq_json(type_, json_value) + got = bff_template.convert_from_bq_json(type_, json_value) assert got == expected @@ -76,7 +76,7 @@ def test_convert_from_bq_json(type_, json_value, expected): ], ) def test_convert_from_bq_json_none(type_): - got = remote_function_template.convert_from_bq_json(type_, None) + got = bff_template.convert_from_bq_json(type_, None) assert got is None @@ -113,7 +113,7 @@ def test_convert_from_bq_json_none(type_): ), ) def test_convert_to_bq_json(type_, value, expected): - got = remote_function_template.convert_to_bq_json(type_, value) + got = bff_template.convert_to_bq_json(type_, value) assert got == expected @@ -130,7 +130,7 @@ def test_convert_to_bq_json(type_, value, expected): ], ) def test_convert_to_bq_json_none(type_): - got = remote_function_template.convert_to_bq_json(type_, None) + got = bff_template.convert_to_bq_json(type_, None) assert got is None @@ -176,7 +176,7 @@ def test_convert_to_bq_json_none(type_): ), ) def test_get_pd_series(row_json, expected): - got = remote_function_template.get_pd_series(row_json) + got = bff_template.get_pd_series(row_json) pandas.testing.assert_series_equal(got, expected) diff --git a/tests/unit/polars_session.py b/tests/unit/polars_session.py index dfb1f5bfa6..cffd8ff7ca 100644 --- a/tests/unit/polars_session.py +++ b/tests/unit/polars_session.py @@ -82,7 +82,7 @@ def __init__(self): self._allow_ambiguity = False # type: ignore self._default_index_type = bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64 self._metrics = bigframes.session.metrics.ExecutionMetrics() - self._remote_function_session = None # type: ignore + self._function_session = None # type: ignore self._temp_storage_manager = None # type: ignore self._executor = TestExecutor() self._loader = None # type: ignore diff --git a/tests/unit/test_remote_function.py b/tests/unit/test_remote_function.py index a8c4f2ac2e..413a694680 100644 --- a/tests/unit/test_remote_function.py +++ b/tests/unit/test_remote_function.py @@ -21,7 +21,7 @@ import bigframes.core.compile.ibis_types import bigframes.dtypes -import bigframes.functions.remote_function +import bigframes.functions.function as bff import bigframes.series from tests.unit import resources @@ -42,9 +42,7 @@ def test_series_input_types_to_str(series_type): """Check that is_row_processor=True uses str as the input type to serialize a row.""" session = resources.create_bigquery_session() - remote_function_decorator = bigframes.functions.remote_function.remote_function( - session=session - ) + remote_function_decorator = bff.remote_function(session=session) with pytest.warns( bigframes.exceptions.PreviewWarning, @@ -75,9 +73,7 @@ def test_supported_types_correspond(): def test_missing_input_types(): session = resources.create_bigquery_session() - remote_function_decorator = bigframes.functions.remote_function.remote_function( - session=session - ) + remote_function_decorator = bff.remote_function(session=session) def function_without_parameter_annotations(myparam) -> str: return str(myparam) @@ -93,9 +89,7 @@ def function_without_parameter_annotations(myparam) -> str: def test_missing_output_type(): session = resources.create_bigquery_session() - remote_function_decorator = bigframes.functions.remote_function.remote_function( - session=session - ) + remote_function_decorator = bff.remote_function(session=session) def function_without_return_annotation(myparam: int): return str(myparam)