feat: add bigframes.bigquery.to_json (#2078)

Author: chelsea-lin

bigframes/bigquery/__init__.py

@@ -51,6 +51,7 @@
     json_value,
     json_value_array,
     parse_json,
+    to_json,
     to_json_string,
 )
 from bigframes.bigquery._operations.search import create_vector_index, vector_search
@@ -89,6 +90,7 @@
     json_value,
     json_value_array,
     parse_json,
+    to_json,
     to_json_string,
     # search ops
     create_vector_index,

bigframes/bigquery/_operations/json.py

@@ -430,6 +430,40 @@ def json_value_array(
     return input._apply_unary_op(ops.JSONValueArray(json_path=json_path))
 
 
+def to_json(
+    input: series.Series,
+) -> series.Series:
+    """Converts a series with a JSON value to a JSON-formatted STRING value.
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> bpd.options.display.progress_bar = None
+
+        >>> s = bpd.Series([1, 2, 3])
+        >>> bbq.to_json(s)
+        0    1
+        1    2
+        2    3
+        dtype: extension<dbjson<JSONArrowType>>[pyarrow]
+
+        >>> s = bpd.Series([{"int": 1, "str": "pandas"}, {"int": 2, "str": "numpy"}])
+        >>> bbq.to_json(s)
+        0    {"int":1,"str":"pandas"}
+        1     {"int":2,"str":"numpy"}
+        dtype: extension<dbjson<JSONArrowType>>[pyarrow]
+
+    Args:
+        input (bigframes.series.Series):
+            The Series containing JSON or JSON-formatted string values.
+
+    Returns:
+        bigframes.series.Series: A new Series with the JSON value.
+    """
+    return input._apply_unary_op(ops.ToJSON())
+
+
 def to_json_string(
     input: series.Series,
 ) -> series.Series:

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1302,6 +1302,11 @@ def parse_json_op_impl(x: ibis_types.Value, op: ops.ParseJSON):
     return parse_json(json_str=x)
 
 
+@scalar_op_compiler.register_unary_op(ops.ToJSON)
+def to_json_op_impl(json_obj: ibis_types.Value):
+    return to_json(json_obj=json_obj)
+
+
 @scalar_op_compiler.register_unary_op(ops.ToJSONString)
 def to_json_string_op_impl(x: ibis_types.Value):
     return to_json_string(value=x)
@@ -2093,6 +2098,11 @@ def json_extract_string_array(  # type: ignore[empty-body]
     """Extracts a JSON array and converts it to a SQL ARRAY of STRINGs."""
 
 
+@ibis_udf.scalar.builtin(name="to_json")
+def to_json(json_obj) -> ibis_dtypes.JSON:  # type: ignore[empty-body]
+    """Convert to JSON."""
+
+
 @ibis_udf.scalar.builtin(name="to_json_string")
 def to_json_string(value) -> ibis_dtypes.String:  # type: ignore[empty-body]
     """Convert value to JSON-formatted string."""

bigframes/operations/__init__.py

@@ -124,6 +124,7 @@
     JSONValue,
     JSONValueArray,
     ParseJSON,
+    ToJSON,
     ToJSONString,
 )
 from bigframes.operations.numeric_ops import (
@@ -376,6 +377,7 @@
     "JSONValue",
     "JSONValueArray",
     "ParseJSON",
+    "ToJSON",
     "ToJSONString",
     # Bool ops
     "and_op",

bigframes/operations/json_ops.py

@@ -102,6 +102,20 @@ def output_type(self, *input_types):
         return dtypes.JSON_DTYPE
 
 
+@dataclasses.dataclass(frozen=True)
+class ToJSON(base_ops.UnaryOp):
+    name: typing.ClassVar[str] = "to_json"
+
+    def output_type(self, *input_types):
+        input_type = input_types[0]
+        if not dtypes.is_json_encoding_type(input_type):
+            raise TypeError(
+                "The value to be assigned must be a type that can be encoded as JSON."
+                + f"Received type: {input_type}"
+            )
+        return dtypes.JSON_DTYPE
+
+
 @dataclasses.dataclass(frozen=True)
 class ToJSONString(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "to_json_string"

tests/system/small/bigquery/test_json.py

@@ -386,6 +386,31 @@ def test_parse_json_w_invalid_series_type():
         bbq.parse_json(s)
 
 
+def test_to_json_from_int():
+    s = bpd.Series([1, 2, None, 3])
+    actual = bbq.to_json(s)
+    expected = bpd.Series(["1.0", "2.0", "null", "3.0"], dtype=dtypes.JSON_DTYPE)
+    pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
+def test_to_json_from_struct():
+    s = bpd.Series(
+        [
+            {"version": 1, "project": "pandas"},
+            {"version": 2, "project": "numpy"},
+        ]
+    )
+    assert dtypes.is_struct_like(s.dtype)
+
+    actual = bbq.to_json(s)
+    expected = bpd.Series(
+        ['{"project":"pandas","version":1}', '{"project":"numpy","version":2}'],
+        dtype=dtypes.JSON_DTYPE,
+    )
+
+    pd.testing.assert_series_equal(actual.to_pandas(), expected.to_pandas())
+
+
 def test_to_json_string_from_int():
     s = bpd.Series([1, 2, None, 3])
     actual = bbq.to_json_string(s)