Skip to content

Commit be86de3

Browse files
feat(bigquery): expose start index parameter for query result (#121)
* feat(bigquery): expose start index parameter for query result * feat(bigquery): nit
1 parent cfaa901 commit be86de3

File tree

3 files changed

+61
-1
lines changed

3 files changed

+61
-1
lines changed

google/cloud/bigquery/job.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3160,7 +3160,12 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None):
31603160
raise
31613161

31623162
def result(
3163-
self, page_size=None, max_results=None, retry=DEFAULT_RETRY, timeout=None
3163+
self,
3164+
page_size=None,
3165+
max_results=None,
3166+
retry=DEFAULT_RETRY,
3167+
timeout=None,
3168+
start_index=None,
31643169
):
31653170
"""Start the job and wait for it to complete and get the result.
31663171
@@ -3177,6 +3182,8 @@ def result(
31773182
before using ``retry``.
31783183
If multiple requests are made under the hood, ``timeout``
31793184
applies to each individual request.
3185+
start_index (Optional[int]):
3186+
The zero-based index of the starting row to read.
31803187
31813188
Returns:
31823189
google.cloud.bigquery.table.RowIterator:
@@ -3230,6 +3237,7 @@ def result(
32303237
dest_table,
32313238
page_size=page_size,
32323239
max_results=max_results,
3240+
start_index=start_index,
32333241
retry=retry,
32343242
timeout=timeout,
32353243
)

tests/system.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1577,6 +1577,18 @@ def test_query_w_page_size(self):
15771577
iterator = query_job.result(page_size=page_size)
15781578
self.assertEqual(next(iterator.pages).num_items, page_size)
15791579

1580+
def test_query_w_start_index(self):
1581+
start_index = 164652
1582+
query_job = Config.CLIENT.query(
1583+
"SELECT word FROM `bigquery-public-data.samples.shakespeare`;",
1584+
job_id_prefix="test_query_w_start_index_",
1585+
)
1586+
result1 = query_job.result(start_index=start_index)
1587+
total_rows = result1.total_rows
1588+
1589+
self.assertEqual(result1.extra_params["startIndex"], start_index)
1590+
self.assertEqual(len(list(result1)), total_rows - start_index)
1591+
15801592
def test_query_statistics(self):
15811593
"""
15821594
A system test to exercise some of the extended query statistics.

tests/unit/test_job.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4666,6 +4666,46 @@ def test_result_w_page_size(self):
46664666
]
46674667
)
46684668

4669+
def test_result_with_start_index(self):
4670+
from google.cloud.bigquery.table import RowIterator
4671+
4672+
query_resource = {
4673+
"jobComplete": True,
4674+
"jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
4675+
"schema": {"fields": [{"name": "col1", "type": "STRING"}]},
4676+
"totalRows": "5",
4677+
}
4678+
tabledata_resource = {
4679+
"totalRows": "5",
4680+
"pageToken": None,
4681+
"rows": [
4682+
{"f": [{"v": "abc"}]},
4683+
{"f": [{"v": "def"}]},
4684+
{"f": [{"v": "ghi"}]},
4685+
{"f": [{"v": "jkl"}]},
4686+
],
4687+
}
4688+
connection = _make_connection(query_resource, tabledata_resource)
4689+
client = _make_client(self.PROJECT, connection=connection)
4690+
resource = self._make_resource(ended=True)
4691+
job = self._get_target_class().from_api_repr(resource, client)
4692+
4693+
start_index = 1
4694+
4695+
result = job.result(start_index=start_index)
4696+
4697+
self.assertIsInstance(result, RowIterator)
4698+
self.assertEqual(result.total_rows, 5)
4699+
4700+
rows = list(result)
4701+
4702+
self.assertEqual(len(rows), 4)
4703+
self.assertEqual(len(connection.api_request.call_args_list), 2)
4704+
tabledata_list_request = connection.api_request.call_args_list[1]
4705+
self.assertEqual(
4706+
tabledata_list_request[1]["query_params"]["startIndex"], start_index
4707+
)
4708+
46694709
def test_result_error(self):
46704710
from google.cloud import exceptions
46714711

0 commit comments

Comments
 (0)