add more testcase

Author: shuoweil

bigframes/display/anywidget.py

@@ -231,6 +231,8 @@ def _set_table_html(self) -> None:
                 cached_data = self._cached_data
             else:
                 break
+
+        # Get the data for the current page
         page_data = cached_data.iloc[start:end]
 
         # Generate HTML table

notebooks/dataframes/anywidget_mode.ipynb

@@ -142,7 +142,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1d718cdbafcb42898120637cdb3fa267",
+       "model_id": "93dd10072d564a02a0278817d14855a9",
        "version_major": 2,
        "version_minor": 0
       },
@@ -205,7 +205,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "519297c3ad19403aa844cbeabcd5eb44",
+       "model_id": "6e2538d446e344ac8505e4706730243e",
        "version_major": 2,
        "version_minor": 0
       },
@@ -304,7 +304,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "37ba207603aa40a38c9786a210e712fd",
+       "model_id": "d6faf367ea5d44ad9d275506d870557a",
        "version_major": 2,
        "version_minor": 0
       },
@@ -333,6 +333,14 @@
     "The `AI.GENERATE` function in BigQuery returns results in a JSON column. While BigQuery's JSON type is not natively supported by the underlying Arrow `to_pandas_batches()` method used in anywidget mode ([Apache Arrow issue #45262](https://github.com/apache/arrow/issues/45262)), BigQuery Dataframes automatically converts JSON columns to strings for display. This allows you to view the results of generative AI functions seamlessly."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fdadcad6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": 10,
@@ -343,7 +351,7 @@
      "data": {
       "text/html": [
        "✅ Completed. \n",
-       "    Query processed 85.9 kB in 23 seconds of slot time.\n",
+       "    Query processed 85.9 kB in 24 seconds of slot time.\n",
        "    "
       ],
       "text/plain": [
@@ -382,7 +390,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "379998ea9a744e7b8afd9c1bcb36548d",
+       "model_id": "b6d6f3bacc2c43fc9a335e6039db12a5",
        "version_major": 2,
        "version_minor": 0
       },

tests/system/small/test_anywidget.py

@@ -527,6 +527,137 @@ def test_json_column_anywidget_mode(mock_display, json_df: bf.dataframe.DataFram
             assert result == ""
 
 
+def mock_execute_result_with_params(
+    self, schema, total_rows_val, arrow_batches_val, *args, **kwargs
+):
+    """
+    Mocks an execution result with configurable total_rows and arrow_batches.
+    """
+    from bigframes.session.executor import ExecuteResult
+
+    return ExecuteResult(
+        iter(arrow_batches_val),
+        schema=schema,
+        query_job=None,
+        total_bytes=None,
+        total_rows=total_rows_val,
+    )
+
+
+def test_widget_row_count_should_be_immutable_after_creation(
+    paginated_bf_df: bf.dataframe.DataFrame,
+):
+    """
+    Given a widget created with a specific configuration when global display
+    options are changed later, the widget's original row_count should remain
+    unchanged.
+    """
+    from bigframes.display.anywidget import TableWidget
+
+    # Use a context manager to ensure the option is reset
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+        widget = TableWidget(paginated_bf_df)
+        initial_row_count = widget.row_count
+
+    # Change a global option that could influence row count
+    bf.options.display.max_rows = 10
+
+    # Verify the row count remains immutable.
+    assert widget.row_count == initial_row_count
+
+
+class FaultyIterator:
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        raise ValueError("Simulated read error")
+
+
+def test_widget_should_fallback_to_zero_rows_with_invalid_total_rows(
+    paginated_bf_df: bf.dataframe.DataFrame,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """
+    Given an internal component fails to return valid execution data,
+    when the TableWidget is created, its error_message should be set and displayed.
+    """
+    # Patch the executor's 'execute' method to simulate an error.
+    monkeypatch.setattr(
+        "bigframes.session.bq_caching_executor.BigQueryCachingExecutor.execute",
+        lambda self, *args, **kwargs: mock_execute_result_with_params(
+            self, paginated_bf_df._block.expr.schema, None, [], *args, **kwargs
+        ),
+    )
+
+    # Create the TableWidget under the error condition.
+    with bf.option_context("display.repr_mode", "anywidget"):
+        from bigframes.display.anywidget import TableWidget
+
+        # The widget should handle the faulty data from the mock without crashing.
+        widget = TableWidget(paginated_bf_df)
+
+    # The widget should have an error message and display it in the HTML.
+    assert widget.row_count == 0
+    assert widget._error_message is not None
+    assert "Could not determine total row count" in widget._error_message
+    assert widget._error_message in widget.table_html
+
+
+def test_widget_row_count_reflects_actual_data_available(
+    paginated_bf_df: bf.dataframe.DataFrame,
+):
+    """
+    Test that widget row_count reflects the actual data available,
+    regardless of theoretical limits.
+    """
+    from bigframes.display.anywidget import TableWidget
+
+    # Set up display options that define a page size.
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+        widget = TableWidget(paginated_bf_df)
+
+        # The widget should report the total rows in the DataFrame,
+        # not limited by page_size (which only affects pagination)
+        assert widget.row_count == EXPECTED_ROW_COUNT
+        assert widget.page_size == 2  # Respects the display option
+
+
 # TODO(shuowei): Add tests for custom index and multiindex
 # This may not be necessary for the SQL Cell use case but should be
 # considered for completeness.
+
+
+@pytest.fixture(scope="module")
+def empty_json_df(session: bf.Session) -> bf.dataframe.DataFrame:
+    """Create an empty DataFrame with a JSON column for testing."""
+    import bigframes.dtypes
+
+    pandas_df = pd.DataFrame(
+        {
+            "a": pd.Series(dtype="int64"),
+            "b": pd.Series(dtype=bigframes.dtypes.JSON_DTYPE),
+        }
+    )
+    return session.read_pandas(pandas_df)
+
+
+def test_empty_widget_with_json_column(empty_json_df: bf.dataframe.DataFrame):
+    """Given an empty DataFrame with a JSON column, the widget should render table headers."""
+    with bf.option_context("display.repr_mode", "anywidget"):
+        from bigframes.display.anywidget import TableWidget
+
+        widget = TableWidget(empty_json_df)
+        html = widget.table_html
+
+        assert widget.row_count == 0
+        assert "<table" in html
+        assert "a" in html
+        assert "b" in html
+
+
+def test_json_column_conversion_warning(json_df: bf.dataframe.DataFrame):
+    """Test that a warning is shown when converting JSON columns."""
+    with bf.option_context("display.repr_mode", "anywidget"):
+        with pytest.warns(UserWarning, match="Converting JSON columns to strings"):
+            json_df._repr_html_()

tests/unit/test_dataframe.py

@@ -181,3 +181,26 @@ def test_dataframe_ai_property_future_warning(
 
     with pytest.warns(FutureWarning):
         dataframe.ai
+
+
+@pytest.fixture()
+def json_df(polars_session: bigframes.session.Session) -> bigframes.dataframe.DataFrame:
+    """Create a DataFrame with a JSON column for testing."""
+    import bigframes.dtypes
+
+    pandas_df = pd.DataFrame(
+        {
+            "a": [1],
+            "b": ['{"c": 2, "d": 3}'],
+        }
+    )
+    pandas_df["b"] = pandas_df["b"].astype(bigframes.dtypes.JSON_DTYPE)
+    return polars_session.read_pandas(pandas_df)
+
+
+def test_to_pandas_batches_with_json_column(json_df: bigframes.dataframe.DataFrame):
+    """Test that JSON columns are converted to strings in to_pandas_batches."""
+    batches = list(json_df._to_pandas_batches(page_size=10))
+    assert len(batches) > 0
+    # Verify the JSON column is now string type
+    assert batches[0]["b"].dtype == pd.StringDtype(storage="pyarrow")

tests/unit/test_polars_compiler.py

@@ -0,0 +1,86 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import bigframes as bf
+import bigframes.core.compile.polars.compiler as polars_compiler
+import bigframes.core.nodes as nodes
+import bigframes.operations.json_ops as json_ops
+
+
+def test_polars_to_json_string():
+    """Test ToJSONString operation in Polars compiler."""
+    compiler = polars_compiler.PolarsExpressionCompiler()
+    op = json_ops.ToJSONString()
+    # Polars doesn't have a native JSON type, it uses strings.
+    # The operation is a cast to string.
+    input_expr = pl.lit('{"b": 2}', dtype=pl.String)
+    result = compiler.compile_op(op, input_expr)
+
+    df = pl.DataFrame({"a": ['{"b": 2}']}).lazy()
+    result_df = df.with_columns(result.alias("b")).collect()
+    assert result_df["b"][0] == '{"b": 2}'
+    assert result_df["b"].dtype == pl.String
+
+
+def test_polars_parse_json():
+    """Test ParseJSON operation in Polars compiler."""
+    compiler = polars_compiler.PolarsExpressionCompiler()
+    op = json_ops.ParseJSON()
+    input_expr = pl.lit('{"b": 2}', dtype=pl.String)
+    result = compiler.compile_op(op, input_expr)
+
+    df = pl.DataFrame({"a": ['{"b": 2}']}).lazy()
+    result_df = df.with_columns(result.alias("b")).collect()
+    # The result of json_decode is a struct
+    assert isinstance(result_df["b"][0], dict)
+    assert result_df["b"][0] == {"b": 2}
+
+
+@pytest.mark.skip(reason="Polars does not have json_extract on string expressions")
+def test_polars_json_extract():
+    """Test JSONExtract operation in Polars compiler."""
+    compiler = polars_compiler.PolarsExpressionCompiler()
+    op = json_ops.JSONExtract(json_path="$.b")
+    input_expr = pl.lit('{"a": 1, "b": "hello"}', dtype=pl.String)
+    result = compiler.compile_op(op, input_expr)
+
+    df = pl.DataFrame({"a": ['{"b": "world"}']}).lazy()
+    result_df = df.with_columns(result.alias("b")).collect()
+    # json_extract returns a JSON encoded string
+    assert result_df["b"][0] == '"world"'
+
+
+def test_readlocal_with_json_column(polars_session):
+    """Test ReadLocalNode compilation with JSON columns."""
+    pandas_df = pd.DataFrame({"data": ['{"key": "value"}']})
+    pandas_df["data"] = pandas_df["data"].astype(bf.dtypes.JSON_DTYPE)
+    bf_df = polars_session.read_pandas(pandas_df)
+
+    node = bf_df._block.expr.node
+    # Traverse the node tree to find the ReadLocalNode
+    while not isinstance(node, nodes.ReadLocalNode):
+        node = node.child
+    assert isinstance(node, nodes.ReadLocalNode)
+
+    compiler = polars_compiler.PolarsCompiler()
+    lazy_frame = compiler.compile_node(node)
+    result_df = lazy_frame.collect()
+
+    # The compiler should have converted the JSON column to string.
+    assert result_df.schema["column_0"] == pl.String
+    assert result_df["column_0"][0] == '{"key":"value"}'