Added logs

Author: jprakash-db

src/databricks/sql/client.py

@@ -1515,7 +1515,10 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
+        print("Server side has more rows", self.has_more_rows)
+        
         while not self.has_been_closed_server_side and self.has_more_rows:
+            print(f"RESULT SIZE TOTAL {results.num_rows}")
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             if isinstance(results, ColumnTable) and isinstance(

src/databricks/sql/cloudfetch/download_manager.py

@@ -78,13 +78,13 @@ def get_next_downloaded_file(
                     next_row_offset, file.start_row_offset, file.row_count
                 )
             )
-
         return file
 
     def _schedule_downloads(self):
         """
         While download queue has a capacity, peek pending links and submit them to thread pool.
         """
+        print("Schedule_downloads")
         logger.debug("ResultFileDownloadManager: schedule downloads")
         while (len(self._download_tasks) < self._max_download_threads) and (
             len(self._pending_links) > 0

src/databricks/sql/cloudfetch/downloader.py

@@ -95,6 +95,10 @@ def run(self) -> DownloadedFile:
         session.mount("https://", HTTPAdapter(max_retries=retryPolicy))
 
         try:
+            print_text = [
+
+            ]
+            start_time = time.time()
             # Get the file via HTTP request
             response = session.get(
                 self.link.fileLink,
@@ -104,7 +108,8 @@ def run(self) -> DownloadedFile:
                 # TODO: Pass cert from `self._ssl_options`
             )
             response.raise_for_status()
-
+            end_time = time.time()
+            print_text.append(f"Downloaded file in {end_time - start_time} seconds")
             # Save (and decompress if needed) the downloaded file
             compressed_data = response.content
             decompressed_data = (
@@ -127,6 +132,13 @@ def run(self) -> DownloadedFile:
                 )
             )
 
+            print_text.append(
+                f"Downloaded file startRowOffset - {self.link.startRowOffset} - rowCount - {self.link.rowCount}"
+            )
+
+            for text in print_text:
+                print(text)
+
             return DownloadedFile(
                 decompressed_data,
                 self.link.startRowOffset,

src/databricks/sql/utils.py

@@ -235,6 +235,8 @@ def __init__(
                         result_link.startRowOffset, result_link.rowCount
                     )
                 )
+        print("Initial Setup Cloudfetch Queue")
+        print(f"No of result links - {len(result_links)}")
         self.download_manager = ResultFileDownloadManager(
             links=result_links or [],
             max_download_threads=self.max_download_threads,
@@ -288,6 +290,9 @@ def remaining_rows(self) -> "pyarrow.Table":
             # Return empty pyarrow table to cause retry of fetch
             return self._create_empty_table()
         results = self.table.slice(0, 0)
+        
+        print("remaining_rows call")
+        print(f"self.table.num_rows - {self.table.num_rows}")
         while self.table:
             table_slice = self.table.slice(
                 self.table_row_index, self.table.num_rows - self.table_row_index
@@ -296,6 +301,7 @@ def remaining_rows(self) -> "pyarrow.Table":
             self.table_row_index += table_slice.num_rows
             self.table = self._create_next_table()
             self.table_row_index = 0
+        print(f"results.num_rows - {results.num_rows}")
         return results
 
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
@@ -334,7 +340,9 @@ def _create_next_table(self) -> Union["pyarrow.Table", None]:
                 arrow_table.num_rows, self.start_row_index
             )
         )
-
+        
+        print("_create_next_table")
+        print(f"arrow_table.num_rows - {arrow_table.num_rows}")
         return arrow_table
 
     def _create_empty_table(self) -> "pyarrow.Table":