Add StreamingChunkProvider for result fetching #1111

jayantsing-db · 2025-11-27T10:45:02Z

Description

Implement streaming chunk provider that fetches results without dependency on total chunk count. Key components:

StreamingChunkProvider: Memory-bounded parallel downloads with proactive link prefetching
ChunkLinkFetcher interface with SeaChunkLinkFetcher for SEA API
StreamingChunkDownloadTask: Simplified download task with retry logic
Support for initial links from ResultData to avoid extra fetch calls

Enable via URL parameter: EnableStreamingChunkProvider=1

Next: Implement ThriftChunkLinkFetcher to unify SEA and Thrift code paths

Testing

Manual testing

Additional Notes to the Reviewer

Implement streaming chunk provider that fetches results without dependency on total chunk count. Key components: - StreamingChunkProvider: Memory-bounded parallel downloads with proactive link prefetching - ChunkLinkFetcher interface with SeaChunkLinkFetcher for SEA API - StreamingChunkDownloadTask: Simplified download task with retry logic - Support for initial links from ResultData to avoid extra fetch calls Enable via URL parameter: EnableStreamingChunkProvider=1 Next: Implement ThriftChunkLinkFetcher to unify SEA and Thrift code paths

tejassp-db

[Big PR] Review is not completed yet.

src/main/java/com/databricks/jdbc/model/core/ChunkLinkFetchResult.java

tejassp-db · 2025-12-03T10:25:42Z

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

+  private final ReentrantLock prefetchLock = new ReentrantLock();
+  private final Condition consumerAdvanced = prefetchLock.newCondition();
+  private final Condition chunkCreated = prefetchLock.newCondition();


Can these be replaced with a BlockingQueue<ArrowResultChunk>?

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

tejassp-db · 2025-12-03T10:45:46Z

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

+   * @return The refreshed ExternalLink with a new expiration time
+   * @throws DatabricksSQLException if the refetch operation fails
+   */
+  ExternalLink refetchLink(long chunkIndex, long rowOffset) throws DatabricksSQLException;


Does the server endpoint API return multiple links? If so aren't we making unnecessary endpoint API calls?

Yes, we do get multiple links, and likely subsequent links also need to be refreshed if they are not yet downloaded. Due to parallelism, it could be possible that some later links got downloaded earlier.

Does the server endpoint API return multiple links? If so aren't we making unnecessary endpoint API calls?

Yes I need to update this: Basically when we refresh for a chunk index, we should very well refresh the links for all the chunks we get in a single RPC

Note that with this consumer approach signalling link fetch, it would be rare to hit link expiry issue. I will also introduce a param to configure the link fetch window that can be adjusted based on the expected time spent by the main thread to consume the chunks

Then the code should accommodate fetching multiple links with one call and reduce the number of calls. This API should change.

Yes the return type will change to a collection and corresponding changes.

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

tejassp-db · 2025-12-03T10:57:14Z

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

+    if (endOfStreamReached) {
+      return highestKnownChunkIndex + 1;
+    }
+    return -1; // Unknown


This needs to be documented in the API.

tejassp-db · 2025-12-03T11:01:51Z

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

+          LOGGER.debug("Successfully downloaded chunk {}", chunk.getChunkIndex());
+
+        } catch (IOException | DatabricksSQLException e) {
+          retries++;


Does the httpClient also retry internally? Then this accounting is incorrect.

Need to check this. I have not updated this code in the refactor.

There are some ongoing efforts to unify the http retries in a separate branch https://github.com/databricks/databricks-jdbc/tree/retry-unification. These are yet to be merged.

This is the current behaviour:

What HTTP Client Does NOT Retry

Error Type HTTP Client Retries?

400 Bad Request No

401 Unauthorized No

403 Forbidden No

404 Not Found No

408 Request Timeout No

500 Internal Server Error No

502 Bad Gateway No

504 Gateway Timeout No

Network errors (IOException) No

Connection timeout No

SSL/TLS errors No

DNS resolution failures No

503 without Retry-After header No

429 without Retry-After header No

Given that chunk downloads from cloud storage (S3, Azure Blob, GCS):

Typically return 403 for expired links (not retried by HTTP client)

May return 500/502/504 for transient errors (not retried by HTTP client)

Rarely return 503/429 with Retry-After header (this is specific to Sql gateway)

For most practical scenarios, the HTTP client will NOT retry chunk downloads and therefore these task level retries were introduced. But the unify retry work https://github.com/databricks/databricks-jdbc/tree/retry-unification should address this later.

Ideally, post that PR, we will only manually handle refetching an expired link once in the task and rest of the scenarios like IOException, error codes, retry delay will be handled by the HTTP client.

tejassp-db · 2025-12-03T11:04:04Z

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

+                "Retry {} for chunk {}: {}", retries, chunk.getChunkIndex(), e.getMessage());
+            chunk.setStatus(ChunkStatus.DOWNLOAD_RETRY);
+            try {
+              Thread.sleep(RETRY_DELAY_MS);


If the httpClient retries with a delay, then this is incorrect accounting. Also why are we adding delays to a retry here?

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

gopalldb · 2025-12-05T03:15:06Z

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

+    }
+
+    List<ExternalLink> linkList =
+        externalLinks instanceof List


why converting to ArrayList explicitly? can't you do collection.stream directly?

gopalldb · 2025-12-05T03:16:36Z

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

+    int lastIndex = resultLinks.size() - 1;
+    boolean hasMoreRows = resultsResp.hasMoreRows;
+
+    for (int i = 0; i < resultLinks.size(); i++) {


instead of i, use readable names

gopalldb · 2025-12-05T03:22:44Z

src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java

+    }
+
+    List<ExternalLink> linkList =
+        links instanceof List ? (List<ExternalLink>) links : new ArrayList<>(links);


why this conversion to list? It is already a collection

gopalldb · 2025-12-05T03:30:32Z

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

+    long nextRowOffset = rowOffset;
+    long nextFetchIndex = chunkIndex;
+
+    for (int i = 0; i < resultLinks.size(); i++) {


better readable iterator name

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

src/main/java/com/databricks/jdbc/model/core/ChunkLinkFetchResult.java

gopalldb · 2025-12-09T03:27:06Z

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

+    }
+
+    // Exact match not found - this indicates a server bug
+    throw new DatabricksSQLException(


can we log this?

@gopalldb Throwing an exception is part of the API contract, and exceptions are expected to be handled at the place they are caught. Logging an exception is one way of handling it at the catch site.

Logging and throwing can create multiple/duplicate log entries for the same error if the place where the exception is caught also logs the error. Is there a specific reason to log and throw?

gopalldb · 2025-12-09T03:28:01Z

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

+        }
+
+        // Exact match not found - this indicates a server bug
+        throw new DatabricksSQLException(


add logging

gopalldb · 2025-12-09T03:28:13Z

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

+      // No links returned - check if we should retry
+      if (!result.hasMore()) {
+        // No more data and no links - this is unexpected for a refetch
+        throw new DatabricksSQLException(


gopalldb · 2025-12-09T03:28:33Z

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

+          maxRetries);
+    }
+
+    throw new DatabricksSQLException(


gopalldb · 2025-12-09T03:30:03Z

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

+    }
+
+    if (chunk == null) {
+      throw new DatabricksSQLException(


add logging

gopalldb · 2025-12-09T03:30:35Z

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

+
+    if (chunk == null) {
+      throw new DatabricksSQLException(
+          "Chunk " + currentChunkIndex + " not found after waiting",


nit: can use String.format for concating string

String.format throws an IllegalFormatException on any formatting errors including incorrect format specifiers. It should be used when very specific formatting is required.

Generally it is safer to concatenate using + operator since it calls toString() on all the objects and handles null cases safely as well.

gopalldb · 2025-12-09T03:30:50Z

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

+          e,
+          DatabricksDriverErrorCode.THREAD_INTERRUPTED_ERROR);
+    } catch (ExecutionException e) {
+      throw new DatabricksSQLException(


add logging for all exceptions

jayantsing-db · 2025-12-09T09:31:44Z

Created a test PR for this: #1125

jayantsing-db added 4 commits November 27, 2025 10:39

fmt

51b0dee

Add Thrift implementation

cd14529

fmt

d23eed4

tejassp-db assigned jayantsing-db Dec 3, 2025

tejassp-db reviewed Dec 3, 2025

View reviewed changes

gopalldb reviewed Dec 5, 2025

View reviewed changes

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java Show resolved Hide resolved

gopalldb reviewed Dec 5, 2025

View reviewed changes

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java Show resolved Hide resolved

gopalldb reviewed Dec 5, 2025

View reviewed changes

src/main/java/com/databricks/jdbc/model/core/ChunkLinkFetchResult.java Outdated Show resolved Hide resolved

Address review comments

5bffdb5

gopalldb reviewed Dec 9, 2025

View reviewed changes

jayantsing-db closed this Dec 9, 2025

jayantsing-db deleted the jayantsing-db/chunk-download branch December 9, 2025 08:28

jayantsing-db restored the jayantsing-db/chunk-download branch December 9, 2025 08:29

jayantsing-db reopened this Dec 9, 2025

jayantsing-db and others added 2 commits December 9, 2025 14:14

Merge branch 'main' into jayantsing-db/chunk-download

5818265

fmt

22fb391

Error Type	HTTP Client Retries?
400 Bad Request	No
401 Unauthorized	No
403 Forbidden	No
404 Not Found	No
408 Request Timeout	No
500 Internal Server Error	No
502 Bad Gateway	No
504 Gateway Timeout	No
Network errors (IOException)	No
Connection timeout	No
SSL/TLS errors	No
DNS resolution failures	No
503 without Retry-After header	No
429 without Retry-After header	No

Add StreamingChunkProvider for result fetching #1111

Are you sure you want to change the base?

Add StreamingChunkProvider for result fetching #1111

Conversation

jayantsing-db commented Nov 27, 2025

Description

Testing

Additional Notes to the Reviewer

Uh oh!

tejassp-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gopalldb Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tejassp-db Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayantsing-db commented Dec 9, 2025

Uh oh!

Reviewers

Assignees

Labels

gopalldb Dec 5, 2025 •

edited

Loading

tejassp-db Dec 9, 2025 •

edited

Loading