feat: add bbq.json_query_array and warn bbq.json_extract_array deprecated (#1811)

* feat: add bbq.json_query_array and warn bbq.json_extract_array deprecated * complete features --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: chelsea-lin

bigframes/bigquery/__init__.py

@@ -40,6 +40,7 @@
  json_extract_array,
  json_extract_string_array,
  json_query,
+ json_query_array,
  json_set,
  json_value,
  parse_json,
@@ -67,6 +68,7 @@
  "json_extract_array",
  "json_extract_string_array",
  "json_query",
+ "json_query_array",
  "json_set",
  "json_value",
  "parse_json",

bigframes/bigquery/_operations/json.py

@@ -133,6 +133,10 @@ def json_extract_array(
  `STRING` or `JSON` values. This function uses single quotes and brackets to
  escape invalid JSONPath characters in JSON keys.
 
+ .. deprecated:: 2.5.0
+ The ``json_extract_array`` is deprecated and will be removed in a future version.
+ Use ``json_query_array`` instead.
+
  **Examples:**
 
  >>> import bigframes.pandas as bpd
@@ -172,6 +176,11 @@ def json_extract_array(
  Returns:
  bigframes.series.Series: A new Series with the parsed arrays from the input.
  """
+ msg = (
+ "The `json_extract_array` is deprecated and will be removed in a future version. "
+ "Use `json_query_array` instead."
+ )
+ warnings.warn(bfe.format_message(msg), category=UserWarning)
  return input._apply_unary_op(ops.JSONExtractArray(json_path=json_path))
 
 
@@ -273,6 +282,56 @@ def json_query(
  return input._apply_unary_op(ops.JSONQuery(json_path=json_path))
 
 
+def json_query_array(
+ 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 double quotes to escape invalid
+ JSONPath characters in JSON keys. For example: `"a.b"`.
+
+ **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_query_array(s)
+ 0 ['1' '2' '3']
+ 1 ['4' '5']
+ dtype: list<item: string>[pyarrow]
+
+ >>> s = bpd.Series([
+ ... '{"fruits": [{"name": "apple"}, {"name": "cherry"}]}',
+ ... '{"fruits": [{"name": "guava"}, {"name": "grapes"}]}'
+ ... ])
+ >>> bbq.json_query_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_query_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.
+
+ Returns:
+ bigframes.series.Series: A new Series with the parsed arrays from the input.
+ """
+ return input._apply_unary_op(ops.JSONQueryArray(json_path=json_path))
+
+
 def json_value(
  input: series.Series,
  json_path: str,

bigframes/core/compile/scalar_op_compiler.py

@@ -1379,6 +1379,19 @@ def json_query(json_or_json_string, json_path: ibis_dtypes.str): # type: ignore
  return json_query_op(json_or_json_string=x, json_path=op.json_path)
 
 
+@scalar_op_compiler.register_unary_op(ops.JSONQueryArray, pass_op=True)
+def json_query_array_op_impl(x: ibis_types.Value, op: ops.JSONQueryArray):
+ # Define a user-defined function whose returned type is dynamically matching the input.
+ def json_query_array(json_or_json_string, json_path: ibis_dtypes.str): # type: ignore
+ """Extracts a JSON value and converts it to a SQL JSON-formatted STRING or JSON value."""
+ ...
+
+ return_type = x.type()
+ json_query_array.__annotations__["return"] = ibis_dtypes.Array[return_type] # type: ignore
+ json_query_op = ibis_udf.scalar.builtin(json_query_array)
+ return json_query_op(json_or_json_string=x, json_path=op.json_path)
+
+
 @scalar_op_compiler.register_unary_op(ops.ParseJSON, pass_op=True)
 def parse_json_op_impl(x: ibis_types.Value, op: ops.ParseJSON):
  return parse_json(json_str=x)

bigframes/operations/__init__.py

@@ -109,6 +109,7 @@
  JSONExtractArray,
  JSONExtractStringArray,
  JSONQuery,
+ JSONQueryArray,
  JSONSet,
  JSONValue,
  ParseJSON,
@@ -359,6 +360,7 @@
  "JSONExtractArray",
  "JSONExtractStringArray",
  "JSONQuery",
+ "JSONQueryArray",
  "JSONSet",
  "JSONValue",
  "ParseJSON",

bigframes/operations/json_ops.py

@@ -37,6 +37,23 @@ def output_type(self, *input_types):
  return input_type
 
 
+@dataclasses.dataclass(frozen=True)
+class JSONQueryArray(base_ops.UnaryOp):
+ name: typing.ClassVar[str] = "json_query_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 a valid JSON object or JSON-formatted string type."
+ + f" Received type: {input_type}"
+ )
+ return pd.ArrowDtype(
+ pa.list_(dtypes.bigframes_dtype_to_arrow_dtype(input_type))
+ )
+
+
 @dataclasses.dataclass(frozen=True)
 class JSONExtractArray(base_ops.UnaryOp):
  name: typing.ClassVar[str] = "json_extract_array"

tests/system/small/bigquery/test_json.py

@@ -128,7 +128,8 @@ def test_json_extract_array_from_json():
  ['{"a": ["ab", "2", "3 xy"]}', '{"a": []}', '{"a": ["4", "5"]}', "{}"],
  dtype=dtypes.JSON_DTYPE,
  )
- actual = bbq.json_extract_array(s, "$.a")
+ with pytest.warns(UserWarning, match="The `json_extract_array` is deprecated"):
+ actual = bbq.json_extract_array(s, "$.a")
 
  # This code provides a workaround for issue https://github.com/apache/arrow/issues/45262,
  # which currently prevents constructing a series using the pa.list_(db_types.JSONArrrowType())
@@ -241,6 +242,66 @@ def test_json_query_w_invalid_series_type():
  bbq.json_query(s, "$.a")
 
 
+def test_json_query_array_from_json():
+ s = bpd.Series(
+ ['{"a": ["ab", "2", "3 xy"]}', '{"a": []}', '{"a": ["4", "5"]}', "{}"],
+ dtype=dtypes.JSON_DTYPE,
+ )
+ actual = bbq.json_query_array(s, "$.a")
+
+ # This code provides a workaround for issue https://github.com/apache/arrow/issues/45262,
+ # which currently prevents constructing a series using the pa.list_(db_types.JSONArrrowType())
+ sql = """
+ SELECT 0 AS id, [JSON '"ab"', JSON '"2"', JSON '"3 xy"'] AS data,
+ UNION ALL
+ SELECT 1, [],
+ UNION ALL
+ SELECT 2, [JSON '"4"', JSON '"5"'],
+ UNION ALL
+ SELECT 3, null,
+ """
+ df = bpd.read_gbq(sql).set_index("id").sort_index()
+ expected = df["data"]
+ expected.index.name = None
+ expected.name = None
+
+ pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_query_array_from_json_strings():
+ s = bpd.Series(
+ ['{"a": ["ab", "2", "3 xy"]}', '{"a": []}', '{"a": ["4","5"]}', "{}"],
+ dtype=pd.StringDtype(storage="pyarrow"),
+ )
+ actual = bbq.json_query_array(s, "$.a")
+ expected = bpd.Series(
+ [['"ab"', '"2"', '"3 xy"'], [], ['"4"', '"5"'], None],
+ dtype=pd.ArrowDtype(pa.list_(pa.string())),
+ )
+
+ pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_query_array_from_json_array_strings():
+ s = bpd.Series(
+ ["[1, 2, 3]", "[]", "[4,5]"],
+ dtype=pd.StringDtype(storage="pyarrow"),
+ )
+ actual = bbq.json_query_array(s)
+ expected = bpd.Series(
+ [["1", "2", "3"], [], ["4", "5"]],
+ dtype=pd.ArrowDtype(pa.list_(pa.string())),
+ )
+
+ pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_json_query_array_w_invalid_series_type():
+ s = bpd.Series([1, 2])
+ with pytest.raises(TypeError):
+ bbq.json_query_array(s)
+
+
 def test_json_value_from_json():
  s = bpd.Series(
  ['{"a": {"b": [1, 2]}}', '{"a": {"c": 1}}', '{"a": {"b": 0}}'],