make download manager less defensive

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

Author: varun-edachali-dbx

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

@@ -357,14 +357,14 @@ def __init__(
         # Initialize table and position
         self.table = self._create_next_table()
 
-    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+    def _create_next_table(self) -> "pyarrow.Table":
         """Create next table by retrieving the logical next downloaded file."""
         if self.link_fetcher is None:
-            return None
+            return self._create_empty_table()
 
         chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index)
         if chunk_link is None:
-            return None
+            return self._create_empty_table()
 
         row_offset = chunk_link.row_offset
         # NOTE: link has already been submitted to download manager at this point

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,13 +1,15 @@
 import logging
 
 from concurrent.futures import ThreadPoolExecutor, Future
+import threading
 from typing import List, Union
 
 from databricks.sql.cloudfetch.downloader import (
     ResultSetDownloadHandler,
     DownloadableResultSettings,
     DownloadedFile,
 )
+from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
@@ -34,16 +36,16 @@ def __init__(
             )
             self._pending_links.append(link)
 
-        self._download_tasks: List[Future[DownloadedFile]] = []
         self._max_download_threads: int = max_download_threads
+
+        self._download_condition = threading.Condition()
+        self._download_tasks: List[Future[DownloadedFile]] = []
         self._thread_pool = ThreadPoolExecutor(max_workers=self._max_download_threads)
 
         self._downloadable_result_settings = DownloadableResultSettings(lz4_compressed)
         self._ssl_options = ssl_options
 
-    def get_next_downloaded_file(
-        self, next_row_offset: int
-    ) -> Union[DownloadedFile, None]:
+    def get_next_downloaded_file(self, next_row_offset: int) -> DownloadedFile:
         """
         Get next file that starts at given offset.
 
@@ -62,8 +64,10 @@ def get_next_downloaded_file(
 
         # No more files to download from this batch of links
         if len(self._download_tasks) == 0:
-            self._shutdown_manager()
-            return None
+            if self._thread_pool._shutdown:
+                raise Error("download manager shut down before file was ready")
+            with self._download_condition:
+                self._download_condition.wait()
 
         task = self._download_tasks.pop(0)
         # Future's `result()` method will wait for the call to complete, and return
@@ -124,3 +128,4 @@ def _shutdown_manager(self):
         self._pending_links = []
         self._download_tasks = []
         self._thread_pool.shutdown(wait=False)
+        self._download_condition.notify_all()

src/databricks/sql/utils.py

@@ -236,7 +236,7 @@ def __init__(
         self._ssl_options = ssl_options
 
         # Table state
-        self.table = None
+        self.table = self._create_empty_table()
         self.table_row_index = 0
 
         # Initialize download manager
@@ -256,23 +256,20 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
         Returns:
             pyarrow.Table
         """
-        if not self.table:
-            logger.debug("CloudFetchQueue: no more rows available")
-            # Return empty pyarrow table to cause retry of fetch
-            return self._create_empty_table()
+
         logger.debug("CloudFetchQueue: trying to get {} next rows".format(num_rows))
         results = self.table.slice(0, 0)
-        while num_rows > 0 and self.table:
+        while num_rows > 0 and self.table.num_rows > 0:
+            # Replace current table with the next table if we are at the end of the current table
+            if self.table_row_index == self.table.num_rows:
+                self.table = self._create_next_table()
+                self.table_row_index = 0
+
             # Get remaining of num_rows or the rest of the current table, whichever is smaller
             length = min(num_rows, self.table.num_rows - self.table_row_index)
             table_slice = self.table.slice(self.table_row_index, length)
             results = pyarrow.concat_tables([results, table_slice])
             self.table_row_index += table_slice.num_rows
-
-            # Replace current table with the next table if we are at the end of the current table
-            if self.table_row_index == self.table.num_rows:
-                self.table = self._create_next_table()
-                self.table_row_index = 0
             num_rows -= table_slice.num_rows
 
         logger.debug("CloudFetchQueue: collected {} next rows".format(results.num_rows))
@@ -286,11 +283,8 @@ def remaining_rows(self) -> "pyarrow.Table":
             pyarrow.Table
         """
 
-        if not self.table:
-            # Return empty pyarrow table to cause retry of fetch
-            return self._create_empty_table()
         results = self.table.slice(0, 0)
-        while self.table:
+        while self.table.num_rows > 0:
             table_slice = self.table.slice(
                 self.table_row_index, self.table.num_rows - self.table_row_index
             )
@@ -300,17 +294,11 @@ def remaining_rows(self) -> "pyarrow.Table":
             self.table_row_index = 0
         return results
 
-    def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
+    def _create_table_at_offset(self, offset: int) -> "pyarrow.Table":
         """Create next table at the given row offset"""
 
         # Create next table by retrieving the logical next downloaded file, or return None to signal end of queue
         downloaded_file = self.download_manager.get_next_downloaded_file(offset)
-        if not downloaded_file:
-            logger.debug(
-                "CloudFetchQueue: Cannot find downloaded file for row {}".format(offset)
-            )
-            # None signals no more Arrow tables can be built from the remaining handlers if any remain
-            return None
         arrow_table = create_arrow_table_from_arrow_file(
             downloaded_file.file_bytes, self.description
         )
@@ -326,7 +314,7 @@ def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
         return arrow_table
 
     @abstractmethod
-    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+    def _create_next_table(self) -> "pyarrow.Table":
         """Create next table by retrieving the logical next downloaded file."""
         pass
 
@@ -345,7 +333,7 @@ class ThriftCloudFetchQueue(CloudFetchQueue):
 
     def __init__(
         self,
-        schema_bytes,
+        schema_bytes: Optional[bytes],
         max_download_threads: int,
         ssl_options: SSLOptions,
         start_row_offset: int = 0,
@@ -390,23 +378,24 @@ def __init__(
                 )
                 self.download_manager.add_link(result_link)
 
-        # Initialize table and position
-        self.table = self._create_next_table()
+            # Initialize table and position
+            self.table = self._create_next_table()
+        else:
+            self.table = self._create_empty_table()
 
-    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+    def _create_next_table(self) -> "pyarrow.Table":
         logger.debug(
             "ThriftCloudFetchQueue: Trying to get downloaded file for row {}".format(
                 self.start_row_index
             )
         )
         arrow_table = self._create_table_at_offset(self.start_row_index)
-        if arrow_table:
-            self.start_row_index += arrow_table.num_rows
-            logger.debug(
-                "ThriftCloudFetchQueue: Found downloaded file, row count: {}, new start offset: {}".format(
-                    arrow_table.num_rows, self.start_row_index
-                )
+        self.start_row_index += arrow_table.num_rows
+        logger.debug(
+            "ThriftCloudFetchQueue: Found downloaded file, row count: {}, new start offset: {}".format(
+                arrow_table.num_rows, self.start_row_index
             )
+        )
         return arrow_table