Repository: googleapis/google-cloud-python

Issue:
BigQuery: add page_size parameter for list_rows and use in DB-API for arraysize
With https://github.com/GoogleCloudPlatform/google-cloud-python/pull/4144 the arraysize parameter no longer affects how much data the library actually requests from the API.

We should add a `page_size` parameter to `Client.list_rows` and use that parameter in the DB-API implementation.

See also: https://github.com/GoogleCloudPlatform/google-cloud-python/issues/4134


Discussion:

Resolving patch:
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
--- a/bigquery/google/cloud/bigquery/client.py
+++ b/bigquery/google/cloud/bigquery/client.py
@@ -1247,7 +1247,8 @@ def create_rows_json(self, *args, **kwargs):
         return self.insert_rows_json(*args, **kwargs)
 
     def list_rows(self, table, selected_fields=None, max_results=None,
-                  page_token=None, start_index=None, retry=DEFAULT_RETRY):
+                  page_token=None, start_index=None, page_size=None,
+                  retry=DEFAULT_RETRY):
         """List the rows of the table.
 
         See
@@ -1287,6 +1288,10 @@ def list_rows(self, table, selected_fields=None, max_results=None,
         :param start_index: (Optional) The zero-based index of the starting
                            row to read.
 
+        :type page_size: int
+        :param page_size: (Optional) The maximum number of items to return
+                          per page in the iterator.
+
         :type retry: :class:`google.api_core.retry.Retry`
         :param retry: (Optional) How to retry the RPC.
 
@@ -1314,7 +1319,6 @@ def list_rows(self, table, selected_fields=None, max_results=None,
         if selected_fields is not None:
             params['selectedFields'] = ','.join(
                 field.name for field in selected_fields)
-
         if start_index is not None:
             params['startIndex'] = start_index
 
@@ -1325,6 +1329,7 @@ def list_rows(self, table, selected_fields=None, max_results=None,
             schema=schema,
             page_token=page_token,
             max_results=max_results,
+            page_size=page_size,
             extra_params=params)
         return row_iterator
 
diff --git a/bigquery/google/cloud/bigquery/dbapi/cursor.py b/bigquery/google/cloud/bigquery/dbapi/cursor.py
--- a/bigquery/google/cloud/bigquery/dbapi/cursor.py
+++ b/bigquery/google/cloud/bigquery/dbapi/cursor.py
@@ -190,10 +190,11 @@ def _try_fetch(self, size=None):
 
         if self._query_data is None:
             client = self.connection._client
-            # TODO(tswast): pass in page size to list_rows based on arraysize
             rows_iter = client.list_rows(
                 self._query_job.destination,
-                selected_fields=self._query_job._query_results.schema)
+                selected_fields=self._query_job._query_results.schema,
+                page_size=self.arraysize
+            )
             self._query_data = iter(rows_iter)
 
     def fetchone(self):
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
--- a/bigquery/google/cloud/bigquery/table.py
+++ b/bigquery/google/cloud/bigquery/table.py
@@ -1148,14 +1148,15 @@ class RowIterator(HTTPIterator):
         path (str): The method path to query for the list of items.
         page_token (str): A token identifying a page in a result set to start
             fetching results from.
-        max_results (int): The maximum number of results to fetch.
+        max_results (int): (Optional) The maximum number of results to fetch.
+        page_size (int): (Optional) The number of items to return per page.
         extra_params (dict): Extra query string parameters for the API call.
 
     .. autoattribute:: pages
     """
 
     def __init__(self, client, api_request, path, schema, page_token=None,
-                 max_results=None, extra_params=None):
+                 max_results=None, page_size=None, extra_params=None):
         super(RowIterator, self).__init__(
             client, api_request, path, item_to_value=_item_to_row,
             items_key='rows', page_token=page_token, max_results=max_results,
@@ -1164,6 +1165,21 @@ def __init__(self, client, api_request, path, schema, page_token=None,
         self._schema = schema
         self._field_to_index = _field_to_index_mapping(schema)
         self._total_rows = None
+        self._page_size = page_size
+
+    def _get_next_page_response(self):
+        """Requests the next page from the path provided.
+
+        Returns:
+            dict: The parsed JSON response of the next page's contents.
+        """
+        params = self._get_query_params()
+        if self._page_size is not None:
+            params['maxResults'] = self._page_size
+        return self.api_request(
+            method=self._HTTP_METHOD,
+            path=self.path,
+            query_params=params)
 
     @property
     def schema(self):


Test patch:
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
--- a/bigquery/tests/system.py
+++ b/bigquery/tests/system.py
@@ -1737,6 +1737,42 @@ def test_nested_table_to_dataframe(self):
             row['record_col']['nested_record']['nested_nested_string'],
             'some deep insight')
 
+    def test_list_rows_page_size(self):
+        from google.cloud.bigquery.job import SourceFormat
+        from google.cloud.bigquery.job import WriteDisposition
+
+        num_items = 7
+        page_size = 3
+        num_pages, num_last_page = divmod(num_items, page_size)
+
+        SF = bigquery.SchemaField
+        schema = [SF('string_col', 'STRING', mode='NULLABLE')]
+        to_insert = [{'string_col': 'item%d' % i} for i in range(num_items)]
+        rows = [json.dumps(row) for row in to_insert]
+        body = six.BytesIO('{}\n'.format('\n'.join(rows)))
+
+        table_id = 'test_table'
+        dataset = self.temp_dataset(_make_dataset_id('nested_df'))
+        table = dataset.table(table_id)
+        self.to_delete.insert(0, table)
+        job_config = bigquery.LoadJobConfig()
+        job_config.write_disposition = WriteDisposition.WRITE_TRUNCATE
+        job_config.source_format = SourceFormat.NEWLINE_DELIMITED_JSON
+        job_config.schema = schema
+        # Load a table using a local JSON file from memory.
+        Config.CLIENT.load_table_from_file(
+            body, table, job_config=job_config).result()
+
+        df = Config.CLIENT.list_rows(
+            table, selected_fields=schema, page_size=page_size)
+        pages = df.pages
+
+        for i in range(num_pages):
+            page = next(pages)
+            self.assertEqual(page.num_items, page_size)
+        page = next(pages)
+        self.assertEqual(page.num_items, num_last_page)
+
     def temp_dataset(self, dataset_id, location=None):
         dataset = Dataset(Config.CLIENT.dataset(dataset_id))
         if location:
diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py
--- a/bigquery/tests/unit/test_table.py
+++ b/bigquery/tests/unit/test_table.py
@@ -1071,6 +1071,29 @@ def test_iterate(self):
         api_request.assert_called_once_with(
             method='GET', path=path, query_params={})
 
+    def test_page_size(self):
+        from google.cloud.bigquery.table import RowIterator
+        from google.cloud.bigquery.table import SchemaField
+
+        schema = [
+            SchemaField('name', 'STRING', mode='REQUIRED'),
+            SchemaField('age', 'INTEGER', mode='REQUIRED')
+        ]
+        rows = [
+            {'f': [{'v': 'Phred Phlyntstone'}, {'v': '32'}]},
+            {'f': [{'v': 'Bharney Rhubble'}, {'v': '33'}]},
+        ]
+        path = '/foo'
+        api_request = mock.Mock(return_value={'rows': rows})
+
+        row_iterator = RowIterator(
+            mock.sentinel.client, api_request, path, schema, page_size=4)
+        row_iterator._get_next_page_response()
+
+        api_request.assert_called_once_with(
+            method='GET', path=path, query_params={
+                'maxResults': row_iterator._page_size})
+
     @unittest.skipIf(pandas is None, 'Requires `pandas`')
     def test_to_dataframe(self):
         from google.cloud.bigquery.table import RowIterator