Merge branch 'ext-links-sea' into sea-hybrid

Author: varun-edachali-dbx

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

@@ -177,10 +177,10 @@ def __init__(
         )
 
         # Track the current chunk we're processing
-        self._current_chunk_link: Optional["ExternalLink"] = initial_link
+        self._current_chunk_link = initial_link
 
         # Initialize table and position
-        self.table = self._create_next_table()
+        self.table = self._create_table_from_link(self._current_chunk_link)
 
     def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
@@ -225,22 +225,30 @@ def _progress_chunk_link(self):
             f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_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
+    def _create_table_from_link(
+        self, link: "ExternalLink"
+    ) -> Union["pyarrow.Table", None]:
+        """Create a table from a link."""
 
         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)
+        thrift_link = self._convert_to_thrift_link(link)
         self.download_manager.add_link(thrift_link)
 
-        row_offset = self._current_chunk_link.row_offset
+        row_offset = link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)
 
+        return arrow_table
+
+    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+        """Create next table by retrieving the logical next downloaded file."""
+
         self._progress_chunk_link()
 
-        return arrow_table
+        if not self._current_chunk_link:
+            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
+            return None
+
+        return self._create_table_from_link(self._current_chunk_link)

tests/e2e/common/large_queries_mixin.py

@@ -2,6 +2,8 @@
 import math
 import time
 
+import pytest
+
 log = logging.getLogger(__name__)
 
 
@@ -42,7 +44,14 @@ def fetch_rows(self, cursor, row_count, fetchmany_size):
             + "assuming 10K fetch size."
         )
 
-    def test_query_with_large_wide_result_set(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_query_with_large_wide_result_set(self, extra_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8192  # B
         rows = resultSize // width
@@ -52,7 +61,7 @@ def test_query_with_large_wide_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 1000
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             for lz4_compression in [False, True]:
                 cursor.connection.lz4_compression = lz4_compression
                 uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
@@ -68,7 +77,14 @@ def test_query_with_large_wide_result_set(self):
                     assert row[0] == row_id  # Verify no rows are dropped in the middle.
                     assert len(row[1]) == 36
 
-    def test_query_with_large_narrow_result_set(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_query_with_large_narrow_result_set(self, extra_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8  # sizeof(long)
         rows = resultSize / width
@@ -77,12 +93,19 @@ def test_query_with_large_narrow_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 10000000
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             cursor.execute("SELECT * FROM RANGE({rows})".format(rows=rows))
             for row_id, row in enumerate(self.fetch_rows(cursor, rows, fetchmany_size)):
                 assert row[0] == row_id
 
-    def test_long_running_query(self):
+    @pytest.mark.parametrize(
+        "extra_params",
+        [
+            {},
+            {"use_sea": True},
+        ],
+    )
+    def test_long_running_query(self, extra_params):
         """Incrementally increase query size until it takes at least 3 minutes,
         and asserts that the query completes successfully.
         """
@@ -92,7 +115,7 @@ def test_long_running_query(self):
         duration = -1
         scale0 = 10000
         scale_factor = 1
-        with self.cursor() as cursor:
+        with self.cursor(extra_params) as cursor:
             while duration < min_duration:
                 assert scale_factor < 1024, "Detected infinite loop"
                 start = time.time()

tests/unit/test_sea_queue.py

@@ -213,7 +213,9 @@ def test_build_queue_arrow_stream(
 
         with patch(
             "databricks.sql.backend.sea.queue.ResultFileDownloadManager"
-        ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
+        ), patch.object(
+            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
+        ):
             queue = SeaResultSetQueueFactory.build_queue(
                 result_data=result_data,
                 manifest=arrow_manifest,
@@ -342,7 +344,9 @@ def test_init_with_valid_initial_link(
         mock_download_manager_class.return_value = mock_download_manager
 
         # Create a queue with valid initial link
-        with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
+        with patch.object(
+            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
+        ):
             queue = SeaCloudFetchQueue(
                 initial_links=[sample_external_link],
                 max_download_threads=5,
@@ -528,13 +532,14 @@ def test_create_next_table_success(self, mock_logger):
         # Mock the dependencies
         mock_table = Mock()
         queue._create_table_at_offset = Mock(return_value=mock_table)
+        queue._create_table_from_link = Mock(return_value=mock_table)
         queue._progress_chunk_link = Mock()
 
         # Call the method directly
         result = SeaCloudFetchQueue._create_next_table(queue)
 
         # Verify the table was created
-        queue._create_table_at_offset.assert_called_once_with(50)
+        queue._create_table_from_link.assert_called_once_with(queue._current_chunk_link)
 
         # Verify progress was called
         queue._progress_chunk_link.assert_called_once()