feat: allow functions returned from bpd.read_gbq_function to execute outside of apply (#706)

* feat: allow functions decorated with `@bpd.remote_function` to execute locally

* fix read_gbq_function

* feat: allow functions returned from `bpd.read_gbq_function` to execute outside of `apply`

* fix for rare case where re-deploy exact same function object

* feat: support type annotations to supply input and output types to `@remote_function` decorator

* make tests robust to cloud function listing failures too

* remove unused bigquery_client argument

* add test that function can be called directly

Author: tswast

bigframes/functions/remote_function.py

@@ -1146,21 +1146,14 @@ def try_delattr(attr):
 
 def read_gbq_function(
     function_name: str,
-    session: Optional[Session] = None,
-    bigquery_client: Optional[bigquery.Client] = None,
+    *,
+    session: Session,
 ):
     """
     Read an existing BigQuery function and prepare it for use in future queries.
     """
-
-    # A BigQuery client is required to perform BQ operations
-    if not bigquery_client and session:
-        bigquery_client = session.bqclient
-    if not bigquery_client:
-        raise ValueError(
-            "A bigquery client must be provided, either directly or via session. "
-            f"{constants.FEEDBACK_LINK}"
-        )
+    bigquery_client = session.bqclient
+    ibis_client = session.ibis_client
 
     try:
         routine_ref = get_routine_reference(function_name, bigquery_client, session)
@@ -1192,8 +1185,10 @@ def read_gbq_function(
     # non-standard names for the arguments here.
     def func(*ignored_args, **ignored_kwargs):
         f"""Remote function {str(routine_ref)}."""
-        # TODO(swast): Construct an ibis client from bigquery_client and
-        # execute node via a query.
+        nonlocal node  # type: ignore
+
+        expr = node(*ignored_args, **ignored_kwargs)  # type: ignore
+        return ibis_client.execute(expr)
 
     # TODO: Move ibis logic to compiler step
     func.__name__ = routine_ref.routine_id

tests/system/small/test_remote_function.py

@@ -537,19 +537,20 @@ def add_one(x):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_read_gbq_function_detects_invalid_function(bigquery_client, dataset_id):
+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(
             str(dataset_ref.routine("not_a_function")),
-            bigquery_client=bigquery_client,
+            session=session,
         )
 
     assert "Unknown function" in str(e.value)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
 def test_read_gbq_function_like_original(
+    session,
     bigquery_client,
     bigqueryconnection_client,
     cloudfunctions_client,
@@ -577,7 +578,7 @@ def square1(x):
 
     square2 = rf.read_gbq_function(
         function_name=square1.bigframes_remote_function,
-        bigquery_client=bigquery_client,
+        session=session,
     )
 
     # The newly-created function (square1) should have a remote function AND a
@@ -607,7 +608,14 @@ def square1(x):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id):
+def test_read_gbq_function_runs_existing_udf(session, bigquery_client, dataset_id):
+    func = session.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only")
+    got = func("AURÉLIE")
+    assert got == "aurÉlie"
+
+
+@pytest.mark.flaky(retries=2, delay=120)
+def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id):
     dataset_ref = bigquery.DatasetReference.from_string(dataset_id)
     arg = bigquery.RoutineArgument(
         name="x",
@@ -633,7 +641,8 @@ def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id):
         # 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(
-            str(routine.reference), bigquery_client=bigquery_client
+            str(routine.reference),
+            session=session,
         )
 
         # It should point to the named routine and yield the expected results.
@@ -658,7 +667,9 @@ def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_read_gbq_function_enforces_explicit_types(bigquery_client, dataset_id):
+def test_read_gbq_function_enforces_explicit_types(
+    session, bigquery_client, dataset_id
+):
     dataset_ref = bigquery.DatasetReference.from_string(dataset_id)
     typed_arg = bigquery.RoutineArgument(
         name="x",
@@ -702,18 +713,22 @@ def test_read_gbq_function_enforces_explicit_types(bigquery_client, dataset_id):
     bigquery_client.create_routine(neither_type_specified, exists_ok=True)
 
     rf.read_gbq_function(
-        str(both_types_specified.reference), bigquery_client=bigquery_client
+        str(both_types_specified.reference),
+        session=session,
     )
     rf.read_gbq_function(
-        str(only_return_type_specified.reference), bigquery_client=bigquery_client
+        str(only_return_type_specified.reference),
+        session=session,
     )
     with pytest.raises(ValueError):
         rf.read_gbq_function(
-            str(only_arg_type_specified.reference), bigquery_client=bigquery_client
+            str(only_arg_type_specified.reference),
+            session=session,
         )
     with pytest.raises(ValueError):
         rf.read_gbq_function(
-            str(neither_type_specified.reference), bigquery_client=bigquery_client
+            str(neither_type_specified.reference),
+            session=session,
         )