handle post-processing virtualized array type

Author: TrevorBergeron

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1156,6 +1156,13 @@ def array_reduce_op_impl(x: ibis_types.Value, op: ops.ArrayReduceOp):
         )
 
 
+@scalar_op_compiler.register_unary_op(ops.ArrayMapOp, pass_op=True)
+def array_map_op_impl(x: ibis_types.Value, op: ops.ArrayMapOp):
+    return typing.cast(ibis_types.ArrayValue, x).map(
+        lambda arr_vals: scalar_op_compiler.compile_row_op(op.map_op, (arr_vals,))
+    )
+
+
 # JSON Ops
 @scalar_op_compiler.register_binary_op(ops.JSONSet, pass_op=True)
 def json_set_op_impl(x: ibis_types.Value, y: ibis_types.Value, op: ops.JSONSet):

bigframes/core/compile/sqlglot/expressions/array_ops.py

@@ -73,6 +73,28 @@ def _(expr: TypedExpr, op: ops.ArrayReduceOp) -> sge.Expression:
     )
 
 
+@register_unary_op(ops.ArrayMapOp, pass_op=True)
+def _(expr: TypedExpr, op: ops.ArrayMapOp) -> sge.Expression:
+    sub_expr = sg.to_identifier("bf_arr_map_uid")
+    sub_type = dtypes.get_array_inner_type(expr.dtype)
+
+    # TODO: Expression should be provided instead of invoking compiler manually
+    map_expr = expression_compiler.expression_compiler.compile_row_op(
+        op.map_op, (TypedExpr(sub_expr, sub_type),)
+    )
+
+    return sge.array(
+        sge.select(map_expr)
+        .from_(
+            sge.Unnest(
+                expressions=[expr.expr],
+                alias=sge.TableAlias(columns=[sub_expr]),
+            )
+        )
+        .subquery()
+    )
+
+
 @register_unary_op(ops.ArraySliceOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.ArraySliceOp) -> sge.Expression:
     if expr.dtype == dtypes.STRING_DTYPE:

bigframes/functions/_utils.py

@@ -292,20 +292,6 @@ def get_python_version(is_compat: bool = False) -> str:
     return f"python{major}{minor}" if is_compat else f"python-{major}.{minor}"
 
 
-def build_unnest_post_routine(py_list_type: type[list]):
-    sdk_type = function_typing.sdk_array_output_type_from_python_type(py_list_type)
-    assert sdk_type.array_element_type is not None
-    inner_sdk_type = sdk_type.array_element_type
-    result_dtype = function_typing.sdk_type_to_bf_type(inner_sdk_type)
-
-    def post_process(input):
-        import bigframes.bigquery as bbq
-
-        return bbq.json_extract_string_array(input, value_dtype=result_dtype)
-
-    return post_process
-
-
 def has_conflict_input_type(
     signature: inspect.Signature,
     input_types: Sequence[Any],

bigframes/functions/udf_def.py

@@ -158,10 +158,19 @@ def emulating_type(self) -> DirectScalarType:
     def out_expr(
         self, expr: bigframes.core.expression.Expression
     ) -> bigframes.core.expression.Expression:
+        # essentially we are undoing json.dumps in sql
         import bigframes.operations as ops
 
-        # convert json string to array of underlying type
-        return ops.JSONValueArray(json_path="$").as_expr(expr)
+        as_str_list = ops.JSONValueArray(json_path="$").as_expr(expr)
+        if self.inner_dtype.py_type is str:
+            return as_str_list
+        elif self.inner_dtype.py_type is bool:
+            # TODO: hack so we don't need to make ArrayMap support general expressions yet
+            return ops.ArrayMapOp(ops.IsInOp(values=("true",))).as_expr(as_str_list)
+        else:
+            return ops.ArrayMapOp(ops.AsTypeOp(self.inner_dtype.bf_type)).as_expr(
+                as_str_list
+            )
 
     @property
     def sql_type(self) -> str:

bigframes/operations/__init__.py

@@ -25,6 +25,7 @@
 )
 from bigframes.operations.array_ops import (
     ArrayIndexOp,
+    ArrayMapOp,
     ArrayReduceOp,
     ArraySliceOp,
     ArrayToStringOp,
@@ -440,4 +441,5 @@
     "NUMPY_TO_OP",
     "ToArrayOp",
     "ArrayReduceOp",
+    "ArrayMapOp",
 ]

bigframes/operations/array_ops.py

@@ -88,3 +88,17 @@ def output_type(self, *input_types):
         assert dtypes.is_array_like(input_type)
         inner_type = dtypes.get_array_inner_type(input_type)
         return self.aggregation.output_type(inner_type)
+
+
+@dataclasses.dataclass(frozen=True)
+class ArrayMapOp(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "array_map"
+    # TODO: Generalize to chained expressions
+    map_op: base_ops.UnaryOp
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        assert dtypes.is_array_like(input_type)
+        inner_type = dtypes.get_array_inner_type(input_type)
+        out_inner_type = self.map_op.output_type(inner_type)
+        return dtypes.list_type(out_inner_type)