feat: create deploy_remote_function and deploy_udf functions to immediately deploy functions to BigQuery

bigframes/functions/_function_session.py

@@ -668,6 +668,75 @@ def wrapper(func):
 
         return wrapper
 
+    def deploy_remote_function(
+        self,
+        *,
+        input_types: Union[None, type, Sequence[type]] = None,
+        output_type: Optional[type] = None,
+        session: Optional[Session] = None,
+        bigquery_client: Optional[bigquery.Client] = None,
+        bigquery_connection_client: Optional[
+            bigquery_connection_v1.ConnectionServiceClient
+        ] = None,
+        cloud_functions_client: Optional[functions_v2.FunctionServiceClient] = None,
+        resource_manager_client: Optional[resourcemanager_v3.ProjectsClient] = None,
+        dataset: Optional[str] = None,
+        bigquery_connection: Optional[str] = None,
+        reuse: bool = True,
+        name: Optional[str] = None,
+        packages: Optional[Sequence[str]] = None,
+        cloud_function_service_account: str,
+        cloud_function_kms_key_name: Optional[str] = None,
+        cloud_function_docker_repository: Optional[str] = None,
+        max_batching_rows: Optional[int] = 1000,
+        cloud_function_timeout: Optional[int] = 600,
+        cloud_function_max_instances: Optional[int] = None,
+        cloud_function_vpc_connector: Optional[str] = None,
+        cloud_function_memory_mib: Optional[int] = 1024,
+        cloud_function_ingress_settings: Literal[
+            "all", "internal-only", "internal-and-gclb"
+        ] = "internal-only",
+        cloud_build_service_account: Optional[str] = None,
+    ):
+        """Orchestrates the creation of a BigQuery remote function that deploys immediately.
+
+        This method ensures that the remote function is created and available for
+        use in BigQuery as soon as this call is made. It achieves this by calling
+        the standard :meth:`~bigframes.functions._function_session.FunctionSession.remote_function`
+        method, which currently handles immediate deployment.
+
+        All arguments are passed directly to
+        :meth:`~bigframes.functions._function_session.FunctionSession.remote_function`.
+        Please see its docstring for parameter details.
+
+        Returns:
+            The same result as :meth:`~bigframes.functions._function_session.FunctionSession.remote_function`.
+        """
+        return self.remote_function(  # type: ignore
+            input_types=input_types,
+            output_type=output_type,
+            session=session,
+            bigquery_client=bigquery_client,
+            bigquery_connection_client=bigquery_connection_client,
+            cloud_functions_client=cloud_functions_client,
+            resource_manager_client=resource_manager_client,
+            dataset=dataset,
+            bigquery_connection=bigquery_connection,
+            reuse=reuse,
+            name=name,
+            packages=packages,
+            cloud_function_service_account=cloud_function_service_account,
+            cloud_function_kms_key_name=cloud_function_kms_key_name,
+            cloud_function_docker_repository=cloud_function_docker_repository,
+            max_batching_rows=max_batching_rows,
+            cloud_function_timeout=cloud_function_timeout,
+            cloud_function_max_instances=cloud_function_max_instances,
+            cloud_function_vpc_connector=cloud_function_vpc_connector,
+            cloud_function_memory_mib=cloud_function_memory_mib,
+            cloud_function_ingress_settings=cloud_function_ingress_settings,
+            cloud_build_service_account=cloud_build_service_account,
+        )
+
     def udf(
         self,
         input_types: Union[None, type, Sequence[type]] = None,
@@ -866,6 +935,42 @@ def wrapper(func):
 
         return wrapper
 
+    def deploy_udf(
+        self,
+        input_types: Union[None, type, Sequence[type]] = None,
+        output_type: Optional[type] = None,
+        session: Optional[Session] = None,
+        bigquery_client: Optional[bigquery.Client] = None,
+        dataset: Optional[str] = None,
+        bigquery_connection: Optional[str] = None,
+        name: Optional[str] = None,
+        packages: Optional[Sequence[str]] = None,
+    ):
+        """Orchestrates the creation of a BigQuery UDF that deploys immediately.
+
+        This method ensures that the UDF is created and available for
+        use in BigQuery as soon as this call is made. It achieves this by calling
+        the standard :meth:`~bigframes.functions._function_session.FunctionSession.udf`
+        method, which currently handles immediate deployment.
+
+        All arguments are passed directly to
+        :meth:`~bigframes.functions._function_session.FunctionSession.udf`.
+        Please see its docstring for parameter details.
+
+        Returns:
+            The same result as :meth:`~bigframes.functions._function_session.FunctionSession.udf`.
+        """
+        return self.udf(  # type: ignore
+            input_types=input_types,
+            output_type=output_type,
+            session=session,
+            bigquery_client=bigquery_client,
+            dataset=dataset,
+            bigquery_connection=bigquery_connection,
+            name=name,
+            packages=packages,
+        )
+
 
 def _convert_row_processor_sig(
     signature: inspect.Signature,

bigframes/pandas/__init__.py

@@ -116,6 +116,65 @@ def remote_function(
 remote_function.__doc__ = inspect.getdoc(bigframes.session.Session.remote_function)
 
 
+def deploy_remote_function(
+    # Make sure that the input/output types, and dataset can be used
+    # positionally. This avoids the worst of the breaking change from 1.x to
+    # 2.x while still preventing possible mixups between consecutive str
+    # parameters.
+    input_types: Union[None, type, Sequence[type]] = None,
+    output_type: Optional[type] = None,
+    dataset: Optional[str] = None,
+    *,
+    bigquery_connection: Optional[str] = None,
+    reuse: bool = True,
+    name: Optional[str] = None,
+    packages: Optional[Sequence[str]] = None,
+    cloud_function_service_account: str,
+    cloud_function_kms_key_name: Optional[str] = None,
+    cloud_function_docker_repository: Optional[str] = None,
+    max_batching_rows: Optional[int] = 1000,
+    cloud_function_timeout: Optional[int] = 600,
+    cloud_function_max_instances: Optional[int] = None,
+    cloud_function_vpc_connector: Optional[str] = None,
+    cloud_function_memory_mib: Optional[int] = 1024,
+    cloud_function_ingress_settings: Literal[
+        "all", "internal-only", "internal-and-gclb"
+    ] = "internal-only",
+    cloud_build_service_account: Optional[str] = None,
+):
+    """Deploys a remote function immediately.
+
+    This function takes the same arguments as :func:`~bigframes.pandas.remote_function`.
+    The key difference is that this function ensures the remote function
+    is created and deployed in BigQuery at the time of this call.
+
+    For detailed argument descriptions, see :func:`~bigframes.pandas.remote_function`.
+    """
+    return global_session.with_default_session(
+        bigframes.session.Session.deploy_remote_function,
+        input_types=input_types,
+        output_type=output_type,
+        dataset=dataset,
+        bigquery_connection=bigquery_connection,
+        reuse=reuse,
+        name=name,
+        packages=packages,
+        cloud_function_service_account=cloud_function_service_account,
+        cloud_function_kms_key_name=cloud_function_kms_key_name,
+        cloud_function_docker_repository=cloud_function_docker_repository,
+        max_batching_rows=max_batching_rows,
+        cloud_function_timeout=cloud_function_timeout,
+        cloud_function_max_instances=cloud_function_max_instances,
+        cloud_function_vpc_connector=cloud_function_vpc_connector,
+        cloud_function_memory_mib=cloud_function_memory_mib,
+        cloud_function_ingress_settings=cloud_function_ingress_settings,
+        cloud_build_service_account=cloud_build_service_account,
+    )
+
+
+# deploy_remote_function.__doc__ is now set directly in the SEARCH block.
+
+
 def udf(
     *,
     input_types: Union[None, type, Sequence[type]] = None,
@@ -139,6 +198,37 @@ def udf(
 udf.__doc__ = inspect.getdoc(bigframes.session.Session.udf)
 
 
+def deploy_udf(
+    *,
+    input_types: Union[None, type, Sequence[type]] = None,
+    output_type: Optional[type] = None,
+    dataset: str,
+    bigquery_connection: Optional[str] = None,
+    name: str,
+    packages: Optional[Sequence[str]] = None,
+):
+    """Deploys a UDF immediately.
+
+    This function takes the same arguments as :func:`~bigframes.pandas.udf`.
+    The key difference is that this function ensures the UDF
+    is created and deployed in BigQuery at the time of this call.
+
+    For detailed argument descriptions, see :func:`~bigframes.pandas.udf`.
+    """
+    return global_session.with_default_session(
+        bigframes.session.Session.deploy_udf,
+        input_types=input_types,
+        output_type=output_type,
+        dataset=dataset,
+        bigquery_connection=bigquery_connection,
+        name=name,
+        packages=packages,
+    )
+
+
+# deploy_udf.__doc__ is now set directly in the SEARCH block.
+
+
 @typing.overload
 def to_datetime(
     arg: Union[
@@ -329,6 +419,8 @@ def reset_session():
     clean_up_by_session_id,
     concat,
     cut,
+    deploy_remote_function,
+    deploy_udf,
     get_default_session_id,
     get_dummies,
     merge,

tests/unit/functions/test_remote_function.py

@@ -89,3 +89,165 @@ def function_without_return_annotation(myparam: int):
         match="'output_type' was not set .* missing a return type annotation",
     ):
         remote_function_decorator(function_without_return_annotation)
+
+
+# --- Tests for bpd.deploy_remote_function ---
+@mock.patch("bigframes.session.Session.deploy_remote_function")
+def test_bpd_deploy_remote_function_calls_session_deploy_method(mock_session_deploy_rf_method):
+    mock_pandas_session = mocks.create_bigquery_session()
+    # bpd.deploy_remote_function uses global_session.with_default_session,
+    # which will call the 'deploy_remote_function' method on the provided/global session instance.
+    # The @mock.patch above targets the class method on Session.
+
+    @bpd.deploy_remote_function(session=mock_pandas_session, cloud_function_service_account="test_sa@example.com")
+    def my_remote_func(x: int) -> int:
+        return x * 2
+
+    mock_session_deploy_rf_method.assert_called_once()
+    # Check some key args passed to the session's deploy_remote_function method
+    args, kwargs = mock_session_deploy_rf_method.call_args
+    # The first arg to Session.deploy_remote_function is the actual python function
+    assert callable(args[0])
+    assert args[0].__name__ == "my_remote_func"
+    # Other args are passed as kwargs by with_default_session
+    assert kwargs.get("cloud_function_service_account") == "test_sa@example.com"
+    assert kwargs.get("reuse") is True  # Default reuse is True
+    assert kwargs.get("name") is None # Default name is None
+
+    # Test that the function is still callable locally
+    assert my_remote_func(10) == 20
+
+
+@mock.patch("bigframes.session.Session.deploy_remote_function")
+def test_bpd_deploy_remote_function_no_reuse(mock_session_deploy_rf_method):
+    mock_pandas_session = mocks.create_bigquery_session()
+
+    @bpd.deploy_remote_function(session=mock_pandas_session, cloud_function_service_account="test_sa@example.com", reuse=False)
+    def my_remote_func_no_reuse(x: int) -> int:
+        return x * 3
+
+    mock_session_deploy_rf_method.assert_called_once()
+    args, kwargs = mock_session_deploy_rf_method.call_args
+    assert kwargs.get("reuse") is False
+
+    assert my_remote_func_no_reuse(5) == 15
+
+
+@mock.patch("bigframes.session.Session.deploy_remote_function")
+def test_bpd_deploy_remote_function_with_name(mock_session_deploy_rf_method):
+    mock_pandas_session = mocks.create_bigquery_session()
+
+    @bpd.deploy_remote_function(
+        session=mock_pandas_session,
+        cloud_function_service_account="test_sa@example.com",
+        name="custom_name"
+    )
+    def my_named_remote_func(x: int) -> int:
+        return x * 4
+
+    mock_session_deploy_rf_method.assert_called_once()
+    args, kwargs = mock_session_deploy_rf_method.call_args
+    assert kwargs.get("name") == "custom_name"
+
+    assert my_named_remote_func(3) == 12
+
+
+# --- Tests for bpd.remote_function ---
+@mock.patch("bigframes.session.Session.remote_function")
+def test_bpd_remote_function_calls_session_method(mock_session_rf_method):
+    mock_pandas_session = mocks.create_bigquery_session()
+
+    @bpd.remote_function(session=mock_pandas_session, cloud_function_service_account="test_sa@example.com")
+    def my_std_remote_func(x: int) -> int:
+        return x * 5
+
+    mock_session_rf_method.assert_called_once()
+    args, kwargs = mock_session_rf_method.call_args
+    assert callable(args[0])
+    assert args[0].__name__ == "my_std_remote_func"
+    assert kwargs.get("cloud_function_service_account") == "test_sa@example.com"
+    # No deploy_immediately flag to check here by design
+
+    assert my_std_remote_func(6) == 30
+
+
+# --- Tests for bpd.deploy_udf ---
+@mock.patch("bigframes.session.Session.deploy_udf")
+def test_bpd_deploy_udf_calls_session_deploy_method(mock_session_deploy_udf_method):
+    mock_pandas_session = mocks.create_bigquery_session(
+        default_project="test-project",
+        default_location="us-central1"
+    )
+
+    @bpd.deploy_udf(session=mock_pandas_session, dataset="my_dataset", name="my_udf_1")
+    def my_udf(y: str) -> str:
+        return f"hello {y}"
+
+    mock_session_deploy_udf_method.assert_called_once()
+    args, kwargs = mock_session_deploy_udf_method.call_args
+    assert callable(args[0])
+    assert args[0].__name__ == "my_udf"
+    assert kwargs.get("name") == "my_udf_1"
+    assert kwargs.get("dataset") == "my_dataset"
+
+    assert my_udf("world") == "hello world"
+
+
+@mock.patch("bigframes.session.Session.deploy_udf")
+def test_bpd_deploy_udf_no_name(mock_session_deploy_udf_method):
+    mock_pandas_session = mocks.create_bigquery_session(
+        default_project="test-project",
+        default_location="us-central1"
+    )
+
+    @bpd.deploy_udf(session=mock_pandas_session, dataset="my_dataset") # No explicit name
+    def my_anon_udf(val: float) -> float:
+        return val + 1.0
+
+    mock_session_deploy_udf_method.assert_called_once()
+    args, kwargs = mock_session_deploy_udf_method.call_args
+    assert kwargs.get("name") is None # Name should be None, letting system generate
+
+    assert my_anon_udf(1.5) == 2.5
+
+
+@mock.patch("bigframes.session.Session.deploy_udf")
+def test_bpd_deploy_udf_with_default_dataset(mock_session_deploy_udf_method):
+    mock_pandas_session = mocks.create_bigquery_session(
+        default_project="test-project",
+        default_location="us-central1"
+    )
+
+    @bpd.deploy_udf(session=mock_pandas_session, name="my_udf_2") # dataset should come from session
+    def my_udf_default_ds(z: bytes) -> bytes:
+        return z + b" extra"
+
+    mock_session_deploy_udf_method.assert_called_once()
+    args, kwargs = mock_session_deploy_udf_method.call_args
+    assert kwargs.get("name") == "my_udf_2"
+    assert kwargs.get("dataset") is None # bpd wrapper passes None if not specified by user
+
+    assert my_udf_default_ds(b"data") == b"data extra"
+
+
+# --- Tests for bpd.udf ---
+@mock.patch("bigframes.session.Session.udf")
+def test_bpd_udf_calls_session_method(mock_session_udf_method):
+    mock_pandas_session = mocks.create_bigquery_session(
+        default_project="test-project",
+        default_location="us-central1"
+    )
+
+    @bpd.udf(session=mock_pandas_session, dataset="my_dataset_std", name="my_std_udf")
+    def my_std_udf(y: str) -> str:
+        return f"standard hello {y}"
+
+    mock_session_udf_method.assert_called_once()
+    args, kwargs = mock_session_udf_method.call_args
+    assert callable(args[0])
+    assert args[0].__name__ == "my_std_udf"
+    assert kwargs.get("name") == "my_std_udf"
+    assert kwargs.get("dataset") == "my_dataset_std"
+    # No deploy_immediately flag to check here by design
+
+    assert my_std_udf("dev") == "standard hello dev"