update to binary operator for series supports

Author: chelsea-lin

bigframes/bigquery/__init__.py

@@ -177,8 +177,11 @@ def json_set(
         >>> bpd.options.display.progress_bar = None
 
         >>> s = bpd.read_gbq("SELECT JSON '{\\\"a\\\": 1}' AS data")["data"]
-        >>> bbq.json_set(s, json_path_value_pairs=[("$.a", 100), ("$.b", "hi")])
-            0    {"a":100,"b":"hi"}
+        >>> bbq.json_set(s, json_path_value_pairs=[("$.a", 100)])
+            0    {"a":100}
+            Name: data, dtype: string
+        >>> bbq.json_set(s, json_path_value_pairs=[("$.b", "hi")])
+            0    {"a":1,"b":"hi"}
             Name: data, dtype: string
 
     Args:
@@ -193,19 +196,23 @@ def json_set(
     """
     # SQLGlot parser does not support the "create_if_missing => true" syntax, so
     # create_if_missing is not currently implemented.
-    json_path_value_tuples = []
-    for json_path_value_pair in json_path_value_pairs:
-        if len(json_path_value_pair) != 2:
-            raise ValueError(
-                "Incorrect format: Expected (<json_path>, <json_value>), but found: "
-                + f"{json_path_value_pair}"
-            )
-        json_path_value_tuples.append(tuple(json_path_value_pair))
-
-    return series._apply_unary_op(
-        ops.JSONSet(
-            json_path_value_pairs=tuple(json_path_value_tuples),
+
+    # Currently limited to single JSON path/value pairs (binary operations only).
+    if len(json_path_value_pairs) != 1:
+        raise ValueError(
+            "Expected exactly one JSON path and value pair but found "
+            + f"{len(json_path_value_pairs)} pairs."
         )
+
+    if len(json_path_value_pairs[0]) != 2:
+        raise ValueError(
+            "Incorrect format: Expected (<json_path>, <json_value>), but found: "
+            + f"{json_path_value_pairs[0]}"
+        )
+
+    json_path, json_value = json_path_value_pairs[0]
+    return series._apply_binary_op(
+        json_value, ops.JSONSet(json_path=json_path), alignment="left"
     )
 
 

bigframes/core/compile/scalar_op_compiler.py

@@ -895,11 +895,12 @@ def array_to_string_op_impl(x: ibis_types.Value, op: ops.ArrayToStringOp):
 
 
 # JSON Ops
-@scalar_op_compiler.register_unary_op(ops.JSONSet, pass_op=True)
-def json_set_op_impl(x: ibis_types.Value, op: ops.JSONSet):
+@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):
     return vendored_ibis_ops.JSONSet(
         x,
-        json_path_value_pairs=op.json_path_value_pairs,
+        json_value=y,
+        json_path=op.json_path,
     ).to_expr()
 
 

bigframes/operations/__init__.py

@@ -602,19 +602,6 @@ def output_type(self, *input_types):
         return dtypes.STRING_DTYPE
 
 
-## JSON Ops
-@dataclasses.dataclass(frozen=True)
-class JSONSet(UnaryOp):
-    name: typing.ClassVar[str] = "json_set"
-    json_path_value_pairs: typing.Tuple[typing.Tuple[str, typing.Any], ...]
-
-    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 JSON or JSON-formatted string type.")
-        return input_type
-
-
 # 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)
@@ -720,6 +707,19 @@ def output_type(self, *input_types):
 strconcat_op = StrConcatOp()
 
 
+## JSON Ops
+@dataclasses.dataclass(frozen=True)
+class JSONSet(BinaryOp):
+    name: typing.ClassVar[str] = "json_set"
+    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 JSON or JSON-formatted string type.")
+        return input_type
+
+
 # Ternary Ops
 @dataclasses.dataclass(frozen=True)
 class WhereOp(TernaryOp):

tests/system/small/bigquery/test_json.py

@@ -15,6 +15,7 @@
 import json
 
 import pandas as pd
+import pytest
 
 import bigframes.bigquery as bbq
 import bigframes.pandas as bpd
@@ -27,54 +28,59 @@ def _get_series_from_json(json_data):
     return bpd.read_gbq(sql)["data"]
 
 
-def test_json_set():
-    init_json = [
-        {"a": 1},
-    ]
-    s = _get_series_from_json(init_json)
-    actual = bbq.json_set(s, json_path_value_pairs=[("$.a", 10)])
+@pytest.mark.parametrize(
+    ("json_path", "expected_json"),
+    [
+        pytest.param("$.a", [{"a": 10}], id="simple"),
+        pytest.param("$.a.b.c", [{"a": {"b": {"c": 10, "d": []}}}], id="nested"),
+    ],
+)
+def test_json_set_at_json_path(json_path, expected_json):
+    s = _get_series_from_json([{"a": {"b": {"c": "tester", "d": []}}}])
+    actual = bbq.json_set(s, json_path_value_pairs=[(json_path, 10)])
 
-    expected_json = [
-        {"a": 10},
-    ]
     expected = _get_series_from_json(expected_json)
     pd.testing.assert_series_equal(
         actual.to_pandas(),
         expected.to_pandas(),
     )
 
 
-def test_json_set_w_nested_json():
-    init_json = [
-        {"a": {"b": {"c": "tester", "d": []}}},
-    ]
-    s = _get_series_from_json(init_json)
-    actual = bbq.json_set(s, json_path_value_pairs=[("$.a.b.c", "user")])
+@pytest.mark.parametrize(
+    ("json_value", "expected_json"),
+    [
+        pytest.param(10, [{"a": {"b": 10}}, {"a": {"b": 10}}], id="int"),
+        pytest.param(0.333, [{"a": {"b": 0.333}}, {"a": {"b": 0.333}}], id="float"),
+        pytest.param("eng", [{"a": {"b": "eng"}}, {"a": {"b": "eng"}}], id="string"),
+        pytest.param([1, 1], [{"a": {"b": 1}}, {"a": {"b": 1}}], id="series"),
+    ],
+)
+def test_json_set_at_json_value_type(json_value, expected_json):
+    s = _get_series_from_json([{"a": {"b": "dev"}}, {"a": {"b": [1, 2]}}])
+    actual = bbq.json_set(s, json_path_value_pairs=[("$.a.b", json_value)])
 
-    expected_json = [
-        {"a": {"b": {"c": "user", "d": []}}},
-    ]
     expected = _get_series_from_json(expected_json)
     pd.testing.assert_series_equal(
         actual.to_pandas(),
         expected.to_pandas(),
     )
 
 
-def test_json_set_w_ordered_pairs():
-    init_json: object = [
-        {"a": {"b": {"c": {}}}},
-    ]
-    s = _get_series_from_json(init_json)
-    actual = bbq.json_set(
-        s, json_path_value_pairs=[("$.a.b.e", "user"), ("$.a.b.e", "dev")]
-    )
-
-    expected_json: object = [
-        {"a": {"b": {"c": {}, "e": "dev"}}},
-    ]
-    expected = _get_series_from_json(expected_json)
-    pd.testing.assert_series_equal(
-        actual.to_pandas(),
-        expected.to_pandas(),
-    )
+@pytest.mark.parametrize(
+    ("json_path_value_pairs"),
+    [
+        pytest.param(
+            [("$.a", 1), ("$.b", 2)],
+            id="two_pairs",
+            marks=pytest.mark.xfail(raises=ValueError),
+        ),
+        pytest.param(
+            [("$.a", 1, 100)],
+            id="invalid_pair",
+            marks=pytest.mark.xfail(raises=ValueError),
+        ),
+    ],
+)
+def test_json_set_w_invalid_param(json_path_value_pairs):
+    s = _get_series_from_json([{"a": {"b": {"c": {}, "e": "dev"}}}])
+    bbq.json_set(s, json_path_value_pairs=json_path_value_pairs)

third_party/bigframes_vendored/ibis/backends/bigquery/registry.py

@@ -60,10 +60,9 @@ def _array_aggregate(translator, op: vendored_ibis_ops.ArrayAggregate):
 
 def _json_set(translator, op: vendored_ibis_ops.JSONSet):
     arg = translator.translate(op.arg)
-    json_path_value_pairs_list = [
-        translator.translate(item) for pair in op.json_path_value_pairs for item in pair
-    ]
-    return f"JSON_SET(PARSE_JSON({arg}), {', '.join(json_path_value_pairs_list)})"
+    json_value = translator.translate(op.json_value)
+    json_path = translator.translate(op.json_path)
+    return f"JSON_SET(PARSE_JSON({arg}), {json_path}, {json_value})"
 
 
 patched_ops = {

third_party/bigframes_vendored/ibis/expr/operations/json.py

@@ -1,7 +1,6 @@
 # Contains code from https://github.com/ibis-project/ibis/blob/master/ibis/expr/operations/json.py
 from __future__ import annotations
 
-import ibis.common.typing as ibis_typing
 import ibis.expr.datatypes as dt
 import ibis.expr.operations.core as ibis_ops_core
 import ibis.expr.rules as rlz
@@ -12,9 +11,8 @@ class ToJsonString(ibis_ops_core.Unary):
 
 
 class JSONSet(ibis_ops_core.Unary):
-    json_path_value_pairs: ibis_typing.VarTuple[
-        ibis_typing.VarTuple[ibis_ops_core.Value[dt.Any]]
-    ]
+    json_value: ibis_ops_core.Value[dt.Any]
+    json_path: ibis_ops_core.Value[dt.String]
 
     shape = rlz.shape_like("arg")
     dtype = rlz.dtype_like("arg")