move queue and result set into SEA specific dir

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

Author: varun-edachali-dbx

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

@@ -1,13 +1,30 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
+
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
+import dateutil
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
-from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+from databricks.sql.backend.sea.models.base import (
+    ExternalLink,
+    ResultData,
+    ResultManifest,
+)
 from databricks.sql.backend.sea.utils.constants import ResultFormat
 from databricks.sql.exc import ProgrammingError
-from databricks.sql.utils import ResultSetQueue
+from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
+from databricks.sql.types import SSLOptions
+from databricks.sql.utils import CloudFetchQueue, ResultSetQueue
+
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class SeaResultSetQueueFactory(ABC):
@@ -42,8 +59,30 @@ def build_queue(
             return JsonQueue(sea_result_data.data)
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             # EXTERNAL_LINKS disposition
-            raise NotImplementedError(
-                "EXTERNAL_LINKS disposition is not implemented for SEA backend"
+            if not max_download_threads:
+                raise ValueError(
+                    "Max download threads is required for EXTERNAL_LINKS disposition"
+                )
+            if not ssl_options:
+                raise ValueError(
+                    "SSL options are required for EXTERNAL_LINKS disposition"
+                )
+            if not sea_client:
+                raise ValueError(
+                    "SEA client is required for EXTERNAL_LINKS disposition"
+                )
+            if not manifest:
+                raise ValueError("Manifest is required for EXTERNAL_LINKS disposition")
+
+            return SeaCloudFetchQueue(
+                initial_links=sea_result_data.external_links,
+                max_download_threads=max_download_threads,
+                ssl_options=ssl_options,
+                sea_client=sea_client,
+                statement_id=statement_id,
+                total_chunk_count=manifest.total_chunk_count,
+                lz4_compressed=lz4_compressed,
+                description=description,
             )
         raise ProgrammingError("Invalid result format")
 
@@ -69,3 +108,134 @@ def remaining_rows(self) -> List[List[str]]:
         slice = self.data_array[self.cur_row_index :]
         self.cur_row_index += len(slice)
         return slice
+
+
+class SeaCloudFetchQueue(CloudFetchQueue):
+    """Queue implementation for EXTERNAL_LINKS disposition with ARROW format for SEA backend."""
+
+    def __init__(
+        self,
+        initial_links: List["ExternalLink"],
+        max_download_threads: int,
+        ssl_options: SSLOptions,
+        sea_client: "SeaDatabricksClient",
+        statement_id: str,
+        total_chunk_count: int,
+        lz4_compressed: bool = False,
+        description: Optional[List[Tuple]] = None,
+    ):
+        """
+        Initialize the SEA CloudFetchQueue.
+
+        Args:
+            initial_links: Initial list of external links to download
+            schema_bytes: Arrow schema bytes
+            max_download_threads: Maximum number of download threads
+            ssl_options: SSL options for downloads
+            sea_client: SEA client for fetching additional links
+            statement_id: Statement ID for the query
+            total_chunk_count: Total number of chunks in the result set
+            lz4_compressed: Whether the data is LZ4 compressed
+            description: Column descriptions
+        """
+
+        super().__init__(
+            max_download_threads=max_download_threads,
+            ssl_options=ssl_options,
+            schema_bytes=None,
+            lz4_compressed=lz4_compressed,
+            description=description,
+        )
+
+        self._sea_client = sea_client
+        self._statement_id = statement_id
+
+        logger.debug(
+            "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
+                statement_id, total_chunk_count
+            )
+        )
+
+        initial_link = next((l for l in initial_links if l.chunk_index == 0), None)
+        if not initial_link:
+            raise ValueError("No initial link found for chunk index 0")
+
+        self.download_manager = ResultFileDownloadManager(
+            links=[],
+            max_download_threads=max_download_threads,
+            lz4_compressed=lz4_compressed,
+            ssl_options=ssl_options,
+        )
+
+        # Track the current chunk we're processing
+        self._current_chunk_link: Optional["ExternalLink"] = initial_link
+        self._download_current_link()
+
+        # Initialize table and position
+        self.table = self._create_next_table()
+
+    def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
+        """Convert SEA external links to Thrift format for compatibility with existing download manager."""
+        # Parse the ISO format expiration time
+        expiry_time = int(dateutil.parser.parse(link.expiration).timestamp())
+        return TSparkArrowResultLink(
+            fileLink=link.external_link,
+            expiryTime=expiry_time,
+            rowCount=link.row_count,
+            bytesNum=link.byte_count,
+            startRowOffset=link.row_offset,
+            httpHeaders=link.http_headers or {},
+        )
+
+    def _download_current_link(self):
+        """Download the current chunk link."""
+        if not self._current_chunk_link:
+            return None
+
+        if not self.download_manager:
+            logger.debug("SeaCloudFetchQueue: No download manager, returning")
+            return None
+
+        thrift_link = self._convert_to_thrift_link(self._current_chunk_link)
+        self.download_manager.add_link(thrift_link)
+
+    def _progress_chunk_link(self):
+        """Progress to the next chunk link."""
+        if not self._current_chunk_link:
+            return None
+
+        next_chunk_index = self._current_chunk_link.next_chunk_index
+
+        if next_chunk_index is None:
+            self._current_chunk_link = None
+            return None
+
+        try:
+            self._current_chunk_link = self._sea_client.get_chunk_link(
+                self._statement_id, next_chunk_index
+            )
+        except Exception as e:
+            logger.error(
+                "SeaCloudFetchQueue: Error fetching link for chunk {}: {}".format(
+                    next_chunk_index, e
+                )
+            )
+            return None
+
+        logger.debug(
+            f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
+        )
+        self._download_current_link()
+
+    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+        """Create next table by retrieving the logical next downloaded file."""
+        if not self._current_chunk_link:
+            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
+            return None
+
+        row_offset = self._current_chunk_link.row_offset
+        arrow_table = self._create_table_at_offset(row_offset)
+
+        self._progress_chunk_link()
+
+        return arrow_table

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

@@ -196,10 +196,10 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
 
-        if not isinstance(self.results, JsonQueue):
-            raise NotImplementedError("fetchmany_arrow only supported for JSON data")
+        results = self.results.next_n_rows(size)
+        if isinstance(self.results, JsonQueue):
+            results = self._convert_json_to_arrow_table(results)
 
-        results = self._convert_json_to_arrow_table(self.results.next_n_rows(size))
         self._next_row_index += results.num_rows
 
         return results
@@ -209,10 +209,10 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         Fetch all remaining rows as an Arrow table.
         """
 
-        if not isinstance(self.results, JsonQueue):
-            raise NotImplementedError("fetchall_arrow only supported for JSON data")
+        results = self.results.remaining_rows()
+        if isinstance(self.results, JsonQueue):
+            results = self._convert_json_to_arrow_table(results)
 
-        results = self._convert_json_to_arrow_table(self.results.remaining_rows())
         self._next_row_index += results.num_rows
 
         return results
@@ -229,7 +229,7 @@ def fetchone(self) -> Optional[Row]:
         if isinstance(self.results, JsonQueue):
             res = self._create_json_table(self.fetchmany_json(1))
         else:
-            raise NotImplementedError("fetchone only supported for JSON data")
+            res = self._convert_arrow_table(self.fetchmany_arrow(1))
 
         return res[0] if res else None
 
@@ -250,7 +250,7 @@ def fetchmany(self, size: int) -> List[Row]:
         if isinstance(self.results, JsonQueue):
             return self._create_json_table(self.fetchmany_json(size))
         else:
-            raise NotImplementedError("fetchmany only supported for JSON data")
+            return self._convert_arrow_table(self.fetchmany_arrow(size))
 
     def fetchall(self) -> List[Row]:
         """
@@ -263,4 +263,4 @@ def fetchall(self) -> List[Row]:
         if isinstance(self.results, JsonQueue):
             return self._create_json_table(self.fetchall_json())
         else:
-            raise NotImplementedError("fetchall only supported for JSON data")
+            return self._convert_arrow_table(self.fetchall_arrow())

src/databricks/sql/result_set.py

@@ -25,8 +25,6 @@
 from databricks.sql.utils import (
     ColumnTable,
     ColumnQueue,
-    JsonQueue,
-    SeaResultSetQueueFactory,
 )
 from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse