fix: Address PR review comments for streaming prefetch feature

- 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>

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).
 
 ### Updated
 - Implemented lazy loading for inline Arrow results, fetching arrow batches on demand instead of all at once. This improves memory usage and initial response time for large result sets when using the Thrift protocol with Arrow format.

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

@@ -26,6 +26,9 @@ class ExecutionResultFactory {
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(ExecutionResultFactory.class);
 
+  /** Default timeout in seconds for waiting for a batch to be ready in streaming mode. */
+  private static final int DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS = 300;
+
   static IExecutionResult getResultSet(
       ResultData data,
       ResultManifest manifest,
@@ -129,7 +132,11 @@ private static IExecutionResult createThriftColumnarResult(
       LOGGER.info("Using StreamingColumnarResult for improved throughput (default)");
       int maxBatchesInMemory = connectionContext.getThriftMaxBatchesInMemory();
       return new StreamingColumnarResult(
-          resultsResp, parentStatement, session, maxBatchesInMemory, 300); // 5 minute timeout
+          resultsResp,
+          parentStatement,
+          session,
+          maxBatchesInMemory,
+          DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS);
     } else {
       LOGGER.info("Using LazyThriftResult (streaming explicitly disabled)");
       return new LazyThriftResult(resultsResp, parentStatement, session);
@@ -153,7 +160,11 @@ private static IExecutionResult createInlineArrowResult(
       LOGGER.info("Using StreamingInlineArrowResult for improved throughput (default)");
       int maxBatchesInMemory = connectionContext.getThriftMaxBatchesInMemory();
       return new StreamingInlineArrowResult(
-          resultsResp, parentStatement, session, maxBatchesInMemory, 300); // 5 minute timeout
+          resultsResp,
+          parentStatement,
+          session,
+          maxBatchesInMemory,
+          DEFAULT_STREAMING_BATCH_TIMEOUT_SECONDS);
     } else {
       LOGGER.info("Using LazyThriftInlineArrowResult (streaming explicitly disabled)");
       return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);

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

@@ -167,18 +167,23 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
   /** Validates state before getting an object. */
   private void validateGetObjectState(int columnIndex) throws DatabricksSQLException {
     if (isClosed) {
+      LOGGER.debug("[STREAMING] Attempted to access closed result");
       throw new DatabricksSQLException(
           "Result is closed", DatabricksDriverErrorCode.STATEMENT_CLOSED);
     }
     if (globalRowIndex == -1) {
+      LOGGER.debug("[STREAMING] Attempted to access data before first row");
       throw new DatabricksSQLException(
           "Cursor is before first row", DatabricksDriverErrorCode.INVALID_STATE);
     }
     if (currentChunkIterator == null) {
+      LOGGER.debug("[STREAMING] No current chunk available at row {}", globalRowIndex);
       throw new DatabricksSQLException(
           "No current chunk available", DatabricksDriverErrorCode.INVALID_STATE);
     }
     if (columnIndex < 0 || columnIndex >= columnInfos.size()) {
+      LOGGER.debug(
+          "[STREAMING] Column index {} out of bounds (0-{})", columnIndex, columnInfos.size() - 1);
       throw new DatabricksSQLException(
           "Column index out of bounds: " + columnIndex, DatabricksDriverErrorCode.INVALID_STATE);
     }

src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java

@@ -101,6 +101,7 @@ public StreamingBatch<ArrowResultChunk> processResponse(
       return batch;
 
     } catch (DatabricksParsingException e) {
+      LOGGER.error("Failed to process inline Arrow batch {}: {}", batchIndex, e.getMessage(), e);
       batch.setError(e);
       throw new DatabricksSQLException(
           "Failed to process Arrow data", e, DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
@@ -141,6 +142,7 @@ private ByteArrayInputStream createArrowByteStream(TFetchResultsResp response)
 
       return new ByteArrayInputStream(baos.toByteArray());
     } catch (DatabricksSQLException | IOException e) {
+      LOGGER.error("Failed to create Arrow byte stream: {}", e.getMessage(), e);
       throw new DatabricksParsingException(
           "Failed to create Arrow byte stream: " + e.getMessage(),
           e,
@@ -204,13 +206,17 @@ private Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema)
       throws DatabricksParsingException {
     List<Field> fields = new ArrayList<>();
     if (hiveSchema == null) {
+      LOGGER.debug("Hive schema is null, returning empty Arrow schema");
       return new Schema(fields);
     }
     try {
+      LOGGER.debug(
+          "Converting Hive schema to Arrow schema with {} columns", hiveSchema.getColumnsSize());
       for (TColumnDesc columnDesc : hiveSchema.getColumns()) {
         fields.add(getArrowField(columnDesc));
       }
     } catch (SQLException e) {
+      LOGGER.error("Failed to convert Hive schema to Arrow: {}", e.getMessage(), e);
       throw new DatabricksParsingException(
           "Failed to convert Hive schema to Arrow: " + e.getMessage(),
           e,

src/main/java/com/databricks/jdbc/api/impl/streaming/ThriftStreamingProvider.java

@@ -174,6 +174,17 @@ private ThriftStreamingProvider(
       int timeoutSeconds)
       throws DatabricksSQLException {
 
+    // Validate required parameters
+    if (initialResponse == null) {
+      throw new IllegalArgumentException("initialResponse cannot be null");
+    }
+    if (fetcher == null) {
+      throw new IllegalArgumentException("fetcher cannot be null");
+    }
+    if (processor == null) {
+      throw new IllegalArgumentException("processor cannot be null");
+    }
+
     this.batchFetcher = fetcher;
     this.processor = processor;
     this.maxBatchesInMemory = Math.max(2, maxBatchesInMemory);
@@ -355,13 +366,22 @@ public void close() {
     }
 
     // Release all batches using type-safe release action
+    // Continue releasing subsequent batches even if one fails
     for (StreamingBatch<T> batch : batches.values()) {
-      batch.release();
+      try {
+        batch.release();
+      } catch (Exception e) {
+        LOGGER.warn("Error releasing batch during close: {}", e.getMessage(), e);
+      }
     }
     batches.clear();
 
     if (batchFetcher != null) {
-      batchFetcher.close();
+      try {
+        batchFetcher.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error closing batchFetcher: {}", e.getMessage(), e);
+      }
     }
   }
 
@@ -459,6 +479,9 @@ private void releaseBatch(long batchIndex) {
   private void waitForBatchCreation(long batchIndex) throws DatabricksSQLException {
     prefetchLock.lock();
     try {
+      long waitStartTime = System.currentTimeMillis();
+      long timeoutMillis = batchReadyTimeoutSeconds * 1000L;
+
       while (!closed && !batches.containsKey(batchIndex)) {
         checkPrefetchError();
         if (endOfStreamReached && batchIndex > highestFetchedBatchIndex.get()) {
@@ -470,8 +493,22 @@ private void waitForBatchCreation(long batchIndex) throws DatabricksSQLException
                   + ")",
               DatabricksDriverErrorCode.CHUNK_READY_ERROR);
         }
+
+        // Check for timeout
+        long elapsedMillis = System.currentTimeMillis() - waitStartTime;
+        if (elapsedMillis >= timeoutMillis) {
+          throw new DatabricksSQLException(
+              "Timeout waiting for batch "
+                  + batchIndex
+                  + " to be created (timeout: "
+                  + batchReadyTimeoutSeconds
+                  + "s)",
+              DatabricksDriverErrorCode.CHUNK_READY_ERROR);
+        }
+
         try {
-          batchAvailable.await();
+          long remainingMillis = timeoutMillis - elapsedMillis;
+          batchAvailable.await(remainingMillis, java.util.concurrent.TimeUnit.MILLISECONDS);
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
           throw new DatabricksSQLException(