feat: support routines with ARRAY return type in read_gbq_function (#1412)

shobsi · web-flow · commit 4b60049e8362 · 2025-02-21T15:31:24.000-08:00
* feat: support routines with ARRAY return type in `read_gbq_function`

* remove commented out code, fix test to work for all python versions

* support array output in binary and nary applications of read_gbq_function
diff --git a/bigframes/core/compile/ibis_types.py b/bigframes/core/compile/ibis_types.py
@@ -463,10 +463,19 @@ def ibis_array_output_type_from_python_type(t: type) -> ibis_dtypes.DataType:
     return python_type_to_ibis_type(t)
 
 
-def ibis_type_from_type_kind(tk: bigquery.StandardSqlTypeNames) -> ibis_dtypes.DataType:
+def ibis_type_from_bigquery_type(
+    type_: bigquery.StandardSqlDataType,
+) -> ibis_dtypes.DataType:
     """Convert bq type to ibis. Only to be used for remote functions, does not handle all types."""
-    if tk not in bigframes.dtypes.RF_SUPPORTED_IO_BIGQUERY_TYPEKINDS:
+    if type_.type_kind not in bigframes.dtypes.RF_SUPPORTED_IO_BIGQUERY_TYPEKINDS:
         raise UnsupportedTypeError(
-            tk, bigframes.dtypes.RF_SUPPORTED_IO_BIGQUERY_TYPEKINDS
+            type_.type_kind, bigframes.dtypes.RF_SUPPORTED_IO_BIGQUERY_TYPEKINDS
+        )
+    elif type_.type_kind == "ARRAY":
+        return ibis_dtypes.Array(
+            value_type=ibis_type_from_bigquery_type(
+                typing.cast(bigquery.StandardSqlDataType, type_.array_element_type)
+            )
         )
-    return third_party_ibis_bqtypes.BigQueryType.to_ibis(tk)
+    else:
+        return third_party_ibis_bqtypes.BigQueryType.to_ibis(type_.type_kind)
diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py
@@ -4088,9 +4088,12 @@ def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
                 )
             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):
+            # If the result type is string but the function output is intended
+            # to be an array, reconstruct the array from the string assuming it
+            # is a json serialized form of the array.
+            if bigframes.dtypes.is_string_like(
+                result_series.dtype
+            ) and bigframes.dtypes.is_array_like(func.output_dtype):
                 import bigframes.bigquery as bbq
 
                 result_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py
@@ -874,4 +874,5 @@ def lcd_type_or_throw(dtype1: Dtype, dtype2: Dtype) -> Dtype:
     "INT64",
     "INTEGER",
     "STRING",
+    "ARRAY",
 }
diff --git a/bigframes/functions/_function_session.py b/bigframes/functions/_function_session.py
@@ -501,6 +501,7 @@ def try_delattr(attr):
             try_delattr("bigframes_remote_function")
             try_delattr("input_dtypes")
             try_delattr("output_dtype")
+            try_delattr("bigframes_bigquery_function_output_dtype")
             try_delattr("is_row_processor")
             try_delattr("ibis_node")
 
@@ -589,6 +590,11 @@ def try_delattr(attr):
                     ibis_signature.output_type
                 )
             )
+            func.bigframes_bigquery_function_output_dtype = (
+                bigframes.core.compile.ibis_types.ibis_dtype_to_bigframes_dtype(
+                    ibis_output_type_for_bqrf
+                )
+            )
             func.is_row_processor = is_row_processor
             func.ibis_node = node
 
diff --git a/bigframes/functions/function.py b/bigframes/functions/function.py
@@ -56,8 +56,10 @@ class ReturnTypeMissingError(ValueError):
 # TODO: Move this to compile folder
 def ibis_signature_from_routine(routine: bigquery.Routine) -> _utils.IbisSignature:
     if routine.return_type:
-        ibis_output_type = bigframes.core.compile.ibis_types.ibis_type_from_type_kind(
-            routine.return_type.type_kind
+        ibis_output_type = (
+            bigframes.core.compile.ibis_types.ibis_type_from_bigquery_type(
+                routine.return_type
+            )
         )
     else:
         raise ReturnTypeMissingError
@@ -82,8 +84,8 @@ def ibis_signature_from_routine(routine: bigquery.Routine) -> _utils.IbisSignatu
     return _utils.IbisSignature(
         parameter_names=[arg.name for arg in routine.arguments],
         input_types=[
-            bigframes.core.compile.ibis_types.ibis_type_from_type_kind(
-                arg.data_type.type_kind
+            bigframes.core.compile.ibis_types.ibis_type_from_bigquery_type(
+                arg.data_type
             )
             if arg.data_type
             else None
@@ -233,6 +235,8 @@ def func(*bigframes_args, **bigframes_kwargs):
         else ibis_signature.output_type
     )
 
+    func.bigframes_bigquery_function_output_dtype = bigframes.core.compile.ibis_types.ibis_dtype_to_bigframes_dtype(ibis_signature.output_type)  # type: ignore
+
     func.is_row_processor = is_row_processor  # type: ignore
     func.ibis_node = node  # type: ignore
     return func
diff --git a/bigframes/operations/remote_function_ops.py b/bigframes/operations/remote_function_ops.py
@@ -15,7 +15,6 @@
 import dataclasses
 import typing
 
-from bigframes import dtypes
 from bigframes.operations import base_ops
 
 
@@ -31,17 +30,10 @@ def expensive(self) -> bool:
 
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
-        if hasattr(self.func, "output_dtype"):
-            if dtypes.is_array_like(self.func.output_dtype):
-                # 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 remote function
-                # level, and parse that to the intended output type at the
-                # bigframes level.
-                return dtypes.STRING_DTYPE
-            return self.func.output_dtype
+        if hasattr(self.func, "bigframes_bigquery_function_output_dtype"):
+            return self.func.bigframes_bigquery_function_output_dtype
         else:
-            raise AttributeError("output_dtype not defined")
+            raise AttributeError("bigframes_bigquery_function_output_dtype not defined")
 
 
 @dataclasses.dataclass(frozen=True)
@@ -55,17 +47,10 @@ def expensive(self) -> bool:
 
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
-        if hasattr(self.func, "output_dtype"):
-            if dtypes.is_array_like(self.func.output_dtype):
-                # 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 remote function
-                # level, and parse that to the intended output type at the
-                # bigframes level.
-                return dtypes.STRING_DTYPE
-            return self.func.output_dtype
+        if hasattr(self.func, "bigframes_bigquery_function_output_dtype"):
+            return self.func.bigframes_bigquery_function_output_dtype
         else:
-            raise AttributeError("output_dtype not defined")
+            raise AttributeError("bigframes_bigquery_function_output_dtype not defined")
 
 
 @dataclasses.dataclass(frozen=True)
@@ -79,14 +64,7 @@ def expensive(self) -> bool:
 
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
-        if hasattr(self.func, "output_dtype"):
-            if dtypes.is_array_like(self.func.output_dtype):
-                # 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 remote function
-                # level, and parse that to the intended output type at the
-                # bigframes level.
-                return dtypes.STRING_DTYPE
-            return self.func.output_dtype
+        if hasattr(self.func, "bigframes_bigquery_function_output_dtype"):
+            return self.func.bigframes_bigquery_function_output_dtype
         else:
-            raise AttributeError("output_dtype not defined")
+            raise AttributeError("bigframes_bigquery_function_output_dtype not defined")
diff --git a/bigframes/series.py b/bigframes/series.py
@@ -1545,9 +1545,12 @@ def apply(
             ops.RemoteFunctionOp(func=func, apply_on_null=True)
         )
 
-        # if the output is an array, reconstruct it from the json serialized
-        # string form
-        if bigframes.dtypes.is_array_like(func.output_dtype):
+        # If the result type is string but the function output is intended to
+        # be an array, reconstruct the array from the string assuming it is a
+        # json serialized form of the array.
+        if bigframes.dtypes.is_string_like(
+            result_series.dtype
+        ) and bigframes.dtypes.is_array_like(func.output_dtype):
             import bigframes.bigquery as bbq
 
             result_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
@@ -1585,9 +1588,12 @@ def combine(
             other, ops.BinaryRemoteFunctionOp(func=func)
         )
 
-        # if the output is an array, reconstruct it from the json serialized
-        # string form
-        if bigframes.dtypes.is_array_like(func.output_dtype):
+        # If the result type is string but the function output is intended to
+        # be an array, reconstruct the array from the string assuming it is a
+        # json serialized form of the array.
+        if bigframes.dtypes.is_string_like(
+            result_series.dtype
+        ) and bigframes.dtypes.is_array_like(func.output_dtype):
             import bigframes.bigquery as bbq
 
             result_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
diff --git a/tests/system/conftest.py b/tests/system/conftest.py
@@ -251,6 +251,11 @@ def table_id_unique(dataset_id: str):
     return f"{dataset_id}.{prefixer.create_prefix()}"
 
 
+@pytest.fixture(scope="function")
+def routine_id_unique(dataset_id: str):
+    return f"{dataset_id}.{prefixer.create_prefix()}"
+
+
 @pytest.fixture(scope="session")
 def scalars_schema(bigquery_client: bigquery.Client):
     # TODO(swast): Add missing scalar data types such as BIGNUMERIC.
diff --git a/tests/system/large/functions/test_remote_function.py b/tests/system/large/functions/test_remote_function.py
@@ -2193,6 +2193,10 @@ def foo(x, y, z):
                 )
             )
         )
+        assert (
+            getattr(foo, "bigframes_bigquery_function_output_dtype")
+            == bigframes.dtypes.STRING_DTYPE
+        )
 
         # Fails to apply on dataframe with incompatible number of columns
         with pytest.raises(
diff --git a/tests/system/small/functions/test_remote_function.py b/tests/system/small/functions/test_remote_function.py
diff --git a/tests/unit/functions/test_remote_function.py b/tests/unit/functions/test_remote_function.py

Original file line number	Diff line number	Diff line change
`@@ -874,4 +874,5 @@ def lcd_type_or_throw(dtype1: Dtype, dtype2: Dtype) -> Dtype:`
`874`	`874`	`"INT64",`
`875`	`875`	`"INTEGER",`
`876`	`876`	`"STRING",`
	`877`	`+ "ARRAY",`
`877`	`878`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2193,6 +2193,10 @@ def foo(x, y, z):`
`2193`	`2193`	`)`
`2194`	`2194`	`)`
`2195`	`2195`	`)`
	`2196`	`+ assert (`
	`2197`	`+ getattr(foo, "bigframes_bigquery_function_output_dtype")`
	`2198`	`+ == bigframes.dtypes.STRING_DTYPE`
	`2199`	`+ )`
`2196`	`2200`
`2197`	`2201`	`# Fails to apply on dataframe with incompatible number of columns`
`2198`	`2202`	`with pytest.raises(`