supports more than two pairs, defines as ibis udf

Author: chelsea-lin

bigframes/bigquery/__init__.py

@@ -197,23 +197,18 @@ def json_set(
     # SQLGlot parser does not support the "create_if_missing => true" syntax, so
     # create_if_missing is not currently implemented.
 
-    # 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]}"
+    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, json_value = json_path_value_pair
+        series = series._apply_binary_op(
+            json_value, ops.JSONSet(json_path=json_path), alignment="left"
         )
-
-    json_path, json_value = json_path_value_pairs[0]
-    return series._apply_binary_op(
-        json_value, ops.JSONSet(json_path=json_path), alignment="left"
-    )
+    return series
 
 
 def vector_search(

bigframes/core/compile/scalar_op_compiler.py

@@ -897,11 +897,21 @@ def array_to_string_op_impl(x: ibis_types.Value, op: ops.ArrayToStringOp):
 # 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):
-    return vendored_ibis_ops.JSONSet(
-        x,
-        json_value=y,
-        json_path=op.json_path,
-    ).to_expr()
+    if x.type().is_json():
+        return json_set(
+            json_obj=x,
+            json_path=op.json_path,
+            json_value=y,
+        ).to_expr()
+    else:
+        # Enabling JSON type eliminates the need for less efficient string conversions.
+        return vendored_ibis_ops.ToJsonString(
+            json_set(
+                json_obj=parse_json(x),
+                json_path=op.json_path,
+                json_value=y,
+            )
+        ).to_expr()
 
 
 ### Binary Ops
@@ -1479,3 +1489,15 @@ def float_floor(a: float) -> float:
 def float_ceil(a: float) -> float:
     """Convert string to timestamp."""
     return 0  # pragma: NO COVER
+
+
+@ibis.udf.scalar.builtin(name="parse_json")
+def parse_json(a: str) -> ibis_dtypes.JSON:
+    """Converts a JSON-formatted STRING value to a JSON value."""
+
+
+@ibis.udf.scalar.builtin(name="json_set")
+def json_set(
+    json_obj: ibis_dtypes.JSON, json_path: ibis_dtypes.str, json_value
+) -> ibis_dtypes.JSON:
+    """Produces a new SQL JSON value with the specified JSON data inserted or replaced."""

bigframes/dtypes.py

@@ -245,6 +245,12 @@ def is_json_like(type: ExpressionType) -> bool:
     return type == STRING_DTYPE
 
 
+def is_json_encoding_type(type: ExpressionType) -> bool:
+    # Types can be converted into JSON.
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#json_encodings
+    return type != GEO_DTYPE
+
+
 def is_numeric(type: ExpressionType) -> bool:
     return type in NUMERIC_BIGFRAMES_TYPES_PERMISSIVE
 

bigframes/operations/__init__.py

@@ -714,10 +714,21 @@ class JSONSet(BinaryOp):
     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
+        left_type = input_types[0]
+        right_type = input_types[1]
+        if not dtypes.is_json_like(left_type):
+            raise TypeError(
+                "Input type must be an valid JSON object or JSON-formatted string type."
+                + f" Received type: {left_type}"
+            )
+        if not dtypes.is_json_encoding_type(right_type):
+            raise TypeError(
+                "The value to be assigned must be a type that can be encoded as JSON."
+                + f"Received type: {right_type}"
+            )
+
+        # After JSON type implementation, ONLY return JSON data.
+        return left_type
 
 
 # Ternary Ops

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

@@ -58,13 +58,6 @@ def _array_aggregate(translator, op: vendored_ibis_ops.ArrayAggregate):
     return f"ARRAY_AGG({arg} IGNORE NULLS {order_by_sql})"
 
 
-def _json_set(translator, op: vendored_ibis_ops.JSONSet):
-    arg = translator.translate(op.arg)
-    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 = {
     vendored_ibis_ops.ApproximateMultiQuantile: _approx_quantiles,  # type:ignore
     vendored_ibis_ops.FirstNonNullValue: _first_non_null_value,  # type:ignore
@@ -74,7 +67,6 @@ def _json_set(translator, op: vendored_ibis_ops.JSONSet):
     vendored_ibis_ops.SafeCastToDatetime: _safe_cast_to_datetime,  # type:ignore
     ibis_reductions.Quantile: _quantile,  # type:ignore
     vendored_ibis_ops.ArrayAggregate: _array_aggregate,  # type:ignore
-    vendored_ibis_ops.JSONSet: _json_set,  # type:ignore
 }
 
 OPERATION_REGISTRY.update(patched_ops)

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

@@ -3,16 +3,7 @@
 
 import ibis.expr.datatypes as dt
 import ibis.expr.operations.core as ibis_ops_core
-import ibis.expr.rules as rlz
 
 
 class ToJsonString(ibis_ops_core.Unary):
     dtype = dt.string
-
-
-class JSONSet(ibis_ops_core.Unary):
-    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")