feat: allow functions decorated with @bpd.remote_function to execute locally (#704)

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

* fix read_gbq_function

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

Author: tswast

bigframes/core/compile/scalar_op_compiler.py

@@ -856,11 +856,12 @@ def to_timestamp_op_impl(x: ibis_types.Value, op: ops.ToTimestampOp):
 
 @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True)
 def remote_function_op_impl(x: ibis_types.Value, op: ops.RemoteFunctionOp):
-    if not hasattr(op.func, "bigframes_remote_function"):
+    ibis_node = getattr(op.func, "ibis_node", None)
+    if ibis_node is None:
         raise TypeError(
             f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}"
         )
-    x_transformed = op.func(x)
+    x_transformed = ibis_node(x)
     if not op.apply_on_null:
         x_transformed = ibis.case().when(x.isnull(), x).else_(x_transformed).end()
     return x_transformed
@@ -1342,11 +1343,12 @@ def minimum_impl(
 def binary_remote_function_op_impl(
     x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp
 ):
-    if not hasattr(op.func, "bigframes_remote_function"):
+    ibis_node = getattr(op.func, "ibis_node", None)
+    if ibis_node is None:
         raise TypeError(
             f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}"
         )
-    x_transformed = op.func(x, y)
+    x_transformed = ibis_node(x, y)
     return x_transformed
 
 

bigframes/functions/remote_function.py

@@ -1013,11 +1013,11 @@ def remote_function(
 
     bq_connection_manager = None if session is None else session.bqconnectionmanager
 
-    def wrapper(f):
+    def wrapper(func):
         nonlocal input_types, output_type
 
-        if not callable(f):
-            raise TypeError("f must be callable, got {}".format(f))
+        if not callable(func):
+            raise TypeError("f must be callable, got {}".format(func))
 
         if sys.version_info >= (3, 10):
             # Add `eval_str = True` so that deferred annotations are turned into their
@@ -1028,7 +1028,7 @@ def wrapper(f):
             signature_kwargs = {}
 
         signature = inspect.signature(
-            f,
+            func,
             **signature_kwargs,
         )
 
@@ -1089,8 +1089,23 @@ def wrapper(f):
             session=session,  # type: ignore
         )
 
+        # In the unlikely case where the user is trying to re-deploy the same
+        # function, cleanup the attributes we add below, first. This prevents
+        # the pickle from having dependencies that might not otherwise be
+        # present such as ibis or pandas.
+        def try_delattr(attr):
+            try:
+                delattr(func, attr)
+            except AttributeError:
+                pass
+
+        try_delattr("bigframes_cloud_function")
+        try_delattr("bigframes_remote_function")
+        try_delattr("output_dtype")
+        try_delattr("ibis_node")
+
         rf_name, cf_name = remote_function_client.provision_bq_remote_function(
-            f,
+            func,
             ibis_signature.input_types,
             ibis_signature.output_type,
             reuse,
@@ -1105,19 +1120,20 @@ def wrapper(f):
 
         # TODO: Move ibis logic to compiler step
         node = ibis.udf.scalar.builtin(
-            f,
+            func,
             name=rf_name,
             schema=f"{dataset_ref.project}.{dataset_ref.dataset_id}",
             signature=(ibis_signature.input_types, ibis_signature.output_type),
         )
-        node.bigframes_cloud_function = (
+        func.bigframes_cloud_function = (
             remote_function_client.get_cloud_function_fully_qualified_name(cf_name)
         )
-        node.bigframes_remote_function = str(dataset_ref.routine(rf_name))  # type: ignore
-        node.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(
+        func.bigframes_remote_function = str(dataset_ref.routine(rf_name))  # type: ignore
+        func.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(
             ibis_signature.output_type
         )
-        return node
+        func.ibis_node = node
+        return func
 
     return wrapper
 
@@ -1168,19 +1184,23 @@ def read_gbq_function(
 
     # The name "args" conflicts with the Ibis operator, so we use
     # non-standard names for the arguments here.
-    def node(*ignored_args, **ignored_kwargs):
+    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.
 
     # TODO: Move ibis logic to compiler step
-    node.__name__ = routine_ref.routine_id
+    func.__name__ = routine_ref.routine_id
+
     node = ibis.udf.scalar.builtin(
-        node,
+        func,
         name=routine_ref.routine_id,
         schema=f"{routine_ref.project}.{routine_ref.dataset_id}",
         signature=(ibis_signature.input_types, ibis_signature.output_type),
     )
-    node.bigframes_remote_function = str(routine_ref)  # type: ignore
-    node.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(  # type: ignore
+    func.bigframes_remote_function = str(routine_ref)  # type: ignore
+    func.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype(  # type: ignore
         ibis_signature.output_type
     )
-    return node
+    func.ibis_node = node  # type: ignore
+    return func

tests/system/large/test_remote_function.py

@@ -194,6 +194,9 @@ def test_remote_function_stringify_with_ibis(
         def stringify(x):
             return f"I got {x}"
 
+        # Function should work locally.
+        assert stringify(42) == "I got 42"
+
         _, dataset_name, table_name = scalars_table_id.split(".")
         if not ibis_client.dataset:
             ibis_client.dataset = dataset_name
@@ -205,7 +208,7 @@ def stringify(x):
         pandas_df_orig = bigquery_client.query(sql).to_dataframe()
 
         col = table[col_name]
-        col_2x = stringify(col).name("int64_str_col")
+        col_2x = stringify.ibis_node(col).name("int64_str_col")
         table = table.mutate([col_2x])
         sql = table.compile()
         pandas_df_new = bigquery_client.query(sql).to_dataframe()

tests/system/small/test_remote_function.py

@@ -67,7 +67,7 @@ def bq_cf_connection_location_project(bigquery_client) -> str:
 
 @pytest.fixture(scope="module")
 def bq_cf_connection_location_project_mismatched() -> str:
-    """Pre-created BQ connection in the migframes-metrics project in US location,
+    """Pre-created BQ connection in the bigframes-metrics project in US location,
     in format PROJECT_ID.LOCATION.CONNECTION_NAME, used to invoke cloud function.
 
     $ bq show --connection --location=us --project_id=PROJECT_ID bigframes-rf-conn
@@ -108,11 +108,15 @@ def test_remote_function_direct_no_session_param(
         reuse=True,
     )
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
 
-    assert square.bigframes_remote_function
-    assert square.bigframes_cloud_function
+    # Function should still work normally.
+    assert square(2) == 4
+
+    # Function should have extra metadata attached for remote execution.
+    assert hasattr(square, "bigframes_remote_function")
+    assert hasattr(square, "bigframes_cloud_function")
+    assert hasattr(square, "ibis_node")
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -161,8 +165,10 @@ def test_remote_function_direct_no_session_param_location_specified(
         reuse=True,
     )
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square(2) == 4
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -197,7 +203,10 @@ def test_remote_function_direct_no_session_param_location_mismatched(
     dataset_id_permanent,
     bq_cf_connection_location_mismatched,
 ):
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match=re.escape("The location does not match BigQuery connection location:"),
+    ):
 
         @rf.remote_function(
             [int],
@@ -212,7 +221,8 @@ def test_remote_function_direct_no_session_param_location_mismatched(
             reuse=True,
         )
         def square(x):
-            # This executes on a remote function, where coverage isn't tracked.
+            # Not expected to reach this code, as the location of the
+            # connection doesn't match the location of the dataset.
             return x * x  # pragma: NO COVER
 
 
@@ -239,8 +249,10 @@ def test_remote_function_direct_no_session_param_location_project_specified(
         reuse=True,
     )
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square(2) == 4
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -275,7 +287,12 @@ def test_remote_function_direct_no_session_param_project_mismatched(
     dataset_id_permanent,
     bq_cf_connection_location_project_mismatched,
 ):
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "The project_id does not match BigQuery connection gcp_project_id:"
+        ),
+    ):
 
         @rf.remote_function(
             [int],
@@ -290,7 +307,8 @@ def test_remote_function_direct_no_session_param_project_mismatched(
             reuse=True,
         )
         def square(x):
-            # This executes on a remote function, where coverage isn't tracked.
+            # Not expected to reach this code, as the project of the
+            # connection doesn't match the project of the dataset.
             return x * x  # pragma: NO COVER
 
 
@@ -302,8 +320,10 @@ def test_remote_function_direct_session_param(session_with_bq_connection, scalar
         session=session_with_bq_connection,
     )
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square(2) == 4
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -340,8 +360,10 @@ def test_remote_function_via_session_default(session_with_bq_connection, scalars
     # cloud function would be common and quickly reused.
     @session_with_bq_connection.remote_function([int], int)
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square(2) == 4
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -380,8 +402,10 @@ def test_remote_function_via_session_with_overrides(
         reuse=True,
     )
     def square(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square(2) == 4
 
     scalars_df, scalars_pandas_df = scalars_dfs
 
@@ -508,7 +532,7 @@ def test_skip_bq_connection_check(dataset_id_permanent):
 
         @session.remote_function([int], int, dataset=dataset_id_permanent)
         def add_one(x):
-            # This executes on a remote function, where coverage isn't tracked.
+            # Not expected to reach this code, as the connection doesn't exist.
             return x + 1  # pragma: NO COVER
 
 
@@ -546,8 +570,10 @@ def test_read_gbq_function_like_original(
         reuse=True,
     )
     def square1(x):
-        # This executes on a remote function, where coverage isn't tracked.
-        return x * x  # pragma: NO COVER
+        return x * x
+
+    # Function should still work normally.
+    assert square1(2) == 4
 
     square2 = rf.read_gbq_function(
         function_name=square1.bigframes_remote_function,