[PECOBLR-4216] Cache inline Arrow compression metadata

NEXT_CHANGELOG.md

@@ -12,7 +12,7 @@
 ### Fixed
 - Invalid or incomplete Databricks JDBC URLs now fail with a descriptive `DatabricksSQLException`
   instead of leaking a `NullPointerException` when required connection parameters are missing.
-
+- Fixed inline Arrow result processing when later Thrift batches omit compression metadata.
 - Fixed connections failing when the same parameter is provided in both the JDBC URL and the connection properties, with the JDBC URL taking precedence.
 - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails.
 

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

@@ -14,6 +14,7 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 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.TSparkRowSetType;
 import com.databricks.jdbc.model.core.ResultData;
 import com.databricks.jdbc.model.core.ResultManifest;
@@ -94,7 +95,14 @@ private static IExecutionResult getResultHandler(
       IDatabricksStatementInternal parentStatement,
       IDatabricksSession session)
       throws SQLException {
-    TSparkRowSetType resultFormat = resultsResp.getResultSetMetadata().getResultFormat();
+    TGetResultSetMetadataResp metadata = resultsResp.getResultSetMetadata();
+    if (metadata == null) {
+      throw new DatabricksSQLException(
+          "Missing result set metadata",
+          DatabricksDriverErrorCode.INVALID_STATE.name(),
+          DatabricksDriverErrorCode.INVALID_STATE);
+    }
+    TSparkRowSetType resultFormat = metadata.getResultFormat();
     TelemetryHelper.setResultFormat(session.getConnectionContext(), parentStatement, resultFormat);
     LOGGER.info("Processing result of format {} from Thrift server", resultFormat);
     switch (resultFormat) {

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

@@ -8,11 +8,13 @@
 import com.databricks.jdbc.api.impl.IExecutionResult;
 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.TFetchResultsResp;
+import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp;
 import com.databricks.jdbc.model.core.ColumnInfo;
 import com.databricks.jdbc.model.core.ColumnInfoTypeName;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
@@ -39,6 +41,7 @@ public class LazyThriftInlineArrowResult implements IExecutionResult {
   private long totalRowsFetched;
   private List<ColumnInfo> columnInfos;
   private byte[] cachedSchema; // Cache schema from first response for subsequent batches
+  private final CompressionCodec cachedCompressionCodec;
 
   /**
    * Creates a new LazyThriftInlineArrowResult that lazily fetches arrow data on demand.
@@ -61,12 +64,19 @@ public LazyThriftInlineArrowResult(
     this.isClosed = false;
     this.totalRowsFetched = 0;
 
-    // Initialize column info from metadata
-    this.columnInfos = getColumnInfoList(initialResponse.getResultSetMetadata());
+    TGetResultSetMetadataResp metadata = initialResponse.getResultSetMetadata();
+    if (metadata == null) {
+      throw new DatabricksSQLException(
+          "Initial inline Arrow response is missing result set metadata",
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(),
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+    }
+    this.columnInfos = getColumnInfoList(metadata);
+    this.cachedCompressionCodec = CompressionCodec.getCompressionMapping(metadata);
 
     // Cache the schema from the first response for use in subsequent batches
     try {
-      this.cachedSchema = getSerializedSchema(initialResponse.getResultSetMetadata());
+      this.cachedSchema = getSerializedSchema(metadata);
     } catch (DatabricksParsingException e) {
       LOGGER.error("Failed to cache Arrow schema: {}", e.getMessage(), e);
       throw new DatabricksSQLException(
@@ -258,7 +268,7 @@ public List<String> getArrowMetadata() throws DatabricksSQLException {
   private void loadCurrentChunk() throws DatabricksSQLException {
     try {
       ByteArrayInputStream byteStream =
-          createArrowByteStream(cachedSchema, currentResponse, getClass());
+          createArrowByteStream(cachedSchema, currentResponse, cachedCompressionCodec, getClass());
       long rowCount = getTotalRowsInResponse(currentResponse);
 
       ArrowResultChunk.Builder builder =

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

@@ -14,6 +14,7 @@
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp;
+import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp;
 import com.databricks.jdbc.model.core.ColumnInfo;
 import com.databricks.jdbc.model.core.ColumnInfoTypeName;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
@@ -80,8 +81,14 @@ public StreamingInlineArrowResult(
     this.hasReachedEnd = false;
     this.isClosed = false;
 
-    // Initialize column info from metadata
-    this.columnInfos = getColumnInfoList(initialResponse.getResultSetMetadata());
+    TGetResultSetMetadataResp metadata = initialResponse.getResultSetMetadata();
+    if (metadata == null) {
+      throw new DatabricksSQLException(
+          "Initial inline Arrow response is missing result set metadata",
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(),
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+    }
+    this.columnInfos = getColumnInfoList(metadata);
 
     // Create batch fetcher and type-safe generic provider for Arrow
     ThriftBatchFetcher fetcher = new ThriftBatchFetcherImpl(session, statement);

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

@@ -5,12 +5,14 @@
 import static com.databricks.jdbc.common.util.ArrowUtil.getTotalRowsInResponse;
 
 import com.databricks.jdbc.api.impl.arrow.ArrowResultChunk;
+import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 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.TFetchResultsResp;
+import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import java.io.ByteArrayInputStream;
 import java.util.function.Consumer;
@@ -20,7 +22,8 @@
  *
  * <p>This processor converts {@link TFetchResultsResp} into {@link ArrowResultChunk} for inline
  * Arrow result handling. It caches the Arrow schema from the first response for use in subsequent
- * batches.
+ * batches. The compression codec is also cached because later Thrift responses may omit result set
+ * metadata.
  */
 public class InlineArrowResponseProcessor implements ThriftResponseProcessor<ArrowResultChunk> {
 
@@ -30,6 +33,7 @@ public class InlineArrowResponseProcessor implements ThriftResponseProcessor<Arr
   private final StatementId statementId;
   private volatile byte[]
       cachedSchema; // Cache schema from first response, volatile for visibility across threads
+  private volatile CompressionCodec cachedCompressionCodec;
 
   /**
    * Creates a new inline Arrow response processor.
@@ -44,9 +48,17 @@ public InlineArrowResponseProcessor(StatementId statementId) {
   public StreamingBatch<ArrowResultChunk> processInitialResponse(TFetchResultsResp response)
       throws DatabricksSQLException {
     LOGGER.debug("Processing initial inline Arrow response");
-    // Cache the schema for subsequent batches
+    TGetResultSetMetadataResp metadata = response.getResultSetMetadata();
+    if (metadata == null) {
+      throw new DatabricksSQLException(
+          "Initial inline Arrow response is missing result set metadata",
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(),
+          DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR);
+    }
+
     try {
-      this.cachedSchema = getSerializedSchema(response.getResultSetMetadata());
+      this.cachedSchema = getSerializedSchema(metadata);
+      this.cachedCompressionCodec = CompressionCodec.getCompressionMapping(metadata);
     } catch (DatabricksParsingException e) {
       LOGGER.error("Failed to serialize Arrow schema: {}", e.getMessage(), e);
       throw new DatabricksSQLException(
@@ -65,7 +77,8 @@ public StreamingBatch<ArrowResultChunk> processResponse(
         new StreamingBatch<>(batchIndex, rowOffset, getReleaseAction());
 
     try {
-      ByteArrayInputStream byteStream = createArrowByteStream(cachedSchema, response, getClass());
+      ByteArrayInputStream byteStream =
+          createArrowByteStream(cachedSchema, response, cachedCompressionCodec, getClass());
       long rowCount = getTotalRowsInResponse(response);
 
       ArrowResultChunk.Builder builder =

src/main/java/com/databricks/jdbc/common/CompressionCodec.java

@@ -31,6 +31,9 @@ public static CompressionCodec parseCompressionType(String compressionType) {
   }
 
   public static CompressionCodec getCompressionMapping(TGetResultSetMetadataResp metadataResp) {
+    if (metadataResp == null) {
+      return CompressionCodec.NONE;
+    }
     if (!metadataResp.isSetLz4Compressed()) {
       return CompressionCodec.NONE;
     }

src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java

@@ -151,10 +151,31 @@ public static Field columnDescToArrowField(TColumnDesc columnDesc) throws SQLExc
   public static ByteArrayInputStream createArrowByteStream(
       byte[] cachedSchema, TFetchResultsResp response, Class<?> callerClass)
       throws DatabricksParsingException {
+    return createArrowByteStream(
+        cachedSchema,
+        response,
+        CompressionCodec.getCompressionMapping(response.getResultSetMetadata()),
+        callerClass);
+  }
+
+  /**
+   * Creates an Arrow IPC byte stream using a previously resolved compression codec.
+   *
+   * @param cachedSchema The serialized Arrow schema bytes
+   * @param response The Thrift fetch response containing Arrow batches
+   * @param compressionCodec The compression codec resolved from the initial response
+   * @param callerClass The calling class for logging context
+   * @return ByteArrayInputStream containing the Arrow IPC data
+   * @throws DatabricksParsingException if processing fails
+   */
+  public static ByteArrayInputStream createArrowByteStream(
+      byte[] cachedSchema,
+      TFetchResultsResp response,
+      CompressionCodec compressionCodec,
+      Class<?> callerClass)
+      throws DatabricksParsingException {
     String context = callerClass.getSimpleName();
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
-    CompressionCodec compressionCodec =
-        CompressionCodec.getCompressionMapping(response.getResultSetMetadata());
 
     try {
       // Write schema if available

src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java

@@ -12,11 +12,13 @@
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
+import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException;
 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.core.ResultSchema;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import com.databricks.sdk.service.sql.Format;
 import java.io.ByteArrayOutputStream;
 import java.sql.SQLException;
@@ -127,6 +129,17 @@ public void testGetResultSet_thriftRow() {
         () -> ExecutionResultFactory.getResultSet(fetchResultsResp, session, parentStatement));
   }
 
+  @Test
+  public void testGetResultSet_thriftMissingMetadata() {
+    DatabricksSQLException thrown =
+        assertThrows(
+            DatabricksSQLException.class,
+            () -> ExecutionResultFactory.getResultSet(fetchResultsResp, session, parentStatement));
+
+    assertEquals(DatabricksDriverErrorCode.INVALID_STATE.name(), thrown.getSQLState());
+    assertEquals(DatabricksDriverErrorCode.INVALID_STATE.getCode(), thrown.getErrorCode());
+  }
+
   @Test
   public void testGetResultSet_thriftURL() throws SQLException {
     when(connectionContext.getConnectionUuid()).thenReturn("sample-uuid");

src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java

@@ -13,11 +13,13 @@
 import com.databricks.jdbc.model.client.thrift.generated.*;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
+import net.jpountz.lz4.LZ4FrameOutputStream;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -140,6 +142,39 @@ private TFetchResultsResp createFetchResultsResp(
     return response;
   }
 
+  private TFetchResultsResp createFetchResultsRespWithoutMetadata(
+      byte[] arrowData, int rowCount, boolean hasMoreRows) {
+    TSparkArrowBatch arrowBatch = new TSparkArrowBatch().setRowCount(rowCount).setBatch(arrowData);
+    TRowSet rowSet = new TRowSet().setArrowBatches(Collections.singletonList(arrowBatch));
+    TFetchResultsResp response = new TFetchResultsResp().setResults(rowSet);
+    response.hasMoreRows = hasMoreRows;
+    return response;
+  }
+
+  private static byte[] compressLz4(byte[] data) {
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    try (LZ4FrameOutputStream lz4 = new LZ4FrameOutputStream(out)) {
+      lz4.write(data);
+    } catch (IOException e) {
+      throw new RuntimeException("Failed to create compressed Arrow data", e);
+    }
+    return out.toByteArray();
+  }
+
+  @Test
+  void testMissingInitialMetadataThrowsTypedParsingError() {
+    TFetchResultsResp response = new TFetchResultsResp();
+
+    DatabricksSQLException thrown =
+        assertThrows(
+            DatabricksSQLException.class,
+            () -> new LazyThriftInlineArrowResult(response, statement, session));
+
+    assertEquals(DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(), thrown.getSQLState());
+    assertEquals(
+        DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.getCode(), thrown.getErrorCode());
+  }
+
   @Test
   void testEmptyResultSet() throws SQLException {
     byte[] arrowData = createValidArrowData(1, 0);
@@ -378,6 +413,31 @@ void testFetchNextChunkFromServer() throws SQLException {
     verify(databricksClient).getMoreResults(statement);
   }
 
+  @Test
+  void testCompressedNextChunkWithoutMetadataUsesInitialCompression() throws SQLException {
+    int rowsPerChunk = 2;
+    byte[] firstArrowData = compressLz4(createValidArrowData(1, rowsPerChunk));
+    byte[] secondArrowData = compressLz4(createValidArrowData(1, rowsPerChunk));
+    TFetchResultsResp initialResponse = createFetchResultsResp(firstArrowData, rowsPerChunk, true);
+    initialResponse.getResultSetMetadata().setLz4Compressed(true);
+    TFetchResultsResp secondResponse =
+        createFetchResultsRespWithoutMetadata(secondArrowData, rowsPerChunk, false);
+
+    when(statement.getStatementId()).thenReturn(STATEMENT_ID);
+    when(session.getDatabricksClient()).thenReturn(databricksClient);
+    when(databricksClient.getMoreResults(statement)).thenReturn(secondResponse);
+
+    LazyThriftInlineArrowResult result =
+        new LazyThriftInlineArrowResult(initialResponse, statement, session);
+
+    assertTrue(result.next());
+    assertTrue(result.next());
+    assertTrue(result.next());
+    assertTrue(result.next());
+    assertFalse(result.next());
+    assertEquals(rowsPerChunk * 2, result.getTotalRowsFetched());
+  }
+
   @Test
   void testGetRowCountReturnsCurrentChunkRowCount() throws SQLException {
     int rowCount = 5;

src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java

@@ -59,6 +59,20 @@ void setUp() throws Exception {
     lenient().when(statement.getStatementId()).thenReturn(STATEMENT_ID);
   }
 
+  @Test
+  void testMissingInitialMetadataThrowsTypedParsingError() {
+    TFetchResultsResp response = new TFetchResultsResp();
+
+    DatabricksSQLException thrown =
+        assertThrows(
+            DatabricksSQLException.class,
+            () -> new StreamingInlineArrowResult(response, statement, session));
+
+    assertEquals(DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(), thrown.getSQLState());
+    assertEquals(
+        DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.getCode(), thrown.getErrorCode());
+  }
+
   @Test
   void testBasicIteration() throws SQLException {
     int rowCount = 5;