Merge branch 'exec-resp-norm' into fetch-json-inline

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -3,22 +3,20 @@
 import logging
 import math
 import time
-import uuid
 import threading
 from typing import List, Union, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
-from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
     CommandState,
     SessionId,
     CommandId,
-    BackendType,
-    guid_to_hex_id,
     ExecuteResponse,
 )
+from databricks.sql.backend.utils import guid_to_hex_id
+
 
 try:
     import pyarrow
@@ -759,11 +757,13 @@ def _results_message_to_execute_response(self, resp, operation_state):
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
+
         has_more_rows = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
         )
+
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
@@ -779,43 +779,25 @@ 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 = ThriftResultSetQueueFactory.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)
         if status is None:
             raise ValueError(f"Unknown command state: {operation_state}")
 
-        return (
-            ExecuteResponse(
-                command_id=command_id,
-                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,
-            ),
-            schema_bytes,
+        execute_response = ExecuteResponse(
+            command_id=command_id,
+            status=status,
+            description=description,
+            has_been_closed_server_side=has_been_closed_server_side,
+            lz4_compressed=lz4_compressed,
+            is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
+            arrow_schema_bytes=schema_bytes,
+            result_format=t_result_set_metadata_resp.resultFormat,
         )
 
+        return execute_response, has_more_rows
+
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
     ) -> "ResultSet":
@@ -840,9 +822,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
         )
@@ -857,27 +836,21 @@ def get_execution_result(
         else:
             schema_bytes = None
 
-        queue = ThriftResultSetQueueFactory.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)
 
         execute_response = ExecuteResponse(
             command_id=command_id,
             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(
@@ -887,7 +860,10 @@ def get_execution_result(
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=schema_bytes,
+            t_row_set=resp.results,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -918,7 +894,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         self._check_command_not_in_error_or_closed_state(thrift_handle, poll_resp)
         state = CommandState.from_thrift_state(operation_state)
         if state is None:
-            raise ValueError(f"Invalid operation state: {operation_state}")
+            raise ValueError(f"Unknown command state: {operation_state}")
         return state
 
     @staticmethod
@@ -1000,18 +976,25 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, arrow_schema_bytes = self._handle_execute_response(
+            execute_response, has_more_rows = 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,
-                arrow_schema_bytes=arrow_schema_bytes,
+                t_row_set=t_row_set,
+                max_download_threads=self.max_download_threads,
+                ssl_options=self._ssl_options,
+                has_more_rows=has_more_rows,
             )
 
     def get_catalogs(
@@ -1033,9 +1016,11 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = 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,
@@ -1044,7 +1029,10 @@ def get_catalogs(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_schemas(
@@ -1070,9 +1058,11 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = 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,
@@ -1081,7 +1071,10 @@ def get_schemas(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_tables(
@@ -1111,9 +1104,11 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = 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,
@@ -1122,7 +1117,10 @@ def get_tables(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_columns(
@@ -1152,9 +1150,11 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = 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,
@@ -1163,7 +1163,10 @@ def get_columns(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            arrow_schema_bytes=arrow_schema_bytes,
+            t_row_set=t_row_set,
+            max_download_threads=self.max_download_threads,
+            ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1177,11 +1180,7 @@ def _handle_execute_response(self, resp, cursor):
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        (
-            execute_response,
-            arrow_schema_bytes,
-        ) = self._results_message_to_execute_response(resp, final_operation_state)
-        return execute_response, arrow_schema_bytes
+        return self._results_message_to_execute_response(resp, final_operation_state)
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
@@ -1225,7 +1224,9 @@ def fetch_results(
                 )
             )
 
-        queue = ThriftResultSetQueueFactory.build_queue(
+        from databricks.sql.utils import ResultSetQueueFactory
+
+        queue = ResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,
             t_row_set=resp.results,
             arrow_schema_bytes=arrow_schema_bytes,

src/databricks/sql/backend/types.py

@@ -423,11 +423,9 @@ class ExecuteResponse:
 
     command_id: CommandId
     status: CommandState
-    description: Optional[
-        List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]]
-    ] = None
-    has_more_rows: bool = False
-    results_queue: Optional[Any] = None
+    description: Optional[List[Tuple]] = 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

@@ -158,7 +158,10 @@ def __init__(
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
-        arrow_schema_bytes: Optional[bytes] = None,
+        t_row_set=None,
+        max_download_threads: int = 10,
+        ssl_options=None,
+        has_more_rows: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -170,13 +173,32 @@ 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
-            arrow_schema_bytes: Arrow schema bytes for the result set
+            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
+            has_more_rows: Whether there are more rows to fetch
         """
         # Initialize ThriftResultSet-specific attributes
-        self._arrow_schema_bytes = arrow_schema_bytes
+        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,
@@ -186,8 +208,8 @@ 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=execute_response.has_more_rows,
-            results_queue=execute_response.results_queue,
+            has_more_rows=has_more_rows,
+            results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
         )