fix: json in struct destination type (#1187)

* fix: json in struct destination type

* patch output json_str to json

* fix type hint

* fix type hint

* remove draft

* use pd.ArrowDtype(pa.large_string) as bf dtype

* wording

* fix tests

* fix tests

* remove json conversions

* fix remote func

Author: GarrettWu

bigframes/bigquery/_operations/json.py

@@ -47,7 +47,7 @@ def json_set(
         >>> 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"}
-            Name: data, dtype: string
+            Name: data, dtype: large_string[pyarrow]
 
     Args:
         input (bigframes.series.Series):

bigframes/core/__init__.py

@@ -107,7 +107,7 @@ def from_table(
             raise ValueError("must set at most one of 'offests', 'primary_key'")
         if any(i.field_type == "JSON" for i in table.schema if i.name in schema.names):
             warnings.warn(
-                "Interpreting JSON column(s) as StringDtype. This behavior may change in future versions.",
+                "Interpreting JSON column(s) as StringDtype and pyarrow.large_string. This behavior may change in future versions.",
                 bigframes.exceptions.PreviewWarning,
             )
         # define data source only for needed columns, this makes row-hashing cheaper

bigframes/core/blocks.py

@@ -2564,13 +2564,13 @@ def _get_rows_as_json_values(self) -> Block:
 ),
 T1 AS (
     SELECT *,
-           JSON_OBJECT(
+           TO_JSON_STRING(JSON_OBJECT(
                "names", [{column_names_csv}],
                "types", [{column_types_csv}],
                "values", [{column_references_csv}],
                "indexlength", {index_columns_count},
                "dtype", {pandas_row_dtype}
-           ) AS {googlesql.identifier(row_json_column_name)} FROM T0
+           )) AS {googlesql.identifier(row_json_column_name)} FROM T0
 )
 SELECT {select_columns_csv} FROM T1
 """

bigframes/core/compile/compiled.py

@@ -464,7 +464,7 @@ def _set_or_replace_by_id(
             builder.columns = [*self.columns, new_value.name(id)]
         return builder.build()
 
-    def _select(self, values: typing.Tuple[ibis_types.Value]) -> UnorderedIR:
+    def _select(self, values: typing.Tuple[ibis_types.Value, ...]) -> UnorderedIR:
         builder = self.builder()
         builder.columns = values
         return builder.build()
@@ -1129,7 +1129,7 @@ def _set_or_replace_by_id(self, id: str, new_value: ibis_types.Value) -> Ordered
             builder.columns = [*self.columns, new_value.name(id)]
         return builder.build()
 
-    def _select(self, values: typing.Tuple[ibis_types.Value]) -> OrderedIR:
+    def _select(self, values: typing.Tuple[ibis_types.Value, ...]) -> OrderedIR:
         """Safely assign by id while maintaining ordering integrity."""
         # TODO: Split into explicit set and replace methods
         ordering_col_ids = set(

bigframes/core/compile/compiler.py

@@ -59,11 +59,11 @@ def compile_sql(
         node = self.set_output_names(node, output_ids)
         if ordered:
             node, limit = rewrites.pullup_limit_from_slice(node)
-            return self.compile_ordered_ir(self._preprocess(node)).to_sql(
-                ordered=True, limit=limit
-            )
+            ir = self.compile_ordered_ir(self._preprocess(node))
+            return ir.to_sql(ordered=True, limit=limit)
         else:
-            return self.compile_unordered_ir(self._preprocess(node)).to_sql()
+            ir = self.compile_unordered_ir(self._preprocess(node))  # type: ignore
+            return ir.to_sql()
 
     def compile_peek_sql(self, node: nodes.BigFrameNode, n_rows: int) -> str:
         return self.compile_unordered_ir(self._preprocess(node)).peek_sql(n_rows)

bigframes/core/compile/ibis_types.py

@@ -24,7 +24,6 @@
 from bigframes_vendored.ibis.expr.datatypes.core import (
     dtype as python_type_to_bigquery_type,
 )
-import bigframes_vendored.ibis.expr.operations as ibis_ops
 import bigframes_vendored.ibis.expr.types as ibis_types
 import geopandas as gpd  # type: ignore
 import google.cloud.bigquery as bigquery
@@ -46,6 +45,7 @@
     ibis_dtypes.Binary,
     ibis_dtypes.Decimal,
     ibis_dtypes.GeoSpatial,
+    ibis_dtypes.JSON,
 ]
 
 
@@ -74,6 +74,7 @@
         ibis_dtypes.GeoSpatial(geotype="geography", srid=4326, nullable=True),
         gpd.array.GeometryDtype(),
     ),
+    (ibis_dtypes.json, pd.ArrowDtype(pa.large_string())),
 )
 
 BIGFRAMES_TO_IBIS: Dict[bigframes.dtypes.Dtype, ibis_dtypes.DataType] = {
@@ -219,12 +220,6 @@ def ibis_value_to_canonical_type(value: ibis_types.Value) -> ibis_types.Value:
     """
     ibis_type = value.type()
     name = value.get_name()
-    if ibis_type.is_json():
-        value = ibis_ops.ToJsonString(value).to_expr()  # type: ignore
-        value = (
-            value.case().when("null", bigframes_vendored.ibis.null()).else_(value).end()
-        )
-        return value.name(name)
     # Allow REQUIRED fields to be joined with NULLABLE fields.
     nullable_type = ibis_type.copy(nullable=True)
     return value.cast(nullable_type).name(name)
@@ -314,7 +309,7 @@ def ibis_dtype_to_bigframes_dtype(
             "Interpreting JSON as string. This behavior may change in future versions.",
             bigframes.exceptions.PreviewWarning,
         )
-        return bigframes.dtypes.STRING_DTYPE
+        return bigframes.dtypes.JSON_DTYPE
 
     if ibis_dtype in IBIS_TO_BIGFRAMES:
         return IBIS_TO_BIGFRAMES[ibis_dtype]

bigframes/core/compile/scalar_op_compiler.py

@@ -1181,7 +1181,10 @@ def json_set_op_impl(x: ibis_types.Value, y: ibis_types.Value, op: ops.JSONSet):
 
 @scalar_op_compiler.register_unary_op(ops.JSONExtract, pass_op=True)
 def json_extract_op_impl(x: ibis_types.Value, op: ops.JSONExtract):
-    return json_extract(json_obj=x, json_path=op.json_path)
+    if x.type().is_json():
+        return json_extract(json_obj=x, json_path=op.json_path)
+    # json string
+    return json_extract_string(json_obj=x, json_path=op.json_path)
 
 
 @scalar_op_compiler.register_unary_op(ops.JSONExtractArray, pass_op=True)
@@ -1845,7 +1848,7 @@ def float_ceil(a: float) -> float:
 
 
 @ibis_udf.scalar.builtin(name="parse_json")
-def parse_json(a: str) -> ibis_dtypes.JSON:  # type: ignore[empty-body]
+def parse_json(json_str: str) -> ibis_dtypes.JSON:  # type: ignore[empty-body]
     """Converts a JSON-formatted STRING value to a JSON value."""
 
 
@@ -1860,7 +1863,14 @@ def json_set(  # type: ignore[empty-body]
 def json_extract(  # type: ignore[empty-body]
     json_obj: ibis_dtypes.JSON, json_path: ibis_dtypes.String
 ) -> ibis_dtypes.JSON:
-    """Extracts a JSON value and converts it to a SQL JSON-formatted STRING or JSON value."""
+    """Extracts a JSON value and converts it to a JSON value."""
+
+
+@ibis_udf.scalar.builtin(name="json_extract")
+def json_extract_string(  # type: ignore[empty-body]
+    json_obj: ibis_dtypes.String, json_path: ibis_dtypes.String
+) -> ibis_dtypes.String:
+    """Extracts a JSON SRING value and converts it to a SQL JSON-formatted STRING."""
 
 
 @ibis_udf.scalar.builtin(name="json_extract_array")

bigframes/dtypes.py

@@ -55,6 +55,8 @@
 BIGNUMERIC_DTYPE = pd.ArrowDtype(pa.decimal256(76, 38))
 # No arrow equivalent
 GEO_DTYPE = gpd.array.GeometryDtype()
+# JSON
+JSON_DTYPE = pd.ArrowDtype(pa.large_string())
 
 # Used when storing Null expressions
 DEFAULT_DTYPE = FLOAT_DTYPE
@@ -132,6 +134,13 @@ class SimpleDtypeInfo:
         orderable=True,
         clusterable=True,
     ),
+    SimpleDtypeInfo(
+        dtype=JSON_DTYPE,
+        arrow_dtype=pa.large_string(),
+        type_kind=("JSON",),
+        orderable=False,
+        clusterable=False,
+    ),
     SimpleDtypeInfo(
         dtype=DATE_DTYPE,
         arrow_dtype=pa.date32(),
@@ -281,7 +290,7 @@ def is_struct_like(type_: ExpressionType) -> bool:
 
 def is_json_like(type_: ExpressionType) -> bool:
     # TODO: Add JSON type support
-    return type_ == STRING_DTYPE
+    return type_ == JSON_DTYPE or type_ == STRING_DTYPE  # Including JSON string
 
 
 def is_json_encoding_type(type_: ExpressionType) -> bool:
@@ -455,8 +464,6 @@ def infer_literal_arrow_type(literal) -> typing.Optional[pa.DataType]:
     return bigframes_dtype_to_arrow_dtype(infer_literal_type(literal))
 
 
-# Don't have dtype for json, so just end up interpreting as STRING
-_REMAPPED_TYPEKINDS = {"JSON": "STRING"}
 _TK_TO_BIGFRAMES = {
     type_kind: mapping.dtype
     for mapping in SIMPLE_TYPES
@@ -480,16 +487,13 @@ def convert_schema_field(
         pa_struct = pa.struct(fields)
         pa_type = pa.list_(pa_struct) if is_repeated else pa_struct
         return field.name, pd.ArrowDtype(pa_type)
-    elif (
-        field.field_type in _TK_TO_BIGFRAMES or field.field_type in _REMAPPED_TYPEKINDS
-    ):
-        singular_type = _TK_TO_BIGFRAMES[
-            _REMAPPED_TYPEKINDS.get(field.field_type, field.field_type)
-        ]
+    elif field.field_type in _TK_TO_BIGFRAMES:
         if is_repeated:
-            pa_type = pa.list_(bigframes_dtype_to_arrow_dtype(singular_type))
+            pa_type = pa.list_(
+                bigframes_dtype_to_arrow_dtype(_TK_TO_BIGFRAMES[field.field_type])
+            )
             return field.name, pd.ArrowDtype(pa_type)
-        return field.name, singular_type
+        return field.name, _TK_TO_BIGFRAMES[field.field_type]
     else:
         raise ValueError(f"Cannot handle type: {field.field_type}")
 
@@ -639,7 +643,7 @@ def can_coerce(source_type: ExpressionType, target_type: ExpressionType) -> bool
         return True  # None can be coerced to any supported type
     else:
         return (source_type == STRING_DTYPE) and (
-            target_type in TEMPORAL_BIGFRAMES_TYPES
+            target_type in TEMPORAL_BIGFRAMES_TYPES + [JSON_DTYPE]
         )
 
 

bigframes/session/executor.py

@@ -393,6 +393,7 @@ def peek(
     def head(
         self, array_value: bigframes.core.ArrayValue, n_rows: int
     ) -> ExecuteResult:
+
         maybe_row_count = self._local_get_row_count(array_value)
         if (maybe_row_count is not None) and (maybe_row_count <= n_rows):
             return self.execute(array_value, ordered=True)
@@ -452,7 +453,7 @@ def cached(
         # use a heuristic for whether something needs to be cached
         if (not force) and self._is_trivially_executable(array_value):
             return
-        elif use_session:
+        if use_session:
             self._cache_with_session_awareness(array_value)
         else:
             self._cache_with_cluster_cols(array_value, cluster_cols=cluster_cols)
@@ -658,7 +659,7 @@ def _sql_as_cached_temp_table(
     def _validate_result_schema(
         self,
         array_value: bigframes.core.ArrayValue,
-        bq_schema: list[bigquery.schema.SchemaField],
+        bq_schema: list[bigquery.SchemaField],
     ):
         actual_schema = tuple(bq_schema)
         ibis_schema = bigframes.core.compile.test_only_ibis_inferred_schema(
@@ -667,6 +668,7 @@ def _validate_result_schema(
         internal_schema = array_value.schema
         if not bigframes.features.PANDAS_VERSIONS.is_arrow_list_dtype_usable:
             return
+
         if internal_schema.to_bigquery() != actual_schema:
             raise ValueError(
                 f"This error should only occur while testing. BigFrames internal schema: {internal_schema.to_bigquery()} does not match actual schema: {actual_schema}"

tests/system/small/bigquery/test_json.py

@@ -134,6 +134,7 @@ def test_json_extract_from_string():
         actual.to_pandas(),
         expected.to_pandas(),
         check_names=False,
+        check_dtype=False,  # json_extract returns string type. While _get_series_from_json gives a JSON series (pa.large_string).
     )
 
 
@@ -200,3 +201,11 @@ def test_json_extract_string_array_as_float_array_from_array_strings():
 def test_json_extract_string_array_w_invalid_series_type():
     with pytest.raises(TypeError):
         bbq.json_extract_string_array(bpd.Series([1, 2]))
+
+
+# b/381148539
+def test_json_in_struct():
+    df = bpd.read_gbq(
+        "SELECT STRUCT(JSON '{\\\"a\\\": 1}' AS data, 1 AS number) as struct_col"
+    )
+    assert df["struct_col"].struct.field("data")[0] == '{"a":1}'