Skip to content

Commit ef73c01

Browse files
authored
Revert "feat: Disable progress bars in Anywidget mode to reduce notebook clut…"
This reverts commit 853240d.
1 parent 853240d commit ef73c01

File tree

8 files changed

+253
-195
lines changed

8 files changed

+253
-195
lines changed

bigframes/core/compile/polars/compiler.py

Lines changed: 1 addition & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import dataclasses
1717
import functools
1818
import itertools
19-
import json
2019
from typing import cast, Literal, Optional, Sequence, Tuple, Type, TYPE_CHECKING
2120

2221
import pandas as pd
@@ -430,68 +429,7 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
430429
@compile_op.register(json_ops.JSONDecode)
431430
def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
432431
assert isinstance(op, json_ops.JSONDecode)
433-
target_dtype = _bigframes_dtype_to_polars_dtype(op.to_type)
434-
if op.safe:
435-
# Polars does not support safe JSON decoding (returning null on failure).
436-
# We use map_elements to provide safe JSON decoding.
437-
def safe_decode(val):
438-
if val is None:
439-
return None
440-
try:
441-
decoded = json.loads(val)
442-
except Exception:
443-
return None
444-
445-
if decoded is None:
446-
return None
447-
448-
if op.to_type == bigframes.dtypes.INT_DTYPE:
449-
if type(decoded) is bool:
450-
return None
451-
if isinstance(decoded, int):
452-
return decoded
453-
if isinstance(decoded, float):
454-
if decoded.is_integer():
455-
return int(decoded)
456-
if isinstance(decoded, str):
457-
try:
458-
return int(decoded)
459-
except Exception:
460-
pass
461-
return None
462-
463-
if op.to_type == bigframes.dtypes.FLOAT_DTYPE:
464-
if type(decoded) is bool:
465-
return None
466-
if isinstance(decoded, (int, float)):
467-
return float(decoded)
468-
if isinstance(decoded, str):
469-
try:
470-
return float(decoded)
471-
except Exception:
472-
pass
473-
return None
474-
475-
if op.to_type == bigframes.dtypes.BOOL_DTYPE:
476-
if isinstance(decoded, bool):
477-
return decoded
478-
if isinstance(decoded, str):
479-
if decoded.lower() == "true":
480-
return True
481-
if decoded.lower() == "false":
482-
return False
483-
return None
484-
485-
if op.to_type == bigframes.dtypes.STRING_DTYPE:
486-
if isinstance(decoded, str):
487-
return decoded
488-
return None
489-
490-
return decoded
491-
492-
return input.map_elements(safe_decode, return_dtype=target_dtype)
493-
494-
return input.str.json_decode(target_dtype)
432+
return input.str.json_decode(_DTYPE_MAPPING[op.to_type])
495433

496434
@compile_op.register(arr_ops.ToArrayOp)
497435
def _(self, op: ops.ToArrayOp, *inputs: pl.Expr) -> pl.Expr:

bigframes/core/compile/polars/lowering.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
391391
return arg
392392

393393
if arg.output_type == dtypes.JSON_DTYPE:
394-
return json_ops.JSONDecode(cast_op.to_type, safe=cast_op.safe).as_expr(arg)
394+
return json_ops.JSONDecode(cast_op.to_type).as_expr(arg)
395395
if (
396396
arg.output_type == dtypes.STRING_DTYPE
397397
and cast_op.to_type == dtypes.DATETIME_DTYPE

bigframes/display/anywidget.py

Lines changed: 18 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -133,26 +133,25 @@ def _initial_load(self) -> None:
133133
# obtain the row counts
134134
# TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
135135
# before we get here so that the count might already be cached.
136-
with bigframes.option_context("display.progress_bar", None):
137-
self._reset_batches_for_new_page_size()
136+
self._reset_batches_for_new_page_size()
138137

139-
if self._batches is None:
140-
self._error_message = (
141-
"Could not retrieve data batches. Data might be unavailable or "
142-
"an error occurred."
143-
)
144-
self.row_count = None
145-
elif self._batches.total_rows is None:
146-
# Total rows is unknown, this is an expected state.
147-
# TODO(b/461536343): Cheaply discover if we have exactly 1 page.
148-
# There are cases where total rows is not set, but there are no additional
149-
# pages. We could disable the "next" button in these cases.
150-
self.row_count = None
151-
else:
152-
self.row_count = self._batches.total_rows
138+
if self._batches is None:
139+
self._error_message = (
140+
"Could not retrieve data batches. Data might be unavailable or "
141+
"an error occurred."
142+
)
143+
self.row_count = None
144+
elif self._batches.total_rows is None:
145+
# Total rows is unknown, this is an expected state.
146+
# TODO(b/461536343): Cheaply discover if we have exactly 1 page.
147+
# There are cases where total rows is not set, but there are no additional
148+
# pages. We could disable the "next" button in these cases.
149+
self.row_count = None
150+
else:
151+
self.row_count = self._batches.total_rows
153152

154-
# get the initial page
155-
self._set_table_html()
153+
# get the initial page
154+
self._set_table_html()
156155

157156
@traitlets.observe("_initial_load_complete")
158157
def _on_initial_load_complete(self, change: dict[str, Any]):
@@ -282,9 +281,7 @@ def _reset_batches_for_new_page_size(self) -> None:
282281
def _set_table_html(self) -> None:
283282
"""Sets the current html data based on the current page and page size."""
284283
new_page = None
285-
with self._setting_html_lock, bigframes.option_context(
286-
"display.progress_bar", None
287-
):
284+
with self._setting_html_lock:
288285
if self._error_message:
289286
self.table_html = (
290287
f"<div class='bigframes-error-message'>"

bigframes/display/html.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -363,8 +363,7 @@ def repr_mimebundle(
363363

364364
if opts.repr_mode == "anywidget":
365365
try:
366-
with bigframes.option_context("display.progress_bar", None):
367-
return get_anywidget_bundle(obj, include=include, exclude=exclude)
366+
return get_anywidget_bundle(obj, include=include, exclude=exclude)
368367
except ImportError:
369368
# Anywidget is an optional dependency, so warn rather than fail.
370369
# TODO(shuowei): When Anywidget becomes the default for all repr modes,

bigframes/operations/json_ops.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,6 @@ def output_type(self, *input_types):
220220
class JSONDecode(base_ops.UnaryOp):
221221
name: typing.ClassVar[str] = "json_decode"
222222
to_type: dtypes.Dtype
223-
safe: bool = False
224223

225224
def output_type(self, *input_types):
226225
input_type = input_types[0]

bigframes/session/polars_executor.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
numeric_ops,
3535
string_ops,
3636
)
37-
import bigframes.operations.json_ops as json_ops
3837
from bigframes.session import executor, semi_executor
3938

4039
if TYPE_CHECKING:
@@ -95,7 +94,6 @@
9594
string_ops.EndsWithOp,
9695
string_ops.StrContainsOp,
9796
string_ops.StrContainsRegexOp,
98-
json_ops.JSONDecode,
9997
)
10098
_COMPATIBLE_AGG_OPS = (
10199
agg_ops.SizeOp,

0 commit comments

Comments
 (0)