feat: Implement streaming prefetch for Thrift inline results (databricks#1184)

## Summary

Implements proactive prefetching with a sliding window for both Thrift
columnar and inline Arrow results, eliminating blocking at batch
boundaries and improving throughput.

## Key Components

### New Streaming Infrastructure
- **`ThriftStreamingProvider<T>`**: Generic type-safe streaming provider
with background prefetch thread and configurable sliding window
- **`StreamingBatch<T>`**: Type-safe batch container with lifecycle
management and error handling
- **`ThriftResponseProcessor<T>`**: Interface for pluggable response
processors
  - `ColumnarResponseProcessor`: Processes Thrift columnar results
- `InlineArrowResponseProcessor`: Processes inline Arrow results with
schema caching

### Result Implementations
- **`StreamingInlineArrowResult`**: High-throughput streaming
implementation for inline Arrow results with background prefetching
- **`StreamingColumnarResult`**: Streaming implementation for Thrift
columnar results with prefetch

### Supporting Classes
- **`ThriftBatchFetcher`** / **`ThriftBatchFetcherImpl`**: Abstraction
for fetching batches from the Thrift server

<img width="1792" height="1234" alt="streaming inline"
src="https://github.com/user-attachments/assets/66ea9b83-a16b-42d5-9280-cb1fb81dadeb"
/>

## Configuration

| Parameter | Description | Default |
|-----------|-------------|---------|
| `EnableInlineStreaming` | Toggle streaming mode for inline results |
`1` (enabled) |
| `ThriftMaxBatchesInMemory` | Sliding window size (max batches kept in
memory) | `3` |

## Key Features

1. **Background Prefetching**: Dedicated thread fetches batches ahead of
consumption
2. **Sliding Window**: Configurable memory limit prevents unbounded
memory growth
3. **Type Safety**: Generic `ThriftStreamingProvider<T>` eliminates
unsafe casting
4. **Graceful Error Handling**: 
   - Try-catch around resource cleanup to prevent cascading failures
   - Timeout on batch creation wait to prevent indefinite blocking
5. **Comprehensive Logging**: Debug/error logging for troubleshooting

## Testing

- Updated `ExecutionResultFactoryTest` for new factory logic
- Updated `DatabricksThriftServiceClientTest` for CloudFetch control
- Existing integration tests cover streaming behavior

## Usage

Streaming is enabled by default. To disable and use lazy loading
instead:

```
jdbc:databricks://host:port/default;EnableInlineStreaming=0;...
```

To adjust the sliding window size:

```
jdbc:databricks://host:port/default;ThriftMaxBatchesInMemory=5;...
```

---------

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

NEXT_CHANGELOG.md

@@ -3,6 +3,9 @@
 ## [Unreleased]
 
 ### Added
+- Added streaming prefetch mode for Thrift inline results (columnar and Arrow) with background batch prefetching and configurable sliding window for improved throughput.
+- Added `EnableInlineStreaming` connection parameter to enable/disable streaming mode (default: enabled).
+- Added `ThriftMaxBatchesInMemory` connection parameter to control the sliding window size for streaming (default: 3).
 - Added support for disabling CloudFetch via `EnableQueryResultDownload=0` to use inline Arrow results instead.
 
 ### Updated

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

@@ -1198,11 +1198,27 @@ public boolean isStreamingChunkProviderEnabled() {
     return getParameter(DatabricksJdbcUrlParams.ENABLE_STREAMING_CHUNK_PROVIDER).equals("1");
   }
 
+  @Override
+  public boolean isInlineStreamingEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_INLINE_STREAMING).equals("1");
+  }
+
   @Override
   public boolean isCloudFetchEnabled() {
     return getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH).equals("1");
   }
 
+  @Override
+  public int getThriftMaxBatchesInMemory() {
+    try {
+      return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.THRIFT_MAX_BATCHES_IN_MEMORY));
+    } catch (NumberFormatException e) {
+      LOGGER.warn("Invalid value for ThriftMaxBatchesInMemory, using default value");
+      return Integer.parseInt(
+          DatabricksJdbcUrlParams.THRIFT_MAX_BATCHES_IN_MEMORY.getDefaultValue());
+    }
+  }
+
   @Override
   public int getLinkPrefetchWindow() {
     return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.LINK_PREFETCH_WINDOW));

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

@@ -8,8 +8,11 @@
 import com.databricks.jdbc.api.IExecutionStatus;
 import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
 import com.databricks.jdbc.api.impl.arrow.ChunkProvider;
+import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult;
+import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult;
 import com.databricks.jdbc.api.impl.converters.ConverterHelper;
 import com.databricks.jdbc.api.impl.converters.ObjectConverter;
+import com.databricks.jdbc.api.impl.thrift.StreamingColumnarResult;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
 import com.databricks.jdbc.api.internal.IDatabricksResultSetInternal;
 import com.databricks.jdbc.api.internal.IDatabricksSession;
@@ -592,12 +595,11 @@ public boolean isBeforeFirst() throws SQLException {
   /**
    * {@inheritDoc}
    *
-   * <p><b>Limitation:</b> For lazy-loaded result sets ({@link LazyThriftResult}), particularly
-   * those using {@link
-   * com.databricks.jdbc.model.client.thrift.generated.TSparkRowSetType#COLUMN_BASED_SET}, this
-   * method cannot reliably determine the cursor position. The total row count remains unknown until
-   * all rows are fetched, preventing accurate detection of whether the cursor is after the last
-   * row. This is specific to Databricks JDBC dialect.
+   * <p><b>Limitation:</b> For lazy/streaming result sets ({@link LazyThriftResult}, {@link
+   * StreamingColumnarResult}, {@link LazyThriftInlineArrowResult}, {@link
+   * StreamingInlineArrowResult}), this method cannot reliably determine the cursor position. The
+   * total row count remains unknown until all rows are fetched, preventing accurate detection of
+   * whether the cursor is after the last row. This is specific to Databricks JDBC dialect.
    */
   @Override
   public boolean isAfterLast() throws SQLException {
@@ -617,8 +619,10 @@ public boolean isFirst() throws SQLException {
    * <p>This method uses different strategies based on the result set type:
    *
    * <ul>
-   *   <li>For {@link LazyThriftResult} instances: Checks if there are no more rows available (using
-   *       {@code hasNext()}), since the total row count is unknown until all rows are fetched.
+   *   <li>For lazy/streaming result types ({@link LazyThriftResult}, {@link
+   *       StreamingColumnarResult}, {@link LazyThriftInlineArrowResult}, {@link
+   *       StreamingInlineArrowResult}): Checks if there are no more rows available (using {@code
+   *       hasNext()}), since the total row count is unknown until all rows are fetched.
    *   <li>For other result types: Compares the current row position against the known total row
    *       count.
    * </ul>
@@ -629,7 +633,10 @@ public boolean isFirst() throws SQLException {
   @Override
   public boolean isLast() throws SQLException {
     checkIfClosed();
-    if (executionResult instanceof LazyThriftResult) {
+    if (executionResult instanceof LazyThriftResult
+        || executionResult instanceof StreamingColumnarResult
+        || executionResult instanceof LazyThriftInlineArrowResult
+        || executionResult instanceof StreamingInlineArrowResult) {
       return executionResult.getCurrentRow() >= 0 && !executionResult.hasNext();
     }
     return executionResult.getCurrentRow() == resultSetMetaData.getTotalRows() - 1;

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

@@ -2,7 +2,10 @@
 
 import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
 import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult;
+import com.databricks.jdbc.api.impl.arrow.StreamingInlineArrowResult;
+import com.databricks.jdbc.api.impl.thrift.StreamingColumnarResult;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
+import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.api.internal.IDatabricksSession;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.common.util.DatabricksThriftUtil;
@@ -96,9 +99,9 @@ private static IExecutionResult getResultHandler(
     LOGGER.info("Processing result of format {} from Thrift server", resultFormat);
     switch (resultFormat) {
       case COLUMN_BASED_SET:
-        return new LazyThriftResult(resultsResp, parentStatement, session);
+        return createThriftColumnarResult(resultsResp, parentStatement, session);
       case ARROW_BASED_SET:
-        return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);
+        return createInlineArrowResult(resultsResp, parentStatement, session);
       case URL_BASED_SET:
         return new ArrowStreamResult(resultsResp, parentStatement, session);
       case ROW_BASED_SET:
@@ -110,6 +113,50 @@ private static IExecutionResult getResultHandler(
     }
   }
 
+  /**
+   * Creates the appropriate result handler for Thrift columnar results. Uses
+   * StreamingColumnarResult by default for improved throughput, otherwise falls back to
+   * LazyThriftResult if disabled.
+   */
+  private static IExecutionResult createThriftColumnarResult(
+      TFetchResultsResp resultsResp,
+      IDatabricksStatementInternal parentStatement,
+      IDatabricksSession session)
+      throws DatabricksSQLException {
+    IDatabricksConnectionContext connectionContext = session.getConnectionContext();
+
+    // Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1")
+    if (connectionContext.isInlineStreamingEnabled()) {
+      LOGGER.info("Using StreamingColumnarResult for improved throughput (default)");
+      return new StreamingColumnarResult(resultsResp, parentStatement, session);
+    } else {
+      LOGGER.info("Using LazyThriftResult (streaming explicitly disabled)");
+      return new LazyThriftResult(resultsResp, parentStatement, session);
+    }
+  }
+
+  /**
+   * Creates the appropriate result handler for inline Arrow results. Uses
+   * StreamingInlineArrowResult by default for improved throughput, otherwise falls back to
+   * LazyThriftInlineArrowResult.
+   */
+  private static IExecutionResult createInlineArrowResult(
+      TFetchResultsResp resultsResp,
+      IDatabricksStatementInternal parentStatement,
+      IDatabricksSession session)
+      throws DatabricksSQLException {
+    IDatabricksConnectionContext connectionContext = session.getConnectionContext();
+
+    // Streaming is enabled by default (ENABLE_INLINE_STREAMING defaults to "1")
+    if (connectionContext.isInlineStreamingEnabled()) {
+      LOGGER.info("Using StreamingInlineArrowResult for improved throughput (default)");
+      return new StreamingInlineArrowResult(resultsResp, parentStatement, session);
+    } else {
+      LOGGER.info("Using LazyThriftInlineArrowResult (streaming explicitly disabled)");
+      return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);
+    }
+  }
+
   static IExecutionResult getResultSet(Object[][] rows) {
     return new InlineJsonResult(rows);
   }

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

@@ -1,24 +1,21 @@
 package com.databricks.jdbc.api.impl.arrow;
 
+import static com.databricks.jdbc.common.util.ArrowUtil.getColumnInfoList;
 import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createExternalLink;
-import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getColumnInfoFromTColumnDesc;
 
 import com.databricks.jdbc.api.impl.ComplexDataTypeParser;
 import com.databricks.jdbc.api.impl.IExecutionResult;
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.api.internal.IDatabricksSession;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.common.CompressionCodec;
-import com.databricks.jdbc.common.util.DatabricksThriftUtil;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
-import com.databricks.jdbc.model.client.thrift.generated.TColumnDesc;
 import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp;
-import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp;
 import com.databricks.jdbc.model.client.thrift.generated.TSparkArrowResultLink;
 import com.databricks.jdbc.model.core.ChunkLinkFetchResult;
 import com.databricks.jdbc.model.core.ColumnInfo;
@@ -163,7 +160,7 @@ public ArrowStreamResult(
       IDatabricksHttpClient httpClient)
       throws DatabricksSQLException {
     this.session = session;
-    setColumnInfo(resultsResp.getResultSetMetadata());
+    this.columnInfos = getColumnInfoList(resultsResp.getResultSetMetadata());
     this.chunkProvider =
         createThriftRemoteChunkProvider(resultsResp, parentStatement, session, httpClient);
   }
@@ -337,22 +334,6 @@ public ChunkProvider getChunkProvider() {
     return chunkProvider;
   }
 
-  private void setColumnInfo(TGetResultSetMetadataResp resultManifest)
-      throws DatabricksSQLException {
-    columnInfos = new ArrayList<>();
-    List<String> arrowMetadataList = DatabricksThriftUtil.getArrowMetadata(resultManifest);
-    if (resultManifest.getSchema() == null) {
-      return;
-    }
-    List<TColumnDesc> columns = resultManifest.getSchema().getColumns();
-    for (int columnIndex = 0; columnIndex < columns.size(); columnIndex++) {
-      TColumnDesc tColumnDesc = columns.get(columnIndex);
-      String columnArrowMetadata =
-          arrowMetadataList != null ? arrowMetadataList.get(columnIndex) : null;
-      columnInfos.add(getColumnInfoFromTColumnDesc(tColumnDesc, columnArrowMetadata));
-    }
-  }
-
   /**
    * Helper method to handle complex type and geospatial type conversion when support is disabled.
    *
@@ -397,7 +378,6 @@ protected static Object getObjectWithComplexTypeHandling(
 
     if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) {
       LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING");
-
       Object result =
           chunkIterator.getColumnObjectAtCurrentRow(
               columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfo);