Support and tests for large queries

Added support for large queries:
- Still fetching results if there are more even if the fast path is taken
- Fixed a bug: we should actually always specify the offset we want to fetch from
- Setting the gRPC message size properly

New DriverLocal test for large queries

Author: Niall Egan <niall.egan@databricks.com>

Author: Niall Egan

cmdexec/clients/python/command_exec_client.py

@@ -13,7 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_BUFFER_SIZE_ROWS = 1000
+DEFAULT_RESULT_BUFFER_SIZE_BYTES = 10485760
 
 
 def connect(**kwargs):
@@ -48,10 +48,11 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
-    def cursor(self, buffer_size_rows=DEFAULT_BUFFER_SIZE_ROWS):
+    def cursor(self, buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES):
         if not self.open:
             raise Error("Cannot create cursor from closed connection")
-        cursor = Cursor(self, buffer_size_rows)
+
+        cursor = Cursor(self, buffer_size_bytes)
         self._cursors.append(cursor)
         return cursor
 
@@ -65,12 +66,12 @@ def close(self):
 
 
 class Cursor:
-    def __init__(self, connection, buffer_size_rows=DEFAULT_BUFFER_SIZE_ROWS):
+    def __init__(self, connection, result_buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES):
         self.connection = connection
         self.description = None
         self.rowcount = -1
         self.arraysize = 1
-        self.buffer_size_rows = buffer_size_rows
+        self.buffer_size_bytes = result_buffer_size_bytes
         self.active_result_set = None
         # Note that Cursor closed => active result set closed, but not vice versa
         self.open = True
@@ -92,10 +93,11 @@ def _response_to_result_set(self, execute_command_response, status):
         command_id = execute_command_response.command_id
         arrow_results = execute_command_response.results.arrow_ipc_stream
         has_been_closed_server_side = execute_command_response.closed
+        has_more_rows = execute_command_response.results.has_more_rows
         num_valid_rows = execute_command_response.results.num_valid_rows
 
         return ResultSet(self.connection, command_id, status, has_been_closed_server_side,
-                         arrow_results, num_valid_rows, self.buffer_size_rows)
+                         has_more_rows, arrow_results, num_valid_rows, self.buffer_size_bytes)
 
     def _close_and_clear_active_result_set(self):
         try:
@@ -127,15 +129,13 @@ def _poll_for_state(self, command_id):
 
     def _wait_until_command_done(self, command_id, initial_status):
         status = initial_status
-        print("initial status: %s" % status)
         while status in [command_pb2.PENDING, command_pb2.RUNNING]:
             resp = self._poll_for_state(command_id)
             status = resp.status.state
             self._check_response_for_error(resp, command_id)
-            print("status is: %s" % status)
 
             # TODO: Remove this sleep once we have long-polling on the server (SC-77653)
-            time.sleep(1)
+            time.sleep(0.1)
         return status
 
     def execute(self, operation, query_params=None, metadata=None):
@@ -197,23 +197,25 @@ def __init__(self,
                  command_id,
                  status,
                  has_been_closed_server_side,
+                 has_more_rows,
                  arrow_ipc_stream=None,
                  num_valid_rows=None,
-                 buffer_size_rows=DEFAULT_BUFFER_SIZE_ROWS):
+                 result_buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES):
         self.connection = connection
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.buffer_size_rows = buffer_size_rows
+        self.has_more_rows = has_more_rows
+        self.buffer_size_bytes = result_buffer_size_bytes
+        self._row_index = 0
 
         assert (self.status not in [command_pb2.PENDING, command_pb2.RUNNING])
 
         if arrow_ipc_stream:
             # In the case we are passed in an initial result set, the server has taken the
             # fast path and has no more rows to send
             self.results = ArrowQueue(
-                pyarrow.ipc.open_stream(arrow_ipc_stream).read_all(), num_valid_rows)
-            self.has_more_rows = False
+                pyarrow.ipc.open_stream(arrow_ipc_stream).read_all(), num_valid_rows, 0)
         else:
             # In this case, there are results waiting on the server so we fetch now for simplicity
             self._fill_results_buffer()
@@ -230,15 +232,17 @@ def _fetch_and_deserialize_results(self):
         fetch_results_request = command_pb2.FetchCommandResultsRequest(
             id=self.command_id,
             options=command_pb2.CommandResultOptions(
-                max_rows=self.buffer_size_rows,
+                max_bytes=self.buffer_size_bytes,
+                row_offset=self._row_index,
                 include_metadata=True,
             ))
 
         result_message = self.connection.base_client.make_request(
             self.connection.base_client.stub.FetchCommandResults, fetch_results_request).results
         num_valid_rows = result_message.num_valid_rows
         arrow_table = pyarrow.ipc.open_stream(result_message.arrow_ipc_stream).read_all()
-        results = ArrowQueue(arrow_table, num_valid_rows)
+        results = ArrowQueue(arrow_table, num_valid_rows,
+                             self._row_index - result_message.start_row_offset)
         return results, result_message.has_more_rows
 
     def _fill_results_buffer(self):
@@ -268,23 +272,29 @@ def fetchmany_arrow(self, n_rows):
             raise ValueError("n_rows argument for fetchmany is %s but must be >= 0", n_rows)
         results = self.results.next_n_rows(n_rows)
         n_remaining_rows = n_rows - results.num_rows
+        self._row_index += results.num_rows
 
         while n_remaining_rows > 0 and not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
             results = pyarrow.concat_tables([results, partial_results])
             n_remaining_rows -= partial_results.num_rows
+            self._row_index += partial_results.num_rows
 
         return results
 
     def fetchall_arrow(self):
         """
          Fetch all (remaining) rows of a query result, returning them as a PyArrow table.
         """
-        results = self.fetchmany_arrow(self.buffer_size_rows)
-        while self.has_more_rows:
-            # TODO: What's the optimal sequence of sizes to fetch?
-            results = pyarrow.concat_tables([results, self.fetchmany_arrow(self.buffer_size_rows)])
+        results = self.results.remaining_rows()
+        self._row_index += results.num_rows
+
+        while not self.has_been_closed_server_side and self.has_more_rows:
+            self._fill_results_buffer()
+            partial_results = self.results.remaining_rows()
+            results = pyarrow.concat_tables([results, partial_results])
+            self._row_index += partial_results.num_rows
 
         return results
 
@@ -293,6 +303,7 @@ def fetchone(self):
         Fetch the next row of a query result set, returning a single sequence,
         or None when no more data is available.
         """
+        self._row_index += 1
         res = self._convert_arrow_table(self.fetchmany_arrow(1))
         if len(res) > 0:
             return res[0]
@@ -329,8 +340,8 @@ def close(self):
 
 
 class ArrowQueue:
-    def __init__(self, arrow_table, n_valid_rows):
-        self.cur_row_index = 0
+    def __init__(self, arrow_table, n_valid_rows, start_row_index):
+        self.cur_row_index = start_row_index
         self.arrow_table = arrow_table
         self.n_valid_rows = n_valid_rows
 
@@ -343,6 +354,11 @@ def next_n_rows(self, num_rows):
         self.cur_row_index += slice.num_rows
         return slice
 
+    def remaining_rows(self):
+        slice = self.arrow_table.slice(self.cur_row_index, self.n_valid_rows - self.cur_row_index)
+        self.cur_row_index += slice.num_rows
+        return slice
+
 
 class CmdExecBaseHttpClient:
     """
@@ -352,7 +368,8 @@ class CmdExecBaseHttpClient:
     def __init__(self, host: str, port: int, http_headers: List[Tuple[str, str]]):
         self.host_url = host + ":" + str(port)
         self.http_headers = [(k.lower(), v) for (k, v) in http_headers]
-        self.channel = grpc.insecure_channel(self.host_url)
+        self.channel = grpc.insecure_channel(
+            self.host_url, options=[('grpc.max_receive_message_length', -1)])
         self.stub = SqlCommandServiceStub(self.channel)
 
     def make_request(self, method, request):

cmdexec/clients/python/tests/test_fetches.py

@@ -32,7 +32,7 @@ def make_arrow_ipc_stream(batch):
     @staticmethod
     def make_arrow_queue(batch):
         table = FetchTests.make_arrow_table(batch)
-        queue = command_exec_client.ArrowQueue(table, len(batch))
+        queue = command_exec_client.ArrowQueue(table, len(batch), 0)
         return queue
 
     @staticmethod
@@ -45,7 +45,8 @@ def make_dummy_result_set_from_initial_results(initial_results):
             None,
             True,
             arrow_ipc_stream=arrow_ipc_stream,
-            num_valid_rows=len(initial_results))
+            num_valid_rows=len(initial_results),
+            has_more_rows=False)
 
     @staticmethod
     def make_dummy_result_set_from_batch_list(batch_list):
@@ -59,7 +60,7 @@ def _fetch_and_deserialize_results(self):
 
                 return results, batch_index < len(batch_list)
 
-        return SemiFakeResultSet(None, None, None, False)
+        return SemiFakeResultSet(None, None, None, False, False)
 
     def test_fetchmany_with_initial_results(self):
         # Fetch all in one go

cmdexec/clients/python/tests/tests.py

@@ -100,7 +100,8 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(
             command_pb2.SUCCESS,
             False,
             arrow_ipc_stream=Mock(),
-            num_valid_rows=0)
+            num_valid_rows=0,
+            has_more_rows=False)
         mock_connection.open = False
 
         result_set.close()
@@ -115,9 +116,10 @@ def test_closing_result_set_hard_closes_commands(self, pyarrow_ipc_open_stream):
         mock_connection = Mock()
         mock_response = Mock()
         mock_response.id = b'\x22'
+        mock_response.results.start_row_offset = 0
         mock_connection.base_client.make_request.return_value = mock_response
-        result_set = command_exec_client.ResultSet(mock_connection, b'\x10', command_pb2.SUCCESS,
-                                                   False)
+        result_set = command_exec_client.ResultSet(
+            mock_connection, b'\x10', command_pb2.SUCCESS, False, has_more_rows=False)
         mock_connection.open = True
 
         result_set.close()
@@ -171,7 +173,15 @@ def test_closed_cursor_doesnt_allow_operations(self):
 
     @patch("pyarrow.ipc.open_stream")
     def test_negative_fetch_throws_exception(self, pyarrow_ipc_open_stream_mock):
-        result_set = command_exec_client.ResultSet(Mock(), b'\x22', command_pb2.SUCCESS, Mock())
+        mock_connection = Mock()
+        mock_response = Mock()
+        mock_response.id = b'\x22'
+        mock_response.results.start_row_offset = 0
+        mock_response.status.state = command_pb2.SUCCESS
+        mock_connection.base_client.make_request.return_value = mock_response
+
+        result_set = command_exec_client.ResultSet(
+            mock_connection, b'\x22', command_pb2.SUCCESS, Mock(), has_more_rows=False)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)