introduce SeaResultSetQueueFactory

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

Author: varun-edachali-dbx

src/databricks/sql/backend/sea_result_set.py

@@ -18,7 +18,7 @@
 from databricks.sql.types import Row
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.exc import Error
-from databricks.sql.utils import ResultSetQueueFactory, JsonQueue
+from databricks.sql.utils import SeaResultSetQueueFactory, JsonQueue
 
 from databricks.sql.backend.models import (
     StatementStatus,
@@ -117,9 +117,9 @@ def __init__(
 
         # Initialize queue for result data
         if self.result:
-            self.results = ResultSetQueueFactory.build_queue(
+            self.results = SeaResultSetQueueFactory.build_queue(
                 sea_result_data=self.result,
-                description=cast(Optional[List[List[Any]]], self.description)
+                description=cast(Optional[List[List[Any]]], self.description),
             )
             self._has_more_rows = True if self.result.data else False
         else:
@@ -168,7 +168,7 @@ def _fill_results_buffer(self) -> None:
         if not self.result or not self.result.data:
             self._has_more_rows = False
             return
-        
+
         # For INLINE disposition, we already have all the data
         # No need to fetch more data from the backend
         self._has_more_rows = False  # No more rows to fetch for INLINE
@@ -177,21 +177,21 @@ def _convert_rows_to_arrow_table(self, rows):
         """Convert rows to Arrow table."""
         if not self.description:
             return pyarrow.Table.from_pylist([])
-        
+
         # Create dict of column data
         column_data = {}
         column_names = [col[0] for col in self.description]
-        
+
         for i, name in enumerate(column_names):
             column_data[name] = [row[i] for row in rows]
-        
+
         return pyarrow.Table.from_pydict(column_data)
 
     def _create_empty_arrow_table(self):
         """Create an empty Arrow table with the correct schema."""
         if not self.description:
             return pyarrow.Table.from_pylist([])
-        
+
         column_names = [col[0] for col in self.description]
         return pyarrow.Table.from_pydict({name: [] for name in column_names})
 
@@ -201,9 +201,9 @@ def fetchone(self) -> Optional[Row]:
             rows = self.results.next_n_rows(1)
             if not rows:
                 return None
-            
+
             row = rows[0]
-            
+
             # Convert to Row object
             if self.description:
                 column_names = [col[0] for col in self.description]
@@ -219,13 +219,13 @@ def fetchmany(self, size: Optional[int] = None) -> List[Row]:
         """Fetch the next set of rows of a query result."""
         if size is None:
             size = self.arraysize
-            
+
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
-        
+
         if isinstance(self.results, JsonQueue):
             rows = self.results.next_n_rows(size)
-            
+
             # Convert to Row objects
             if self.description:
                 column_names = [col[0] for col in self.description]
@@ -241,7 +241,7 @@ def fetchall(self) -> List[Row]:
         """Fetch all remaining rows of a query result."""
         if isinstance(self.results, JsonQueue):
             rows = self.results.remaining_rows()
-            
+
             # Convert to Row objects
             if self.description:
                 column_names = [col[0] for col in self.description]
@@ -257,25 +257,25 @@ def fetchmany_arrow(self, size: int) -> Any:
         """Fetch the next set of rows as an Arrow table."""
         if not pyarrow:
             raise ImportError("PyArrow is required for Arrow support")
-        
+
         rows = self.fetchmany(size)
         if not rows:
             # Return empty Arrow table with schema
             return self._create_empty_arrow_table()
-        
+
         # Convert rows to Arrow table
         return self._convert_rows_to_arrow_table(rows)
 
     def fetchall_arrow(self) -> Any:
         """Fetch all remaining rows as an Arrow table."""
         if not pyarrow:
             raise ImportError("PyArrow is required for Arrow support")
-        
+
         rows = self.fetchall()
         if not rows:
             # Return empty Arrow table with schema
             return self._create_empty_arrow_table()
-        
+
         # Convert rows to Arrow table
         return self._convert_rows_to_arrow_table(rows)
 
@@ -288,4 +288,4 @@ def close(self) -> None:
                     CommandId.from_sea_statement_id(self.statement_id)
                 )
             except Exception as e:
-                logger.warning(f"Error closing SEA statement: {e}")
+                logger.warning(f"Error closing SEA statement: {e}")

src/databricks/sql/backend/thrift_backend.py

@@ -45,7 +45,7 @@
     _bound,
     RequestErrorInfo,
     NoRetryReason,
-    ResultSetQueueFactory,
+    ThriftResultSetQueueFactory,
     convert_arrow_based_set_to_arrow_table,
     convert_decimals_in_arrow_table,
     convert_column_based_set_to_arrow_table,
@@ -792,7 +792,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             assert direct_results.resultSet.results.startRowOffset == 0
             assert direct_results.resultSetMetadata
 
-            arrow_queue_opt = ResultSetQueueFactory.build_queue(
+            arrow_queue_opt = ThriftResultSetQueueFactory.build_queue(
                 row_set_type=t_result_set_metadata_resp.resultFormat,
                 t_row_set=direct_results.resultSet.results,
                 arrow_schema_bytes=schema_bytes,
@@ -859,7 +859,7 @@ def get_execution_result(
         else:
             schema_bytes = None
 
-        queue = ResultSetQueueFactory.build_queue(
+        queue = ThriftResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,
             t_row_set=resp.results,
             arrow_schema_bytes=schema_bytes,
@@ -1207,7 +1207,7 @@ def fetch_results(
                 )
             )
 
-        queue = ResultSetQueueFactory.build_queue(
+        queue = ThriftResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,
             t_row_set=resp.results,
             arrow_schema_bytes=arrow_schema_bytes,

src/databricks/sql/utils.py

@@ -48,30 +48,7 @@ def remaining_rows(self):
         pass
 
 
-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.n_valid_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.n_valid_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 ResultSetQueueFactory(ABC):
+class ThriftResultSetQueueFactory(ABC):
     @staticmethod
     def build_queue(
         row_set_type: Optional[TSparkRowSetType] = None,
@@ -81,57 +58,38 @@ def build_queue(
         ssl_options: Optional[SSLOptions] = None,
         lz4_compressed: bool = True,
         description: Optional[List[List[Any]]] = None,
-        # SEA specific parameters
-        sea_result_data: Optional[Any] = None,
     ) -> ResultSetQueue:
         """
-        Factory method to build a result set queue.
-        
-        This method is extended to handle both Thrift and SEA result formats.
-        For SEA, the sea_result_data parameter is used instead of row_set_type and t_row_set.
-        
+        Factory method to build a result set queue for Thrift backend.
+
         Args:
-            # Thrift parameters
             row_set_type (enum): Row set type (Arrow, Column, or URL).
             t_row_set (TRowSet): Result containing arrow batches, columns, or cloud fetch links.
-            
-            # Common parameters
             arrow_schema_bytes (bytes): Bytes representing the arrow schema.
             lz4_compressed (bool): Whether result data has been lz4 compressed.
             description (List[List[Any]]): Hive table schema description.
             max_download_threads (int): Maximum number of downloader thread pool threads.
             ssl_options (SSLOptions): SSLOptions object for CloudFetchQueue
-            
-            # SEA parameters
-            sea_result_data (ResultData): Result data from SEA response
-            
+
         Returns:
             ResultSetQueue
         """
-        # Handle SEA result data
-        if sea_result_data is not None:
-            if sea_result_data.data:
-                # INLINE disposition with JSON_ARRAY format
-                return JsonQueue(sea_result_data.data)
-            elif sea_result_data.external_links:
-                # EXTERNAL_LINKS disposition (not implemented yet)
-                raise NotImplementedError(
-                    "EXTERNAL_LINKS disposition is not supported yet"
-                )
-            else:
-                # Empty result set
-                return JsonQueue([])
-        
-        # Handle Thrift result data (existing implementation)
-        if row_set_type == TSparkRowSetType.ARROW_BASED_SET and t_row_set is not None and arrow_schema_bytes is not None:
+        # Handle Thrift result data
+        if (
+            row_set_type == TSparkRowSetType.ARROW_BASED_SET
+            and t_row_set is not None
+            and arrow_schema_bytes is not None
+        ):
             arrow_table, n_valid_rows = convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, arrow_schema_bytes
             )
             converted_arrow_table = convert_decimals_in_arrow_table(
                 arrow_table, description
             )
             return ArrowQueue(converted_arrow_table, n_valid_rows)
-        elif row_set_type == TSparkRowSetType.COLUMN_BASED_SET and t_row_set is not None:
+        elif (
+            row_set_type == TSparkRowSetType.COLUMN_BASED_SET and t_row_set is not None
+        ):
             column_table, column_names = convert_column_based_set_to_column_table(
                 t_row_set.columns, description
             )
@@ -141,7 +99,13 @@ def build_queue(
             )
 
             return ColumnQueue(ColumnTable(converted_column_table, column_names))
-        elif row_set_type == TSparkRowSetType.URL_BASED_SET and t_row_set is not None and arrow_schema_bytes is not None and max_download_threads is not None and ssl_options is not None:
+        elif (
+            row_set_type == TSparkRowSetType.URL_BASED_SET
+            and t_row_set is not None
+            and arrow_schema_bytes is not None
+            and max_download_threads is not None
+            and ssl_options is not None
+        ):
             return CloudFetchQueue(
                 schema_bytes=arrow_schema_bytes,
                 start_row_offset=t_row_set.startRowOffset,
@@ -155,6 +119,56 @@ def build_queue(
             raise AssertionError("Row set type is not valid")
 
 
+class SeaResultSetQueueFactory(ABC):
+    @staticmethod
+    def build_queue(
+        sea_result_data: Any,
+        description: Optional[List[List[Any]]] = None,
+    ) -> ResultSetQueue:
+        """
+        Factory method to build a result set queue for SEA backend.
+
+        Args:
+            sea_result_data (ResultData): Result data from SEA response
+            description (List[List[Any]]): Column descriptions
+
+        Returns:
+            ResultSetQueue: The appropriate queue for the result data
+        """
+        if sea_result_data.data:
+            # INLINE disposition with JSON_ARRAY format
+            return JsonQueue(sea_result_data.data)
+        elif sea_result_data.external_links:
+            # EXTERNAL_LINKS disposition (not implemented yet)
+            raise NotImplementedError("EXTERNAL_LINKS disposition is not supported yet")
+        else:
+            # Empty result set
+            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.n_valid_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.n_valid_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