feat: Display JSON columns in anywidget mode

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1036,7 +1036,7 @@ def astype_op_impl(x: ibis_types.Value, op: ops.AsTypeOp):
         if to_type == ibis_dtypes.bool:
             return cast_json_to_bool_in_safe(x) if op.safe else cast_json_to_bool(x)
         if to_type == ibis_dtypes.string:
-            return cast_json_to_string_in_safe(x) if op.safe else cast_json_to_string(x)
+            return to_json_string(x)
 
     # TODO: either inline this function, or push rest of this op into the function
     return bigframes.core.compile.ibis_types.cast_ibis_value(x, to_type, safe=op.safe)

bigframes/core/compile/polars/compiler.py

@@ -407,6 +407,19 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, json_ops.JSONDecode)
             return input.str.json_decode(_DTYPE_MAPPING[op.to_type])
 
+        @compile_op.register(json_ops.ToJSONString)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            return input.str.json_decode(pl.String())
+
+        @compile_op.register(json_ops.ParseJSON)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            return input.str.json_decode(pl.String())
+
+        @compile_op.register(json_ops.JSONExtract)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, json_ops.JSONExtract)
+            return input.str.json_extract(json_path=op.json_path)
+
         @compile_op.register(arr_ops.ToArrayOp)
         def _(self, op: ops.ToArrayOp, *inputs: pl.Expr) -> pl.Expr:
             return pl.concat_list(*inputs)

bigframes/dataframe.py

@@ -783,7 +783,8 @@ def __repr__(self) -> str:
 
         opts = bigframes.options.display
         max_results = opts.max_rows
-        # anywdiget mode uses the same display logic as the "deferred" mode
+
+        # anywidget mode uses the same display logic as the "deferred" mode
         # for faster execution
         if opts.repr_mode in ("deferred", "anywidget"):
             return formatter.repr_query_job(self._compute_dry_run())
@@ -851,18 +852,34 @@ def _repr_html_(self) -> str:
 
         if opts.repr_mode == "anywidget":
             try:
+                import anywidget  # noqa: F401
                 from IPython.display import display as ipython_display
+                import traitlets  # noqa: F401
 
                 from bigframes import display
 
+                # The anywidget frontend doesn't support the db_dtypes JSON type, so
+                # convert to strings for display.
+                json_cols = [
+                    series_name
+                    for series_name, series in df.items()
+                    if bigframes.dtypes.contains_db_dtypes_json_dtype(series.dtype)
+                ]
+                if json_cols:
+                    warnings.warn(
+                        "Converting JSON columns to strings for display. "
+                        "This is temporary and will be removed when the frontend supports JSON types."
+                    )
+                    for col in json_cols:
+                        df[col] = df[col]._apply_unary_op(ops.json_ops.ToJSONString())
+
                 # Always create a new widget instance for each display call
                 # This ensures that each cell gets its own widget and prevents
                 # unintended sharing between cells
                 widget = display.TableWidget(df.copy())
 
                 ipython_display(widget)
                 return ""  # Return empty string since we used display()
-
             except (AttributeError, ValueError, ImportError):
                 # Fallback if anywidget is not available
                 warnings.warn(

bigframes/display/anywidget.py

@@ -62,7 +62,20 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
             )
 
         super().__init__()
-        self._dataframe = dataframe
+        # Workaround for Arrow bug https://github.com/apache/arrow/issues/45262
+        # JSON columns are not supported in `to_pandas_batches` and will be converted to string.
+        json_cols = [
+            col
+            for col, dtype in dataframe.dtypes.items()
+            if dtype == bigframes.dtypes.JSON_DTYPE
+        ]
+        if json_cols:
+            df_copy = dataframe.copy()
+            for col in json_cols:
+                df_copy[str(col)] = df_copy[str(col)].astype("string")
+            self._dataframe = df_copy
+        else:
+            self._dataframe = dataframe
 
         # Initialize attributes that might be needed by observers FIRST
         self._table_id = str(uuid.uuid4())
@@ -76,9 +89,6 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # Initialize data fetching attributes.
         self._batches = dataframe._to_pandas_batches(page_size=initial_page_size)
 
-        # set traitlets properties that trigger observers
-        self.page_size = initial_page_size
-
         # len(dataframe) is expensive, since it will trigger a
         # SELECT COUNT(*) query. It is a must have however.
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
@@ -87,6 +97,9 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # there are multiple pages and show "page 1 of many" in this case.
         self.row_count = self._batches.total_rows or 0
 
+        # set traitlets properties that trigger observers
+        self.page_size = initial_page_size
+
         # get the initial page
         self._set_table_html()
 
@@ -221,8 +234,5 @@ def _page_size_changed(self, _change: Dict[str, Any]):
         # Reset the page to 0 when page size changes to avoid invalid page states
         self.page = 0
 
-        # Reset batches to use new page size for future data fetching
-        self._reset_batches_for_new_page_size()
-
         # Update the table display
         self._set_table_html()

bigframes/series.py

@@ -608,6 +608,15 @@ def astype(
         if errors not in ["raise", "null"]:
             raise ValueError("Argument 'errors' must be one of 'raise' or 'null'")
         dtype = bigframes.dtypes.bigframes_type(dtype)
+
+        # BigQuery doesn't support CAST(json_col AS STRING), but it does support
+        # TO_JSON_STRING(json_col).
+        if (
+            self.dtype == bigframes.dtypes.JSON_DTYPE
+            and dtype == bigframes.dtypes.STRING_DTYPE
+        ):
+            return self._apply_unary_op(ops.json_ops.ToJSONString())
+
         return self._apply_unary_op(
             bigframes.operations.AsTypeOp(to_type=dtype, safe=(errors == "null"))
         )