move Queue construction to ResultSert

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -3,7 +3,6 @@
 import logging
 import math
 import time
-import uuid
 import threading
 from typing import List, Union, Any, TYPE_CHECKING
 
@@ -728,7 +727,7 @@ def _col_to_description(col):
         else:
             precision, scale = None, None
 
-        return col.columnName, cleaned_type, None, None, precision, scale, None
+        return [col.columnName, cleaned_type, None, None, precision, scale, None]
 
     @staticmethod
     def _hive_schema_to_description(t_table_schema):
@@ -778,23 +777,6 @@ def _results_message_to_execute_response(self, resp, operation_state):
             schema_bytes = None
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
-        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        if direct_results and direct_results.resultSet:
-            assert direct_results.resultSet.results.startRowOffset == 0
-            assert direct_results.resultSetMetadata
-
-            arrow_queue_opt = ResultSetQueueFactory.build_queue(
-                row_set_type=t_result_set_metadata_resp.resultFormat,
-                t_row_set=direct_results.resultSet.results,
-                arrow_schema_bytes=schema_bytes,
-                max_download_threads=self.max_download_threads,
-                lz4_compressed=lz4_compressed,
-                description=description,
-                ssl_options=self._ssl_options,
-            )
-        else:
-            arrow_queue_opt = None
-
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
 
         status = CommandState.from_thrift_state(operation_state)
@@ -806,11 +788,11 @@ def _results_message_to_execute_response(self, resp, operation_state):
             status=status,
             description=description,
             has_more_rows=has_more_rows,
-            results_queue=arrow_queue_opt,
             has_been_closed_server_side=has_been_closed_server_side,
             lz4_compressed=lz4_compressed,
-            is_staging_operation=is_staging_operation,
+            is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
             arrow_schema_bytes=schema_bytes,
+            result_format=t_result_set_metadata_resp.resultFormat,
         )
 
     def get_execution_result(
@@ -837,9 +819,6 @@ def get_execution_result(
 
         t_result_set_metadata_resp = resp.resultSetMetadata
 
-        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
-        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        has_more_rows = resp.hasMoreRows
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
@@ -854,15 +833,9 @@ def get_execution_result(
         else:
             schema_bytes = None
 
-        queue = ResultSetQueueFactory.build_queue(
-            row_set_type=resp.resultSetMetadata.resultFormat,
-            t_row_set=resp.results,
-            arrow_schema_bytes=schema_bytes,
-            max_download_threads=self.max_download_threads,
-            lz4_compressed=lz4_compressed,
-            description=description,
-            ssl_options=self._ssl_options,
-        )
+        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
+        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
+        has_more_rows = resp.hasMoreRows
 
         status = self.get_query_state(command_id)
 
@@ -871,11 +844,11 @@ def get_execution_result(
             status=status,
             description=description,
             has_more_rows=has_more_rows,
-            results_queue=queue,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
             arrow_schema_bytes=schema_bytes,
+            result_format=t_result_set_metadata_resp.resultFormat,
         )
 
         return ThriftResultSet(
@@ -885,6 +858,9 @@ def get_execution_result(
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            t_row_set=resp.results,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -999,13 +975,20 @@ def execute_command(
         else:
             execute_response = self._handle_execute_response(resp, cursor)
 
+            t_row_set = None
+            if resp.directResults and resp.directResults.resultSet:
+                t_row_set = resp.directResults.resultSet.results
+
             return ThriftResultSet(
                 connection=cursor.connection,
                 execute_response=execute_response,
                 thrift_client=self,
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
                 use_cloud_fetch=use_cloud_fetch,
+                t_row_set=t_row_set,
+                max_download_threads=self.max_download_threads,
+                ssl_options=self._ssl_options,
             )
 
     def get_catalogs(
@@ -1029,13 +1012,20 @@ def get_catalogs(
 
         execute_response = self._handle_execute_response(resp, cursor)
 
+        t_row_set = None
+        if resp.directResults and resp.directResults.resultSet:
+            t_row_set = resp.directResults.resultSet.results
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
         )
 
     def get_schemas(
@@ -1063,13 +1053,20 @@ def get_schemas(
 
         execute_response = self._handle_execute_response(resp, cursor)
 
+        t_row_set = None
+        if resp.directResults and resp.directResults.resultSet:
+            t_row_set = resp.directResults.resultSet.results
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
         )
 
     def get_tables(
@@ -1101,13 +1098,20 @@ def get_tables(
 
         execute_response = self._handle_execute_response(resp, cursor)
 
+        t_row_set = None
+        if resp.directResults and resp.directResults.resultSet:
+            t_row_set = resp.directResults.resultSet.results
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
         )
 
     def get_columns(
@@ -1139,13 +1143,20 @@ def get_columns(
 
         execute_response = self._handle_execute_response(resp, cursor)
 
+        t_row_set = None
+        if resp.directResults and resp.directResults.resultSet:
+            t_row_set = resp.directResults.resultSet.results
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1203,6 +1214,8 @@ def fetch_results(
                 )
             )
 
+        from databricks.sql.utils import ResultSetQueueFactory
+
         queue = ResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,
             t_row_set=resp.results,

src/databricks/sql/backend/types.py

@@ -423,12 +423,10 @@ class ExecuteResponse:
 
     command_id: CommandId
     status: CommandState
-    description: Optional[
-        List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]]
-    ] = None
+    description: Optional[List[List[Any]]] = None
     has_more_rows: bool = False
-    results_queue: Optional[Any] = None
     has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
     is_staging_operation: bool = False
     arrow_schema_bytes: Optional[bytes] = None
+    result_format: Optional[Any] = None

src/databricks/sql/result_set.py

@@ -157,6 +157,9 @@ def __init__(
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
+        t_row_set=None,
+        max_download_threads: int = 10,
+        ssl_options=None,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -168,12 +171,31 @@ def __init__(
             buffer_size_bytes: Buffer size for fetching results
             arraysize: Default number of rows to fetch
             use_cloud_fetch: Whether to use cloud fetch for retrieving results
+            t_row_set: The TRowSet containing result data (if available)
+            max_download_threads: Maximum number of download threads for cloud fetch
+            ssl_options: SSL options for cloud fetch
         """
         # Initialize ThriftResultSet-specific attributes
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._use_cloud_fetch = use_cloud_fetch
         self.lz4_compressed = execute_response.lz4_compressed
 
+        # Build the results queue if t_row_set is provided
+        results_queue = None
+        if t_row_set and execute_response.result_format is not None:
+            from databricks.sql.utils import ResultSetQueueFactory
+            
+            # Create the results queue using the provided format
+            results_queue = ResultSetQueueFactory.build_queue(
+                row_set_type=execute_response.result_format,
+                t_row_set=t_row_set,
+                arrow_schema_bytes=execute_response.arrow_schema_bytes or b"",
+                max_download_threads=max_download_threads,
+                lz4_compressed=execute_response.lz4_compressed,
+                description=execute_response.description,
+                ssl_options=ssl_options,
+            )
+
         # Call parent constructor with common attributes
         super().__init__(
             connection=connection,
@@ -184,7 +206,7 @@ def __init__(
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
             has_more_rows=execute_response.has_more_rows,
-            results_queue=execute_response.results_queue,
+            results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
         )

tests/unit/test_client.py

@@ -104,6 +104,7 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 # Mock the backend that will be used by the real ThriftResultSet
                 mock_backend = Mock(spec=ThriftDatabricksClient)
                 mock_backend.staging_allowed_local_path = None
+                mock_backend.fetch_results.return_value = (Mock(), False)
 
                 # Configure the decorator's mock to return our specific mock_backend
                 mock_thrift_client_class.return_value = mock_backend
@@ -184,6 +185,7 @@ def test_arraysize_buffer_size_passthrough(
     def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_connection = Mock()
         mock_backend = Mock()
+        mock_backend.fetch_results.return_value = (Mock(), False)
 
         result_set = ThriftResultSet(
             connection=mock_connection,
@@ -210,6 +212,8 @@ def test_closing_result_set_hard_closes_commands(self):
         mock_session.open = True
         type(mock_connection).session = PropertyMock(return_value=mock_session)
 
+        mock_thrift_backend.fetch_results.return_value = (Mock(), False)
+
         result_set = ThriftResultSet(
             mock_connection, mock_results_response, mock_thrift_backend
         )
@@ -254,7 +258,10 @@ def test_closed_cursor_doesnt_allow_operations(self):
             self.assertIn("closed", e.msg)
 
     def test_negative_fetch_throws_exception(self):
-        result_set = ThriftResultSet(Mock(), Mock(), Mock())
+        mock_backend = Mock()
+        mock_backend.fetch_results.return_value = (Mock(), False)
+
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)

tests/unit/test_fetches.py

@@ -40,25 +40,34 @@ def make_dummy_result_set_from_initial_results(initial_results):
         # If the initial results have been set, then we should never try and fetch more
         schema, arrow_table = FetchTests.make_arrow_table(initial_results)
         arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
+
+        # Create a mock backend that will return the queue when _fill_results_buffer is called
+        mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
+        mock_thrift_backend.fetch_results.return_value = (arrow_queue, False)
+
+        num_cols = len(initial_results[0]) if initial_results else 0
+        description = [
+            (f"col{col_id}", "integer", None, None, None, None, None)
+            for col_id in range(num_cols)
+        ]
+
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=True,
                 has_more_rows=False,
-                description=Mock(),
-                lz4_compressed=Mock(),
-                results_queue=arrow_queue,
+                description=description,
+                lz4_compressed=True,
                 is_staging_operation=False,
             ),
-            thrift_client=None,
+            thrift_client=mock_thrift_backend,
+            t_row_set=None,
         )
-        num_cols = len(initial_results[0]) if initial_results else 0
-        rs.description = [
-            (f"col{col_id}", "integer", None, None, None, None, None)
-            for col_id in range(num_cols)
-        ]
+
+        # Replace the results queue with our arrow_queue
+        rs.results = arrow_queue
         return rs
 
     @staticmethod
@@ -85,19 +94,20 @@ def fetch_results(
         mock_thrift_backend.fetch_results = fetch_results
         num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
 
+        description = [
+            (f"col{col_id}", "integer", None, None, None, None, None)
+            for col_id in range(num_cols)
+        ]
+
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=False,
                 has_more_rows=True,
-                description=[
-                    (f"col{col_id}", "integer", None, None, None, None, None)
-                    for col_id in range(num_cols)
-                ],
-                lz4_compressed=Mock(),
-                results_queue=None,
+                description=description,
+                lz4_compressed=True,
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,