Skip to content

Commit 9e7b62b

Browse files
chemelnucfinchemelnucfin
authored
Bigquery: adds page_size parameter for list_rows and use in DB-API for arraysize (googleapis#4931)
1 parent b5404a3 commit 9e7b62b

5 files changed

Lines changed: 87 additions & 6 deletions

File tree

‎bigquery/google/cloud/bigquery/client.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1247,7 +1247,8 @@ def create_rows_json(self, *args, **kwargs):
12471247
return self.insert_rows_json(*args, **kwargs)
12481248

12491249
def list_rows(self, table, selected_fields=None, max_results=None,
1250-
page_token=None, start_index=None, retry=DEFAULT_RETRY):
1250+
page_token=None, start_index=None, page_size=None,
1251+
retry=DEFAULT_RETRY):
12511252
"""List the rows of the table.
12521253
12531254
See
@@ -1287,6 +1288,10 @@ def list_rows(self, table, selected_fields=None, max_results=None,
12871288
:param start_index: (Optional) The zero-based index of the starting
12881289
row to read.
12891290
1291+
:type page_size: int
1292+
:param page_size: (Optional) The maximum number of items to return
1293+
per page in the iterator.
1294+
12901295
:type retry: :class:`google.api_core.retry.Retry`
12911296
:param retry: (Optional) How to retry the RPC.
12921297
@@ -1314,7 +1319,6 @@ def list_rows(self, table, selected_fields=None, max_results=None,
13141319
if selected_fields is not None:
13151320
params['selectedFields'] = ','.join(
13161321
field.name for field in selected_fields)
1317-
13181322
if start_index is not None:
13191323
params['startIndex'] = start_index
13201324

@@ -1325,6 +1329,7 @@ def list_rows(self, table, selected_fields=None, max_results=None,
13251329
schema=schema,
13261330
page_token=page_token,
13271331
max_results=max_results,
1332+
page_size=page_size,
13281333
extra_params=params)
13291334
return row_iterator
13301335

‎bigquery/google/cloud/bigquery/dbapi/cursor.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,10 +190,11 @@ def _try_fetch(self, size=None):
190190

191191
if self._query_data is None:
192192
client = self.connection._client
193-
# TODO(tswast): pass in page size to list_rows based on arraysize
194193
rows_iter = client.list_rows(
195194
self._query_job.destination,
196-
selected_fields=self._query_job._query_results.schema)
195+
selected_fields=self._query_job._query_results.schema,
196+
page_size=self.arraysize
197+
)
197198
self._query_data = iter(rows_iter)
198199

199200
def fetchone(self):

‎bigquery/google/cloud/bigquery/table.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1148,14 +1148,15 @@ class RowIterator(HTTPIterator):
11481148
path (str): The method path to query for the list of items.
11491149
page_token (str): A token identifying a page in a result set to start
11501150
fetching results from.
1151-
max_results (int): The maximum number of results to fetch.
1151+
max_results (int): (Optional) The maximum number of results to fetch.
1152+
page_size (int): (Optional) The number of items to return per page.
11521153
extra_params (dict): Extra query string parameters for the API call.
11531154
11541155
.. autoattribute:: pages
11551156
"""
11561157

11571158
def __init__(self, client, api_request, path, schema, page_token=None,
1158-
max_results=None, extra_params=None):
1159+
max_results=None, page_size=None, extra_params=None):
11591160
super(RowIterator, self).__init__(
11601161
client, api_request, path, item_to_value=_item_to_row,
11611162
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,
11641165
self._schema = schema
11651166
self._field_to_index = _field_to_index_mapping(schema)
11661167
self._total_rows = None
1168+
self._page_size = page_size
1169+
1170+
def _get_next_page_response(self):
1171+
"""Requests the next page from the path provided.
1172+
1173+
Returns:
1174+
dict: The parsed JSON response of the next page's contents.
1175+
"""
1176+
params = self._get_query_params()
1177+
if self._page_size is not None:
1178+
params['maxResults'] = self._page_size
1179+
return self.api_request(
1180+
method=self._HTTP_METHOD,
1181+
path=self.path,
1182+
query_params=params)
11671183

11681184
@property
11691185
def schema(self):

‎bigquery/tests/system.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1737,6 +1737,42 @@ def test_nested_table_to_dataframe(self):
17371737
row['record_col']['nested_record']['nested_nested_string'],
17381738
'some deep insight')
17391739

1740+
def test_list_rows_page_size(self):
1741+
from google.cloud.bigquery.job import SourceFormat
1742+
from google.cloud.bigquery.job import WriteDisposition
1743+
1744+
num_items = 7
1745+
page_size = 3
1746+
num_pages, num_last_page = divmod(num_items, page_size)
1747+
1748+
SF = bigquery.SchemaField
1749+
schema = [SF('string_col', 'STRING', mode='NULLABLE')]
1750+
to_insert = [{'string_col': 'item%d' % i} for i in range(num_items)]
1751+
rows = [json.dumps(row) for row in to_insert]
1752+
body = six.BytesIO('{}\n'.format('\n'.join(rows)))
1753+
1754+
table_id = 'test_table'
1755+
dataset = self.temp_dataset(_make_dataset_id('nested_df'))
1756+
table = dataset.table(table_id)
1757+
self.to_delete.insert(0, table)
1758+
job_config = bigquery.LoadJobConfig()
1759+
job_config.write_disposition = WriteDisposition.WRITE_TRUNCATE
1760+
job_config.source_format = SourceFormat.NEWLINE_DELIMITED_JSON
1761+
job_config.schema = schema
1762+
# Load a table using a local JSON file from memory.
1763+
Config.CLIENT.load_table_from_file(
1764+
body, table, job_config=job_config).result()
1765+
1766+
df = Config.CLIENT.list_rows(
1767+
table, selected_fields=schema, page_size=page_size)
1768+
pages = df.pages
1769+
1770+
for i in range(num_pages):
1771+
page = next(pages)
1772+
self.assertEqual(page.num_items, page_size)
1773+
page = next(pages)
1774+
self.assertEqual(page.num_items, num_last_page)
1775+
17401776
def temp_dataset(self, dataset_id, location=None):
17411777
dataset = Dataset(Config.CLIENT.dataset(dataset_id))
17421778
if location:

‎bigquery/tests/unit/test_table.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1071,6 +1071,29 @@ def test_iterate(self):
10711071
api_request.assert_called_once_with(
10721072
method='GET', path=path, query_params={})
10731073

1074+
def test_page_size(self):
1075+
from google.cloud.bigquery.table import RowIterator
1076+
from google.cloud.bigquery.table import SchemaField
1077+
1078+
schema = [
1079+
SchemaField('name', 'STRING', mode='REQUIRED'),
1080+
SchemaField('age', 'INTEGER', mode='REQUIRED')
1081+
]
1082+
rows = [
1083+
{'f': [{'v': 'Phred Phlyntstone'}, {'v': '32'}]},
1084+
{'f': [{'v': 'Bharney Rhubble'}, {'v': '33'}]},
1085+
]
1086+
path = '/foo'
1087+
api_request = mock.Mock(return_value={'rows': rows})
1088+
1089+
row_iterator = RowIterator(
1090+
mock.sentinel.client, api_request, path, schema, page_size=4)
1091+
row_iterator._get_next_page_response()
1092+
1093+
api_request.assert_called_once_with(
1094+
method='GET', path=path, query_params={
1095+
'maxResults': row_iterator._page_size})
1096+
10741097
@unittest.skipIf(pandas is None, 'Requires `pandas`')
10751098
def test_to_dataframe(self):
10761099
from google.cloud.bigquery.table import RowIterator

0 commit comments

Comments
 (0)