From 7276964a5949a20be518363ee62149d89aeda0e2 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Wed, 8 Oct 2025 01:56:44 +0000 Subject: [PATCH] fix: Yield row count from read session if otherwise unknown --- bigframes/session/read_api_execution.py | 2 +- tests/system/small/test_dataframe.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/bigframes/session/read_api_execution.py b/bigframes/session/read_api_execution.py index 037fde011f..2530a1dc8d 100644 --- a/bigframes/session/read_api_execution.py +++ b/bigframes/session/read_api_execution.py @@ -102,7 +102,7 @@ def process_page(page): if peek: batches = pyarrow_utils.truncate_pyarrow_iterable(batches, max_results=peek) - rows = node.source.n_rows + rows = node.source.n_rows or session.estimated_row_count if peek and rows: rows = min(peek, rows) diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 1a942a023e..851c934838 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -993,6 +993,12 @@ def test_filter_df(scalars_dfs): assert_pandas_df_equal(bf_result, pd_result) +def test_read_gbq_direct_to_batches_row_count(unordered_session): + df = unordered_session.read_gbq("bigquery-public-data.usa_names.usa_1910_2013") + iter = df.to_pandas_batches() + assert iter.total_rows == 5552452 + + def test_df_to_pandas_batches(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs