feat: support array output in remote_function (#1057)

* feat: support array output in `remote_function`

This is feature request to support use cases like creating custom
feature vectors, embeddings etc.

* add multiindex test

* move array type conversion to bigquery module, test multiindex

* add `bigframes.bigquery.json_extract_string_array`, support int and str array outputs

* increase cleanup rate

* update input and output types doc

* support array output in DataFrame.apply

* support read_gbq_function on a remote function created for array output

* fix the json_set after variable renaming

* add tests for output_type in read_gbq_function

* temporarily exclude system 3.9 tests and include 3.10 and 3.11

* Revert "temporarily exclude system 3.9 tests and include 3.10 and 3.11"

This reverts commit 2485aa3.

* add more info in the unexpected exception

* more debug info

* use unique routine name across tests

* Revert "more debug info"

This reverts commit 86fe316.

* Revert "add more info in the unexpected exception"

This reverts commit fe010cb.

* support array output in binary remote function operations

* support array output in nary remote function operations

* preserve array output type in function description to avoid explit output_type in read_gbq_function

* fix one failing read_gbq_function test

* make test parameterization order deterministic

* fix sorting of types for mypy

* remove test parameterization with sorting inside

* include partial ordering mode testing for read_gbq_function

* add remote function array out test in partial ordering mode

* avoid repr-eval for output type serialization/deserialization

* remove unsupported scenarios system tests, use common exception for unsupported

Author: shobsi

bigframes/clients.py

@@ -21,6 +21,7 @@
 from typing import cast, Optional
 
 import google.api_core.exceptions
+import google.api_core.retry
 from google.cloud import bigquery_connection_v1, resourcemanager_v3
 from google.iam.v1 import iam_policy_pb2, policy_pb2
 

bigframes/core/compile/ibis_types.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import textwrap
+import typing
 from typing import Any, cast, Dict, Iterable, Optional, Tuple, Union
 import warnings
 
@@ -22,7 +23,7 @@
 import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes
 import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
 from bigframes_vendored.ibis.expr.datatypes.core import (
-    dtype as python_type_to_bigquery_type,
+    dtype as python_type_to_ibis_type,
 )
 import bigframes_vendored.ibis.expr.types as ibis_types
 import geopandas as gpd  # type: ignore
@@ -472,12 +473,24 @@ class UnsupportedTypeError(ValueError):
     def __init__(self, type_, supported_types):
         self.type = type_
         self.supported_types = supported_types
+        super().__init__(
+            f"'{type_}' is not one of the supported types {supported_types}"
+        )
 
 
 def ibis_type_from_python_type(t: type) -> ibis_dtypes.DataType:
     if t not in bigframes.dtypes.RF_SUPPORTED_IO_PYTHON_TYPES:
         raise UnsupportedTypeError(t, bigframes.dtypes.RF_SUPPORTED_IO_PYTHON_TYPES)
-    return python_type_to_bigquery_type(t)
+    return python_type_to_ibis_type(t)
+
+
+def ibis_array_output_type_from_python_type(t: type) -> ibis_dtypes.DataType:
+    array_of = typing.get_args(t)[0]
+    if array_of not in bigframes.dtypes.RF_SUPPORTED_ARRAY_OUTPUT_PYTHON_TYPES:
+        raise UnsupportedTypeError(
+            array_of, bigframes.dtypes.RF_SUPPORTED_ARRAY_OUTPUT_PYTHON_TYPES
+        )
+    return python_type_to_ibis_type(t)
 
 
 def ibis_type_from_type_kind(tk: bigquery.StandardSqlTypeNames) -> ibis_dtypes.DataType:

bigframes/dataframe.py

@@ -4017,6 +4017,19 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
                     ops.NaryRemoteFunctionOp(func=func), series_list[1:]
                 )
             result_series.name = None
+
+            # if the output is an array, reconstruct it from the json serialized
+            # string form
+            if bigframes.dtypes.is_array_like(func.output_dtype):
+                import bigframes.bigquery as bbq
+
+                result_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
+                    func.output_dtype.pyarrow_dtype.value_type
+                )
+                result_series = bbq.json_extract_string_array(
+                    result_series, value_dtype=result_dtype
+                )
+
             return result_series
 
         # Per-column apply

bigframes/dtypes.py

@@ -701,6 +701,13 @@ def lcd_type_or_throw(dtype1: Dtype, dtype2: Dtype) -> Dtype:
 # https://cloud.google.com/bigquery/docs/remote-functions#limitations
 RF_SUPPORTED_IO_PYTHON_TYPES = {bool, bytes, float, int, str}
 
+# Support array output types in BigQuery DataFrames remote functions even though
+# it is not currently (2024-10-06) supported in BigQuery remote functions.
+# https://cloud.google.com/bigquery/docs/remote-functions#limitations
+# TODO(b/284515241): remove this special handling when BigQuery remote functions
+# support array.
+RF_SUPPORTED_ARRAY_OUTPUT_PYTHON_TYPES = {bool, float, int, str}
+
 RF_SUPPORTED_IO_BIGQUERY_TYPEKINDS = {
     "BOOLEAN",
     "BOOL",

bigframes/functions/_remote_function_client.py

@@ -95,6 +95,7 @@ def create_bq_remote_function(
         endpoint,
         bq_function_name,
         max_batching_rows,
+        metadata,
     ):
         """Create a BigQuery remote function given the artifacts of a user defined
         function and the http endpoint of a corresponding cloud function."""
@@ -120,9 +121,14 @@ def create_bq_remote_function(
             "max_batching_rows": max_batching_rows,
         }
 
+        if metadata:
+            # We are using the description field to store this structured
+            # bigframes specific metadata for the lack of a better option
+            remote_function_options["description"] = metadata
+
         remote_function_options_str = ", ".join(
             [
-                f'{key}="{val}"' if isinstance(val, str) else f"{key}={val}"
+                f"{key}='{val}'" if isinstance(val, str) else f"{key}={val}"
                 for key, val in remote_function_options.items()
                 if val is not None
             ]
@@ -200,14 +206,7 @@ def generate_cloud_function_code(
         package_requirements=None,
         is_row_processor=False,
     ):
-        """Generate the cloud function code for a given user defined function.
-
-        Args:
-            input_types (tuple[str]):
-                Types of the input arguments in BigQuery SQL data type names.
-            output_type (str):
-                Types of the output scalar as a BigQuery SQL data type name.
-        """
+        """Generate the cloud function code for a given user defined function."""
 
         # requirements.txt
         if package_requirements:
@@ -240,14 +239,7 @@ def create_cloud_function(
         memory_mib=1024,
         ingress_settings="all",
     ):
-        """Create a cloud function from the given user defined function.
-
-        Args:
-            input_types (tuple[str]):
-                Types of the input arguments in BigQuery SQL data type names.
-            output_type (str):
-                Types of the output scalar as a BigQuery SQL data type name.
-        """
+        """Create a cloud function from the given user defined function."""
 
         # Build and deploy folder structure containing cloud function
         with tempfile.TemporaryDirectory() as directory:
@@ -394,6 +386,7 @@ def provision_bq_remote_function(
         cloud_function_vpc_connector,
         cloud_function_memory_mib,
         cloud_function_ingress_settings,
+        bq_metadata,
     ):
         """Provision a BigQuery remote function."""
         # Augment user package requirements with any internal package
@@ -473,6 +466,7 @@ def provision_bq_remote_function(
                 cf_endpoint,
                 remote_function_name,
                 max_batching_rows,
+                bq_metadata,
             )
 
             created_new = True

bigframes/functions/_remote_function_session.py

@@ -34,6 +34,7 @@
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes
+import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
 import bigframes_vendored.ibis.expr.operations.udf as ibis_udf
 import cloudpickle
 import google.api_core.exceptions
@@ -167,12 +168,19 @@ def remote_function(
                 `$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`.
 
         Args:
-            input_types (None, type, or sequence(type)):
+            input_types (type or sequence(type), Optional):
                 For scalar user defined function it should be the input type or
-                sequence of input types. For row processing user defined function,
-                type `Series` should be specified.
-            output_type (Optional[type]):
-                Data type of the output in the user defined function.
+                sequence of input types. The supported scalar input types are
+                `bool`, `bytes`, `float`, `int`, `str`. For row processing user
+                defined function (i.e. functions that receive a single input
+                representing a row in form of a Series), type `Series` should be
+                specified.
+            output_type (type, Optional):
+                Data type of the output in the user defined function. If the
+                user defined function returns an array, then `list[type]` should
+                be specified. The supported output types are `bool`, `bytes`,
+                `float`, `int`, `str`, `list[bool]`, `list[float]`, `list[int]`
+                and `list[str]`.
             session (bigframes.Session, Optional):
                 BigQuery DataFrames session to use for getting default project,
                 dataset and BigQuery connection.
@@ -497,6 +505,24 @@ def try_delattr(attr):
             try_delattr("is_row_processor")
             try_delattr("ibis_node")
 
+            # resolve the output type that can be supported in the bigframes,
+            # ibis, BQ remote functions and cloud functions integration
+            ibis_output_type_for_bqrf = ibis_signature.output_type
+            bqrf_metadata = None
+            if isinstance(ibis_signature.output_type, ibis_dtypes.Array):
+                # TODO(b/284515241): remove this special handling to support
+                # array output types once BQ remote functions support ARRAY.
+                # Until then, use json serialized strings at the cloud function
+                # and BQ level, and parse that to the intended output type at
+                # the bigframes level.
+                ibis_output_type_for_bqrf = ibis_dtypes.String()
+                bqrf_metadata = _utils.get_bigframes_metadata(
+                    python_output_type=output_type
+                )
+            bqrf_output_type = third_party_ibis_bqtypes.BigQueryType.from_ibis(
+                ibis_output_type_for_bqrf
+            )
+
             (
                 rf_name,
                 cf_name,
@@ -508,9 +534,7 @@ def try_delattr(attr):
                     for type_ in ibis_signature.input_types
                     if type_ is not None
                 ),
-                output_type=third_party_ibis_bqtypes.BigQueryType.from_ibis(
-                    ibis_signature.output_type
-                ),
+                output_type=bqrf_output_type,
                 reuse=reuse,
                 name=name,
                 package_requirements=packages,
@@ -521,6 +545,7 @@ def try_delattr(attr):
                 cloud_function_vpc_connector=cloud_function_vpc_connector,
                 cloud_function_memory_mib=cloud_function_memory_mib,
                 cloud_function_ingress_settings=cloud_function_ingress_settings,
+                bq_metadata=bqrf_metadata,
             )
 
             # TODO(shobs): Find a better way to support udfs with param named "name".
@@ -541,7 +566,7 @@ def try_delattr(attr):
                 name=rf_name,
                 catalog=dataset_ref.project,
                 database=dataset_ref.dataset_id,
-                signature=(ibis_signature.input_types, ibis_signature.output_type),
+                signature=(ibis_signature.input_types, ibis_output_type_for_bqrf),
             )  # type: ignore
             func.bigframes_cloud_function = (
                 remote_function_client.get_cloud_function_fully_qualified_name(cf_name)

bigframes/functions/_utils.py

@@ -15,6 +15,8 @@
 
 import hashlib
 import inspect
+import json
+import typing
 from typing import cast, List, NamedTuple, Optional, Sequence, Set
 
 import bigframes_vendored.ibis.expr.datatypes.core as ibis_dtypes
@@ -26,6 +28,7 @@
 import pyarrow
 
 import bigframes.core.compile.ibis_types
+import bigframes.dtypes
 
 # Naming convention for the remote function artifacts
 _BIGFRAMES_REMOTE_FUNCTION_PREFIX = "bigframes"
@@ -194,6 +197,7 @@ class IbisSignature(NamedTuple):
     parameter_names: List[str]
     input_types: List[Optional[ibis_dtypes.DataType]]
     output_type: ibis_dtypes.DataType
+    output_type_override: Optional[ibis_dtypes.DataType] = None
 
 
 def ibis_signature_from_python_signature(
@@ -202,13 +206,77 @@ def ibis_signature_from_python_signature(
     output_type: type,
 ) -> IbisSignature:
 
+    ibis_input_types: List[Optional[ibis_dtypes.DataType]] = [
+        bigframes.core.compile.ibis_types.ibis_type_from_python_type(t)
+        for t in input_types
+    ]
+
+    if typing.get_origin(output_type) is list:
+        ibis_output_type = (
+            bigframes.core.compile.ibis_types.ibis_array_output_type_from_python_type(
+                output_type
+            )
+        )
+    else:
+        ibis_output_type = bigframes.core.compile.ibis_types.ibis_type_from_python_type(
+            output_type
+        )
+
     return IbisSignature(
         parameter_names=list(signature.parameters.keys()),
-        input_types=[
-            bigframes.core.compile.ibis_types.ibis_type_from_python_type(t)
-            for t in input_types
-        ],
-        output_type=bigframes.core.compile.ibis_types.ibis_type_from_python_type(
-            output_type
-        ),
+        input_types=ibis_input_types,
+        output_type=ibis_output_type,
     )
+
+
+def get_python_output_type_from_bigframes_metadata(
+    metadata_text: str,
+) -> Optional[type]:
+    try:
+        metadata_dict = json.loads(metadata_text)
+    except (TypeError, json.decoder.JSONDecodeError):
+        return None
+
+    try:
+        output_type = metadata_dict["value"]["python_array_output_type"]
+    except KeyError:
+        return None
+
+    for (
+        python_output_array_type
+    ) in bigframes.dtypes.RF_SUPPORTED_ARRAY_OUTPUT_PYTHON_TYPES:
+        if python_output_array_type.__name__ == output_type:
+            return list[python_output_array_type]  # type: ignore
+
+    return None
+
+
+def get_bigframes_metadata(*, python_output_type: Optional[type] = None) -> str:
+    # Let's keep the actual metadata inside one level of nesting so that in
+    # future we can use a top level key "version" (parallel to "value"), based
+    # on which "value" can be interpreted according to the "version". The
+    # absence of "version" should be interpreted as default version.
+    inner_metadata = {}
+    if typing.get_origin(python_output_type) is list:
+        python_output_array_type = typing.get_args(python_output_type)[0]
+        if (
+            python_output_array_type
+            in bigframes.dtypes.RF_SUPPORTED_ARRAY_OUTPUT_PYTHON_TYPES
+        ):
+            inner_metadata[
+                "python_array_output_type"
+            ] = python_output_array_type.__name__
+
+    metadata = {"value": inner_metadata}
+    metadata_ser = json.dumps(metadata)
+
+    # let's make sure the serialized value is deserializable
+    if (
+        get_python_output_type_from_bigframes_metadata(metadata_ser)
+        != python_output_type
+    ):
+        raise ValueError(
+            f"python_output_type {python_output_type} is not serializable."
+        )
+
+    return metadata_ser