From e5d629e633815038df620f0b1c5816826bece38f Mon Sep 17 00:00:00 2001
From: shubha-rajan <shubhadayini@google.com>
Date: Tue, 3 Sep 2019 14:01:52 -0700
Subject: [PATCH 1/2] added max_results param to QueryJob.result()

---
 bigquery/google/cloud/bigquery/job.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
index 6e1eb81648f5..155474536d14 100644
--- a/bigquery/google/cloud/bigquery/job.py
+++ b/bigquery/google/cloud/bigquery/job.py
@@ -2905,7 +2905,9 @@ def _begin(self, client=None, retry=DEFAULT_RETRY):
             exc.message += self._format_for_exception(self.query, self.job_id)
             raise
 
-    def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY):
+    def result(
+        self, timeout=None, page_size=None, retry=DEFAULT_RETRY, max_results=None
+    ):
         """Start the job and wait for it to complete and get the result.
 
         Args:
@@ -2956,7 +2958,9 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY):
         dest_table_ref = self.destination
         dest_table = Table(dest_table_ref, schema=schema)
         dest_table._properties["numRows"] = self._query_results.total_rows
-        rows = self._client.list_rows(dest_table, page_size=page_size, retry=retry)
+        rows = self._client.list_rows(
+            dest_table, page_size=page_size, retry=retry, max_results=max_results
+        )
         rows._preserve_order = _contains_order_by(self.query)
         return rows
 

From 8f8bbef892219ad1cb68f25d600d300856b99389 Mon Sep 17 00:00:00 2001
From: shubha-rajan <shubhadayini@google.com>
Date: Tue, 3 Sep 2019 17:02:44 -0700
Subject: [PATCH 2/2] added test to check that max_results param appears in
 tabledata.list call arguments

---
 bigquery/tests/unit/test_job.py | 39 +++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
index 19409d8d43c3..98090a5257fd 100644
--- a/bigquery/tests/unit/test_job.py
+++ b/bigquery/tests/unit/test_job.py
@@ -4133,6 +4133,45 @@ def test_result(self):
         # on the response from tabledata.list.
         self.assertEqual(result.total_rows, 1)
 
+    def test_result_with_max_results(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"}]},
+            ],
+        }
+        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)
+
+        max_results = 3
+
+        result = job.result(max_results=max_results)
+
+        self.assertIsInstance(result, RowIterator)
+        self.assertEqual(result.total_rows, 5)
+
+        rows = list(result)
+
+        self.assertEqual(len(rows), 3)
+        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"]["maxResults"], max_results
+        )
+
     def test_result_w_empty_schema(self):
         from google.cloud.bigquery.table import _EmptyRowIterator