select Queue based on format in manifest

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

Author: varun-edachali-dbx

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

@@ -5,6 +5,8 @@
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+from databricks.sql.backend.sea.utils.constants import ResultFormat
+from databricks.sql.exc import ProgrammingError
 from databricks.sql.utils import ResultSetQueue
 
 
@@ -35,15 +37,15 @@ def build_queue(
             ResultSetQueue: The appropriate queue for the result data
         """
 
-        if sea_result_data.data is not None:
+        if manifest.format == ResultFormat.JSON_ARRAY.value:
             # INLINE disposition with JSON_ARRAY format
             return JsonQueue(sea_result_data.data)
-        elif sea_result_data.external_links is not None:
+        elif manifest.format == ResultFormat.ARROW_STREAM.value:
             # EXTERNAL_LINKS disposition
             raise NotImplementedError(
                 "EXTERNAL_LINKS disposition is not implemented for SEA backend"
             )
-        return JsonQueue([])
+        raise ProgrammingError("Invalid result format")
 
 
 class JsonQueue(ResultSetQueue):
@@ -53,7 +55,7 @@ def __init__(self, data_array):
         """Initialize with JSON array data."""
         self.data_array = data_array
         self.cur_row_index = 0
-        self.num_rows = len(data_array)
+        self.num_rows = len(data_array) if data_array else 0
 
     def next_n_rows(self, num_rows):
         """Get the next n rows from the data array."""

tests/unit/test_sea_queue.py

@@ -9,6 +9,7 @@
 
 from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+from databricks.sql.backend.sea.utils.constants import ResultFormat
 
 
 class TestJsonQueue:
@@ -104,6 +105,15 @@ def mock_description(self):
             ("col3", "boolean", None, None, None, None, None),
         ]
 
+    def _create_empty_manifest(self, format: ResultFormat):
+        return ResultManifest(
+            format=format.value,
+            schema={},
+            total_row_count=-1,
+            total_byte_count=-1,
+            total_chunk_count=-1,
+        )
+
     def test_build_queue_with_inline_data(self, mock_sea_client, mock_description):
         """Test building a queue with inline JSON data."""
         # Create sample data for inline JSON result
@@ -116,7 +126,7 @@ def test_build_queue_with_inline_data(self, mock_sea_client, mock_description):
         result_data = ResultData(data=data, external_links=None, row_count=len(data))
 
         # Create a manifest (not used for inline data)
-        manifest = None
+        manifest = self._create_empty_manifest(ResultFormat.JSON_ARRAY)
 
         # Build the queue
         queue = SeaResultSetQueueFactory.build_queue(
@@ -140,7 +150,7 @@ def test_build_queue_with_empty_data(self, mock_sea_client, mock_description):
         # Build the queue
         queue = SeaResultSetQueueFactory.build_queue(
             result_data,
-            None,
+            self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             "test-statement-123",
             description=mock_description,
             sea_client=mock_sea_client,
@@ -165,7 +175,7 @@ def test_build_queue_with_external_links(self, mock_sea_client, mock_description
         ):
             SeaResultSetQueueFactory.build_queue(
                 result_data,
-                None,
+                self._create_empty_manifest(ResultFormat.ARROW_STREAM),
                 "test-statement-123",
                 description=mock_description,
                 sea_client=mock_sea_client,

tests/unit/test_sea_result_set.py

@@ -10,6 +10,7 @@
 
 from databricks.sql.backend.sea.result_set import SeaResultSet, Row
 from databricks.sql.backend.sea.queue import JsonQueue
+from databricks.sql.backend.sea.utils.constants import ResultFormat
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
 
@@ -59,6 +60,16 @@ def sample_data(self):
             ["value5", "5", "true"],
         ]
 
+    def _create_empty_manifest(self, format: ResultFormat):
+        """Create an empty manifest."""
+        return ResultManifest(
+            format=format.value,
+            schema={},
+            total_row_count=-1,
+            total_byte_count=-1,
+            total_chunk_count=-1,
+        )
+
     @pytest.fixture
     def result_set_with_data(
         self, mock_connection, mock_sea_client, execute_response, sample_data
@@ -75,7 +86,7 @@ def result_set_with_data(
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=result_data,
-            manifest=None,
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -88,18 +99,6 @@ def json_queue(self, sample_data):
         """Create a JsonQueue with sample data."""
         return JsonQueue(sample_data)
 
-    def empty_manifest(self):
-        """Create an empty manifest."""
-        return ResultManifest(
-            format="JSON_ARRAY",
-            schema={},
-            total_row_count=0,
-            total_byte_count=0,
-            total_chunk_count=0,
-            truncated=False,
-            is_volume_operation=False,
-        )
-
     def test_init_with_execute_response(
         self, mock_connection, mock_sea_client, execute_response
     ):
@@ -109,7 +108,7 @@ def test_init_with_execute_response(
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=ResultData(data=[]),
-            manifest=self.empty_manifest(),
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -130,7 +129,7 @@ def test_close(self, mock_connection, mock_sea_client, execute_response):
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=ResultData(data=[]),
-            manifest=self.empty_manifest(),
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -152,7 +151,7 @@ def test_close_when_already_closed_server_side(
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=ResultData(data=[]),
-            manifest=self.empty_manifest(),
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -176,7 +175,7 @@ def test_close_when_connection_closed(
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=ResultData(data=[]),
-            manifest=self.empty_manifest(),
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )
@@ -332,7 +331,7 @@ def test_fetchmany_arrow_not_implemented(
                 execute_response=execute_response,
                 sea_client=mock_sea_client,
                 result_data=ResultData(data=None, external_links=[]),
-                manifest=self.empty_manifest(),
+                manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
                 buffer_size_bytes=1000,
                 arraysize=100,
             )
@@ -352,7 +351,7 @@ def test_fetchall_arrow_not_implemented(
                 execute_response=execute_response,
                 sea_client=mock_sea_client,
                 result_data=ResultData(data=None, external_links=[]),
-                manifest=self.empty_manifest(),
+                manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
                 buffer_size_bytes=1000,
                 arraysize=100,
             )
@@ -370,7 +369,7 @@ def test_is_staging_operation(
             execute_response=execute_response,
             sea_client=mock_sea_client,
             result_data=ResultData(data=[]),
-            manifest=self.empty_manifest(),
+            manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
             buffer_size_bytes=1000,
             arraysize=100,
         )