Reapply "Merge branch 'sea-res-set' into fetch-json-inline"

This reverts commit 597f116.

Author: varun-edachali-dbx

src/databricks/sql/result_set.py

@@ -48,10 +48,11 @@ def __init__(
         command_id: CommandId,
         status: CommandState,
         has_been_closed_server_side: bool = False,
-        has_more_rows: bool = False,
         results_queue=None,
         description=None,
         is_staging_operation: bool = False,
+        lz4_compressed: bool = False,
+        arrow_schema_bytes: bytes = b"",
     ):
         """
         A ResultSet manages the results of a single command.
@@ -79,9 +80,10 @@ def __init__(
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.has_more_rows = has_more_rows
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
+        self.lz4_compressed = lz4_compressed
+        self._arrow_schema_bytes = arrow_schema_bytes
 
     def __iter__(self):
         while True:
@@ -185,9 +187,24 @@ def __init__(
             has_more_rows: Whether there are more rows to 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
+        self.has_more_rows = has_more_rows
+
+        # 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 ThriftResultSetQueueFactory
+
+            # Create the results queue using the provided format
+            results_queue = ThriftResultSetQueueFactory.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__(
@@ -198,7 +215,6 @@ def __init__(
             command_id=execute_response.command_id,
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
-            has_more_rows=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
@@ -434,7 +450,7 @@ def map_col_type(type_):
 
 
 class SeaResultSet(ResultSet):
-    """ResultSet implementation for the SEA backend."""
+    """ResultSet implementation for SEA backend."""
 
     def __init__(
         self,
@@ -451,11 +467,12 @@ def __init__(
 
         Args:
             connection: The parent connection
+            execute_response: Response from the execute command
             sea_client: The SeaDatabricksClient instance for direct access
             buffer_size_bytes: Buffer size for fetching results
             arraysize: Default number of rows to fetch
-            execute_response: Response from the execute command (new style)
-            sea_response: Direct SEA response (legacy style)
+            result_data: Result data from SEA response (optional)
+            manifest: Manifest from SEA response (optional)
         """
 
         if result_data:
@@ -481,6 +498,8 @@ def __init__(
             results_queue=queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
+            lz4_compressed=execute_response.lz4_compressed,
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
     def _convert_to_row_objects(self, rows):