feat: Implement streaming prefetch for Thrift inline results#1
Open
vikrantpuppala wants to merge 27 commits into
Open
feat: Implement streaming prefetch for Thrift inline results#1vikrantpuppala wants to merge 27 commits into
vikrantpuppala wants to merge 27 commits into
Conversation
## Description <!-- Provide a brief summary of the changes made and the issue they aim to address.--> Added query tags to telemetry ## Testing <!-- Describe how the changes have been tested--> Tested with real workspace in both cases: when query tags are present / not present. Behaviour is working as expected. ## Additional Notes to the Reviewer <!-- Share any additional context or insights that may help the reviewer understand the changes better. This could include challenges faced, limitations, or compromises made during the development process. Also, mention any areas of the code that you would like the reviewer to focus on specifically. --> NO_CHANGELOG=true Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
## Problem Custom user agent from useragententry parameter wasn't included in connector service HTTP requests for feature flag retrieval. ## Root Cause Method execution order issue in UserAgentManager.setUserAgent() - custom user agent was set AFTER the connector service request was made. ## Solution Reordered the method to set custom user agent before calling getClientUserAgent() (which triggers feature flag fetch). ## Testing - Added testCustomUserAgentIncludedBeforeClientTypeEvaluation() test NO_CHANGELOG=true --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
…action preview (databricks#1176) ## Summary - Changes default value of `IgnoreTransactions` parameter from `0` to `1`, making transactions disabled by default - Updates `supportsTransactions()` to respect the `IgnoreTransactions` flag, returning `false` when transactions are ignored (default) and `true` when explicitly enabled - Adds test case for when transactions are explicitly enabled via `IgnoreTransactions=0` ## Background The multi-statement transaction feature is currently in private preview for limited workspaces. When BI tools (Tableau, Power BI, DBeaver) detect transaction support via `supportsTransactions()`, they automatically use transaction methods, causing failures for customers not enrolled in the preview. This change prevents unexpected failures for non-preview customers while allowing preview participants to opt-in by explicitly setting `IgnoreTransactions=0` in their connection string. ## Migration Path - **Non-preview customers**: No action required - transactions are now disabled by default - **Preview participants**: Set `IgnoreTransactions=0` in connection string to enable transaction support - **GA migration**: When multi-statement transactions reach GA, flip the default back to `0` ## Test plan - [ ] Verify existing tests pass - [ ] Verify default connection returns `supportsTransactions() = false` - [ ] Verify connection with `IgnoreTransactions=0` returns `supportsTransactions() = true` - [ ] Verify transaction methods (`setAutoCommit`, `commit`, `rollback`) are no-ops by default 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
## Description Bump version to 3.1.1
…ricks#1177) ### Problem When executing queries that return 0 rows (e.g., `WHERE 1=0`), complex types (ARRAY, MAP, STRUCT) showed only generic type names instead of detailed type information: **Before:** - `ARRAY` instead of `ARRAY<INT>` - `MAP` instead of `MAP<STRING,STRING>` - `STRUCT` instead of `STRUCT<field: TYPE>` **After:** - Detailed type information is correctly preserved for all row counts ### Root Cause In `AbstractArrowResultChunk.java`, Arrow field metadata was only extracted inside the `while(arrowStreamReader.loadNextBatch())` loop. For queries with 0 rows, no batches are loaded, so the loop never executes and metadata is never extracted. **Code location:** `/src/main/java/com/databricks/jdbc/api/impl/arrow/AbstractArrowResultChunk.java:338-359` ### Solution Extract metadata from `VectorSchemaRoot` immediately after obtaining it, **before** the `loadNextBatch()` loop. The Arrow IPC format always sends the schema message first (before any record batches), so field metadata is available even when there are 0 rows. `VectorSchemaRoot` contains field vectors with metadata regardless of row count. **Key changes:** 1. Moved metadata extraction from inside the while loop to before it 2. Added defensive null checks for `VectorSchemaRoot` and field vectors 3. Added debug logging to track metadata extraction ### Testing #### Unit Test Coverage - ✅ Added `testMetadataExtractionWithZeroRows()` to `ArrowResultChunkTest` - ✅ Verifies Arrow field metadata is extracted correctly with 0 rows - ✅ Tests complex types: `ARRAY<INT>`, `MAP<STRING,STRING>` - ✅ All 2,693 unit tests pass #### Manual Verification Tested with queries returning 0 rows: ```sql SELECT array_col, map_col, struct_col FROM table WHERE 1=0 Result: Metadata now correctly shows detailed type information Impact - Scope: Both SQL Exec API and Thrift Server (shared code path) - Risk: Low - backward compatible change, only affects metadata extraction timing - Benefits: - Fixes schema discovery for WHERE 1=0 pattern - Improves metadata availability for empty result sets - Aligns with Arrow IPC specification behavior Additional Context - Arrow IPC specification guarantees schema is sent before record batches - VectorSchemaRoot.getFieldVectors() is available immediately after ArrowStreamReader.getVectorSchemaRoot() - No performance impact: metadata extraction is now done once upfront instead of conditionally on first batch --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
## Description
<!-- Provide a brief summary of the changes made and the issue they aim
to address.-->
This PR introduces lazy loading support for inline Arrow results to
improve memory efficiency when handling large result sets.
Previously, InlineChunkProvider would eagerly fetch all arrow batches
upfront when results had hasMoreRows = true, which could lead to memory
issues with large datasets. This change splits the handling into two
separate paths:
1. Lazy path (new): For Thrift-based inline Arrow results (when
ARROW_BASED_SET is returned), we now use LazyThriftInlineArrowResult
which fetches arrow batches on-demand as the client iterates through
rows. This is similar to how LazyThriftResult works for columnar data.
2. Remote path (existing): For URL-based Arrow results (URL_BASED_SET),
we continue using ArrowStreamResult with RemoteChunkProvider which
downloads chunks from cloud storage.
The InlineChunkProvider is now only used for SEA results with JSON_ARRAY
format and INLINE disposition (contain all data inline {no hasMoreRows
flag set}).
This will reduce memory consumption and improve performance when dealing
with large inline Arrow result sets similar to
databricks#975.
## Testing
<!-- Describe how the changes have been tested-->
- Unit tests
- Integration tests
- Manual testing
## Additional Notes to the Reviewer
<!-- Share any additional context or insights that may help the reviewer
understand the changes better. This could include challenges faced,
limitations, or compromises made during the development process.
Also, mention any areas of the code that you would like the reviewer to
focus on specifically. -->
Bypassing an existing failure on CI/CD because of databricks@3e4f21c
…istency (databricks#1182) ## Summary This PR adds TIMESTAMP_NTZ normalization in the Thrift path to ensure consistent metadata behavior across both SEA and Thrift API paths. ## Background PR databricks#1177 moved Arrow metadata extraction earlier in the processing pipeline, which exposed an inconsistency: the Thrift path started returning the correct "TIMESTAMP_NTZ" from server metadata, while the SEA path was already normalizing it to "TIMESTAMP" for backward compatibility. ## Changes - Added TIMESTAMP_NTZ → TIMESTAMP normalization in `DatabricksResultSetMetaData.java` Thrift constructor (lines 205-208) - This brings Thrift path behavior in line with existing SEA path normalization - Fixes test failure in `PreparedStatementIntegrationTests.testGetMetaData_NoResultSet` ## Testing - ✅ Local test run: `PreparedStatementIntegrationTests.testGetMetaData_NoResultSet` passes - ✅ Metadata now consistent before and after `executeQuery()` for TIMESTAMP_NTZ columns - ✅ Both SEA and Thrift paths return "TIMESTAMP" for TIMESTAMP_NTZ columns ## Related - Builds on PR databricks#1177 (Fix Arrow field metadata not available for queries with 0 rows) - Fixes issue introduced by early metadata extraction in PR databricks#1177 - Maintains backward compatibility with existing behavior 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
vikrantpuppala
force-pushed
the
non-cloud-latency
branch
2 times, most recently
from
January 19, 2026 11:05
4c16fd6 to
a9bd3cf
Compare
…oad parameter (databricks#1183) ## Summary Add support for disabling CloudFetch via `EnableQueryResultDownload=0` connection parameter to use inline Arrow results instead. ## Changes - Add `isCloudFetchEnabled()` method to `IDatabricksConnectionContext` interface - Implement the method in `DatabricksConnectionContext` using existing `EnableQueryResultDownload` parameter - Update `DatabricksThriftServiceClient` to respect this setting when making execute requests - Add unit tests for the new functionality ## Usage To disable CloudFetch and use inline Arrow results: ``` jdbc:databricks://host:port/default;EnableQueryResultDownload=0;... ``` --------- Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Implements proactive prefetching with sliding window for both Thrift columnar and inline Arrow results, eliminating blocking at batch boundaries. Key components: - ThriftStreamingProvider<T>: Generic streaming provider with type-safe batches - StreamingBatch<T>: Type-safe batch container with lifecycle management - ThriftResponseProcessor<T>: Pluggable processors for Columnar and Arrow - StreamingColumnarResult: Streaming variant for Thrift columnar results - StreamingInlineArrowResult: Streaming variant for inline Arrow results Configuration: - EnableInlineStreaming: Toggle streaming (default: enabled) - ThriftMaxBatchesInMemory: Sliding window size (default: 3) Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
- Add null validation for required parameters in ThriftStreamingProvider - Wrap batchFetcher.close() and batch.release() in try-catch blocks - Add timeout to waitForBatchCreation to prevent indefinite waiting - Add logging for error conditions in StreamingInlineArrowResult - Add logging for error conditions in InlineArrowResponseProcessor - Extract timeout constant in ExecutionResultFactory - Update NEXT_CHANGELOG.md to document streaming prefetch feature Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
vikrantpuppala
force-pushed
the
non-cloud-latency
branch
from
January 20, 2026 09:57
f3fc271 to
db7ad68
Compare
Merge latest main branch to incorporate CloudFetch disable feature (PR databricks#1183) Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
…lemetry code audit comments- part1 (databricks#1163) ## Description - 4 things that is improved with respect to telemetry : - Common object mapper across telemetry use-case (This is already thread safe and is expensive to create, i.e., good tor re-use) - Make`flushIntervalMillis` config same across both telemetry clients (un-auth and auth) - Clear connection param cache when connection is closed : this was a memory leak before - Rather than creating a scheduledExecutor for each telemetry client, we share it across a factory. ## Testing - unit tests ## Additional Notes to the Reviewer - When the LAST connection to a host is closed, all pending telemetry events for that host are flushed across all prior connections (since they all shared the same TelemetryClient). i.e., If you have 5 connections to `host-A`, closing connections 1-4 does nothing (just decrements refCount). Only when you close connection 5 (the last one) does the flush occur, sending all accumulated telemetry from all 5 connections. NO_CHANGELOG=true
- Make cachedSchema volatile for thread visibility in InlineArrowResponseProcessor - Add null checks for getData() in StreamingInlineArrowResult - Add null checks for currentBatch and getData() in StreamingColumnarResult Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
vikrantpuppala
force-pushed
the
non-cloud-latency
branch
from
January 21, 2026 08:42
0a2b5b7 to
8d592e4
Compare
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
## Description
This PR enhances geospatial datatype handling to include SRID (Spatial
Reference System Identifier) information in column type names and fixes
multiple issues related to complex datatype handling across different
result formats.
### Key Changes
1. **Geospatial Type Name Enhancement**
- Column type names now include SRID: `GEOMETRY(4326)` instead of
`GEOMETRY`
- Applies to both GEOMETRY and GEOGRAPHY types
- Preserves full type information in metadata for better type
identification
2. **SEA Inline Mode Complex Type Fix**
- Fixed issue where complex types (ARRAY, MAP, STRUCT) were not returned
as complex objects in SEA Inline mode (JSON array result format)
- Now properly converts to complex datatype objects when
`EnableComplexDatatypeSupport=true`
3. **Thrift CloudFetch Metadata Enhancement**
- Fixed error when extracting type details (e.g., `INT` from
`ARRAY<INT>`) in Thrift CloudFetch mode
- Enhanced `getColumnInfoFromTColumnDesc()` to use Arrow schema metadata
alongside `TColumnDesc`
- Arrow schema provides complete type information (e.g., `ARRAY<INT>`)
while `TColumnDesc` only contains base type (e.g., `ARRAY`)
4. **Arrow Metadata Extraction**
- Added `DatabricksThriftUtil.getArrowMetadata()` to deserialize Arrow
schema from `TGetResultSetMetadataResp`
- Fixed null arrow metadata issue in `DatabricksResultSet` constructor
for Thrift CloudFetch mode
## Testing
### Unit Tests
- All existing unit tests pass and additional tests are added for new
methods
### Integration Tests
- `GeospatialTests.java` - Comprehensive E2E integration test
- Tests geospatial types (GEOMETRY and GEOGRAPHY)
- Validates **24 configuration combinations**:
- Protocol: Thrift / SEA
- Serialization: Arrow / Inline
- CloudFetch: Enabled / Disabled (only with Arrow, as CloudFetch
requires Arrow)
- GeoSpatial Support: Enabled / Disabled
- Complex Type Support: Enabled / Disabled
- Validates metadata: column types, type names, class names
- Validates values: WKT representation, SRID
- Validates behavior when geospatial objects are enabled vs. disabled
(STRING fallback)
- **All 24 tests pass** ✅
## Additional Notes to the Reviewer
Other required details are mentioned in comments in the diff
---------
Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
…cloud-latency # Conflicts: # src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java # src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements proactive prefetching with sliding window for both Thrift columnar and inline Arrow results, eliminating blocking at batch boundaries.
Key Components
New Streaming Infrastructure
ThriftStreamingProvider<T>: Generic streaming provider with type-safe batches and background prefetch threadStreamingBatch<T>: Type-safe batch container with lifecycle managementThriftResponseProcessor<T>: Pluggable processors for Columnar (ColumnarResponseProcessor) and Arrow (InlineArrowResponseProcessor)Result Implementations
LazyThriftInlineArrowResult: Lazy loading for inline Arrow results (on-demand fetching)StreamingInlineArrowResult: Streaming variant for inline Arrow with prefetchStreamingThriftResult: Streaming variant for Thrift columnar with prefetchSupporting Classes
ThriftBatch,ThriftBatchFetcher,ThriftBatchFetcherImpl,ThriftBatchProviderConfiguration
EnableInlineStreaming1(enabled)ThriftMaxBatchesInMemory3Changes
Testing
LazyThriftInlineArrowResultTestwith comprehensive unit tests