move SeaResultSetQueueFactory and JsonQueue into separate SEA module

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

Author: varun-edachali-dbx

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

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import List, Optional, Tuple
+
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
+from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+from databricks.sql.utils import ResultSetQueue
+
+
+class SeaResultSetQueueFactory(ABC):
+    @staticmethod
+    def build_queue(
+        sea_result_data: ResultData,
+        manifest: Optional[ResultManifest],
+        statement_id: str,
+        description: List[Tuple] = [],
+        max_download_threads: Optional[int] = None,
+        sea_client: Optional[SeaDatabricksClient] = None,
+        lz4_compressed: bool = False,
+    ) -> ResultSetQueue:
+        """
+        Factory method to build a result set queue for SEA backend.
+
+        Args:
+            sea_result_data (ResultData): Result data from SEA response
+            manifest (ResultManifest): Manifest from SEA response
+            statement_id (str): Statement ID for the query
+            description (List[List[Any]]): Column descriptions
+            schema_bytes (bytes): Arrow schema bytes
+            max_download_threads (int): Maximum number of download threads
+            ssl_options (SSLOptions): SSL options for downloads
+            sea_client (SeaDatabricksClient): SEA client for fetching additional links
+            lz4_compressed (bool): Whether the data is LZ4 compressed
+
+        Returns:
+            ResultSetQueue: The appropriate queue for the result data
+        """
+
+        if sea_result_data.data is not None:
+            # INLINE disposition with JSON_ARRAY format
+            return JsonQueue(sea_result_data.data)
+        elif sea_result_data.external_links is not None:
+            # EXTERNAL_LINKS disposition
+            raise NotImplementedError(
+                "EXTERNAL_LINKS disposition is not implemented for SEA backend"
+            )
+        return JsonQueue([])
+
+
+class JsonQueue(ResultSetQueue):
+    """Queue implementation for JSON_ARRAY format data."""
+
+    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)
+
+    def next_n_rows(self, num_rows):
+        """Get the next n rows from the data array."""
+        length = min(num_rows, self.num_rows - self.cur_row_index)
+        slice = self.data_array[self.cur_row_index : self.cur_row_index + length]
+        self.cur_row_index += length
+        return slice
+
+    def remaining_rows(self):
+        """Get all remaining rows from the data array."""
+        slice = self.data_array[self.cur_row_index :]
+        self.cur_row_index += len(slice)
+        return slice

src/databricks/sql/result_set.py

@@ -24,9 +24,8 @@
 from databricks.sql.utils import (
     ColumnTable,
     ColumnQueue,
-    JsonQueue,
-    SeaResultSetQueueFactory,
 )
+from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
 
 logger = logging.getLogger(__name__)

src/databricks/sql/utils.py

@@ -111,69 +111,6 @@ def build_queue(
             raise AssertionError("Row set type is not valid")
 
 
-class SeaResultSetQueueFactory(ABC):
-    @staticmethod
-    def build_queue(
-        sea_result_data: ResultData,
-        manifest: Optional[ResultManifest],
-        statement_id: str,
-        description: List[Tuple] = [],
-        max_download_threads: Optional[int] = None,
-        sea_client: Optional[SeaDatabricksClient] = None,
-        lz4_compressed: bool = False,
-    ) -> ResultSetQueue:
-        """
-        Factory method to build a result set queue for SEA backend.
-
-        Args:
-            sea_result_data (ResultData): Result data from SEA response
-            manifest (ResultManifest): Manifest from SEA response
-            statement_id (str): Statement ID for the query
-            description (List[List[Any]]): Column descriptions
-            schema_bytes (bytes): Arrow schema bytes
-            max_download_threads (int): Maximum number of download threads
-            ssl_options (SSLOptions): SSL options for downloads
-            sea_client (SeaDatabricksClient): SEA client for fetching additional links
-            lz4_compressed (bool): Whether the data is LZ4 compressed
-
-        Returns:
-            ResultSetQueue: The appropriate queue for the result data
-        """
-
-        if sea_result_data.data is not None:
-            # INLINE disposition with JSON_ARRAY format
-            return JsonQueue(sea_result_data.data)
-        elif sea_result_data.external_links is not None:
-            # EXTERNAL_LINKS disposition
-            raise NotImplementedError(
-                "EXTERNAL_LINKS disposition is not implemented for SEA backend"
-            )
-        return JsonQueue([])
-
-
-class JsonQueue(ResultSetQueue):
-    """Queue implementation for JSON_ARRAY format data."""
-
-    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)
-
-    def next_n_rows(self, num_rows):
-        """Get the next n rows from the data array."""
-        length = min(num_rows, self.num_rows - self.cur_row_index)
-        slice = self.data_array[self.cur_row_index : self.cur_row_index + length]
-        self.cur_row_index += length
-        return slice
-
-    def remaining_rows(self):
-        """Get all remaining rows from the data array."""
-        slice = self.data_array[self.cur_row_index :]
-        self.cur_row_index += len(slice)
-        return slice
-
-
 class ColumnTable:
     def __init__(self, column_table, column_names):
         self.column_table = column_table

tests/unit/test_sea_queue.py

@@ -7,7 +7,7 @@
 import pytest
 from unittest.mock import Mock, MagicMock, patch
 
-from databricks.sql.utils import JsonQueue, SeaResultSetQueueFactory
+from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
 
 

tests/unit/test_sea_result_set.py

@@ -9,7 +9,7 @@
 from unittest.mock import Mock
 
 from databricks.sql.result_set import SeaResultSet, Row
-from databricks.sql.utils import JsonQueue
+from databricks.sql.backend.sea.queue import JsonQueue
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest