Add max_results parameter to QueryJob.result(). (#9167)

* added max_results param to QueryJob.result()

* added test to check that max_results param appears in tabledata.list call arguments

Author: shubha-rajan

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
 

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