chunk download latency

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/backend/thrift_backend.py

@@ -6,6 +6,7 @@
 import time
 import threading
 from typing import List, Optional, Union, Any, TYPE_CHECKING
+from uuid import UUID
 
 from databricks.sql.result_set import ThriftResultSet
 
@@ -1021,7 +1022,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results = self._handle_execute_response(
+            execute_response, is_direct_results, statement_id = self._handle_execute_response(
                 resp, cursor
             )
 
@@ -1040,6 +1041,8 @@ def execute_command(
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
                 is_direct_results=is_direct_results,
+                session_id_hex=self._session_id_hex,
+                statement_id=statement_id,
             )
 
     def get_catalogs(
@@ -1061,7 +1064,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, is_direct_results, _ = self._handle_execute_response(
             resp, cursor
         )
 
@@ -1107,7 +1110,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, is_direct_results, _ = self._handle_execute_response(
             resp, cursor
         )
 
@@ -1157,7 +1160,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, is_direct_results, _ = self._handle_execute_response(
             resp, cursor
         )
 
@@ -1207,7 +1210,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, is_direct_results, _ = self._handle_execute_response(
             resp, cursor
         )
 
@@ -1241,7 +1244,11 @@ def _handle_execute_response(self, resp, cursor):
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        return self._results_message_to_execute_response(resp, final_operation_state)
+        execute_response, is_direct_results = self._results_message_to_execute_response(
+            resp, final_operation_state
+        )
+
+        return execute_response, is_direct_results, cursor.active_command_id.to_hex_guid()
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
@@ -1261,6 +1268,7 @@ def fetch_results(
         arrow_schema_bytes,
         description,
         use_cloud_fetch=True,
+        statement_id=None,
     ):
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
@@ -1297,6 +1305,8 @@ def fetch_results(
             lz4_compressed=lz4_compressed,
             description=description,
             ssl_options=self._ssl_options,
+            session_id_hex=self._session_id_hex,
+            statement_id=statement_id
         )
 
         return queue, resp.hasMoreRows

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,7 +1,7 @@
 import logging
 
 from concurrent.futures import ThreadPoolExecutor, Future
-from typing import List, Union
+from typing import List, Union, Optional, Tuple
 
 from databricks.sql.cloudfetch.downloader import (
     ResultSetDownloadHandler,
@@ -22,24 +22,28 @@ def __init__(
         max_download_threads: int,
         lz4_compressed: bool,
         ssl_options: SSLOptions,
+        session_id_hex: Optional[str] = None,
+        statement_id: Optional[str] = None,
     ):
-        self._pending_links: List[TSparkArrowResultLink] = []
-        for link in links:
+        self._pending_links: List[Tuple[int, TSparkArrowResultLink]] = []
+        for i, link in enumerate(links):
             if link.rowCount <= 0:
                 continue
             logger.debug(
-                "ResultFileDownloadManager: adding file link, start offset {}, row count: {}".format(
-                    link.startRowOffset, link.rowCount
+                "ResultFileDownloadManager: adding file link, chunk id {}, start offset {}, row count: {}".format(
+                    i, link.startRowOffset, link.rowCount
                 )
             )
-            self._pending_links.append(link)
+            self._pending_links.append((i, link))
 
         self._download_tasks: List[Future[DownloadedFile]] = []
         self._max_download_threads: int = max_download_threads
         self._thread_pool = ThreadPoolExecutor(max_workers=self._max_download_threads)
 
         self._downloadable_result_settings = DownloadableResultSettings(lz4_compressed)
         self._ssl_options = ssl_options
+        self.session_id_hex = session_id_hex
+        self.statement_id = statement_id
 
     def get_next_downloaded_file(
         self, next_row_offset: int
@@ -89,14 +93,17 @@ def _schedule_downloads(self):
         while (len(self._download_tasks) < self._max_download_threads) and (
             len(self._pending_links) > 0
         ):
-            link = self._pending_links.pop(0)
+            chunk_id, link = self._pending_links.pop(0)
             logger.debug(
-                "- start: {}, row count: {}".format(link.startRowOffset, link.rowCount)
+                "- chunk: {}, start: {}, row count: {}".format(chunk_id, link.startRowOffset, link.rowCount)
             )
             handler = ResultSetDownloadHandler(
                 settings=self._downloadable_result_settings,
                 link=link,
                 ssl_options=self._ssl_options,
+                chunk_id=chunk_id,
+                session_id_hex=self.session_id_hex,
+                statement_id=self.statement_id
             )
             task = self._thread_pool.submit(handler.run)
             self._download_tasks.append(task)

src/databricks/sql/cloudfetch/downloader.py

@@ -1,5 +1,6 @@
 import logging
 from dataclasses import dataclass
+from typing import Optional
 
 import requests
 from requests.adapters import HTTPAdapter, Retry
@@ -9,6 +10,7 @@
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
+from databricks.sql.telemetry.latency_logger import log_latency
 
 logger = logging.getLogger(__name__)
 
@@ -66,11 +68,18 @@ def __init__(
         settings: DownloadableResultSettings,
         link: TSparkArrowResultLink,
         ssl_options: SSLOptions,
+        chunk_id: int,
+        session_id_hex: Optional[str] = None,
+        statement_id: Optional[str] = None,
     ):
         self.settings = settings
         self.link = link
         self._ssl_options = ssl_options
+        self.chunk_id = chunk_id
+        self.session_id_hex = session_id_hex
+        self.statement_id = statement_id
 
+    @log_latency()
     def run(self) -> DownloadedFile:
         """
         Download the file described in the cloud fetch link.
@@ -80,8 +89,8 @@ def run(self) -> DownloadedFile:
         """
 
         logger.debug(
-            "ResultSetDownloadHandler: starting file download, offset {}, row count {}".format(
-                self.link.startRowOffset, self.link.rowCount
+            "ResultSetDownloadHandler: starting file download, chunk id {}, offset {}, row count {}".format(
+                self.chunk_id, self.link.startRowOffset, self.link.rowCount
             )
         )
 

src/databricks/sql/result_set.py

@@ -198,6 +198,8 @@ def __init__(
         max_download_threads: int = 10,
         ssl_options=None,
         is_direct_results: bool = True,
+        session_id_hex: Optional[str] = None,
+        statement_id: Optional[str] = None,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -233,6 +235,8 @@ def __init__(
                 lz4_compressed=execute_response.lz4_compressed,
                 description=execute_response.description,
                 ssl_options=ssl_options,
+                session_id_hex=session_id_hex,
+                statement_id=statement_id,
             )
 
         # Call parent constructor with common attributes

src/databricks/sql/telemetry/latency_logger.py

@@ -7,7 +7,6 @@
     SqlExecutionEvent,
 )
 from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
-from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
 from uuid import UUID
 
 logger = logging.getLogger(__name__)
@@ -42,6 +41,9 @@ def get_execution_result(self):
     def get_retry_count(self):
         pass
 
+    def get_chunk_id(self):
+        pass
+
 
 class CursorExtractor(TelemetryExtractor):
     """
@@ -63,7 +65,8 @@ def get_is_compressed(self) -> bool:
     def get_execution_result(self) -> ExecutionResultFormat:
         if self.active_result_set is None:
             return ExecutionResultFormat.FORMAT_UNSPECIFIED
-
+        
+        from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
         if isinstance(self.active_result_set.results, ColumnQueue):
             return ExecutionResultFormat.COLUMNAR_INLINE
         elif isinstance(self.active_result_set.results, CloudFetchQueue):
@@ -74,11 +77,14 @@ def get_execution_result(self) -> ExecutionResultFormat:
 
     def get_retry_count(self) -> int:
         if (
-            hasattr(self.thrift_backend, "retry_policy")
-            and self.thrift_backend.retry_policy
+            hasattr(self.backend, "retry_policy")
+            and self.backend.retry_policy
         ):
-            return len(self.thrift_backend.retry_policy.history)
+            return len(self.backend.retry_policy.history)
         return 0
+    
+    def get_chunk_id(self):
+        return None
 
 
 class ResultSetExtractor(TelemetryExtractor):
@@ -101,6 +107,7 @@ def get_is_compressed(self) -> bool:
         return self.lz4_compressed
 
     def get_execution_result(self) -> ExecutionResultFormat:
+        from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
         if isinstance(self.results, ColumnQueue):
             return ExecutionResultFormat.COLUMNAR_INLINE
         elif isinstance(self.results, CloudFetchQueue):
@@ -116,7 +123,34 @@ def get_retry_count(self) -> int:
         ):
             return len(self.thrift_backend.retry_policy.history)
         return 0
+    
+    def get_chunk_id(self):
+        return None
+
+
+class ResultSetDownloadHandlerExtractor(TelemetryExtractor):
+    """
+    Telemetry extractor specialized for ResultSetDownloadHandler objects.
+    """
+    def get_session_id_hex(self) -> Optional[str]:
+        return self._obj.session_id_hex
+
+    def get_statement_id(self) -> Optional[str]:
+        return self._obj.statement_id
+
+    def get_is_compressed(self) -> bool:
+        return self._obj.settings.is_lz4_compressed
+
+    def get_execution_result(self) -> ExecutionResultFormat:
+        return ExecutionResultFormat.EXTERNAL_LINKS
+
+    def get_retry_count(self) -> Optional[int]:
+        # standard requests and urllib3 libraries don't expose retry count
+        return None
 
+    def get_chunk_id(self) -> Optional[int]:
+        return self._obj.chunk_id
+    
 
 def get_extractor(obj):
     """
@@ -133,12 +167,15 @@ def get_extractor(obj):
         TelemetryExtractor: A specialized extractor instance:
             - CursorExtractor for Cursor objects
             - ResultSetExtractor for ResultSet objects
+            - ResultSetDownloadHandlerExtractor for ResultSetDownloadHandler objects
             - None for all other objects
     """
     if obj.__class__.__name__ == "Cursor":
         return CursorExtractor(obj)
     elif obj.__class__.__name__ == "ResultSet":
         return ResultSetExtractor(obj)
+    elif obj.__class__.__name__=="ResultSetDownloadHandler":
+        return ResultSetDownloadHandlerExtractor(obj)
     else:
         logger.debug("No extractor found for %s", obj.__class__.__name__)
         return None
@@ -196,6 +233,7 @@ def _safe_call(func_to_call):
                 duration_ms = int((end_time - start_time) * 1000)
 
                 extractor = get_extractor(self)
+                print("function name", func.__name__, "latency", duration_ms, "session_id_hex", extractor.get_session_id_hex(), "statement_id", extractor.get_statement_id(), flush=True)
 
                 if extractor is not None:
                     session_id_hex = _safe_call(extractor.get_session_id_hex)
@@ -205,7 +243,8 @@ def _safe_call(func_to_call):
                         statement_type=statement_type,
                         is_compressed=_safe_call(extractor.get_is_compressed),
                         execution_result=_safe_call(extractor.get_execution_result),
-                        retry_count=_safe_call(extractor.get_retry_count),
+                        retry_count=extractor.get_retry_count(),
+                        chunk_id=_safe_call(extractor.get_chunk_id),
                     )
 
                     telemetry_client = TelemetryClientFactory.get_telemetry_client(

src/databricks/sql/telemetry/models/event.py

@@ -122,13 +122,14 @@ class SqlExecutionEvent(JsonSerializableMixin):
         is_compressed (bool): Whether the result is compressed
         execution_result (ExecutionResultFormat): Format of the execution result
         retry_count (int): Number of retry attempts made
+        chunk_id (int): ID of the chunk if applicable
     """
 
     statement_type: StatementType
     is_compressed: bool
     execution_result: ExecutionResultFormat
-    retry_count: int
-
+    retry_count: Optional[int]
+    chunk_id: Optional[int]
 
 @dataclass
 class TelemetryEvent(JsonSerializableMixin):