Fix: Anywidget Table Pagination for Small DataFrames & Improve Display Module Docs

Author: shuoweil

bigframes/display/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Interactive display objects for BigQuery DataFrames."""
+
 from __future__ import annotations
 
 try:

bigframes/display/anywidget.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Interactive, paginated table widget for BigFrames DataFrames."""
+
 from __future__ import annotations
 
 from importlib import resources

bigframes/display/table_widget.js

@@ -96,7 +96,7 @@ function render({ model, el }) {
 			// Known total rows
 			const totalPages = Math.ceil(rowCount / pageSize);
 			rowCountLabel.textContent = `${rowCount.toLocaleString()} total rows`;
-			paginationLabel.textContent = `Page ${(currentPage + 1).toLocaleString()} of ${rowCount.toLocaleString()}`;
+			paginationLabel.textContent = `Page ${(currentPage + 1).toLocaleString()} of ${totalPages.toLocaleString()}`;
 			prevPage.disabled = currentPage === 0;
 			nextPage.disabled = currentPage >= totalPages - 1;
 		}

notebooks/dataframes/anywidget_mode.ipynb

@@ -142,7 +142,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8fcad7b7e408422cae71d519cd2d4980",
+       "model_id": "388323f1f2394f44a7199e7ecda7b5d2",
        "version_major": 2,
        "version_minor": 1
       },
@@ -166,7 +166,7 @@
     }
    ],
    "source": [
-    "df.set_index(\"name\")"
+    "df"
    ]
   },
   {
@@ -205,7 +205,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "06cb98c577514d5c9654a7792d93f8e6",
+       "model_id": "24847f9deca94a2abb4fa1a64ec8b616",
        "version_major": 2,
        "version_minor": 1
       },
@@ -305,7 +305,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1672f826f7a347e38539dbb5fb72cd43",
+       "model_id": "8b8a0b2a3572442f8718baa08657bef6",
        "version_major": 2,
        "version_minor": 1
       },
@@ -345,7 +345,7 @@
      "data": {
       "text/html": [
        "✅ Completed. \n",
-       "    Query processed 85.9 kB in 12 seconds of slot time.\n",
+       "    Query processed 85.9 kB in 15 seconds of slot time.\n",
        "    "
       ],
       "text/plain": [
@@ -380,7 +380,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "127a2e356b834c18b6f07c58ee2c4228",
+       "model_id": "1327a3901c194bf3a03f7f3a3358dc19",
        "version_major": 2,
        "version_minor": 1
       },
@@ -415,93 +415,6 @@
     "  LIMIT 5;\n",
     "\"\"\")"
    ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "multi-index-display-markdown",
-   "metadata": {},
-   "source": [
-    "## Display Multi-Index DataFrame in anywidget mode\n",
-    "This section demonstrates how BigFrames can display a DataFrame with multiple levels of indexing (a \"multi-index\") when using the `anywidget` display mode."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "ad7482aa",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "✅ Completed. \n",
-       "    Query processed 483.3 GB in 51 minutes of slot time. [<a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:3eace7c0-7776-48d6-925c-965be33d8738&page=queryresults\">Job bigframes-dev:US.3eace7c0-7776-48d6-925c-965be33d8738 details</a>]\n",
-       "    "
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "✅ Completed. \n",
-       "    Query processed 124.4 MB in 7 seconds of slot time. [<a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:job_UJ5cx4R1jW5cNxq_1H1x-9-ATfqS&page=queryresults\">Job bigframes-dev:US.job_UJ5cx4R1jW5cNxq_1H1x-9-ATfqS details</a>]\n",
-       "    "
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3f9652b5fdc0441eac2b05ab36d571d0",
-       "version_major": 2,
-       "version_minor": 1
-      },
-      "text/plain": [
-       "TableWidget(page_size=10, row_count=3967869, table_html='<table border=\"1\" class=\"dataframe table table-stripe…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [],
-      "text/plain": [
-       "Computation deferred. Computation will process 513.5 GB"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import datetime\n",
-    "\n",
-    " # Read the PyPI downloads dataset\n",
-    "pypi_df = bpd.read_gbq(\"bigquery-public-data.pypi.file_downloads\")\n",
-    "\n",
-    "# Filter for the last 7 days to reduce the data size for this example\n",
-    "seven_days_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=7)\n",
-    "pypi_df_recent = pypi_df[pypi_df[\"timestamp\"] > seven_days_ago]\n",
-    " \n",
-    "# Create a multi-index by grouping by date and project\n",
-    "pypi_df_recent['date'] = pypi_df_recent['timestamp'].dt.date\n",
-    "multi_index_df = pypi_df_recent.groupby([\"date\", \"project\"]).size().to_frame(\"downloads\")\n",
-    " \n",
-    "# Display the DataFrame with the multi-index\n",
-    "multi_index_df"
-   ]
   }
  ],
  "metadata": {

tests/system/small/test_anywidget.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 
+from typing import Any
+from unittest import mock
+
 import pandas as pd
 import pytest
 
@@ -98,6 +101,33 @@ def small_widget(small_bf_df):
         yield TableWidget(small_bf_df)
 
 
+@pytest.fixture
+def unknown_row_count_widget(session):
+    """Fixture to create a TableWidget with an unknown row count."""
+    from bigframes.core import blocks
+    from bigframes.display import TableWidget
+
+    # Create a small DataFrame with known content
+    test_data = pd.DataFrame(
+        {
+            "id": [0, 1, 2, 3, 4],
+            "value": ["row_0", "row_1", "row_2", "row_3", "row_4"],
+        }
+    )
+    bf_df = session.read_pandas(test_data)
+
+    # Simulate a scenario where total_rows is not available from the iterator
+    with mock.patch.object(bf_df, "_to_pandas_batches") as mock_batches:
+        # We need to provide an iterator of DataFrames, not Series
+        batches_iterator = iter([test_data])
+        mock_batches.return_value = blocks.PandasBatches(
+            batches_iterator, total_rows=None
+        )
+        with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+            widget = TableWidget(bf_df)
+            yield widget
+
+
 @pytest.fixture(scope="module")
 def empty_pandas_df() -> pd.DataFrame:
     """Create an empty DataFrame for edge case testing."""
@@ -129,7 +159,7 @@ def execution_metadata(self) -> ExecutionMetadata:
             return ExecutionMetadata()
 
         @property
-        def schema(self):
+        def schema(self) -> Any:
             return schema
 
         def batches(self) -> ResultsIterator:
@@ -180,7 +210,9 @@ def test_widget_initialization_should_calculate_total_row_count(
 def test_widget_initialization_should_set_default_pagination(
     table_widget,
 ):
-    """A TableWidget should initialize with page 0 and the correct page size."""
+    """
+    A TableWidget should initialize with page 0 and the correct page size.
+    """
     # The `table_widget` fixture already creates the widget.
     # Assert its state.
     assert table_widget.page == 0
@@ -304,15 +336,20 @@ def test_widget_with_few_rows_should_display_all_rows(small_widget, small_pandas
 
 def test_widget_with_few_rows_should_have_only_one_page(small_widget):
     """
-    Given a DataFrame smaller than the page size, the widget should
-    clamp page navigation, effectively having only one page.
+    Given a DataFrame with a small number of rows, the widget should
+    report the correct total row count and prevent navigation beyond
+    the first page, ensuring the frontend correctly displays "Page 1 of 1".
     """
+    # For a DataFrame with 2 rows and page_size 5 (from small_widget fixture),
+    # the frontend should calculate 1 total page.
+    assert small_widget.row_count == 2
+
+    # The widget should always be on page 0 for a single-page dataset.
     assert small_widget.page == 0
 
-    # Attempt to navigate past the end
+    # Attempting to navigate to page 1 should be clamped back to page 0,
+    # confirming that only one page is recognized by the backend.
     small_widget.page = 1
-
-    # Should be clamped back to the only valid page
     assert small_widget.page == 0
 
 
@@ -420,8 +457,10 @@ def test_navigation_after_page_size_change_should_use_new_size(
 
 @pytest.mark.parametrize("invalid_size", [0, -5], ids=["zero", "negative"])
 def test_setting_invalid_page_size_should_be_ignored(table_widget, invalid_size: int):
-    """When the page size is set to an invalid number (<=0), the change should
-    be ignored."""
+    """
+    When the page size is set to an invalid number (<=0), the change should
+    be ignored.
+    """
     # Set the initial page to 2.
     initial_size = table_widget.page_size
     assert initial_size == 2