diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index a8173af770..e3b48abf23 100644 --- a/NEXT_CHANGELOG.md +++ b/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. diff --git a/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java b/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java index b9aec3fc1b..56122ed8ae 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ExecutionResultFactory.java +++ b/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) { diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java index 2e63b32892..421dbaf38f 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResult.java +++ b/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 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 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 = diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java index 5f2812284d..347dac8af5 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResult.java +++ b/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); diff --git a/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java b/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java index 477359f2bc..3be56309b4 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessor.java +++ b/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 @@ * *

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 { @@ -30,6 +33,7 @@ public class InlineArrowResponseProcessor implements ThriftResponseProcessor 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 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 = diff --git a/src/main/java/com/databricks/jdbc/common/CompressionCodec.java b/src/main/java/com/databricks/jdbc/common/CompressionCodec.java index 67e7840dff..444bad653b 100644 --- a/src/main/java/com/databricks/jdbc/common/CompressionCodec.java +++ b/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; } diff --git a/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java b/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java index 28f37ae7ea..ce03640b1e 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ArrowUtil.java +++ b/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 diff --git a/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java b/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java index d3a03a26da..d088c94770 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/ExecutionResultFactoryTest.java +++ b/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"); diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java index eb10b7cd1a..e595cbdbc4 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/LazyThriftInlineArrowResultTest.java +++ b/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; diff --git a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java b/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java index 448a02de60..a016110015 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/arrow/StreamingInlineArrowResultTest.java +++ b/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; diff --git a/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java b/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java index f46bb83c62..ecdb031a18 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/streaming/InlineArrowResponseProcessorTest.java @@ -6,10 +6,13 @@ import com.databricks.jdbc.dbclient.impl.common.StatementId; import com.databricks.jdbc.exception.DatabricksSQLException; 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.util.ArrayList; import java.util.Collections; import java.util.List; +import net.jpountz.lz4.LZ4FrameOutputStream; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; @@ -53,7 +56,7 @@ void testProcessSubsequentResponseUsesSchemaCache() throws DatabricksSQLExceptio // Now process a subsequent response (simulating second batch) // The schema should be cached from the initial response - TFetchResultsResp subsequentResponse = createArrowFetchResponseNoSchema(3, false); + TFetchResultsResp subsequentResponse = createArrowFetchResponseWithoutMetadata(3, false, false); StreamingBatch batch2 = processor.processResponse(subsequentResponse, 1, 100); assertNotNull(batch2); @@ -67,6 +70,38 @@ void testProcessSubsequentResponseUsesSchemaCache() throws DatabricksSQLExceptio batch2.release(); } + @Test + void testCompressedSubsequentResponseWithoutMetadataUsesInitialCompression() + throws DatabricksSQLException { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + + StreamingBatch initialBatch = + processor.processInitialResponse(createArrowFetchResponse(2, true, true)); + StreamingBatch subsequentBatch = + processor.processResponse(createArrowFetchResponseWithoutMetadata(3, false, true), 1, 2); + + assertTrue(initialBatch.isReady()); + assertTrue(subsequentBatch.isReady()); + assertEquals(3, subsequentBatch.getRowCount()); + + initialBatch.release(); + subsequentBatch.release(); + } + + @Test + void testInitialResponseWithoutMetadataThrowsTypedParsingError() { + InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); + TFetchResultsResp response = createArrowFetchResponseWithoutMetadata(1, false, false); + + DatabricksSQLException thrown = + assertThrows( + DatabricksSQLException.class, () -> processor.processInitialResponse(response)); + + assertEquals(DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.name(), thrown.getSQLState()); + assertEquals( + DatabricksDriverErrorCode.INLINE_CHUNK_PARSING_ERROR.getCode(), thrown.getErrorCode()); + } + @Test void testReleaseActionReleasesNativeMemory() throws DatabricksSQLException { InlineArrowResponseProcessor processor = new InlineArrowResponseProcessor(STATEMENT_ID); @@ -141,15 +176,24 @@ void testBatchMetadataCorrectlyPropagated() throws DatabricksSQLException { /** Creates an Arrow response with full schema (for initial response). */ private TFetchResultsResp createArrowFetchResponse(int rowCount, boolean hasMoreRows) { + return createArrowFetchResponse(rowCount, hasMoreRows, false); + } + + private TFetchResultsResp createArrowFetchResponse( + int rowCount, boolean hasMoreRows, boolean lz4Compressed) { TFetchResultsResp response = new TFetchResultsResp(); response.hasMoreRows = hasMoreRows; byte[] arrowData = createArrowBytes(rowCount); + if (lz4Compressed) { + arrowData = compressLz4(arrowData); + } TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); metadata.setResultFormat(TSparkRowSetType.ARROW_BASED_SET); metadata.setArrowSchema(new byte[0]); // Empty - schema is in the IPC stream metadata.setSchema(createTableSchema()); + metadata.setLz4Compressed(lz4Compressed); response.setResultSetMetadata(metadata); TSparkArrowBatch arrowBatch = new TSparkArrowBatch(); @@ -163,6 +207,26 @@ private TFetchResultsResp createArrowFetchResponse(int rowCount, boolean hasMore return response; } + private TFetchResultsResp createArrowFetchResponseWithoutMetadata( + int rowCount, boolean hasMoreRows, boolean lz4Compressed) { + TFetchResultsResp response = new TFetchResultsResp(); + response.hasMoreRows = hasMoreRows; + + byte[] arrowData = createArrowBytes(rowCount); + if (lz4Compressed) { + arrowData = compressLz4(arrowData); + } + + TSparkArrowBatch arrowBatch = new TSparkArrowBatch(); + arrowBatch.setRowCount(rowCount); + arrowBatch.setBatch(arrowData); + + TRowSet rowSet = new TRowSet(); + rowSet.setArrowBatches(Collections.singletonList(arrowBatch)); + response.setResults(rowSet); + return response; + } + /** * Creates an Arrow response for subsequent batches (minimal metadata, schema cached from * initial). @@ -215,6 +279,16 @@ private byte[] createArrowBytes(int rowCount) { } } + private 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(); + } + /** Creates a TTableSchema with a single INT column. */ private TTableSchema createTableSchema() { List columns = new ArrayList<>(); diff --git a/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java index 88fa349abc..8fe8dbace7 100644 --- a/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/ArrowUtilTest.java @@ -93,7 +93,8 @@ void testColumnDescToArrowField() throws SQLException { } @Test - void testCreateArrowByteStream() throws DatabricksParsingException { + void testCreateArrowByteStreamWithoutMetadataDefaultsToUncompressed() + throws DatabricksParsingException { byte[] schema = new byte[] {1, 2, 3}; TFetchResultsResp response = new TFetchResultsResp(); @@ -101,10 +102,6 @@ void testCreateArrowByteStream() throws DatabricksParsingException { rowSet.setArrowBatches(Collections.emptyList()); response.setResults(rowSet); - // Set up metadata with no compression - TGetResultSetMetadataResp metadata = new TGetResultSetMetadataResp(); - response.setResultSetMetadata(metadata); - ByteArrayInputStream result = ArrowUtil.createArrowByteStream(schema, response, getClass()); assertNotNull(result);