add more testcases and demo notebook

Author: shuoweil

bigframes/display/__init__.py

@@ -11,6 +11,7 @@
 # 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.
+
 from __future__ import annotations
 
 import warnings

bigframes/display/anywidget.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Iterator
 import uuid
 
 import anywidget  # type: ignore
@@ -26,6 +27,8 @@ class TableWidget(anywidget.AnyWidget):
     An interactive, paginated table widget for BigFrames DataFrames.
     """
 
+    # The _esm variable contains the JavaScript source code for the frontend
+    # component of the widget.
     _esm = """
     function render({ model, el }) {
         const container = document.createElement('div');
@@ -97,20 +100,19 @@ def __init__(self, dataframe):
         Initialize the TableWidget.
 
         Args:
-            dataframe: The Bigframes Dataframe to display
+            dataframe: The Bigframes Dataframe to display.
         """
         super().__init__()
         self._dataframe = dataframe
 
         # respect display options
         self.page_size = bigframes.options.display.max_rows
 
+        # Initialize data fetching attributes.
         self._batches = dataframe.to_pandas_batches(page_size=self.page_size)
         self._cached_data = pd.DataFrame(columns=self._dataframe.columns)
         self._table_id = str(uuid.uuid4())
         self._all_data_loaded = False
-
-        # store the iterator as an instance variable
         self._batch_iterator = None
 
         # len(dataframe) is expensive, since it will trigger a
@@ -120,8 +122,13 @@ def __init__(self, dataframe):
         # get the initial page
         self._set_table_html()
 
-    def _get_next_batch(self):
-        """Gets the next batch of data from the batches generator."""
+    def _get_next_batch(self) -> bool:
+        """
+        Gets the next batch of data from the generator and appends to cache.
+
+        Returns:
+            bool: True if a batch was successfully loaded, False otherwise.
+        """
         if self._all_data_loaded:
             return False
 
@@ -139,8 +146,8 @@ def _get_next_batch(self):
         except Exception as e:
             raise RuntimeError(f"Error during batch processing: {str(e)}") from e
 
-    def _get_batch_iterator(self):
-        """Get batch Iterator."""
+    def _get_batch_iterator(self) -> Iterator[pd.DataFrame]:
+        """Lazily initializes and returns the batch iterator."""
         if self._batch_iterator is None:
             self._batch_iterator = iter(self._batches)
         return self._batch_iterator
@@ -150,12 +157,10 @@ def _set_table_html(self):
         start = self.page * self.page_size
         end = start + self.page_size
 
-        # fetch more dat if the requested page is outside our cache
-        while len(self._cached_data) < end:
-            prev_len = len(self._cached_data)
+        # fetch more data if the requested page is outside our cache
+        while len(self._cached_data) < end and not self._all_data_loaded:
             self._get_next_batch()
-            if len(self._cached_data) == prev_len:
-                break
+
         # Get the data fro the current page
         page_data = self._cached_data.iloc[start:end]
 
@@ -170,5 +175,5 @@ def _set_table_html(self):
 
     @traitlets.observe("page")
     def _page_changed(self, change):
-        """Handler for when the page nubmer is changed from the frontend"""
+        """Handler for when the page nubmer is changed from the frontend."""
         self._set_table_html()

notebooks/dataframes/anywidget_mode.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "d10bfca4",
    "metadata": {},
    "outputs": [],
@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "id": "ca22f059",
    "metadata": {},
    "outputs": [],
@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "1bc5aaf3",
    "metadata": {},
    "outputs": [],
@@ -63,33 +63,26 @@
    "id": "0a354c69",
    "metadata": {},
    "source": [
-    "Display the dataframe in anywidget mode"
+    "Load Sample Data"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "f289d250",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
-       "Query job 4d7057f3-6b68-46b2-b1e3-1dc71cb4682b is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:4d7057f3-6b68-46b2-b1e3-1dc71cb4682b&page=queryresults\">Open Job</a>"
+       "Query job 6fddde2a-5b5f-4920-a0b1-cb38e636bab3 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:US:6fddde2a-5b5f-4920-a0b1-cb38e636bab3&page=queryresults\">Open Job</a>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Computation deferred. Computation will process 171.4 MB\n"
-     ]
     }
    ],
    "source": [
@@ -110,7 +103,18 @@
    "execution_count": null,
    "id": "42bb02ab",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Computation deferred. Computation will process 171.4 MB"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "test_series = df[\"year\"]\n",
     "print(test_series)"
@@ -121,31 +125,113 @@
    "id": "7bcf1bb7",
    "metadata": {},
    "source": [
-    "Interactive BigFrames TableWidget"
+    "Display with Pagination"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 6,
    "id": "ce250157",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/google/home/shuowei/src/python-bigquery-dataframes/bigframes/display/anywidget.py:138: FutureWarning: The behavior of DataFrame concatenation with empty or all-NA entries is deprecated. In a future version, this will no longer exclude empty or all-NA columns when determining the result dtypes. To retain the old behavior, exclude the relevant entries before the concat operation.\n",
+      "  self._cached_data = pd.concat([self._cached_data, batch], ignore_index=True)\n",
+      "/usr/local/google/home/shuowei/src/python-bigquery-dataframes/venv/lib/python3.10/site-packages/IPython/core/formatters.py:429: FormatterWarning: text/html formatter returned invalid type <class 'bigframes.display.anywidget.TableWidget'> (expected <class 'str'>) for object: Computation deferred. Computation will process 171.4 MB\n",
+      "  warnings.warn(\n"
+     ]
+    },
     {
      "data": {
-      "text/html": [
-       "Computation deferred. Computation will process 171.4 MB"
-      ],
       "text/plain": [
        "Computation deferred. Computation will process 171.4 MB"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "df"
+    "df.head(50)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb15bab6",
+   "metadata": {},
+   "source": [
+    "Progarmmatic Navigation Demo"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6920d49b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from bigframes.display.anywidget import TableWidget\n",
+    "import math\n",
+    "  \n",
+    "# Create widget programmatically  \n",
+    "widget = TableWidget(df)  \n",
+    "print(f\"Total pages: {math.ceil(widget.row_count / widget.page_size)}\")  \n",
+    "  \n",
+    "# Display the widget  \n",
+    "widget"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "02cbd1be",
+   "metadata": {},
+   "source": [
+    "Test Navigation Programmatically"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12b68f15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simulate button clicks programmatically  \n",
+    "print(\"Current page:\", widget.page)  \n",
+    "  \n",
+    "# Go to next page  \n",
+    "widget.page = 1  \n",
+    "print(\"After next:\", widget.page)  \n",
+    "  \n",
+    "# Go to previous page  \n",
+    "widget.page = 0  \n",
+    "print(\"After prev:\", widget.page)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9d310138",
+   "metadata": {},
+   "source": [
+    "Edge Case Demonstration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9d5d13a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test with very small dataset  \n",
+    "small_df = df.head(5)  \n",
+    "small_widget = TableWidget(small_df)  \n",
+    "print(f\"Small dataset pages: {math.ceil(small_widget.row_count / small_widget.page_size)}\")  \n",
+    "small_widget"
    ]
   }
  ],

tests/system/small/test_progress_bar.py

@@ -180,16 +180,3 @@ def test_repr_anywidget_index(penguins_df_default_index: bf.dataframe.DataFrame)
         index = penguins_df_default_index.index
         actual_repr = repr(index)
         assert EXPECTED_DRY_RUN_MESSAGE in actual_repr
-
-
-def test_repr_anywidget_pagination_buttons_initial_state(
-    penguins_df_default_index: bf.dataframe.DataFrame,
-):
-    pytest.importorskip("anywidget")
-    with bf.option_context("display.repr_mode", "anywidget"):
-        from bigframes.display import TableWidget
-
-        widget = TableWidget(penguins_df_default_index)
-        assert widget.page == 0
-        assert widget.page_size == bf.options.display.max_rows
-        assert widget.row_count > 0