perf: avoid extra API calls from to_dataframe if all rows are cached (#384)

Follow-up to previous PR, which cached the first page of `getQueryResults`. If the first page is the only page (no `pageToken`), then it is unnecessary to make extra API calls from `to_dataframe` or `to_arrow` to the BigQuery Storage API.

Author: tswast

google/cloud/bigquery/table.py

@@ -1351,6 +1351,41 @@ def __init__(
  self._total_rows = total_rows
  self._first_page_response = first_page_response
 
+ def _is_completely_cached(self):
+ """Check if all results are completely cached.
+
+ This is useful to know, because we can avoid alternative download
+ mechanisms.
+ """
+ if self._first_page_response is None or self.next_page_token:
+ return False
+
+ return self._first_page_response.get(self._next_token) is None
+
+ def _validate_bqstorage(self, bqstorage_client, create_bqstorage_client):
+ """Returns if the BigQuery Storage API can be used.
+
+ Returns:
+ bool
+ True if the BigQuery Storage client can be used or created.
+ """
+ using_bqstorage_api = bqstorage_client or create_bqstorage_client
+ if not using_bqstorage_api:
+ return False
+
+ if self._is_completely_cached():
+ return False
+
+ if self.max_results is not None:
+ warnings.warn(
+ "Cannot use bqstorage_client if max_results is set, "
+ "reverting to fetching data with the REST endpoint.",
+ stacklevel=2,
+ )
+ return False
+
+ return True
+
  def _get_next_page_response(self):
  """Requests the next page from the path provided.
 
@@ -1412,6 +1447,9 @@ def _get_progress_bar(self, progress_bar_type):
  def _to_page_iterable(
  self, bqstorage_download, tabledata_list_download, bqstorage_client=None
  ):
+ if not self._validate_bqstorage(bqstorage_client, False):
+ bqstorage_client = None
+
  if bqstorage_client is not None:
  for item in bqstorage_download():
  yield item
@@ -1503,14 +1541,7 @@ def to_arrow(
  if pyarrow is None:
  raise ValueError(_NO_PYARROW_ERROR)
 
- if (
- bqstorage_client or create_bqstorage_client
- ) and self.max_results is not None:
- warnings.warn(
- "Cannot use bqstorage_client if max_results is set, "
- "reverting to fetching data with the REST endpoint.",
- stacklevel=2,
- )
+ if not self._validate_bqstorage(bqstorage_client, create_bqstorage_client):
  create_bqstorage_client = False
  bqstorage_client = None
 
@@ -1687,14 +1718,7 @@ def to_dataframe(
  if dtypes is None:
  dtypes = {}
 
- if (
- bqstorage_client or create_bqstorage_client
- ) and self.max_results is not None:
- warnings.warn(
- "Cannot use bqstorage_client if max_results is set, "
- "reverting to fetching data with the REST endpoint.",
- stacklevel=2,
- )
+ if not self._validate_bqstorage(bqstorage_client, create_bqstorage_client):
  create_bqstorage_client = False
  bqstorage_client = None
 

tests/unit/job/test_query_pandas.py

@@ -99,6 +99,7 @@ def test_to_dataframe_bqstorage_preserve_order(query):
  ]
  },
  "totalRows": "4",
+ "pageToken": "next-page",
  }
  connection = _make_connection(get_query_results_resource, job_resource)
  client = _make_client(connection=connection)
@@ -133,7 +134,16 @@ def test_to_dataframe_bqstorage_preserve_order(query):
 
 
 @pytest.mark.skipif(pyarrow is None, reason="Requires `pyarrow`")
-def test_to_arrow():
+@pytest.mark.parametrize(
+ "method_kwargs",
+ [
+ {"create_bqstorage_client": False},
+ # Since all rows are contained in the first page of results, the BigQuery
+ # Storage API won't actually be used.
+ {"create_bqstorage_client": True},
+ ],
+)
+def test_to_arrow(method_kwargs):
  from google.cloud.bigquery.job import QueryJob as target_class
 
  begun_resource = _make_job_resource(job_type="query")
@@ -182,7 +192,7 @@ def test_to_arrow():
  client = _make_client(connection=connection)
  job = target_class.from_api_repr(begun_resource, client)
 
- tbl = job.to_arrow(create_bqstorage_client=False)
+ tbl = job.to_arrow(**method_kwargs)
 
  assert isinstance(tbl, pyarrow.Table)
  assert tbl.num_rows == 2
@@ -216,7 +226,16 @@ def test_to_arrow():
 
 
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
-def test_to_dataframe():
+@pytest.mark.parametrize(
+ "method_kwargs",
+ [
+ {"create_bqstorage_client": False},
+ # Since all rows are contained in the first page of results, the BigQuery
+ # Storage API won't actually be used.
+ {"create_bqstorage_client": True},
+ ],
+)
+def test_to_dataframe(method_kwargs):
  from google.cloud.bigquery.job import QueryJob as target_class
 
  begun_resource = _make_job_resource(job_type="query")
@@ -243,7 +262,7 @@ def test_to_dataframe():
  client = _make_client(connection=connection)
  job = target_class.from_api_repr(begun_resource, client)
 
- df = job.to_dataframe(create_bqstorage_client=False)
+ df = job.to_dataframe(**method_kwargs)
 
  assert isinstance(df, pandas.DataFrame)
  assert len(df) == 4 # verify the number of rows
@@ -288,6 +307,7 @@ def test_to_dataframe_bqstorage():
  {"name": "age", "type": "INTEGER", "mode": "NULLABLE"},
  ]
  },
+ "pageToken": "next-page",
  }
  connection = _make_connection(query_resource)
  client = _make_client(connection=connection)