add assert for total_rows

Author: shuoweil

tests/benchmark/read_gbq_colab/aggregate_output.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pathlib
+import typing
 
 import benchmark.utils as utils
 
@@ -27,6 +28,7 @@ def aggregate_output(*, project_id, dataset_id, table_id):
 
     # Simulate getting the first page, since we'll always do that first in the UI.
     batches = df.to_pandas_batches(page_size=PAGE_SIZE)
+    assert typing.cast(typing.Any, batches).total_rows >= 0
     next(iter(batches))
 
     # To simulate very small rows that can only fit a boolean,

tests/benchmark/read_gbq_colab/filter_output.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pathlib
+import typing
 
 import benchmark.utils as utils
 
+import bigframes.core.blocks
 import bigframes.pandas as bpd
 
 PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE
@@ -37,13 +39,14 @@ def filter_output(
     # Simulate the user filtering by a column and visualizing those results
     df_filtered = df[df["col_bool_0"]]
     batches_filtered = df_filtered.to_pandas_batches(page_size=PAGE_SIZE)
-    assert batches_filtered.total_rows >= 0
-
+    batches_filtered = typing.cast(
+        bigframes.core.blocks.PandasBatches, batches_filtered
+    )
+    rows = batches_filtered.total_rows
+    assert rows >= 0
     # It's possible we don't have any pages at all, since we filtered out all
     # matching rows.
     first_page = next(iter(batches_filtered))
-    rows = batches_filtered.total_rows
-    assert rows is not None
     assert len(first_page.index) <= rows
 
 

tests/benchmark/read_gbq_colab/first_page.py

@@ -30,10 +30,9 @@ def first_page(*, project_id, dataset_id, table_id):
 
     # Get number of rows (to calculate number of pages) and the first page.
     batches = df.to_pandas_batches(page_size=PAGE_SIZE)
+    assert typing.cast(typing.Any, batches).total_rows >= 0
     first_page = next(iter(batches))
     assert first_page is not None
-    total_rows = typing.cast(typing.Any, batches).total_rows
-    assert total_rows is not None
 
 
 if __name__ == "__main__":

tests/benchmark/read_gbq_colab/sort_output.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pathlib
+import typing
 
 import benchmark.utils as utils
 
@@ -29,6 +30,7 @@ def sort_output(*, project_id, dataset_id, table_id):
 
     # Simulate getting the first page, since we'll always do that first in the UI.
     batches = df.to_pandas_batches(page_size=PAGE_SIZE)
+    assert typing.cast(typing.Any, batches).total_rows >= 0
     next(iter(batches))
 
     # Simulate the user sorting by a column and visualizing those results
@@ -38,6 +40,7 @@ def sort_output(*, project_id, dataset_id, table_id):
 
     df_sorted = df.sort_values(sort_column)
     batches_sorted = df_sorted.to_pandas_batches(page_size=PAGE_SIZE)
+    assert typing.cast(typing.Any, batches_sorted).total_rows >= 0
     next(iter(batches_sorted))