Add ChunkProvider interface and unify chunk management (#554)

- Added a new ChunkProvider interface to abstract chunk retrieval methods:
  - boolean hasNextChunk();
  - boolean next();
  - ArrowResultChunk getChunk() throws DatabricksSQLException;
  - void close();
- Implemented ChunkProvider in both ChunkDownloader and ChunkExtractor.
- Updated ArrowStreamResult to use ChunkProvider for chunk management.

Author: jayantsing-db

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

@@ -23,10 +23,8 @@
 
 public class ArrowStreamResult implements IExecutionResult {
 
-  private ChunkDownloader chunkDownloader;
-  private ChunkExtractor chunkExtractor;
+  private final ChunkProvider chunkProvider;
   private long currentRowIndex = -1;
-  private boolean isInlineArrow;
   private boolean isClosed;
   private ArrowResultChunk.ArrowResultChunkIterator chunkIterator;
   private List<ColumnInfo> columnInfos;
@@ -53,8 +51,8 @@ public ArrowStreamResult(
       IDatabricksSession session,
       IDatabricksHttpClient httpClient)
       throws DatabricksParsingException {
-    this.chunkDownloader =
-        new ChunkDownloader(
+    this.chunkProvider =
+        new RemoteChunkProvider(
             statementId,
             resultManifest,
             resultData,
@@ -93,22 +91,19 @@ public ArrowStreamResult(
       IDatabricksHttpClient httpClient)
       throws DatabricksParsingException {
     setColumnInfo(resultManifest);
-    this.isInlineArrow = isInlineArrow;
     if (isInlineArrow) {
-      this.chunkExtractor =
-          new ChunkExtractor(resultData.getArrowBatches(), resultManifest, statementId);
-      this.chunkDownloader = null;
+      this.chunkProvider =
+          new InlineChunkProvider(resultData.getArrowBatches(), resultManifest, statementId);
     } else {
       CompressionType compressionType = CompressionType.getCompressionMapping(resultManifest);
-      this.chunkDownloader =
-          new ChunkDownloader(
+      this.chunkProvider =
+          new RemoteChunkProvider(
               statementId,
               resultData,
               session,
               httpClient,
               session.getConnectionContext().getCloudFetchThreadPoolSize(),
               compressionType);
-      this.chunkExtractor = null;
     }
   }
 
@@ -132,20 +127,13 @@ public boolean next() throws DatabricksSQLException {
     if (!hasNext()) {
       return false;
     }
+
     currentRowIndex++;
-    if (isInlineArrow) {
-      if (chunkIterator == null) {
-        chunkIterator = chunkExtractor.next().getChunkIterator();
-      }
-      return chunkIterator.nextRow();
-    }
-    // Either this is first chunk or we are crossing chunk boundary
     if (chunkIterator == null || !chunkIterator.hasNextRow()) {
-      chunkDownloader.next();
-      chunkIterator = chunkDownloader.getChunk().getChunkIterator();
-      return chunkIterator.nextRow();
+      chunkProvider.next();
+      chunkIterator = chunkProvider.getChunk().getChunkIterator();
     }
-    // Traversing within a chunk
+
     return chunkIterator.nextRow();
   }
 
@@ -163,18 +151,14 @@ public boolean hasNext() {
 
     // For inline arrow, check if the chunk extractor has more chunks
     // Otherwise, check the chunk downloader
-    return isInlineArrow ? chunkExtractor.hasNext() : chunkDownloader.hasNextChunk();
+    return chunkProvider.hasNextChunk();
   }
 
   /** {@inheritDoc} */
   @Override
   public void close() {
     isClosed = true;
-    if (isInlineArrow) {
-      chunkExtractor.releaseChunk();
-    } else {
-      chunkDownloader.releaseAllChunks();
-    }
+    chunkProvider.close();
   }
 
   private void setColumnInfo(TGetResultSetMetadataResp resultManifest) {

…pi/impl/arrow/SingleChunkDownloader.java …bc/api/impl/arrow/ChunkDownloadTask.javasrc/main/java/com/databricks/jdbc/api/impl/arrow/SingleChunkDownloader.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java src/main/java/com/databricks/jdbc/api/impl/arrow/SingleChunkDownloader.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloadTask.java

@@ -9,16 +9,16 @@
 import java.util.concurrent.Callable;
 
 /** Task class to manage download for a single chunk. */
-class SingleChunkDownloader implements Callable<Void> {
+class ChunkDownloadTask implements Callable<Void> {
 
-  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(SingleChunkDownloader.class);
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ChunkDownloadTask.class);
   public static final int MAX_RETRIES = 5;
   private static final long RETRY_DELAY_MS = 1500; // 1.5 seconds
   private final ArrowResultChunk chunk;
   private final IDatabricksHttpClient httpClient;
   private final ChunkDownloadCallback chunkDownloader;
 
-  SingleChunkDownloader(
+  ChunkDownloadTask(
       ArrowResultChunk chunk,
       IDatabricksHttpClient httpClient,
       ChunkDownloadCallback chunkDownloader) {

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

@@ -0,0 +1,40 @@
+package com.databricks.jdbc.api.impl.arrow;
+
+import com.databricks.jdbc.exception.DatabricksSQLException;
+
+/**
+ * Implementations of this interface manage the retrieval and iteration over {@link
+ * ArrowResultChunk}s.
+ */
+public interface ChunkProvider {
+
+  /**
+   * Checks if there are more chunks available to iterate over.
+   *
+   * @return {@code true} if there are additional chunks to be retrieved; {@code false} otherwise.
+   */
+  boolean hasNextChunk();
+
+  /**
+   * Advances to the next available chunk. This method should be called before calling {@link
+   * #getChunk()} to retrieve the data from the next chunk.
+   *
+   * @return {@code true} if the next chunk was successfully moved to; {@code false} if there are no
+   *     more chunks.
+   */
+  boolean next();
+
+  /**
+   * Retrieves the current chunk of data after a successful call to {@link #next()}.
+   *
+   * @return The current {@link ArrowResultChunk} containing the data.
+   * @throws DatabricksSQLException if an error occurs while fetching the chunk.
+   */
+  ArrowResultChunk getChunk() throws DatabricksSQLException;
+
+  /**
+   * Closes the chunk provider and releases any resources associated with it. After calling this
+   * method, the chunk provider should not be used again.
+   */
+  void close();
+}

…/jdbc/api/impl/arrow/ChunkExtractor.java …/api/impl/arrow/InlineChunkProvider.javasrc/main/java/com/databricks/jdbc/api/impl/arrow/ChunkExtractor.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkExtractor.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/InlineChunkProvider.java

@@ -23,15 +23,15 @@
 import org.apache.arrow.vector.util.SchemaUtility;
 
 /** Class to manage inline Arrow chunks */
-public class ChunkExtractor {
+public class InlineChunkProvider implements ChunkProvider {
 
-  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ChunkExtractor.class);
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(InlineChunkProvider.class);
   private long totalRows;
   private long currentChunkIndex;
 
   ArrowResultChunk arrowResultChunk; // There is only one packet of data in case of inline arrow
 
-  ChunkExtractor(
+  InlineChunkProvider(
       List<TSparkArrowBatch> arrowBatches, TGetResultSetMetadataResp metadata, String statementId)
       throws DatabricksParsingException {
     this.currentChunkIndex = -1;
@@ -40,18 +40,34 @@ public class ChunkExtractor {
     arrowResultChunk = ArrowResultChunk.builder().withInputStream(byteStream, totalRows).build();
   }
 
-  public boolean hasNext() {
+  /** {@inheritDoc} */
+  @Override
+  public boolean hasNextChunk() {
     return this.currentChunkIndex == -1;
   }
 
-  public ArrowResultChunk next() {
-    if (this.currentChunkIndex != -1) {
-      return null;
+  /** {@inheritDoc} */
+  @Override
+  public boolean next() {
+    if (!hasNextChunk()) {
+      return false;
     }
     this.currentChunkIndex++;
+    return true;
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public ArrowResultChunk getChunk() {
     return arrowResultChunk;
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public void close() {
+    arrowResultChunk.releaseChunk();
+  }
+
   private ByteArrayInputStream initializeByteStream(
       List<TSparkArrowBatch> arrowBatches, TGetResultSetMetadataResp metadata, String statementId)
       throws DatabricksParsingException {
@@ -120,8 +136,7 @@ private static Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema)
 
   private static Field getArrowField(TColumnDesc columnDesc) throws SQLException {
     TTypeId thriftType = getThriftTypeFromTypeDesc(columnDesc.getTypeDesc());
-    ArrowType arrowType = null;
-    arrowType = mapThriftToArrowType(thriftType);
+    ArrowType arrowType = mapThriftToArrowType(thriftType);
     FieldType fieldType = new FieldType(true, arrowType, null);
     return new Field(columnDesc.getColumnName(), fieldType, null);
   }
@@ -132,8 +147,4 @@ static void handleError(Exception e) throws DatabricksParsingException {
     LOGGER.error(errorMessage);
     throw new DatabricksParsingException(errorMessage, e);
   }
-
-  public void releaseChunk() {
-    this.arrowResultChunk.releaseChunk();
-  }
 }

…jdbc/api/impl/arrow/ChunkDownloader.java …/api/impl/arrow/RemoteChunkProvider.javasrc/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloader.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java src/main/java/com/databricks/jdbc/api/impl/arrow/ChunkDownloader.java renamed to src/main/java/com/databricks/jdbc/api/impl/arrow/RemoteChunkProvider.java

@@ -23,17 +23,17 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 /** Class to manage Arrow chunks and fetch them on proactive basis. */
-public class ChunkDownloader implements ChunkDownloadCallback {
+public class RemoteChunkProvider implements ChunkProvider, ChunkDownloadCallback {
 
-  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ChunkDownloader.class);
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(RemoteChunkProvider.class);
   private static final String CHUNKS_DOWNLOADER_THREAD_POOL_PREFIX =
       "databricks-jdbc-chunks-downloader-";
+  private static int chunksDownloaderThreadPoolSize;
   private final IDatabricksSession session;
   private final String statementId;
   private final long totalChunks;
   private final ExecutorService chunkDownloaderExecutorService;
   private final IDatabricksHttpClient httpClient;
-  private static int chunksDownloaderThreadPoolSize;
   private Long currentChunkIndex;
   private long nextChunkToDownload;
   private Long totalChunksInMemory;
@@ -42,15 +42,15 @@ public class ChunkDownloader implements ChunkDownloadCallback {
   private final CompressionType compressionType;
   private final ConcurrentHashMap<Long, ArrowResultChunk> chunkIndexToChunksMap;
 
-  ChunkDownloader(
+  RemoteChunkProvider(
       String statementId,
       ResultManifest resultManifest,
       ResultData resultData,
       IDatabricksSession session,
       IDatabricksHttpClient httpClient,
       int chunksDownloaderThreadPoolSize)
       throws DatabricksParsingException {
-    ChunkDownloader.chunksDownloaderThreadPoolSize = chunksDownloaderThreadPoolSize;
+    RemoteChunkProvider.chunksDownloaderThreadPoolSize = chunksDownloaderThreadPoolSize;
     this.chunkDownloaderExecutorService = createChunksDownloaderExecutorService();
     this.httpClient = httpClient;
     this.session = session;
@@ -61,7 +61,7 @@ public class ChunkDownloader implements ChunkDownloadCallback {
     initializeData();
   }
 
-  ChunkDownloader(
+  RemoteChunkProvider(
       String statementId,
       TRowSet resultData,
       IDatabricksSession session,
@@ -78,15 +78,15 @@ public class ChunkDownloader implements ChunkDownloadCallback {
   }
 
   @VisibleForTesting
-  ChunkDownloader(
+  RemoteChunkProvider(
       String statementId,
       TRowSet resultData,
       IDatabricksSession session,
       IDatabricksHttpClient httpClient,
       int chunksDownloaderThreadPoolSize,
       CompressionType compressionType)
       throws DatabricksParsingException {
-    ChunkDownloader.chunksDownloaderThreadPoolSize = chunksDownloaderThreadPoolSize;
+    RemoteChunkProvider.chunksDownloaderThreadPoolSize = chunksDownloaderThreadPoolSize;
     this.chunkDownloaderExecutorService = createChunksDownloaderExecutorService();
     this.httpClient = httpClient;
     this.compressionType = compressionType;
@@ -117,12 +117,15 @@ public void downloadLinks(long chunkIndexToDownloadLink) throws DatabricksSQLExc
   }
 
   /**
-   * Fetches the chunk for the given index. If chunk is not already downloaded, will download the
+   * {@inheritDoc}
+   *
+   * <p>Fetches the chunk for the given index. If chunk is not already downloaded, will download the
    * chunk first
    *
    * @return the chunk at given index
    */
-  ArrowResultChunk getChunk() throws DatabricksSQLException {
+  @Override
+  public ArrowResultChunk getChunk() throws DatabricksSQLException {
     if (currentChunkIndex < 0) {
       return null;
     }
@@ -153,11 +156,15 @@ public CompressionType getCompressionType() {
     return compressionType;
   }
 
-  boolean hasNextChunk() {
+  /** {@inheritDoc} */
+  @Override
+  public boolean hasNextChunk() {
     return currentChunkIndex < totalChunks - 1;
   }
 
-  boolean next() {
+  /** {@inheritDoc} */
+  @Override
+  public boolean next() {
     if (currentChunkIndex >= 0) {
       // release current chunk
       releaseChunk();
@@ -170,6 +177,19 @@ boolean next() {
     return true;
   }
 
+  /**
+   * {@inheritDoc}
+   *
+   * <p>Release all chunks from memory. This would be called when result-set has been closed.
+   */
+  @Override
+  public void close() {
+    this.isClosed = true;
+    this.chunkDownloaderExecutorService.shutdownNow();
+    this.chunkIndexToChunksMap.values().forEach(ArrowResultChunk::releaseChunk);
+    httpClient.closeExpiredAndIdleConnections();
+  }
+
   /** Release the memory for previous chunk since it is already consumed */
   void releaseChunk() {
     if (chunkIndexToChunksMap.get(currentChunkIndex).releaseChunk()) {
@@ -189,22 +209,13 @@ void setChunkLink(ExternalLink chunkLink) {
     }
   }
 
-  /** Release all chunks from memory. This would be called when result-set has been closed. */
-  void releaseAllChunks() {
-    this.isClosed = true;
-    this.chunkDownloaderExecutorService.shutdownNow();
-    this.chunkIndexToChunksMap.values().forEach(ArrowResultChunk::releaseChunk);
-    httpClient.closeExpiredAndIdleConnections();
-  }
-
   void downloadNextChunks() {
     while (!this.isClosed
         && nextChunkToDownload < totalChunks
         && totalChunksInMemory < allowedChunksInMemory) {
       ArrowResultChunk chunk = chunkIndexToChunksMap.get(nextChunkToDownload);
       if (chunk.getStatus() != ArrowResultChunk.ChunkStatus.DOWNLOAD_SUCCEEDED) {
-        this.chunkDownloaderExecutorService.submit(
-            new SingleChunkDownloader(chunk, httpClient, this));
+        this.chunkDownloaderExecutorService.submit(new ChunkDownloadTask(chunk, httpClient, this));
         totalChunksInMemory++;
       }
       nextChunkToDownload++;