feat: support json_extract_string_array in the bigquery module (#1131)

Author: shobsi

bigframes/bigquery/__init__.py

@@ -25,6 +25,7 @@
 from bigframes.bigquery._operations.json import (
     json_extract,
     json_extract_array,
+    json_extract_string_array,
     json_set,
 )
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
@@ -37,6 +38,7 @@
     "json_set",
     "json_extract",
     "json_extract_array",
+    "json_extract_string_array",
     "approx_top_count",
     "struct",
     "create_vector_index",

bigframes/bigquery/_operations/json.py

@@ -21,14 +21,17 @@
 
 from __future__ import annotations
 
-from typing import Any, Sequence, Tuple
+from typing import Any, cast, Optional, Sequence, Tuple, Union
 
+import bigframes.dtypes
 import bigframes.operations as ops
 import bigframes.series as series
 
+from . import array
+
 
 def json_set(
-    series: series.Series,
+    input: series.Series,
     json_path_value_pairs: Sequence[Tuple[str, Any]],
 ) -> series.Series:
     """Produces a new JSON value within a Series by inserting or replacing values at
@@ -47,7 +50,7 @@ def json_set(
             Name: data, dtype: string
 
     Args:
-        series (bigframes.series.Series):
+        input (bigframes.series.Series):
             The Series containing JSON data (as native JSON objects or JSON-formatted strings).
         json_path_value_pairs (Sequence[Tuple[str, Any]]):
             Pairs of JSON path and the new value to insert/replace.
@@ -59,6 +62,7 @@ def json_set(
     # SQLGlot parser does not support the "create_if_missing => true" syntax, so
     # create_if_missing is not currently implemented.
 
+    result = input
     for json_path_value_pair in json_path_value_pairs:
         if len(json_path_value_pair) != 2:
             raise ValueError(
@@ -67,14 +71,14 @@ def json_set(
             )
 
         json_path, json_value = json_path_value_pair
-        series = series._apply_binary_op(
+        result = result._apply_binary_op(
             json_value, ops.JSONSet(json_path=json_path), alignment="left"
         )
-    return series
+    return result
 
 
 def json_extract(
-    series: series.Series,
+    input: series.Series,
     json_path: str,
 ) -> series.Series:
     """Extracts a JSON value and converts it to a SQL JSON-formatted `STRING` or `JSON`
@@ -93,24 +97,24 @@ def json_extract(
         dtype: string
 
     Args:
-        series (bigframes.series.Series):
+        input (bigframes.series.Series):
             The Series containing JSON data (as native JSON objects or JSON-formatted strings).
         json_path (str):
             The JSON path identifying the data that you want to obtain from the input.
 
     Returns:
         bigframes.series.Series: A new Series with the JSON or JSON-formatted STRING.
     """
-    return series._apply_unary_op(ops.JSONExtract(json_path=json_path))
+    return input._apply_unary_op(ops.JSONExtract(json_path=json_path))
 
 
 def json_extract_array(
-    series: series.Series,
+    input: series.Series,
     json_path: str = "$",
 ) -> series.Series:
-    """Extracts a JSON array and converts it to a SQL array of JSON-formatted `STRING` or `JSON`
-    values. This function uses single quotes and brackets to escape invalid JSONPath
-    characters in JSON keys.
+    """Extracts a JSON array and converts it to a SQL array of JSON-formatted
+    `STRING` or `JSON` values. This function uses single quotes and brackets to
+    escape invalid JSONPath characters in JSON keys.
 
     **Examples:**
 
@@ -124,13 +128,98 @@ def json_extract_array(
         1        ['4' '5']
         dtype: list<item: string>[pyarrow]
 
+        >>> s = bpd.Series([
+        ...   '{"fruits": [{"name": "apple"}, {"name": "cherry"}]}',
+        ...   '{"fruits": [{"name": "guava"}, {"name": "grapes"}]}'
+        ... ])
+        >>> bbq.json_extract_array(s, "$.fruits")
+        0    ['{"name":"apple"}' '{"name":"cherry"}']
+        1    ['{"name":"guava"}' '{"name":"grapes"}']
+        dtype: list<item: string>[pyarrow]
+
+        >>> s = bpd.Series([
+        ...   '{"fruits": {"color": "red",   "names": ["apple","cherry"]}}',
+        ...   '{"fruits": {"color": "green", "names": ["guava", "grapes"]}}'
+        ... ])
+        >>> bbq.json_extract_array(s, "$.fruits.names")
+        0    ['"apple"' '"cherry"']
+        1    ['"guava"' '"grapes"']
+        dtype: list<item: string>[pyarrow]
+
     Args:
-        series (bigframes.series.Series):
+        input (bigframes.series.Series):
             The Series containing JSON data (as native JSON objects or JSON-formatted strings).
         json_path (str):
             The JSON path identifying the data that you want to obtain from the input.
 
     Returns:
-        bigframes.series.Series: A new Series with the JSON or JSON-formatted STRING.
+        bigframes.series.Series: A new Series with the parsed arrays from the input.
     """
-    return series._apply_unary_op(ops.JSONExtractArray(json_path=json_path))
+    return input._apply_unary_op(ops.JSONExtractArray(json_path=json_path))
+
+
+def json_extract_string_array(
+    input: series.Series,
+    json_path: str = "$",
+    value_dtype: Optional[
+        Union[bigframes.dtypes.Dtype, bigframes.dtypes.DtypeString]
+    ] = None,
+) -> series.Series:
+    """Extracts a JSON array and converts it to a SQL array of `STRING` values.
+    A `value_dtype` can be provided to further coerce the data type of the
+    values in the array. This function uses single quotes and brackets to escape
+    invalid JSONPath characters in JSON keys.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> bpd.options.display.progress_bar = None
+
+        >>> s = bpd.Series(['[1, 2, 3]', '[4, 5]'])
+        >>> bbq.json_extract_string_array(s)
+        0    ['1' '2' '3']
+        1        ['4' '5']
+        dtype: list<item: string>[pyarrow]
+
+        >>> bbq.json_extract_string_array(s, value_dtype='Int64')
+        0    [1 2 3]
+        1      [4 5]
+        dtype: list<item: int64>[pyarrow]
+
+        >>> s = bpd.Series([
+        ...   '{"fruits": {"color": "red",   "names": ["apple","cherry"]}}',
+        ...   '{"fruits": {"color": "green", "names": ["guava", "grapes"]}}'
+        ... ])
+        >>> bbq.json_extract_string_array(s, "$.fruits.names")
+        0    ['apple' 'cherry']
+        1    ['guava' 'grapes']
+        dtype: list<item: string>[pyarrow]
+
+    Args:
+        input (bigframes.series.Series):
+            The Series containing JSON data (as native JSON objects or JSON-formatted strings).
+        json_path (str):
+            The JSON path identifying the data that you want to obtain from the input.
+        value_dtype (dtype, Optional):
+            The data type supported by BigFrames DataFrame.
+
+    Returns:
+        bigframes.series.Series: A new Series with the parsed arrays from the input.
+    """
+    array_series = input._apply_unary_op(
+        ops.JSONExtractStringArray(json_path=json_path)
+    )
+    if value_dtype not in [None, bigframes.dtypes.STRING_DTYPE]:
+        array_items_series = array_series.explode()
+        if value_dtype == bigframes.dtypes.BOOL_DTYPE:
+            array_items_series = array_items_series.str.lower() == "true"
+        else:
+            array_items_series = array_items_series.astype(value_dtype)
+        array_series = cast(
+            series.Series,
+            array.array_agg(
+                array_items_series.groupby(level=input.index.names, dropna=False)
+            ),
+        )
+    return array_series

bigframes/core/compile/scalar_op_compiler.py

@@ -1140,6 +1140,13 @@ def json_extract_array_op_impl(x: ibis_types.Value, op: ops.JSONExtractArray):
     return json_extract_array(json_obj=x, json_path=op.json_path)
 
 
+@scalar_op_compiler.register_unary_op(ops.JSONExtractStringArray, pass_op=True)
+def json_extract_string_array_op_impl(
+    x: ibis_types.Value, op: ops.JSONExtractStringArray
+):
+    return json_extract_string_array(json_obj=x, json_path=op.json_path)
+
+
 ### Binary Ops
 def short_circuit_nulls(type_override: typing.Optional[ibis_dtypes.DataType] = None):
     """Wraps a binary operator to generate nulls of the expected type if either input is a null scalar."""
@@ -1801,6 +1808,13 @@ def json_extract_array(
     """Extracts a JSON array and converts it to a SQL ARRAY of JSON-formatted STRINGs or JSON values."""
 
 
+@ibis.udf.scalar.builtin(name="json_extract_string_array")
+def json_extract_string_array(
+    json_obj: ibis_dtypes.JSON, json_path: ibis_dtypes.str
+) -> ibis_dtypes.Array[ibis_dtypes.String]:
+    """Extracts a JSON array and converts it to a SQL ARRAY of STRINGs."""
+
+
 @ibis.udf.scalar.builtin(name="ML.DISTANCE")
 def vector_distance(vector1, vector2, type: str) -> ibis_dtypes.Float64:
     """Computes the distance between two vectors using specified type ("EUCLIDEAN", "MANHATTAN", or "COSINE")"""

bigframes/operations/__init__.py

@@ -25,7 +25,6 @@
 from pandas.tseries.offsets import DateOffset
 import pyarrow as pa
 
-import bigframes.dtypes
 import bigframes.dtypes as dtypes
 import bigframes.operations.type as op_typing
 
@@ -526,6 +525,13 @@ class RemoteFunctionOp(UnaryOp):
     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
         else:
             raise AttributeError("output_dtype not defined")
@@ -548,9 +554,9 @@ class ToDatetimeOp(UnaryOp):
 
     def output_type(self, *input_types):
         if input_types[0] not in (
-            bigframes.dtypes.FLOAT_DTYPE,
-            bigframes.dtypes.INT_DTYPE,
-            bigframes.dtypes.STRING_DTYPE,
+            dtypes.FLOAT_DTYPE,
+            dtypes.INT_DTYPE,
+            dtypes.STRING_DTYPE,
         ):
             raise TypeError("expected string or numeric input")
         return pd.ArrowDtype(pa.timestamp("us", tz=None))
@@ -565,9 +571,9 @@ class ToTimestampOp(UnaryOp):
     def output_type(self, *input_types):
         # Must be numeric or string
         if input_types[0] not in (
-            bigframes.dtypes.FLOAT_DTYPE,
-            bigframes.dtypes.INT_DTYPE,
-            bigframes.dtypes.STRING_DTYPE,
+            dtypes.FLOAT_DTYPE,
+            dtypes.INT_DTYPE,
+            dtypes.STRING_DTYPE,
         ):
             raise TypeError("expected string or numeric input")
         return pd.ArrowDtype(pa.timestamp("us", tz="UTC"))
@@ -699,6 +705,23 @@ def output_type(self, *input_types):
         )
 
 
+@dataclasses.dataclass(frozen=True)
+class JSONExtractStringArray(UnaryOp):
+    name: typing.ClassVar[str] = "json_extract_string_array"
+    json_path: str
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if not dtypes.is_json_like(input_type):
+            raise TypeError(
+                "Input type must be an valid JSON object or JSON-formatted string type."
+                + f" Received type: {input_type}"
+            )
+        return pd.ArrowDtype(
+            pa.list_(dtypes.bigframes_dtype_to_arrow_dtype(dtypes.STRING_DTYPE))
+        )
+
+
 # Binary Ops
 fillna_op = create_binary_op(name="fillna", type_signature=op_typing.COERCE)
 maximum_op = create_binary_op(name="maximum", type_signature=op_typing.COERCE)

tests/system/small/bigquery/test_json.py

@@ -19,6 +19,7 @@
 import pytest
 
 import bigframes.bigquery as bbq
+import bigframes.dtypes
 import bigframes.pandas as bpd
 
 
@@ -142,9 +143,9 @@ def test_json_extract_w_invalid_series_type():
 
 
 def test_json_extract_array_from_json_strings():
-    s = bpd.Series(['{"a": [1, 2, 3]}', '{"a": []}', '{"a": [4,5]}'])
+    s = bpd.Series(['{"a": ["ab", "2", "3 xy"]}', '{"a": []}', '{"a": ["4","5"]}'])
     actual = bbq.json_extract_array(s, "$.a")
-    expected = bpd.Series([["1", "2", "3"], [], ["4", "5"]])
+    expected = bpd.Series([['"ab"', '"2"', '"3 xy"'], [], ['"4"', '"5"']])
     pd.testing.assert_series_equal(
         actual.to_pandas(),
         expected.to_pandas(),
@@ -164,3 +165,38 @@ def test_json_extract_array_from_array_strings():
 def test_json_extract_array_w_invalid_series_type():
     with pytest.raises(TypeError):
         bbq.json_extract_array(bpd.Series([1, 2]))
+
+
+def test_json_extract_string_array_from_json_strings():
+    s = bpd.Series(['{"a": ["ab", "2", "3 xy"]}', '{"a": []}', '{"a": ["4","5"]}'])
+    actual = bbq.json_extract_string_array(s, "$.a")
+    expected = bpd.Series([["ab", "2", "3 xy"], [], ["4", "5"]])
+    pd.testing.assert_series_equal(
+        actual.to_pandas(),
+        expected.to_pandas(),
+    )
+
+
+def test_json_extract_string_array_from_array_strings():
+    s = bpd.Series(["[1, 2, 3]", "[]", "[4,5]"])
+    actual = bbq.json_extract_string_array(s)
+    expected = bpd.Series([["1", "2", "3"], [], ["4", "5"]])
+    pd.testing.assert_series_equal(
+        actual.to_pandas(),
+        expected.to_pandas(),
+    )
+
+
+def test_json_extract_string_array_as_float_array_from_array_strings():
+    s = bpd.Series(["[1, 2.5, 3]", "[]", "[4,5]"])
+    actual = bbq.json_extract_string_array(s, value_dtype=bigframes.dtypes.FLOAT_DTYPE)
+    expected = bpd.Series([[1, 2.5, 3], [], [4, 5]])
+    pd.testing.assert_series_equal(
+        actual.to_pandas(),
+        expected.to_pandas(),
+    )
+
+
+def test_json_extract_string_array_w_invalid_series_type():
+    with pytest.raises(TypeError):
+        bbq.json_extract_string_array(bpd.Series([1, 2]))