Implement lazy loading for inline Arrow results (#1029)

## 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
#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 3e4f21c

Author: jayantsing-db

NEXT_CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Added
 
 ### 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.
 
 ### Fixed
 - Fixed complex data type metadata support when retrieving 0 rows in Arrow format

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

@@ -7,6 +7,7 @@
 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.converters.ConverterHelper;
 import com.databricks.jdbc.api.impl.converters.ObjectConverter;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
@@ -155,6 +156,8 @@ public DatabricksResultSet(
       List<String> arrowMetadata = null;
       if (executionResult instanceof ArrowStreamResult) {
         arrowMetadata = ((ArrowStreamResult) executionResult).getArrowMetadata();
+      } else if (executionResult instanceof LazyThriftInlineArrowResult) {
+        arrowMetadata = ((LazyThriftInlineArrowResult) executionResult).getArrowMetadata();
       }
       this.resultSetMetaData =
           new DatabricksResultSetMetaData(

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

@@ -1,6 +1,7 @@
 package com.databricks.jdbc.api.impl;
 
 import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
+import com.databricks.jdbc.api.impl.arrow.LazyThriftInlineArrowResult;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
 import com.databricks.jdbc.api.internal.IDatabricksSession;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
@@ -96,9 +97,9 @@ private static IExecutionResult getResultHandler(
       case COLUMN_BASED_SET:
         return new LazyThriftResult(resultsResp, parentStatement, session);
       case ARROW_BASED_SET:
-        return new ArrowStreamResult(resultsResp, true, parentStatement, session);
+        return new LazyThriftInlineArrowResult(resultsResp, parentStatement, session);
       case URL_BASED_SET:
-        return new ArrowStreamResult(resultsResp, false, parentStatement, session);
+        return new ArrowStreamResult(resultsResp, parentStatement, session);
       case ROW_BASED_SET:
         throw new DatabricksSQLFeatureNotSupportedException(
             "Invalid state - row based set cannot be received");

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

@@ -143,13 +143,11 @@ private static ChunkProvider createRemoteChunkProvider(
 
   public ArrowStreamResult(
       TFetchResultsResp resultsResp,
-      boolean isInlineArrow,
       IDatabricksStatementInternal parentStatementId,
       IDatabricksSession session)
       throws DatabricksSQLException {
     this(
         resultsResp,
-        isInlineArrow,
         parentStatementId,
         session,
         DatabricksHttpClientFactory.getInstance().getClient(session.getConnectionContext()));
@@ -158,19 +156,14 @@ public ArrowStreamResult(
   @VisibleForTesting
   ArrowStreamResult(
       TFetchResultsResp resultsResp,
-      boolean isInlineArrow,
       IDatabricksStatementInternal parentStatement,
       IDatabricksSession session,
       IDatabricksHttpClient httpClient)
       throws DatabricksSQLException {
     this.session = session;
     setColumnInfo(resultsResp.getResultSetMetadata());
-    if (isInlineArrow) {
-      this.chunkProvider = new InlineChunkProvider(resultsResp, parentStatement, session);
-    } else {
-      this.chunkProvider =
-          createThriftRemoteChunkProvider(resultsResp, parentStatement, session, httpClient);
-    }
+    this.chunkProvider =
+        createThriftRemoteChunkProvider(resultsResp, parentStatement, session, httpClient);
   }
 
   /**
@@ -238,48 +231,15 @@ public List<String> getArrowMetadata() throws DatabricksSQLException {
   /** {@inheritDoc} */
   @Override
   public Object getObject(int columnIndex) throws DatabricksSQLException {
-    ColumnInfoTypeName requiredType = columnInfos.get(columnIndex).getTypeName();
+    ColumnInfo columnInfo = columnInfos.get(columnIndex);
+    ColumnInfoTypeName requiredType = columnInfo.getTypeName();
     String arrowMetadata = chunkIterator.getType(columnIndex);
     if (arrowMetadata == null) {
-      arrowMetadata = columnInfos.get(columnIndex).getTypeText();
+      arrowMetadata = columnInfo.getTypeText();
     }
 
-    // Handle complex type conversion when complex datatype support is disabled
-    boolean isComplexDatatypeSupportEnabled =
-        this.session.getConnectionContext().isComplexDatatypeSupportEnabled();
-    boolean isGeoSpatialSupportEnabled =
-        this.session.getConnectionContext().isGeoSpatialSupportEnabled();
-
-    // Check if we need to convert geospatial types to string when geospatial support is disabled
-    // This check must come before the general complex type check
-    if (!isGeoSpatialSupportEnabled && isGeospatialType(requiredType)) {
-      LOGGER.debug("Geospatial support is disabled, converting {} to STRING", requiredType);
-
-      Object result =
-          chunkIterator.getColumnObjectAtCurrentRow(
-              columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfos.get(columnIndex));
-      if (result == null) {
-        return null;
-      }
-      // Return raw string for geospatial types when support is disabled
-      return result;
-    }
-
-    if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) {
-      LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING");
-
-      Object result =
-          chunkIterator.getColumnObjectAtCurrentRow(
-              columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfos.get(columnIndex));
-      if (result == null) {
-        return null;
-      }
-      ComplexDataTypeParser parser = new ComplexDataTypeParser();
-      return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata);
-    }
-
-    return chunkIterator.getColumnObjectAtCurrentRow(
-        columnIndex, requiredType, arrowMetadata, columnInfos.get(columnIndex));
+    return getObjectWithComplexTypeHandling(
+        session, chunkIterator, columnIndex, requiredType, arrowMetadata, columnInfo);
   }
 
   /**
@@ -384,6 +344,66 @@ private void setColumnInfo(TGetResultSetMetadataResp resultManifest) {
     }
   }
 
+  /**
+   * Helper method to handle complex type and geospatial type conversion when support is disabled.
+   *
+   * <p>This method is also used by LazyThriftInlineArrowResult for consistent type handling.
+   *
+   * @param session The databricks session
+   * @param chunkIterator The chunk iterator
+   * @param columnIndex The column index
+   * @param requiredType The required column type
+   * @param arrowMetadata The arrow metadata
+   * @param columnInfo The column info
+   * @return The object value (converted if complex/geospatial type and support disabled)
+   * @throws DatabricksSQLException if an error occurs
+   */
+  protected static Object getObjectWithComplexTypeHandling(
+      IDatabricksSession session,
+      ArrowResultChunkIterator chunkIterator,
+      int columnIndex,
+      ColumnInfoTypeName requiredType,
+      String arrowMetadata,
+      ColumnInfo columnInfo)
+      throws DatabricksSQLException {
+    boolean isComplexDatatypeSupportEnabled =
+        session.getConnectionContext().isComplexDatatypeSupportEnabled();
+    boolean isGeoSpatialSupportEnabled =
+        session.getConnectionContext().isGeoSpatialSupportEnabled();
+
+    // Check if we need to convert geospatial types to string when geospatial support is disabled
+    // This check must come before the general complex type check
+    if (!isGeoSpatialSupportEnabled && isGeospatialType(requiredType)) {
+      LOGGER.debug("Geospatial support is disabled, converting {} to STRING", requiredType);
+
+      Object result =
+          chunkIterator.getColumnObjectAtCurrentRow(
+              columnIndex, ColumnInfoTypeName.STRING, "STRING", columnInfo);
+      if (result == null) {
+        return null;
+      }
+      // Return raw string for geospatial types when support is disabled
+      return result;
+    }
+
+    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);
+      if (result == null) {
+        return null;
+      }
+      ComplexDataTypeParser parser = new ComplexDataTypeParser();
+
+      return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata);
+    }
+
+    return chunkIterator.getColumnObjectAtCurrentRow(
+        columnIndex, requiredType, arrowMetadata, columnInfo);
+  }
+
   /**
    * Converts a collection of ExternalLinks to a ChunkLinkFetchResult.
    *

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

@@ -1,31 +1,17 @@
 package com.databricks.jdbc.api.impl.arrow;
 
-import static com.databricks.jdbc.common.util.DatabricksTypeUtil.*;
 import static com.databricks.jdbc.common.util.DecompressionUtil.decompress;
 
-import com.databricks.jdbc.api.internal.IDatabricksSession;
-import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 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.*;
 import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.model.core.ResultManifest;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.google.common.annotations.VisibleForTesting;
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.List;
-import org.apache.arrow.vector.types.pojo.ArrowType;
-import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.Schema;
-import org.apache.arrow.vector.util.SchemaUtility;
 
 /** Class to manage inline Arrow chunks */
 public class InlineChunkProvider implements ChunkProvider {
@@ -37,23 +23,6 @@ public class InlineChunkProvider implements ChunkProvider {
   private final ArrowResultChunk
       arrowResultChunk; // There is only one packet of data in case of inline arrow
 
-  InlineChunkProvider(
-      TFetchResultsResp resultsResp,
-      IDatabricksStatementInternal parentStatement,
-      IDatabricksSession session)
-      throws DatabricksParsingException {
-    this.currentChunkIndex = -1;
-    this.totalRows = 0;
-    ByteArrayInputStream byteStream = initializeByteStream(resultsResp, session, parentStatement);
-    ArrowResultChunk.Builder builder =
-        ArrowResultChunk.builder().withInputStream(byteStream, totalRows);
-
-    if (parentStatement != null) {
-      builder.withStatementId(parentStatement.getStatementId());
-    }
-    arrowResultChunk = builder.build();
-  }
-
   /**
    * Constructor for inline arrow chunk provider from {@link ResultData} and {@link ResultManifest}.
    *
@@ -123,97 +92,6 @@ public boolean isClosed() {
     return isClosed;
   }
 
-  private ByteArrayInputStream initializeByteStream(
-      TFetchResultsResp resultsResp,
-      IDatabricksSession session,
-      IDatabricksStatementInternal parentStatement)
-      throws DatabricksParsingException {
-    ByteArrayOutputStream baos = new ByteArrayOutputStream();
-    CompressionCodec compressionType =
-        CompressionCodec.getCompressionMapping(resultsResp.getResultSetMetadata());
-    try {
-      byte[] serializedSchema = getSerializedSchema(resultsResp.getResultSetMetadata());
-      if (serializedSchema != null) {
-        baos.write(serializedSchema);
-      }
-      writeToByteOutputStream(
-          compressionType, parentStatement, resultsResp.getResults().getArrowBatches(), baos);
-      while (resultsResp.hasMoreRows) {
-        resultsResp = session.getDatabricksClient().getMoreResults(parentStatement);
-        writeToByteOutputStream(
-            compressionType, parentStatement, resultsResp.getResults().getArrowBatches(), baos);
-      }
-      return new ByteArrayInputStream(baos.toByteArray());
-    } catch (DatabricksSQLException | IOException e) {
-      handleError(e);
-    }
-    return null;
-  }
-
-  private void writeToByteOutputStream(
-      CompressionCodec compressionCodec,
-      IDatabricksStatementInternal parentStatement,
-      List<TSparkArrowBatch> arrowBatchList,
-      ByteArrayOutputStream baos)
-      throws DatabricksSQLException, IOException {
-    for (TSparkArrowBatch arrowBatch : arrowBatchList) {
-      byte[] decompressedBytes =
-          decompress(
-              arrowBatch.getBatch(),
-              compressionCodec,
-              String.format(
-                  "Data fetch for inline arrow batch [%d] and statement [%s] with decompression algorithm : [%s]",
-                  arrowBatch.getRowCount(), parentStatement, compressionCodec));
-      totalRows += arrowBatch.getRowCount();
-      baos.write(decompressedBytes);
-    }
-  }
-
-  private byte[] getSerializedSchema(TGetResultSetMetadataResp metadata)
-      throws DatabricksSQLException {
-    if (metadata.getArrowSchema() != null) {
-      return metadata.getArrowSchema();
-    }
-    Schema arrowSchema = hiveSchemaToArrowSchema(metadata.getSchema());
-    try {
-      return SchemaUtility.serialize(arrowSchema);
-    } catch (IOException e) {
-      handleError(e);
-    }
-    // should never reach here;
-    return null;
-  }
-
-  private Schema hiveSchemaToArrowSchema(TTableSchema hiveSchema)
-      throws DatabricksParsingException {
-    List<Field> fields = new ArrayList<>();
-    if (hiveSchema == null) {
-      return new Schema(fields);
-    }
-    try {
-      hiveSchema
-          .getColumns()
-          .forEach(
-              columnDesc -> {
-                try {
-                  fields.add(getArrowField(columnDesc));
-                } catch (SQLException e) {
-                  throw new RuntimeException(e);
-                }
-              });
-    } catch (RuntimeException e) {
-      handleError(e);
-    }
-    return new Schema(fields);
-  }
-
-  private Field getArrowField(TColumnDesc columnDesc) throws SQLException {
-    TPrimitiveTypeEntry primitiveTypeEntry = getTPrimitiveTypeOrDefault(columnDesc.getTypeDesc());
-    ArrowType arrowType = mapThriftToArrowType(primitiveTypeEntry.getType());
-    FieldType fieldType = new FieldType(true, arrowType, null);
-    return new Field(columnDesc.getColumnName(), fieldType, null);
-  }
-
   @VisibleForTesting
   void handleError(Exception e) throws DatabricksParsingException {
     String errorMessage =