move arrow_schema_bytes back into ExecuteResult

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -801,18 +801,16 @@ def _results_message_to_execute_response(self, resp, 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,
+        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,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def get_execution_result(
@@ -877,6 +875,7 @@ def get_execution_result(
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
+            arrow_schema_bytes=schema_bytes,
         )
 
         return ThriftResultSet(
@@ -886,7 +885,6 @@ 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,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -999,9 +997,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, arrow_schema_bytes = self._handle_execute_response(
-                resp, cursor
-            )
+            execute_response = self._handle_execute_response(resp, cursor)
 
             return ThriftResultSet(
                 connection=cursor.connection,
@@ -1010,7 +1006,6 @@ def execute_command(
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
                 use_cloud_fetch=use_cloud_fetch,
-                arrow_schema_bytes=arrow_schema_bytes,
             )
 
     def get_catalogs(
@@ -1032,9 +1027,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1043,7 +1036,6 @@ 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,
         )
 
     def get_schemas(
@@ -1069,9 +1061,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1080,7 +1070,6 @@ 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,
         )
 
     def get_tables(
@@ -1110,9 +1099,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1121,7 +1108,6 @@ 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,
         )
 
     def get_columns(
@@ -1151,9 +1137,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, arrow_schema_bytes = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1162,7 +1146,6 @@ 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,
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1176,11 +1159,10 @@ 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
+        execute_response = self._results_message_to_execute_response(
+            resp, final_operation_state
+        )
+        return execute_response
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)

src/databricks/sql/backend/types.py

@@ -431,3 +431,4 @@ class ExecuteResponse:
     has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
     is_staging_operation: bool = False
+    arrow_schema_bytes: Optional[bytes] = None

src/databricks/sql/result_set.py

@@ -157,7 +157,6 @@ def __init__(
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
-        arrow_schema_bytes: Optional[bytes] = None,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -169,10 +168,9 @@ 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
         """
         # 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
 

tests/unit/test_thrift_backend.py

@@ -19,7 +19,13 @@
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.result_set import ResultSet, ThriftResultSet
-from databricks.sql.backend.types import CommandId, CommandState, SessionId, BackendType
+from databricks.sql.backend.types import (
+    CommandId,
+    CommandState,
+    SessionId,
+    BackendType,
+    ExecuteResponse,
+)
 
 
 def retry_policy_factory():
@@ -651,7 +657,7 @@ def test_handle_execute_response_sets_compression_in_direct_results(
                 ssl_options=SSLOptions(),
             )
 
-            execute_response, _ = thrift_backend._handle_execute_response(
+            execute_response = thrift_backend._handle_execute_response(
                 t_execute_resp, Mock()
             )
             self.assertEqual(execute_response.lz4_compressed, lz4Compressed)
@@ -885,7 +891,7 @@ def test_handle_execute_response_can_handle_without_direct_results(
                     auth_provider=AuthProvider(),
                     ssl_options=SSLOptions(),
                 )
-                execute_response, _ = thrift_backend._handle_execute_response(
+                execute_response = thrift_backend._handle_execute_response(
                     execute_resp, Mock()
                 )
 
@@ -963,11 +969,11 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             t_get_result_set_metadata_resp
         )
         thrift_backend = self._make_fake_thrift_backend()
-        execute_response, arrow_schema_bytes = thrift_backend._handle_execute_response(
+        execute_response = thrift_backend._handle_execute_response(
             t_execute_resp, Mock()
         )
 
-        self.assertEqual(arrow_schema_bytes, arrow_schema_mock)
+        self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
@@ -1040,7 +1046,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                 )
                 thrift_backend = self._make_fake_thrift_backend()
 
-                execute_response, _ = thrift_backend._handle_execute_response(
+                execute_response = thrift_backend._handle_execute_response(
                     execute_resp, Mock()
                 )
 
@@ -1172,7 +1178,20 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock(
+            return_value=Mock(
+                spec=ExecuteResponse,
+                command_id=Mock(),
+                status=Mock(),
+                description=Mock(),
+                has_more_rows=Mock(),
+                results_queue=Mock(),
+                has_been_closed_server_side=Mock(),
+                lz4_compressed=Mock(),
+                is_staging_operation=Mock(),
+                arrow_schema_bytes=Mock(),
+            )
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.execute_command(
@@ -1206,7 +1225,20 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock(
+            return_value=Mock(
+                spec=ExecuteResponse,
+                command_id=Mock(),
+                status=Mock(),
+                description=Mock(),
+                has_more_rows=Mock(),
+                results_queue=Mock(),
+                has_been_closed_server_side=Mock(),
+                lz4_compressed=Mock(),
+                is_staging_operation=Mock(),
+                arrow_schema_bytes=Mock(),
+            )
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_catalogs(Mock(), 100, 200, cursor_mock)
@@ -1237,7 +1269,20 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock(
+            return_value=Mock(
+                spec=ExecuteResponse,
+                command_id=Mock(),
+                status=Mock(),
+                description=Mock(),
+                has_more_rows=Mock(),
+                results_queue=Mock(),
+                has_been_closed_server_side=Mock(),
+                lz4_compressed=Mock(),
+                is_staging_operation=Mock(),
+                arrow_schema_bytes=Mock(),
+            )
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_schemas(
@@ -1277,7 +1322,20 @@ def test_get_tables_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock(
+            return_value=Mock(
+                spec=ExecuteResponse,
+                command_id=Mock(),
+                status=Mock(),
+                description=Mock(),
+                has_more_rows=Mock(),
+                results_queue=Mock(),
+                has_been_closed_server_side=Mock(),
+                lz4_compressed=Mock(),
+                is_staging_operation=Mock(),
+                arrow_schema_bytes=Mock(),
+            )
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_tables(
@@ -1321,7 +1379,20 @@ def test_get_columns_calls_client_and_handle_execute_response(
             auth_provider=AuthProvider(),
             ssl_options=SSLOptions(),
         )
-        thrift_backend._handle_execute_response = Mock(return_value=(Mock(), Mock()))
+        thrift_backend._handle_execute_response = Mock(
+            return_value=Mock(
+                spec=ExecuteResponse,
+                command_id=Mock(),
+                status=Mock(),
+                description=Mock(),
+                has_more_rows=Mock(),
+                results_queue=Mock(),
+                has_been_closed_server_side=Mock(),
+                lz4_compressed=Mock(),
+                is_staging_operation=Mock(),
+                arrow_schema_bytes=Mock(),
+            )
+        )
         cursor_mock = Mock()
 
         result = thrift_backend.get_columns(
@@ -2229,7 +2300,18 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     @patch(
         "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response",
-        return_value=(Mock(), Mock()),
+        return_value=Mock(
+            spec=ExecuteResponse,
+            command_id=Mock(),
+            status=Mock(),
+            description=Mock(),
+            has_more_rows=Mock(),
+            results_queue=Mock(),
+            has_been_closed_server_side=Mock(),
+            lz4_compressed=Mock(),
+            is_staging_operation=Mock(),
+            arrow_schema_bytes=Mock(),
+        ),
     )
     def test_execute_command_sets_complex_type_fields_correctly(
         self, mock_handle_execute_response, tcli_service_class