Add chunk latency

Author: samikshya-db

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java

@@ -23,6 +23,7 @@
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.jdbc.telemetry.TelemetryClientFactory;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import com.google.common.annotations.VisibleForTesting;
 import java.sql.*;
 import java.util.*;
@@ -150,6 +151,7 @@ public void close() throws DatabricksSQLException {
     LOGGER.debug("public void close()");
     for (IDatabricksStatementInternal statement : statementSet) {
       statement.close(false);
+      ChunkLatencyHandler.getInstance().clearStatement(statement.getStatementId());
       statementSet.remove(statement);
     }
     this.session.close();

src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java

@@ -1,16 +1,14 @@
 package com.databricks.jdbc.api.impl.arrow;
 
-import static com.databricks.jdbc.telemetry.TelemetryHelper.exportLatencyLog;
-
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.util.DatabricksThreadContextHolder;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
-import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.core.ExternalLink;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.concurrent.ExecutionException;
@@ -67,7 +65,13 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
 
           chunk.downloadData(httpClient, chunkDownloader.getCompressionCodec());
           downloadSuccessful = true;
-        } catch (DatabricksParsingException | IOException e) {
+
+          // Record chunk download latency on successful download
+          long downloadLatency = System.currentTimeMillis() - startTime;
+          ChunkLatencyHandler.getInstance()
+              .recordChunkDownloadLatency(statementId, chunk.getChunkIndex(), downloadLatency);
+
+        } catch (IOException | DatabricksSQLException e) {
           retries++;
           if (retries >= MAX_RETRIES) {
             LOGGER.error(
@@ -110,7 +114,6 @@ public Void call() throws DatabricksSQLException, ExecutionException, Interrupte
         chunk.setStatus(ArrowResultChunk.ChunkStatus.DOWNLOAD_FAILED);
       }
 
-      exportLatencyLog(System.currentTimeMillis() - startTime);
       chunkDownloader.downloadProcessed(chunk.getChunkIndex());
       DatabricksThreadContextHolder.clearAllContext();
     }

src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkLinkDownloadService.java

@@ -10,6 +10,7 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.core.ExternalLink;
+import com.databricks.jdbc.telemetry.latency.ChunkLatencyHandler;
 import java.time.Instant;
 import java.util.Collection;
 import java.util.Map;
@@ -103,6 +104,7 @@ public ChunkLinkDownloadService(
     this.isDownloadInProgress = new AtomicBoolean(false);
     this.isDownloadChainStarted = new AtomicBoolean(false);
     this.isShutdown = false;
+    ChunkLatencyHandler.getInstance().initializeStatement(statementId, totalChunks);
 
     this.chunkIndexToLinkFuture = new ConcurrentHashMap<>();
     for (long i = 0; i < totalChunks; i++) {

src/main/java/com/databricks/jdbc/model/telemetry/DriverConnectionParameters.java

@@ -121,9 +121,6 @@ public class DriverConnectionParameters {
   @JsonProperty("enable_sea_hybrid_results")
   boolean enableSeaHybridResults;
 
-  @JsonProperty("enable_complex_datatype_support")
-  boolean enableComplexSupport;
-
   @JsonProperty("allow_self_signed_support")
   boolean allowSelfSignedSupport;
 
@@ -331,11 +328,6 @@ public DriverConnectionParameters setEnableSeaHybridResults(boolean enableSeaHyb
     return this;
   }
 
-  public DriverConnectionParameters setEnableComplexSupport(boolean enableComplexSupport) {
-    this.enableComplexSupport = enableComplexSupport;
-    return this;
-  }
-
   public DriverConnectionParameters setAllowSelfSignedSupport(boolean allowSelfSignedSupport) {
     this.allowSelfSignedSupport = allowSelfSignedSupport;
     return this;
@@ -397,7 +389,6 @@ public String toString() {
         .add("nonProxyHosts", nonProxyHosts)
         .add("httpConnectionPoolSize", httpConnectionPoolSize)
         .add("enableSeaHybridResults", enableSeaHybridResults)
-        .add("enableComplexSupport", enableComplexSupport)
         .add("allowSelfSignedSupport", allowSelfSignedSupport)
         .add("useSystemTrustStore", useSystemTrustStore)
         .add("rowsFetchedPerBlock", rowsFetchedPerBlock)

src/main/java/com/databricks/jdbc/model/telemetry/latency/ChunkDetails.java

@@ -12,14 +12,20 @@ public class ChunkDetails {
   private Long slowestChunkLatencyMillis;
 
   @JsonProperty("total_chunks_present")
-  private Integer totalChunksPresent;
+  private Long totalChunksPresent;
 
   @JsonProperty("total_chunks_iterated")
-  private Integer totalChunksIterated;
+  private Long totalChunksIterated;
 
   @JsonProperty("sum_chunks_download_time_millis")
   private Long sumChunksDownloadTimeMillis;
 
+  public ChunkDetails(long totalChunks) {
+    this.totalChunksIterated = 0L;
+    this.sumChunksDownloadTimeMillis = 0L;
+    this.totalChunksPresent = totalChunks;
+  }
+
   public ChunkDetails setInitialChunkLatencyMillis(Long initialChunkLatencyMillis) {
     this.initialChunkLatencyMillis = initialChunkLatencyMillis;
     return this;
@@ -30,12 +36,12 @@ public ChunkDetails setSlowestChunkLatencyMillis(Long slowestChunkLatencyMillis)
     return this;
   }
 
-  public ChunkDetails setTotalChunksPresent(Integer totalChunksPresent) {
+  public ChunkDetails setTotalChunksPresent(Long totalChunksPresent) {
     this.totalChunksPresent = totalChunksPresent;
     return this;
   }
 
-  public ChunkDetails setTotalChunksIterated(Integer totalChunksIterated) {
+  public ChunkDetails setTotalChunksIterated(Long totalChunksIterated) {
     this.totalChunksIterated = totalChunksIterated;
     return this;
   }
@@ -45,6 +51,22 @@ public ChunkDetails setSumChunksDownloadTimeMillis(Long sumChunksDownloadTimeMil
     return this;
   }
 
+  public Long getSlowestChunkLatencyMillis() {
+    return slowestChunkLatencyMillis;
+  }
+
+  public Long getTotalChunksPresent() {
+    return totalChunksPresent;
+  }
+
+  public Long getTotalChunksIterated() {
+    return totalChunksIterated;
+  }
+
+  public Long getSumChunksDownloadTimeMillis() {
+    return sumChunksDownloadTimeMillis;
+  }
+
   @Override
   public String toString() {
     return new ToStringer(ChunkDetails.class)

src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java

@@ -231,7 +231,6 @@ private static DriverConnectionParameters buildDriverConnectionParameters(
             .setNonProxyHosts(StringUtil.split(connectionContext.getNonProxyHosts()))
             .setHttpConnectionPoolSize(connectionContext.getHttpConnectionPoolSize())
             .setEnableSeaHybridResults(connectionContext.isSqlExecHybridResultsEnabled())
-            .setEnableComplexSupport(connectionContext.isComplexDatatypeSupportEnabled())
             .setAllowSelfSignedSupport(connectionContext.allowSelfSignedCerts())
             .setUseSystemTrustStore(connectionContext.useSystemTrustStore())
             .setRowsFetchedPerBlock(connectionContext.getRowsFetchedPerBlock())

src/main/java/com/databricks/jdbc/telemetry/latency/ChunkLatencyHandler.java

@@ -0,0 +1,158 @@
+package com.databricks.jdbc.telemetry.latency;
+
+import com.databricks.jdbc.dbclient.impl.common.StatementId;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.jdbc.model.telemetry.latency.ChunkDetails;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Handler for tracking chunk-related latency metrics for Databricks JDBC driver. This class manages
+ * per-statement chunk details and provides logic for data collection.
+ */
+public class ChunkLatencyHandler {
+
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ChunkLatencyHandler.class);
+
+  // Singleton instance for global access
+  private static final ChunkLatencyHandler INSTANCE = new ChunkLatencyHandler();
+
+  // Per-statement latency tracking using ChunkDetails
+  private final ConcurrentHashMap<String, ChunkDetails> statementTrackers =
+      new ConcurrentHashMap<>();
+
+  private ChunkLatencyHandler() {
+    // Private constructor for singleton
+  }
+
+  public static ChunkLatencyHandler getInstance() {
+    return INSTANCE;
+  }
+
+  /**
+   * Initialize tracking for a statement with the total number of chunks.
+   *
+   * @param statementId the statement ID
+   * @param totalChunks the total number of chunks in the result set
+   */
+  public void initializeStatement(StatementId statementId, long totalChunks) {
+    if (statementId == null) {
+      LOGGER.trace("Statement ID is null, skipping initialization");
+      return;
+    }
+    statementTrackers.put(statementId.toString(), new ChunkDetails(totalChunks));
+    LOGGER.trace(
+        "Initialized chunk tracking for statement {} with {} total chunks",
+        statementId.toString(),
+        totalChunks);
+  }
+
+  /**
+   * Records the latency for downloading a chunk and updates metrics.
+   *
+   * @param statementId the statement ID
+   * @param chunkIndex the index of the chunk being downloaded
+   * @param latencyMillis the time taken to download the chunk in milliseconds
+   */
+  public void recordChunkDownloadLatency(String statementId, long chunkIndex, long latencyMillis) {
+    if (statementId == null) {
+      LOGGER.trace("Statement ID is null, skipping chunk latency recording");
+      return;
+    }
+
+    ChunkDetails chunkDetails =
+        statementTrackers.computeIfAbsent(statementId, k -> new ChunkDetails(0));
+
+    // Record initial chunk latency (first chunk downloaded)
+    if (chunkIndex == 0) {
+      chunkDetails.setInitialChunkLatencyMillis(latencyMillis);
+    }
+
+    // Update slowest chunk latency
+    Long currentSlowest = chunkDetails.getSlowestChunkLatencyMillis();
+    if (currentSlowest == null || latencyMillis > currentSlowest) {
+      chunkDetails.setSlowestChunkLatencyMillis(latencyMillis);
+    }
+
+    // Add to sum of all chunk download times
+    Long currentSum = chunkDetails.getSumChunksDownloadTimeMillis();
+    if (currentSum == null) {
+      currentSum = 0L;
+    }
+    chunkDetails.setSumChunksDownloadTimeMillis(currentSum + latencyMillis);
+
+    LOGGER.trace(
+        "Recorded chunk {} latency: {}ms for statement {}", chunkIndex, latencyMillis, statementId);
+  }
+
+  /**
+   * Records when a chunk is iterated/consumed by the result set.
+   *
+   * @param statementId the statement ID
+   * @param chunkIndex the index of the chunk being iterated
+   */
+  public void recordChunkIteration(String statementId, long chunkIndex) {
+    if (statementId == null) {
+      return;
+    }
+
+    ChunkDetails chunkDetails = statementTrackers.get(statementId);
+    if (chunkDetails != null) {
+      Long currentIterated = chunkDetails.getTotalChunksIterated();
+      if (currentIterated == null) {
+        currentIterated = 0L;
+      }
+      chunkDetails.setTotalChunksIterated(currentIterated + 1);
+      LOGGER.trace("Recorded chunk {} iteration for statement {}", chunkIndex, statementId);
+    }
+  }
+
+  /**
+   * Gets the collected chunk details for a statement without removing the tracker.
+   *
+   * @param statementId the statement ID
+   * @return the ChunkDetails object or null if no tracker found
+   */
+  public ChunkDetails getChunkDetails(String statementId) {
+    if (statementId == null) {
+      return null;
+    }
+
+    return statementTrackers.get(statementId);
+  }
+
+  /**
+   * Gets the collected chunk details and removes the tracker from memory (cleanup).
+   *
+   * @param statementId the statement ID
+   * @return the ChunkDetails object or null if no tracker found
+   */
+  public ChunkDetails getChunkDetailsAndCleanup(String statementId) {
+    if (statementId == null) {
+      return null;
+    }
+
+    ChunkDetails chunkDetails = statementTrackers.remove(statementId);
+    if (chunkDetails == null) {
+      LOGGER.trace("No chunk latency telemetry found for statement {}", statementId);
+      return null;
+    }
+    return chunkDetails;
+  }
+
+  /**
+   * Clears tracking data for a statement (useful for cleanup).
+   *
+   * @param statementId the statement ID to clear tracking for
+   */
+  public void clearStatement(StatementId statementId) {
+    if (statementId == null) {
+      LOGGER.trace("Statement ID is null, skipping cleanup");
+      return;
+    }
+
+    String statementIdStr = statementId.toString();
+    statementTrackers.remove(statementIdStr);
+    LOGGER.trace("Cleared tracking for statement {}", statementIdStr);
+  }
+}