mandate ResultData in SeaResultSet constructor

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

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/backend.py

@@ -616,10 +616,10 @@ def get_execution_result(
             connection=cursor.connection,
             execute_response=execute_response,
             sea_client=self,
-            buffer_size_bytes=cursor.buffer_size_bytes,
-            arraysize=cursor.arraysize,
             result_data=response.result,
             manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
         )
 
     # == Metadata Operations ==

src/databricks/sql/backend/sea/utils/filters.py

@@ -77,9 +77,9 @@ def _filter_sea_result_set(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
+            result_data=result_data,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
-            result_data=result_data,
         )
 
         return filtered_result_set

src/databricks/sql/result_set.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import List, Optional, TYPE_CHECKING
 
@@ -450,13 +452,13 @@ class SeaResultSet(ResultSet):
 
     def __init__(
         self,
-        connection: "Connection",
-        execute_response: "ExecuteResponse",
-        sea_client: "SeaDatabricksClient",
+        connection: Connection,
+        execute_response: ExecuteResponse,
+        sea_client: SeaDatabricksClient,
+        result_data: ResultData,
+        manifest: Optional[ResultManifest] = None,
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
-        result_data: Optional["ResultData"] = None,
-        manifest: Optional["ResultManifest"] = None,
     ):
         """
         Initialize a SeaResultSet with the response from a SEA query execution.
@@ -467,21 +469,19 @@ def __init__(
             sea_client: The SeaDatabricksClient instance for direct access
             buffer_size_bytes: Buffer size for fetching results
             arraysize: Default number of rows to fetch
-            result_data: Result data from SEA response (optional)
-            manifest: Manifest from SEA response (optional)
+            result_data: Result data from SEA response
+            manifest: Manifest from SEA response
         """
 
-        results_queue = None
-        if result_data:
-            results_queue = SeaResultSetQueueFactory.build_queue(
-                result_data,
-                manifest,
-                str(execute_response.command_id.to_sea_statement_id()),
-                description=execute_response.description,
-                max_download_threads=sea_client.max_download_threads,
-                sea_client=sea_client,
-                lz4_compressed=execute_response.lz4_compressed,
-            )
+        results_queue = SeaResultSetQueueFactory.build_queue(
+            result_data,
+            manifest,
+            str(execute_response.command_id.to_sea_statement_id()),
+            description=execute_response.description,
+            max_download_threads=sea_client.max_download_threads,
+            sea_client=sea_client,
+            lz4_compressed=execute_response.lz4_compressed,
+        )
 
         # Call parent constructor with common attributes
         super().__init__(
@@ -503,6 +503,8 @@ def _convert_json_to_arrow(self, rows: List) -> "pyarrow.Table":
         """
         Convert raw data rows to Arrow table.
         """
+        if not rows:
+            return pyarrow.Table.from_pydict({})
 
         columns = []
         num_cols = len(rows[0])

src/databricks/sql/utils.py

@@ -21,7 +21,8 @@
 except ImportError:
     pyarrow = None
 
-from databricks.sql import OperationalError, exc
+from databricks.sql import OperationalError
+from databricks.sql.exc import ProgrammingError
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TRowSet,
@@ -148,9 +149,7 @@ def build_queue(
             raise NotImplementedError(
                 "EXTERNAL_LINKS disposition is not implemented for SEA backend"
             )
-        else:
-            # Empty result set
-            return JsonQueue([])
+        raise ProgrammingError("No result data or external links found")
 
 
 class JsonQueue(ResultSetQueue):

tests/unit/test_sea_queue.py

@@ -135,7 +135,7 @@ def test_build_queue_with_inline_data(self, mock_sea_client, mock_description):
     def test_build_queue_with_empty_data(self, mock_sea_client, mock_description):
         """Test building a queue with empty data."""
         # Create a ResultData object with no data
-        result_data = ResultData(data=None, external_links=None, row_count=0)
+        result_data = ResultData(data=[], external_links=None, row_count=0)
 
         # Build the queue
         queue = SeaResultSetQueueFactory.build_queue(

tests/unit/test_sea_result_set.py

@@ -74,10 +74,10 @@ def result_set_with_data(
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
-            buffer_size_bytes=1000,
-            arraysize=100,
             result_data=result_data,
             manifest=None,
+            buffer_size_bytes=1000,
+            arraysize=100,
         )
         result_set.results = JsonQueue(sample_data)
 
@@ -96,6 +96,7 @@ def test_init_with_execute_response(
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
+            result_data=ResultData(data=[]),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -115,6 +116,7 @@ def test_close(self, mock_connection, mock_sea_client, execute_response):
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
+            result_data=ResultData(data=[]),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -135,6 +137,7 @@ def test_close_when_already_closed_server_side(
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
+            result_data=ResultData(data=[]),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -157,6 +160,7 @@ def test_close_when_connection_closed(
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
+            result_data=ResultData(data=[]),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -301,39 +305,40 @@ def test_fetchmany_arrow_not_implemented(
         self, mock_connection, mock_sea_client, execute_response, sample_data
     ):
         """Test that fetchmany_arrow raises NotImplementedError for non-JSON data."""
-        # Create a result set without JSON data
-        result_set = SeaResultSet(
-            connection=mock_connection,
-            execute_response=execute_response,
-            sea_client=mock_sea_client,
-            buffer_size_bytes=1000,
-            arraysize=100,
-        )
 
         # Test that NotImplementedError is raised
         with pytest.raises(
-            NotImplementedError, match="fetchmany_arrow only supported for JSON data"
+            NotImplementedError,
+            match="EXTERNAL_LINKS disposition is not implemented for SEA backend",
         ):
-            result_set.fetchmany_arrow(10)
+            # Create a result set without JSON data
+            result_set = SeaResultSet(
+                connection=mock_connection,
+                execute_response=execute_response,
+                sea_client=mock_sea_client,
+                result_data=ResultData(data=None, external_links=[]),
+                buffer_size_bytes=1000,
+                arraysize=100,
+            )
 
     def test_fetchall_arrow_not_implemented(
         self, mock_connection, mock_sea_client, execute_response, sample_data
     ):
         """Test that fetchall_arrow raises NotImplementedError for non-JSON data."""
-        # Create a result set without JSON data
-        result_set = SeaResultSet(
-            connection=mock_connection,
-            execute_response=execute_response,
-            sea_client=mock_sea_client,
-            buffer_size_bytes=1000,
-            arraysize=100,
-        )
-
         # Test that NotImplementedError is raised
         with pytest.raises(
-            NotImplementedError, match="fetchall_arrow only supported for JSON data"
+            NotImplementedError,
+            match="EXTERNAL_LINKS disposition is not implemented for SEA backend",
         ):
-            result_set.fetchall_arrow()
+            # Create a result set without JSON data
+            result_set = SeaResultSet(
+                connection=mock_connection,
+                execute_response=execute_response,
+                sea_client=mock_sea_client,
+                result_data=ResultData(data=None, external_links=[]),
+                buffer_size_bytes=1000,
+                arraysize=100,
+            )
 
     def test_is_staging_operation(
         self, mock_connection, mock_sea_client, execute_response
@@ -347,6 +352,7 @@ def test_is_staging_operation(
             connection=mock_connection,
             execute_response=execute_response,
             sea_client=mock_sea_client,
+            result_data=ResultData(data=[]),
             buffer_size_bytes=1000,
             arraysize=100,
         )