Merge branch 'main' into tswast-patch-3

Author: tswast

bigframes/core/compile/polars/compiler.py

@@ -16,6 +16,7 @@
 import dataclasses
 import functools
 import itertools
+import json
 from typing import cast, Literal, Optional, Sequence, Tuple, Type, TYPE_CHECKING
 
 import pandas as pd
@@ -429,7 +430,68 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
         @compile_op.register(json_ops.JSONDecode)
         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])
+            target_dtype = _bigframes_dtype_to_polars_dtype(op.to_type)
+            if op.safe:
+                # Polars does not support safe JSON decoding (returning null on failure).
+                # We use map_elements to provide safe JSON decoding.
+                def safe_decode(val):
+                    if val is None:
+                        return None
+                    try:
+                        decoded = json.loads(val)
+                    except Exception:
+                        return None
+
+                    if decoded is None:
+                        return None
+
+                    if op.to_type == bigframes.dtypes.INT_DTYPE:
+                        if type(decoded) is bool:
+                            return None
+                        if isinstance(decoded, int):
+                            return decoded
+                        if isinstance(decoded, float):
+                            if decoded.is_integer():
+                                return int(decoded)
+                        if isinstance(decoded, str):
+                            try:
+                                return int(decoded)
+                            except Exception:
+                                pass
+                        return None
+
+                    if op.to_type == bigframes.dtypes.FLOAT_DTYPE:
+                        if type(decoded) is bool:
+                            return None
+                        if isinstance(decoded, (int, float)):
+                            return float(decoded)
+                        if isinstance(decoded, str):
+                            try:
+                                return float(decoded)
+                            except Exception:
+                                pass
+                        return None
+
+                    if op.to_type == bigframes.dtypes.BOOL_DTYPE:
+                        if isinstance(decoded, bool):
+                            return decoded
+                        if isinstance(decoded, str):
+                            if decoded.lower() == "true":
+                                return True
+                            if decoded.lower() == "false":
+                                return False
+                        return None
+
+                    if op.to_type == bigframes.dtypes.STRING_DTYPE:
+                        if isinstance(decoded, str):
+                            return decoded
+                        return None
+
+                    return decoded
+
+                return input.map_elements(safe_decode, return_dtype=target_dtype)
+
+            return input.str.json_decode(target_dtype)
 
         @compile_op.register(arr_ops.ToArrayOp)
         def _(self, op: ops.ToArrayOp, *inputs: pl.Expr) -> pl.Expr:

bigframes/core/compile/polars/lowering.py

@@ -391,7 +391,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
         return arg
 
     if arg.output_type == dtypes.JSON_DTYPE:
-        return json_ops.JSONDecode(cast_op.to_type).as_expr(arg)
+        return json_ops.JSONDecode(cast_op.to_type, safe=cast_op.safe).as_expr(arg)
     if (
         arg.output_type == dtypes.STRING_DTYPE
         and cast_op.to_type == dtypes.DATETIME_DTYPE

bigframes/display/anywidget.py

@@ -133,25 +133,26 @@ def _initial_load(self) -> None:
         # obtain the row counts
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
         # before we get here so that the count might already be cached.
-        self._reset_batches_for_new_page_size()
+        with bigframes.option_context("display.progress_bar", None):
+            self._reset_batches_for_new_page_size()
 
-        if self._batches is None:
-            self._error_message = (
-                "Could not retrieve data batches. Data might be unavailable or "
-                "an error occurred."
-            )
-            self.row_count = None
-        elif self._batches.total_rows is None:
-            # Total rows is unknown, this is an expected state.
-            # TODO(b/461536343): Cheaply discover if we have exactly 1 page.
-            # There are cases where total rows is not set, but there are no additional
-            # pages. We could disable the "next" button in these cases.
-            self.row_count = None
-        else:
-            self.row_count = self._batches.total_rows
+            if self._batches is None:
+                self._error_message = (
+                    "Could not retrieve data batches. Data might be unavailable or "
+                    "an error occurred."
+                )
+                self.row_count = None
+            elif self._batches.total_rows is None:
+                # Total rows is unknown, this is an expected state.
+                # TODO(b/461536343): Cheaply discover if we have exactly 1 page.
+                # There are cases where total rows is not set, but there are no additional
+                # pages. We could disable the "next" button in these cases.
+                self.row_count = None
+            else:
+                self.row_count = self._batches.total_rows
 
-        # get the initial page
-        self._set_table_html()
+            # get the initial page
+            self._set_table_html()
 
     @traitlets.observe("_initial_load_complete")
     def _on_initial_load_complete(self, change: dict[str, Any]):
@@ -281,7 +282,9 @@ def _reset_batches_for_new_page_size(self) -> None:
     def _set_table_html(self) -> None:
         """Sets the current html data based on the current page and page size."""
         new_page = None
-        with self._setting_html_lock:
+        with self._setting_html_lock, bigframes.option_context(
+            "display.progress_bar", None
+        ):
             if self._error_message:
                 self.table_html = (
                     f"<div class='bigframes-error-message'>"

bigframes/display/html.py

@@ -363,7 +363,8 @@ def repr_mimebundle(
 
     if opts.repr_mode == "anywidget":
         try:
-            return get_anywidget_bundle(obj, include=include, exclude=exclude)
+            with bigframes.option_context("display.progress_bar", None):
+                return get_anywidget_bundle(obj, include=include, exclude=exclude)
         except ImportError:
             # Anywidget is an optional dependency, so warn rather than fail.
             # TODO(shuowei): When Anywidget becomes the default for all repr modes,

bigframes/operations/json_ops.py

@@ -220,6 +220,7 @@ def output_type(self, *input_types):
 class JSONDecode(base_ops.UnaryOp):
     name: typing.ClassVar[str] = "json_decode"
     to_type: dtypes.Dtype
+    safe: bool = False
 
     def output_type(self, *input_types):
         input_type = input_types[0]

bigframes/session/polars_executor.py

@@ -34,6 +34,7 @@
     numeric_ops,
     string_ops,
 )
+import bigframes.operations.json_ops as json_ops
 from bigframes.session import executor, semi_executor
 
 if TYPE_CHECKING:
@@ -94,6 +95,7 @@
     string_ops.EndsWithOp,
     string_ops.StrContainsOp,
     string_ops.StrContainsRegexOp,
+    json_ops.JSONDecode,
 )
 _COMPATIBLE_AGG_OPS = (
     agg_ops.SizeOp,