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,
  )