Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 63 additions & 1 deletion bigframes/core/compile/polars/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion bigframes/core/compile/polars/lowering.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 21 additions & 18 deletions bigframes/display/anywidget.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down Expand Up @@ -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'>"
Expand Down
3 changes: 2 additions & 1 deletion bigframes/display/html.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions bigframes/operations/json_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm curious if this change is related to progress bar. If not, we should separate it out into another PR.

Copy link
Contributor Author

@shuoweil shuoweil Feb 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need this chunk of code for presubmit: screen/56yh654Cnmgft2W


def output_type(self, *input_types):
input_type = input_types[0]
Expand Down
2 changes: 2 additions & 0 deletions bigframes/session/polars_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -94,6 +95,7 @@
string_ops.EndsWithOp,
string_ops.StrContainsOp,
string_ops.StrContainsRegexOp,
json_ops.JSONDecode,
)
_COMPATIBLE_AGG_OPS = (
agg_ops.SizeOp,
Expand Down
Loading
Loading