pre-fetch next chunk link on processing current

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

Author: varun-edachali-dbx

examples/experimental/test_sea_multi_chunk.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 from databricks.sql.client import Connection
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
 
@@ -195,7 +195,7 @@ def main():
         sys.exit(1)
 
     # Get row count from command line or use default
-    requested_row_count = 5000
+    requested_row_count = 10000
 
     if len(sys.argv) > 1:
         try:

src/databricks/sql/cloud_fetch_queue.py

@@ -334,6 +334,7 @@ def __init__(
 
         # 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()
@@ -351,8 +352,22 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
             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
 
@@ -369,24 +384,19 @@ def _progress_chunk_link(self):
                         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 None")
+            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
             return None
 
-        logger.debug(
-            f"SeaCloudFetchQueue: Trying to get downloaded file for chunk {self._current_chunk_link.chunk_index}"
-        )
-
-        if self.download_manager:
-            thrift_link = self._convert_to_thrift_link(self._current_chunk_link)
-            self.download_manager.add_link(thrift_link)
-
         row_offset = self._current_chunk_link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)