From be86de330a3c3801653a0ccef90e3d9bdb3eee7a Mon Sep 17 00:00:00 2001
From: HemangChothani <50404902+HemangChothani@users.noreply.github.com>
Date: Sat, 6 Jun 2020 14:48:46 +0530
Subject: [PATCH] feat(bigquery): expose start index parameter for query result
 (#121)

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

* feat(bigquery): nit
---
 google/cloud/bigquery/job.py | 10 ++++++++-
 tests/system.py              | 12 +++++++++++
 tests/unit/test_job.py       | 40 ++++++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigquery/job.py b/google/cloud/bigquery/job.py
index 52683f234..7a1178a8c 100644
--- a/google/cloud/bigquery/job.py
+++ b/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,
         )
diff --git a/tests/system.py b/tests/system.py
index b47195084..66d7ee259 100644
--- a/tests/system.py
+++ b/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.
diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py
index d97efd946..c89cad749 100644
--- a/tests/unit/test_job.py
+++ b/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