feat(bigquery): expose start index parameter for query result (#121)

* feat(bigquery): expose start index parameter for query result * feat(bigquery): nit

Author: HemangChothani

google/cloud/bigquery/job.py

@@ -3160,7 +3160,12 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None):
  raise
 
  def result(
- self, page_size=None, max_results=None, retry=DEFAULT_RETRY, timeout=None
+ self,
+ page_size=None,
+ max_results=None,
+ retry=DEFAULT_RETRY,
+ timeout=None,
+ start_index=None,
  ):
  """Start the job and wait for it to complete and get the result.
 
@@ -3177,6 +3182,8 @@ def result(
  before using ``retry``.
  If multiple requests are made under the hood, ``timeout``
  applies to each individual request.
+ start_index (Optional[int]):
+ The zero-based index of the starting row to read.
 
  Returns:
  google.cloud.bigquery.table.RowIterator:
@@ -3230,6 +3237,7 @@ def result(
  dest_table,
  page_size=page_size,
  max_results=max_results,
+ start_index=start_index,
  retry=retry,
  timeout=timeout,
  )

tests/system.py

@@ -1577,6 +1577,18 @@ def test_query_w_page_size(self):
  iterator = query_job.result(page_size=page_size)
  self.assertEqual(next(iterator.pages).num_items, page_size)
 
+ def test_query_w_start_index(self):
+ start_index = 164652
+ query_job = Config.CLIENT.query(
+ "SELECT word FROM `bigquery-public-data.samples.shakespeare`;",
+ job_id_prefix="test_query_w_start_index_",
+ )
+ result1 = query_job.result(start_index=start_index)
+ total_rows = result1.total_rows
+
+ self.assertEqual(result1.extra_params["startIndex"], start_index)
+ self.assertEqual(len(list(result1)), total_rows - start_index)
+
  def test_query_statistics(self):
  """
  A system test to exercise some of the extended query statistics.

tests/unit/test_job.py

@@ -4666,6 +4666,46 @@ def test_result_w_page_size(self):
  ]
  )
 
+ def test_result_with_start_index(self):
+ from google.cloud.bigquery.table import RowIterator
+
+ query_resource = {
+ "jobComplete": True,
+ "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+ "schema": {"fields": [{"name": "col1", "type": "STRING"}]},
+ "totalRows": "5",
+ }
+ tabledata_resource = {
+ "totalRows": "5",
+ "pageToken": None,
+ "rows": [
+ {"f": [{"v": "abc"}]},
+ {"f": [{"v": "def"}]},
+ {"f": [{"v": "ghi"}]},
+ {"f": [{"v": "jkl"}]},
+ ],
+ }
+ connection = _make_connection(query_resource, tabledata_resource)
+ client = _make_client(self.PROJECT, connection=connection)
+ resource = self._make_resource(ended=True)
+ job = self._get_target_class().from_api_repr(resource, client)
+
+ start_index = 1
+
+ result = job.result(start_index=start_index)
+
+ self.assertIsInstance(result, RowIterator)
+ self.assertEqual(result.total_rows, 5)
+
+ rows = list(result)
+
+ self.assertEqual(len(rows), 4)
+ self.assertEqual(len(connection.api_request.call_args_list), 2)
+ tabledata_list_request = connection.api_request.call_args_list[1]
+ self.assertEqual(
+ tabledata_list_request[1]["query_params"]["startIndex"], start_index
+ )
+
  def test_result_error(self):
  from google.cloud import exceptions