fix: Improve Anywidget pagination and display for unknown row counts

Author: shuoweil

bigframes/display/anywidget.py

@@ -55,7 +55,11 @@ class TableWidget(WIDGET_BASE):
 
     page = traitlets.Int(0).tag(sync=True)
     page_size = traitlets.Int(0).tag(sync=True)
-    row_count = traitlets.Int(0).tag(sync=True)
+    row_count = traitlets.Union(
+        [traitlets.Int(), traitlets.Instance(type(None))],
+        default_value=None,
+        allow_none=True,
+    ).tag(sync=True)
     table_html = traitlets.Unicode().tag(sync=True)
     _initial_load_complete = traitlets.Bool(False).tag(sync=True)
     _batches: Optional[blocks.PandasBatches] = None
@@ -98,8 +102,10 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # there are multiple pages and show "page 1 of many" in this case
         self._reset_batches_for_new_page_size()
         if self._batches is None or self._batches.total_rows is None:
+            # TODO(b/428238610): We could still end up with a None here if the
+            # underlying execution doesn't produce a total row count.
             self._error_message = "Could not determine total row count. Data might be unavailable or an error occurred."
-            self.row_count = 0
+            self.row_count = None
         else:
             self.row_count = self._batches.total_rows
 
@@ -131,8 +137,13 @@ def _validate_page(self, proposal: Dict[str, Any]) -> int:
         Returns:
             The validated and clamped page number as an integer.
         """
-
         value = proposal["value"]
+
+        # If row count is unknown, allow any non-negative page
+        if self.row_count is None:
+            return max(0, value)
+
+        # If truly empty or invalid page size, stay on page 0
         if self.row_count == 0 or self.page_size == 0:
             return 0
 
@@ -229,6 +240,23 @@ def _set_table_html(self) -> None:
         # Get the data for the current page
         page_data = cached_data.iloc[start:end]
 
+        # Handle case where user navigated beyond available data with unknown row count
+        is_unknown_count = self.row_count is None
+        is_beyond_data = self._all_data_loaded and len(page_data) == 0 and self.page > 0
+        if is_unknown_count and is_beyond_data:
+            # Calculate the last valid page (zero-indexed)
+            total_rows = len(cached_data)
+            if total_rows > 0:
+                last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
+                # Navigate back to the last valid page
+                self.page = last_valid_page
+                # Recursively call to display the correct page
+                return self._set_table_html()
+            else:
+                # If no data at all, stay on page 0 with empty display
+                self.page = 0
+                return self._set_table_html()
+
         # Generate HTML table
         self.table_html = bigframes.display.html.render_html(
             dataframe=page_data,

bigframes/display/table_widget.js

@@ -85,14 +85,23 @@ function render({ model, el }) {
 		const rowCount = model.get(ModelProperty.ROW_COUNT);
 		const pageSize = model.get(ModelProperty.PAGE_SIZE);
 		const currentPage = model.get(ModelProperty.PAGE);
-		const totalPages = Math.ceil(rowCount / pageSize);
-
-		rowCountLabel.textContent = `${rowCount.toLocaleString()} total rows`;
-		paginationLabel.textContent = `Page ${(
-			currentPage + 1
-		).toLocaleString()} of ${(totalPages || 1).toLocaleString()}`;
-		prevPage.disabled = currentPage === 0;
-		nextPage.disabled = currentPage >= totalPages - 1;
+
+		if (rowCount === null) {
+			// Unknown total rows
+			rowCountLabel.textContent = "Total rows unknown";
+			paginationLabel.textContent = `Page ${(currentPage + 1).toLocaleString()}`;
+			prevPage.disabled = currentPage === 0;
+			nextPage.disabled = false; // Allow navigation until we hit the end
+		} else {
+			// Known total rows
+			const totalPages = Math.ceil(rowCount / pageSize);
+			rowCountLabel.textContent = `${rowCount.toLocaleString()} total rows`;
+			paginationLabel.textContent = `Page ${(
+				currentPage + 1
+			).toLocaleString()} of ${(totalPages || 1).toLocaleString()}`;
+			prevPage.disabled = currentPage === 0;
+			nextPage.disabled = currentPage >= totalPages - 1;
+		}
 		pageSizeSelect.value = pageSize;
 	}
 

notebooks/dataframes/anywidget_mode.ipynb

@@ -354,6 +354,15 @@
     "small_widget"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "a9d5d13b",
+   "metadata": {},
+   "source": [
+    "### Handling of Invalid Row Count\n",
+    "In cases where the total number of rows cannot be determined, `row_count` will be `None`. The widget will display 'Total rows unknown', and you can navigate forward until you reach the end of the available data. If you navigate beyond the last page, the widget will automatically return to the last valid page."
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "added-cell-2",

tests/system/small/test_anywidget.py

@@ -549,6 +549,144 @@ def test_widget_row_count_reflects_actual_data_available(
         assert widget.page_size == 2  # Respects the display option
 
 
+def test_widget_with_unknown_row_count_should_auto_navigate_to_last_page(
+    session: bf.Session,
+):
+    """
+    Given a widget with unknown row count (row_count=None), when a user
+    navigates beyond the available data and all data is loaded, then the
+    widget should automatically navigate back to the last valid page.
+    """
+    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)
+
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+        widget = TableWidget(bf_df)
+
+        # Manually set row_count to None to simulate unknown total
+        widget.row_count = None
+
+        # Navigate to a page beyond available data (page 10)
+        # With page_size=2 and 5 rows, valid pages are 0, 1, 2
+        widget.page = 10
+
+        # Force data loading by accessing table_html
+        _ = widget.table_html
+
+        # After all data is loaded, widget should auto-navigate to last valid page
+        # Last valid page = ceil(5 / 2) - 1 = 2
+        assert widget.page == 2
+
+        # Verify the displayed content is the last page
+        html = widget.table_html
+        assert "row_4" in html  # Last row should be visible
+        assert "row_0" not in html  # First row should not be visible
+
+
+def test_widget_with_unknown_row_count_should_display_correct_ui_text(
+    session: bf.Session,
+):
+    """
+    Given a widget with unknown row count, the JavaScript frontend should
+    display appropriate text indicating the total is unknown.
+    """
+    from bigframes.display import TableWidget
+
+    test_data = pd.DataFrame(
+        {
+            "id": [0, 1, 2],
+            "value": ["a", "b", "c"],
+        }
+    )
+    bf_df = session.read_pandas(test_data)
+
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+        widget = TableWidget(bf_df)
+
+        # Set row_count to None
+        widget.row_count = None
+
+        # Verify row_count is None (not 0)
+        assert widget.row_count is None
+
+        # The widget should still function normally
+        assert widget.page == 0
+        assert widget.page_size == 2
+
+
+def test_widget_with_unknown_row_count_should_allow_forward_navigation(
+    session: bf.Session,
+):
+    """
+    Given a widget with unknown row count, users should be able to navigate
+    forward until they reach the end of available data.
+    """
+    from bigframes.display import TableWidget
+
+    test_data = pd.DataFrame(
+        {
+            "id": [0, 1, 2, 3, 4, 5],
+            "value": ["p0_r0", "p0_r1", "p1_r0", "p1_r1", "p2_r0", "p2_r1"],
+        }
+    )
+    bf_df = session.read_pandas(test_data)
+
+    with bf.option_context("display.repr_mode", "anywidget", "display.max_rows", 2):
+        widget = TableWidget(bf_df)
+        widget.row_count = None
+
+        # Navigate to page 1
+        widget.page = 1
+        html = widget.table_html
+        assert "p1_r0" in html
+        assert "p1_r1" in html
+
+        # Navigate to page 2
+        widget.page = 2
+        html = widget.table_html
+        assert "p2_r0" in html
+        assert "p2_r1" in html
+
+        # Navigate beyond available data (page 5)
+        widget.page = 5
+        _ = widget.table_html
+
+        # Should auto-navigate back to last valid page (page 2)
+        assert widget.page == 2
+
+
+def test_widget_with_unknown_row_count_empty_dataframe(
+    session: bf.Session,
+):
+    """
+    Given an empty DataFrame with unknown row count, the widget should
+    stay on page 0 and display empty content.
+    """
+    from bigframes.display import TableWidget
+
+    empty_data = pd.DataFrame(columns=["id", "value"])
+    bf_df = session.read_pandas(empty_data)
+
+    with bf.option_context("display.repr_mode", "anywidget"):
+        widget = TableWidget(bf_df)
+        widget.row_count = None
+
+        # Attempt to navigate to page 5
+        widget.page = 5
+        _ = widget.table_html
+
+        # Should stay on page 0 for empty DataFrame
+        assert widget.page == 0
+
+
 # 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.