Introduce abstract base classes for Arrow result handling by jayantsing-db · Pull Request #850 · databricks/databricks-jdbc

jayantsing-db · 2025-06-11T12:15:52Z

Description

This PR is the first of two and splits the changes originally proposed in #634.

Key changes:

Create AbstractRemoteChunkProvider and AbstractArrowResultChunk as unified base implementations for chunk management
Add state machine (ArrowResultChunkStateMachine) to handle chunk lifecycle transitions
Improve thread synchronization between main thread and IO/download threads

The new async implementation (2nd PR) will provide an alternative path for handling large Arrow datasets with improved scalability, while maintaining the existing synchronous approach for backwards compatibility.

Below is a class diagram that provides a clearer understanding and concise summary of the class structure:

Testing

Fake service tests
Unit tests
Multi-DBR tests
Local testing in a highly concurrent environment

Additional Notes to the Reviewer

Benchmarking document: https://docs.google.com/document/d/1MvKeSnrQVKFGdkuaSPXFCrbb4TDaUqFwHrIW2rxy3zI/edit?usp=sharing

* parser * more complex data type logic * complete complex data types * more tests * refine * integration test * add test to driver test * removed secret * fmt * fix unsupported test * integration test working * fmt * comments * fmt * update method sig * fmt * comments * comments * comments * fmt * fmt * add * more extensive testing * more extensive testing * more extensive testing * more extensive testing * more extensive testing

* Send intial connection log to unauth telemetry endpoint * Add tests + remove debug logs * fmt * Add tests * Add tests * Add tests for classes i didn't add:) * add getters/setters in systemConfiguration + improve coverage

- Run execution queries in `async` mode - Run metadata queries in `async` mode - Refactor result retrieval in Thrift by improving response handling and fixing polling - Refactor credential refreshing by transitioning from Thrift accessor to Thrift transport - Competing driver, SQL Gateway, and SQL Execution API also execute the queries in async mode - Minor code cleanup Running queries in async mode and using polling will be beneficial: - Freeing up server resources under high load instead of holding the connection for ~5 seconds (direct results) - If query takes more than 5 seconds, server will anyway require polling from client - Didn't notice any regressions (slightly better performance) under high concurrent load when running queries in async mode. Some empirical results below: - Running a simple select query on a table with 10-15 columns for 10K rows in a concurrent environment (250 connections executing query in parallel) - Average results when running on `main` ``` Average query time across connections: 73873.58ms Min query time: 22075ms Max query time: 88997ms ``` - Average results when running in `async` mode ``` Average query time across connections: 68382.95ms Min query time: 24718ms Max query time: 95282ms ```

…ch (#622) This adds a workflow to run JDBC comparator in a separate branch for GA readiness. Note that this does not affect the main branch.

Add databricks protected runners

Use databricks runners

Setup PAT correctly

…y paths (#620) * change file paths * change file paths * change file paths * change file paths * change file paths * change file paths * change file paths

Earlier it was PWD which is already set to some path!!

This is because we need the workflow to proceed further and email the differences. The test will instead send the output to STDERR.

Minor changes to github workflow around JDBC comparator

This reverts commit 5124293.

This will send a report around 00:00 UTC. The code is in the branch jdbc-comparator. The reports will be saved as artifacts of github workflow run.

- Add nvdApiKey configuration to OWASP dependency-check-maven plugin - Configure retry settings and API delay for better performance - Update GitHub Actions workflow to pass NVD_API_KEY secret - Resolves warning about missing NVD API key and improves scan performance Without this, release shows a warning: An NVD API Key was not provided - it is highly recommended to use an NVD API key as the update can take a VERY long time without an API Key

- Add `Support for INTERVAL data type.` - This was not added to NEXT_CHANGELOG by authors

… Query (#838) * Basic working * Added tests * nit * nit * PR comments * nit * Added e2e tests * Removed Describe query logs * log fix * Updated error handling

* Update SDK version2 * Fix test

* fixed actions * nit * nit * nit * check * nit

* added application name to track BiTool * add tests * some changes * refactored code; update user agent and telemetry client app * update user agent thrift mode as well * changes fn to package-private. * make parameterized test * spotless * merge UserAgentHelper and UserAgentManager * addressed reviews * update DatabricksHttpClient

* Add statement ID and session ID wherever possible * Add null check tests

* ADded logging * Added logging

* Add support for SSL client certificate authentication - Implemented client certificate authentication via keystore configuration parameters. - Updated `ConfiguratorUtils` to handle loading and managing keystores. - Enhanced unit tests to cover keystore loading and client certificate scenarios. * - Removed redundant null check for the keystore in createKeyManagers method. - Simplified key entry verification by replacing Enumeration with a stream-based approach for better readability and performance. * addresses comments * address comments * fix next changelog * Remove unnecessary copy command for intermediateCA.crt in sslTesting workflow * Update sslTesting workflow to correct client certificate configuration for mutual TLS * Update sslTesting workflow to include intermediate CA certificate and remove obsolete DNS directive * revert to main * Refactor HTTP exception handling to use DatabricksSSLException for SSL/TLS errors. Update relevant classes and tests to replace DatabricksHttpException with DatabricksSSLException, ensuring proper error handling during SSL configuration and handshake processes. * Update error message in DatabricksClientConfiguratorManager to reflect SSL error handling instead of HTTP error. This change enhances clarity in logging and exception management related to SSL configuration failures.

…d in databricks/databricks-jdbc#634. Key changes: - Create AbstractRemoteChunkProvider and AbstractArrowResultChunk as unified base implementations for chunk management - Add state machine (ArrowResultChunkStateMachine) to handle chunk lifecycle transitions - Improve thread synchronization between main thread and IO/download threads The new async implementation (2nd PR) will provide an alternative path for handling large Arrow datasets with improved scalability, while maintaining the existing synchronous approach for backwards compatibility. Below is a class diagram that provides a clearer understanding and concise summary of the class structure: ![class-diagram](https://github.com/user-attachments/assets/e9503d47-5895-439a-9d39-f4963da3e5df) - Fake service tests - Unit tests - Multi-DBR tests - Local testing in a highly concurrent environment Benchmarking document: https://docs.google.com/document/d/1MvKeSnrQVKFGdkuaSPXFCrbb4TDaUqFwHrIW2rxy3zI/edit?usp=sharing

github-actions · 2025-06-11T12:16:24Z

Please ensure that the NEXT_CHANGELOG.md file is updated with any relevant changes.
If this is not necessary for your PR, please include the following in your PR description:
NO_CHANGELOG=true
and rerun the job.

gopalldb · 2025-06-19T04:29:27Z

+      return false;
+    }
+
+    if (getStatus() == ChunkStatus.PROCESSING_SUCCEEDED) {


when is this status set?

This status signifies that the chunk data has been deserialised and is ready to be read. It happens after the download.

gopalldb · 2025-06-19T04:38:07Z

+  protected void downloadData(IDatabricksHttpClient httpClient, CompressionCodec compressionCodec)
      throws DatabricksParsingException, IOException {
-    // Inject error if enabled for testing
-    if (injectError && errorInjectionCount < errorInjectionCountMaxValue) {


we don't need to check for inject errors? Not sure that I understand these errors :|

This was a throw-away code and was added to test the retry behaviour of chunk download using DriverTester

gopalldb · 2025-06-19T04:39:39Z

              response.getEntity().getContent(), compressionCodec, decompressionContext);
      initializeData(uncompressedStream);
-      setStatus(ChunkStatus.DOWNLOAD_SUCCEEDED);
+      setStatus(ChunkStatus.PROCESSING_SUCCEEDED);


shall we do this as part of initializeData itself?

Makes sense will do.

gopalldb · 2025-06-19T04:41:55Z

      this.numRows = rowCount;
      this.inputStream = stream;
-      this.status = ChunkStatus.PENDING;
+      this.status = status == null ? ChunkStatus.DOWNLOAD_SUCCEEDED : status;


if it is pending, then also we can set Download_Succeeded

This is the special case of inline results where we have the results in one go. So the builder method is called with all results downloaded. Hence the status, DOWNLOAD_SUCCEEDED. Note that it will have status PROCESSING_SUCCEEDED only if processing is successful. Another reason to distinguish between DOWNLOAD_SUCCEEDED and PROCESSING_SUCCEEDED :)

gopalldb · 2025-06-19T04:47:34Z

+      return; // No transition needed
+    }
+
+    Set<ChunkStatus> validTargets = VALID_TRANSITIONS.get(currentStatus);


you can call isValidTransition(..) here

gopalldb · 2025-06-19T04:48:04Z

+   *
+   * @return Set of valid target states
+   */
+  public synchronized Set<ChunkStatus> getValidTargetStates() {


where is this used?

this is not used yet in this PR. will be used in later PR for logging purposes. The method returns the valid states from current state.

gopalldb · 2025-06-19T04:53:20Z

    } finally {
-      if (!downloadSuccessful) {
+      if (downloadSuccessful) {
+        chunk.getChunkReadyFuture().complete(null);


add comment on what is null, e.g. null /*.paramName */

gopalldb · 2025-06-19T05:00:12Z

-   * advanced debugging).
-   */
-  @Test
-  void exampleChunkDownloadRetry() throws Exception {


why remove this?

this is not being used anymore. There are separate tests for the retry tests.

vikrantpuppala · 2025-06-19T08:23:01Z

+    return getStatus() == ChunkStatus.PENDING
+        || (!DriverUtil.isRunningAgainstFake()
+            && expiryTime.minusSeconds(SECONDS_BUFFER_FOR_EXPIRY).isBefore(Instant.now()));
+  }


60s seems too long for 22mb download, what's the validity of presigned URLs?

Yes, URL is valid for 15 minutes. Buffer to check expiry is 1 minute. I have retained the old code for now.

vikrantpuppala · 2025-06-19T08:25:04Z

+  /**
+   * Releases all resources associated with this chunk.
+   *
+   * @return true if chunk was released, false if it was already released


why is this contract needed, how is this boolean used? seems like we would need to return false only when release failed, not if already released

The return value is used in

databricks-jdbc/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractRemoteChunkProvider.java

Line 275 in 007aa9f

if (chunkIndexToChunksMap.get(currentChunkIndex).releaseChunk()) {

to decrement the chunk count in memory and queue next chunks for download.

vikrantpuppala · 2025-06-19T08:28:08Z

+   * @return number of record batches
+   */
+  protected int getRecordBatchCountInChunk() {
+    return getStatus() == ChunkStatus.PROCESSING_SUCCEEDED ? recordBatchList.size() : 0;


should we fail here instead of returning 0?

I think it makes sense to fail than to read some junk or throw NPE in chunk iterator. but I kept this processing logic same as before to limit the changes. I will create a separate PR for this.

vikrantpuppala · 2025-06-19T08:31:53Z

+      LOGGER.warn(
+          "Failed to transition to state [%s] from state [%s] for chunk [%d] and statement [%s]. Stack trace: %s",
+          targetStatus, getStatus(), chunkIndex, statementId, ExceptionUtils.getStackTrace(e));
+    }


not sure if it is ok to proceed if transition fails?

Yes initially i was throwing exception for an invalid transitions. But when i introduced this, there were quite a few invalid state transitions (all transparent to user correctness but invalid technically) in the existing code/logic which I fixed. But there could be 1-2 somewhere hidden for which I am throwing the warning for now. I will make this strict sometime later.

vikrantpuppala · 2025-06-19T08:34:59Z

+      chunkReadyFuture.get(CHUNK_READY_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw e;


log this probably?

vikrantpuppala · 2025-06-19T08:40:21Z

+      IDatabricksSession session,
+      IDatabricksHttpClient httpClient,
+      int maxParallelChunkDownloadsPerQuery,
+      CompressionCodec compressionCodec)


one constructor should use the other?

I think you meant we should re-use the constructor.

protected AbstractRemoteChunkProvider( StatementId statementId, ResultManifest resultManifest, ResultData resultData, IDatabricksSession session, IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec)

protected AbstractRemoteChunkProvider( IDatabricksStatementInternal parentStatement, TFetchResultsResp resultsResp, IDatabricksSession session, IDatabricksHttpClient httpClient, int maxParallelChunkDownloadsPerQuery, CompressionCodec compressionCodec)

RemoteChunkProvider (now most of the logic sits in AbstractRemoteChunkProvider) had these two constructors for thrift and SEA flow. Constructor chaining here will look a bit too inflated. Ideally I think we should have separate classes for thrift and SEA in such cases. But I have kept this construction untouched to limit the scope.

oh i think you meant re-using for the common fields. I can do that now.

oh i think you meant re-using for the common fields. I can do that now.

it is getting slightly complex without introducing a third telescope constructor. I will leave this as is for now 😅

vikrantpuppala · 2025-06-19T08:42:37Z

+    return chunkIndexMap;
+  }
+
+  private void populateChunkIndexMap(TRowSet resultData, ConcurrentMap<Long, T> chunkIndexMap)


please add javadoc to these methods

madhav-db and others added 30 commits December 17, 2024 05:30

Send intial connection log to unauth telemetry endpoint (#614)

d6c8f06

* Send intial connection log to unauth telemetry endpoint * Add tests + remove debug logs * fmt * Add tests * Add tests * Add tests for classes i didn't add:) * add getters/setters in systemConfiguration + improve coverage

Run JDBC comparator

2a5ab75

fmt

f2fefb1

Add workflow to run JDBC Driver Comparison in a separate private bran…

8c225c0

…ch (#622) This adds a workflow to run JDBC comparator in a separate branch for GA readiness. Note that this does not affect the main branch.

Create OSS connection first

f99e205

Remove PWD from OSS connection

add14b3

Update runJdbcComparator.yml

e81bf2e

Add databricks protected runners

Update runJdbcComparator.yml

e17cb94

Use databricks runners

Setup PAT correctly

5a98d17

Update runJdbcComparator.yml

c6454ce

Setup PAT correctly

Change integration test cached API call file paths to Windows-friendl…

852d69a

…y paths (#620) * change file paths * change file paths * change file paths * change file paths * change file paths * change file paths * change file paths

Add default to URL

0a888f6

Exclude telemetry

12a1f73

Verify PWD

ecc4c4f

Change env var name

fa376b2

Change token env var name

ce59d0f

Earlier it was PWD which is already set to some path!!

fmt

f6f0e9b

Run comparator on prod benchfood

112cab0

Do not fail the test on finding differences

1cf715d

This is because we need the workflow to proceed further and email the differences. The test will instead send the output to STDERR.

Improve mail body format

00147e9

Improve mail format of JDBC comparator

6a5100d

Minor changes to github workflow around JDBC comparator

Send body as html

5124293

Revert "Send body as html"

d4bef7d

This reverts commit 5124293.

Improved HTML body

1e5315d

Add more details to the body

fd2d5bd

Add a sql query

f632849

Add a workflow to e-mail a jdbc-comparison-report (#625)

6cffd9e

This will send a report around 00:00 UTC. The code is in the branch jdbc-comparator. The reports will be saved as artifacts of github workflow run.

Add telemetry schema to match our frontEndLog proto (#623)

0ac809d

jayantsing-db and others added 12 commits May 30, 2025 13:43

Update CHANGELOG.md

deec7e3

- Add `Support for INTERVAL data type.` - This was not added to NEXT_CHANGELOG by authors

[PECOBLR-414]Prepared Statement getMetaData to work without executing…

8693d24

… Query (#838) * Basic working * Added tests * nit * nit * PR comments * nit * Added e2e tests * Removed Describe query logs * log fix * Updated error handling

Add support for DoD domains and update SDK version to 0.52.0 (#843)

70d51b4

* Update SDK version2 * Fix test

Updated Github actions versions (#845)

cd24356

* fixed actions * nit * nit * nit * check * nit

Add statement ID and session ID wherever possible (#844)

59e524b

* Add statement ID and session ID wherever possible * Add null check tests

Add checks on thrift files based on ID (#846)

0790e02

Logging error, when scale from arrowMetaData is not found (#848)

5e717f5

* ADded logging * Added logging

Merge remote-tracking branch 'databricks/main' into http-client

0a4148b

jayantsing-db temporarily deployed to azure-prod June 11, 2025 12:15 — with GitHub Actions Inactive

jayantsing-db changed the title ~~Jayantsing db/http client~~ Introduce abstract base classes for Arrow result handling Jun 11, 2025

jayantsing-db requested a review from gopalldb June 11, 2025 12:21

gopalldb reviewed Jun 19, 2025

View reviewed changes

gopalldb requested a review from samikshya-db June 19, 2025 04:43

gopalldb reviewed Jun 19, 2025

View reviewed changes

vikrantpuppala reviewed Jun 19, 2025

View reviewed changes

jayantsing-db closed this Jun 20, 2025

jayantsing-db force-pushed the main branch from 33f02ce to cef0a00 Compare June 20, 2025 19:55

jayantsing-db mentioned this pull request Jul 6, 2025

Introduce abstract base classes for Arrow result handling #881

Merged

Uh oh!

Conversation

jayantsing-db commented Jun 11, 2025

Description

Testing

Additional Notes to the Reviewer

Uh oh!

github-actions Bot commented Jun 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayantsing-db Jun 20, 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!

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

jayantsing-db Jun 20, 2025 •

edited

Loading