From 51d71a93c7bdf32059185d7a34fdd7c65ab285f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 1 Oct 2026 00:51:18 +0100 Subject: [PATCH 1/2] feat(flight): negotiate IPC stream compression --- flight/flight-core/pom.xml | 5 + .../org/apache/arrow/flight/CallOptions.java | 23 +++++ .../apache/arrow/flight/DictionaryUtils.java | 17 +++- .../arrow/flight/FlightBindingService.java | 14 ++- .../apache/arrow/flight/FlightConstants.java | 3 + .../org/apache/arrow/flight/FlightServer.java | 23 ++++- .../apache/arrow/flight/FlightService.java | 48 ++++++++-- .../org/apache/arrow/flight/FlightStream.java | 5 + .../apache/arrow/flight/IpcCompression.java | 72 ++++++++++++++ .../flight/OutboundStreamListenerImpl.java | 21 ++++- .../arrow/flight/TestIpcCompression.java | 93 +++++++++++++++++++ flight/flight-sql-jdbc-core/pom.xml | 6 ++ .../driver/jdbc/ArrowFlightConnection.java | 55 ++++++----- .../FlightSqlIpcCompressionExample.java | 62 +++++++++++++ .../ArrowFlightConnectionConfigImpl.java | 30 ++++++ .../jdbc/FlightServerTestExtension.java | 22 ++++- .../arrow/driver/jdbc/IpcCompressionTest.java | 67 +++++++++++++ 17 files changed, 529 insertions(+), 37 deletions(-) create mode 100644 flight/flight-core/src/main/java/org/apache/arrow/flight/IpcCompression.java create mode 100644 flight/flight-core/src/test/java/org/apache/arrow/flight/TestIpcCompression.java create mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/FlightSqlIpcCompressionExample.java create mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/IpcCompressionTest.java diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml index df70007052..fb188c5812 100644 --- a/flight/flight-core/pom.xml +++ b/flight/flight-core/pom.xml @@ -108,6 +108,11 @@ under the License. grpc-services test + + org.apache.arrow + arrow-compression + test + io.grpc grpc-inprocess diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java index b6e052d223..dd16de0167 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java @@ -17,7 +17,11 @@ package org.apache.arrow.flight; import io.grpc.stub.AbstractStub; +import java.util.Arrays; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.apache.arrow.vector.compression.CompressionCodec; +import org.apache.arrow.vector.compression.CompressionUtil; /** Common call options. */ public class CallOptions { @@ -25,6 +29,25 @@ public static CallOption timeout(long duration, TimeUnit unit) { return new Timeout(duration, unit); } + /** + * Advertise support for IPC body compression codecs in preference order. + * + *

Servers that do not support negotiation ignore this option and return uncompressed IPC. + */ + public static CallOption acceptIpcCompression(CompressionUtil.CodecType... codecs) { + if (codecs.length == 0) { + throw new IllegalArgumentException("At least one IPC compression codec is required"); + } + for (CompressionUtil.CodecType codec : codecs) { + CompressionCodec.Factory.INSTANCE.createCodec(codec); + } + final FlightCallHeaders headers = new FlightCallHeaders(); + headers.insert( + FlightConstants.IPC_ACCEPT_COMPRESSION_HEADER, + Arrays.stream(codecs).map(IpcCompression::codecName).collect(Collectors.joining(","))); + return new HeaderCallOption(headers); + } + static > T wrapStub(T stub, CallOption[] options) { for (CallOption option : options) { if (option instanceof GrpcCallOption) { diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/DictionaryUtils.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/DictionaryUtils.java index cecc1b876e..1ca9b424b5 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/DictionaryUtils.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/DictionaryUtils.java @@ -28,6 +28,7 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; @@ -57,6 +58,18 @@ static Schema generateSchemaMessages( final IpcOption option, final Consumer messageCallback) throws Exception { + return generateSchemaMessages( + originalSchema, descriptor, provider, option, null, messageCallback); + } + + static Schema generateSchemaMessages( + final Schema originalSchema, + final FlightDescriptor descriptor, + final DictionaryProvider provider, + final IpcOption option, + final CompressionCodec compressionCodec, + final Consumer messageCallback) + throws Exception { final Set dictionaryIds = new HashSet<>(); final Schema schema = generateSchema(originalSchema, provider, dictionaryIds); MetadataV4UnionChecker.checkForUnion(schema.getFields().iterator(), option.metadataVersion); @@ -77,7 +90,9 @@ static Schema generateSchemaMessages( Collections.singletonList(vector.getField()), Collections.singletonList(vector), count); - final VectorUnloader unloader = new VectorUnloader(dictRoot); + final VectorUnloader unloader = + new VectorUnloader( + dictRoot, /* includeNullCount */ true, compressionCodec, /* alignBuffers */ true); try (final ArrowDictionaryBatch dictionaryBatch = new ArrowDictionaryBatch(id, unloader.getRecordBatch()); final ArrowMessage message = new ArrowMessage(dictionaryBatch, option)) { diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightBindingService.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightBindingService.java index b68f3aa86c..982e2cc75f 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightBindingService.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightBindingService.java @@ -26,6 +26,7 @@ import io.grpc.protobuf.ProtoUtils; import io.grpc.stub.ServerCalls; import io.grpc.stub.StreamObserver; +import java.util.List; import java.util.Set; import java.util.concurrent.ExecutorService; import org.apache.arrow.flight.auth.ServerAuthHandler; @@ -33,6 +34,7 @@ import org.apache.arrow.flight.impl.Flight.PutResult; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.compression.CompressionUtil; /** Extends the basic flight service to override some methods for more efficient implementations. */ class FlightBindingService implements BindableService { @@ -53,8 +55,18 @@ public FlightBindingService( FlightProducer producer, ServerAuthHandler authHandler, ExecutorService executor) { + this(allocator, producer, authHandler, executor, java.util.Collections.emptyList()); + } + + public FlightBindingService( + BufferAllocator allocator, + FlightProducer producer, + ServerAuthHandler authHandler, + ExecutorService executor, + List ipcCompressionCodecs) { this.allocator = allocator; - this.delegate = new FlightService(allocator, producer, authHandler, executor); + this.delegate = + new FlightService(allocator, producer, authHandler, executor, ipcCompressionCodecs); } public static MethodDescriptor getDoGetDescriptor( diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightConstants.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightConstants.java index 6b89c794d6..997a11d950 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightConstants.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightConstants.java @@ -21,6 +21,9 @@ public interface FlightConstants { String SERVICE = "arrow.flight.protocol.FlightService"; + /** Request header listing supported IPC body compression codecs in preference order. */ + String IPC_ACCEPT_COMPRESSION_HEADER = "arrow-ipc-accept-compression"; + FlightServerMiddleware.Key HEADER_KEY = FlightServerMiddleware.Key.of("org.apache.arrow.flight.ServerHeaderMiddleware"); diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java index ac761457f5..ea091a89c4 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java @@ -34,6 +34,8 @@ import java.net.URI; import java.net.URISyntaxException; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -55,6 +57,8 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.util.VisibleForTesting; +import org.apache.arrow.vector.compression.CompressionCodec; +import org.apache.arrow.vector.compression.CompressionUtil; /** * Generic server of flight data that is customized via construction with delegate classes for the @@ -197,6 +201,7 @@ public static final class Builder { private final List> interceptors; // Keep track of inserted interceptors private final Set interceptorKeys; + private List ipcCompressionCodecs = Collections.emptyList(); Builder() { builderOptions = new HashMap<>(); @@ -321,7 +326,7 @@ public FlightServer build() { } final FlightBindingService flightService = - new FlightBindingService(allocator, producer, authHandler, exec); + new FlightBindingService(allocator, producer, authHandler, exec, ipcCompressionCodecs); builder .executor(exec) .maxInboundMessageSize(maxInboundMessageSize) @@ -392,6 +397,22 @@ public Builder backpressureThreshold(int backpressureThreshold) { return this; } + /** + * Enable negotiated IPC body compression for server response streams. + * + *

The client preference order wins. Clients that do not advertise support continue to + * receive uncompressed IPC. + */ + public Builder ipcCompression(CompressionUtil.CodecType... codecs) { + Preconditions.checkArgument(codecs.length > 0, "At least one codec is required"); + for (CompressionUtil.CodecType codec : codecs) { + IpcCompression.codecName(codec); + CompressionCodec.Factory.INSTANCE.createCodec(codec); + } + ipcCompressionCodecs = Collections.unmodifiableList(Arrays.asList(codecs.clone())); + return this; + } + /** * A small utility function to ensure that InputStream attributes. are closed if they are not * null diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java index 9f130463c0..9382f3b634 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightService.java @@ -20,6 +20,7 @@ import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; @@ -38,6 +39,8 @@ import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.compression.CompressionCodec; +import org.apache.arrow.vector.compression.CompressionUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -51,16 +54,27 @@ class FlightService extends FlightServiceImplBase { private final FlightProducer producer; private final ServerAuthHandler authHandler; private final ExecutorService executors; + private final List ipcCompressionCodecs; FlightService( BufferAllocator allocator, FlightProducer producer, ServerAuthHandler authHandler, ExecutorService executors) { + this(allocator, producer, authHandler, executors, Collections.emptyList()); + } + + FlightService( + BufferAllocator allocator, + FlightProducer producer, + ServerAuthHandler authHandler, + ExecutorService executors, + List ipcCompressionCodecs) { this.allocator = allocator; this.producer = producer; this.authHandler = authHandler; this.executors = new ContextPropagatingExecutorService(executors); + this.ipcCompressionCodecs = ipcCompressionCodecs; } private CallContext makeContext(ServerCallStreamObserver responseObserver) { @@ -107,10 +121,14 @@ public void doGetCustom( final ServerCallStreamObserver responseObserver = (ServerCallStreamObserver) responseObserverSimple; + final CallContext context = makeContext(responseObserver); final GetListener listener = - new GetListener(responseObserver, this::handleExceptionWithMiddleware); + new GetListener( + responseObserver, + this::handleExceptionWithMiddleware, + IpcCompression.negotiate(context, ipcCompressionCodecs)); try { - producer.getStream(makeContext(responseObserver), new Ticket(ticket), listener); + producer.getStream(context, new Ticket(ticket), listener); } catch (Exception ex) { listener.error(ex); } @@ -155,7 +173,14 @@ private static class GetListener extends OutboundStreamListenerImpl public GetListener( ServerCallStreamObserver responseObserver, Consumer errorHandler) { - super(null, responseObserver); + this(responseObserver, errorHandler, null); + } + + public GetListener( + ServerCallStreamObserver responseObserver, + Consumer errorHandler, + CompressionCodec compressionCodec) { + super(null, responseObserver, compressionCodec); this.errorHandler = errorHandler; this.completed = false; this.serverCallResponseObserver = responseObserver; @@ -327,7 +352,14 @@ private static class ExchangeListener extends GetListener { public ExchangeListener( ServerCallStreamObserver responseObserver, Consumer errorHandler) { - super(responseObserver, errorHandler); + this(responseObserver, errorHandler, null); + } + + public ExchangeListener( + ServerCallStreamObserver responseObserver, + Consumer errorHandler, + CompressionCodec compressionCodec) { + super(responseObserver, errorHandler, compressionCodec); this.resource = null; super.setOnCancelHandler( () -> { @@ -387,8 +419,12 @@ public StreamObserver doExchangeCustom( StreamObserver responseObserverSimple) { final ServerCallStreamObserver responseObserver = (ServerCallStreamObserver) responseObserverSimple; + final CallContext context = makeContext(responseObserver); final ExchangeListener listener = - new ExchangeListener(responseObserver, this::handleExceptionWithMiddleware); + new ExchangeListener( + responseObserver, + this::handleExceptionWithMiddleware, + IpcCompression.negotiate(context, ipcCompressionCodecs)); final FlightStream fs = new FlightStream( allocator, @@ -405,7 +441,7 @@ public StreamObserver doExchangeCustom( executors.submit( () -> { try { - producer.doExchange(makeContext(responseObserver), fs, listener); + producer.doExchange(context, fs, listener); } catch (Exception ex) { listener.error(ex); } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java index 15cfd6ba85..cdec979e43 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java @@ -38,6 +38,7 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch; @@ -86,6 +87,7 @@ public void close() throws Exception {} private volatile Throwable ex; private volatile ArrowBuf applicationMetadata = null; @VisibleForTesting volatile MetadataVersion metadataVersion = null; + @VisibleForTesting volatile CompressionUtil.CodecType compressionType = null; /** * Constructs a new instance. @@ -272,6 +274,9 @@ public boolean next() { // Ensure we have the root root.get().clear(); try (ArrowRecordBatch arb = msg.asRecordBatch()) { + compressionType = + CompressionUtil.CodecType.fromCompressionType( + arb.getBodyCompression().getCodec()); loader.load(arb); } updateMetadata(msg); diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/IpcCompression.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/IpcCompression.java new file mode 100644 index 0000000000..d38c16eb45 --- /dev/null +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/IpcCompression.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import java.util.List; +import java.util.Locale; +import org.apache.arrow.vector.compression.CompressionCodec; +import org.apache.arrow.vector.compression.CompressionUtil; + +/** Utilities for negotiating IPC body compression. */ +final class IpcCompression { + private IpcCompression() {} + + static String codecName(CompressionUtil.CodecType codec) { + switch (codec) { + case LZ4_FRAME: + return "lz4_frame"; + case ZSTD: + return "zstd"; + default: + throw new IllegalArgumentException("Unsupported IPC compression codec: " + codec); + } + } + + static CompressionCodec negotiate( + FlightProducer.CallContext context, List supportedCodecs) { + if (supportedCodecs.isEmpty()) { + return null; + } + final ServerHeaderMiddleware middleware = context.getMiddleware(FlightConstants.HEADER_KEY); + if (middleware == null) { + return null; + } + final String accepted = middleware.headers().get(FlightConstants.IPC_ACCEPT_COMPRESSION_HEADER); + if (accepted == null) { + return null; + } + for (String name : accepted.split(",")) { + final CompressionUtil.CodecType codec = parseCodec(name); + if (codec != null && supportedCodecs.contains(codec)) { + return CompressionCodec.Factory.INSTANCE.createCodec(codec); + } + } + return null; + } + + private static CompressionUtil.CodecType parseCodec(String name) { + switch (name.trim().toLowerCase(Locale.ROOT)) { + case "lz4": + case "lz4_frame": + return CompressionUtil.CodecType.LZ4_FRAME; + case "zstd": + return CompressionUtil.CodecType.ZSTD; + default: + return null; + } + } +} diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/OutboundStreamListenerImpl.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/OutboundStreamListenerImpl.java index a1bde3a848..c440abb021 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/OutboundStreamListenerImpl.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/OutboundStreamListenerImpl.java @@ -22,6 +22,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.compression.CompressionCodec; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.ipc.message.IpcOption; @@ -32,12 +33,21 @@ abstract class OutboundStreamListenerImpl implements OutboundStreamListener { protected volatile VectorUnloader unloader; // null until stream started protected IpcOption option; // null until stream started protected boolean tryZeroCopy = ArrowMessage.ENABLE_ZERO_COPY_WRITE; + private final CompressionCodec compressionCodec; OutboundStreamListenerImpl( FlightDescriptor descriptor, CallStreamObserver responseObserver) { + this(descriptor, responseObserver, null); + } + + OutboundStreamListenerImpl( + FlightDescriptor descriptor, + CallStreamObserver responseObserver, + CompressionCodec compressionCodec) { Preconditions.checkNotNull(responseObserver, "responseObserver must be provided"); this.descriptor = descriptor; this.responseObserver = responseObserver; + this.compressionCodec = compressionCodec; this.unloader = null; } @@ -56,7 +66,12 @@ public void start(VectorSchemaRoot root, DictionaryProvider dictionaries, IpcOpt this.option = option; try { DictionaryUtils.generateSchemaMessages( - root.getSchema(), descriptor, dictionaries, option, responseObserver::onNext); + root.getSchema(), + descriptor, + dictionaries, + option, + compressionCodec, + responseObserver::onNext); } catch (RuntimeException e) { // Propagate runtime exceptions, like those raised when trying to write unions with V4 // metadata @@ -68,7 +83,9 @@ public void start(VectorSchemaRoot root, DictionaryProvider dictionaries, IpcOpt throw new RuntimeException("Could not generate and send all schema messages", e); } // We include the null count and align buffers to be compatible with Flight/C++ - unloader = new VectorUnloader(root, /* includeNullCount */ true, /* alignBuffers */ true); + unloader = + new VectorUnloader( + root, /* includeNullCount */ true, compressionCodec, /* alignBuffers */ true); } @Override diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestIpcCompression.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestIpcCompression.java new file mode 100644 index 0000000000..42002eae83 --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestIpcCompression.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.charset.StandardCharsets; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.compression.CompressionUtil; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +public class TestIpcCompression { + private static final byte[] VALUE = "compressible-value".getBytes(StandardCharsets.UTF_8); + + @ParameterizedTest + @EnumSource( + value = CompressionUtil.CodecType.class, + names = {"LZ4_FRAME", "ZSTD"}) + public void negotiatesCompression(CompressionUtil.CodecType codec) throws Exception { + assertRoundTrip(codec, true, codec); + } + + @ParameterizedTest + @EnumSource( + value = CompressionUtil.CodecType.class, + names = {"LZ4_FRAME", "ZSTD"}) + public void fallsBackForServersWithoutCompression(CompressionUtil.CodecType codec) + throws Exception { + assertRoundTrip(codec, false, CompressionUtil.CodecType.NO_COMPRESSION); + } + + private static void assertRoundTrip( + CompressionUtil.CodecType requested, + boolean enableServerCompression, + CompressionUtil.CodecType expected) + throws Exception { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + final FlightServer.Builder serverBuilder = + FlightServer.builder( + allocator, + Location.forGrpcInsecure(LOCALHOST, 0), + new NoOpFlightProducer() { + @Override + public void getStream( + CallContext context, Ticket ticket, ServerStreamListener listener) { + try (VarCharVector vector = new VarCharVector("value", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(vector)) { + vector.allocateNew(); + vector.setSafe(0, VALUE); + vector.setValueCount(1); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + listener.completed(); + } + } + }); + if (enableServerCompression) { + serverBuilder.ipcCompression( + CompressionUtil.CodecType.LZ4_FRAME, CompressionUtil.CodecType.ZSTD); + } + try (FlightServer server = serverBuilder.build().start(); + FlightClient client = FlightClient.builder(allocator, server.getLocation()).build(); + FlightStream stream = + client.getStream( + new Ticket(new byte[0]), CallOptions.acceptIpcCompression(requested))) { + assertTrue(stream.next()); + assertEquals("compressible-value", stream.getRoot().getVector(0).getObject(0).toString()); + assertEquals(expected, stream.compressionType); + } + } + } +} diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml index d6fa11688d..f5ba51ae20 100644 --- a/flight/flight-sql-jdbc-core/pom.xml +++ b/flight/flight-sql-jdbc-core/pom.xml @@ -79,6 +79,12 @@ under the License. ${arrow.vector.classifier} + + org.apache.arrow + arrow-compression + runtime + + com.google.guava guava diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 623c2b81be..e10cc17268 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -29,10 +29,12 @@ import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; +import org.apache.arrow.flight.CallOptions; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; import org.apache.calcite.avatica.DriverVersion; @@ -105,30 +107,35 @@ private static ArrowFlightSqlClientHandler createNewClientHandler( final DriverVersion driverVersion) throws SQLException { try { - return new ArrowFlightSqlClientHandler.Builder() - .withHost(config.getHost()) - .withPort(config.getPort()) - .withUsername(config.getUser()) - .withPassword(config.getPassword()) - .withTrustStorePath(config.getTrustStorePath()) - .withTrustStorePassword(config.getTrustStorePassword()) - .withSystemTrustStore(config.useSystemTrustStore()) - .withTlsRootCertificates(config.getTlsRootCertificatesPath()) - .withClientCertificate(config.getClientCertificatePath()) - .withClientKey(config.getClientKeyPath()) - .withBufferAllocator(allocator) - .withEncryption(config.useEncryption()) - .withDisableCertificateVerification(config.getDisableCertificateVerification()) - .withToken(config.getToken()) - .withCallOptions(config.toCallOption()) - .withRetainCookies(config.retainCookies()) - .withRetainAuth(config.retainAuth()) - .withCatalog(config.getCatalog()) - .withClientCache(config.useClientCache() ? new FlightClientCache() : null) - .withConnectTimeout(config.getConnectTimeout()) - .withDriverVersion(driverVersion) - .withOAuthConfiguration(config.getOauthConfiguration()) - .build(); + final ArrowFlightSqlClientHandler.Builder builder = + new ArrowFlightSqlClientHandler.Builder() + .withHost(config.getHost()) + .withPort(config.getPort()) + .withUsername(config.getUser()) + .withPassword(config.getPassword()) + .withTrustStorePath(config.getTrustStorePath()) + .withTrustStorePassword(config.getTrustStorePassword()) + .withSystemTrustStore(config.useSystemTrustStore()) + .withTlsRootCertificates(config.getTlsRootCertificatesPath()) + .withClientCertificate(config.getClientCertificatePath()) + .withClientKey(config.getClientKeyPath()) + .withBufferAllocator(allocator) + .withEncryption(config.useEncryption()) + .withDisableCertificateVerification(config.getDisableCertificateVerification()) + .withToken(config.getToken()) + .withCallOptions(config.toCallOption()) + .withRetainCookies(config.retainCookies()) + .withRetainAuth(config.retainAuth()) + .withCatalog(config.getCatalog()) + .withClientCache(config.useClientCache() ? new FlightClientCache() : null) + .withConnectTimeout(config.getConnectTimeout()) + .withDriverVersion(driverVersion) + .withOAuthConfiguration(config.getOauthConfiguration()); + final CompressionUtil.CodecType[] ipcCompressionCodecs = config.getIpcCompressionCodecs(); + if (ipcCompressionCodecs.length > 0) { + builder.withCallOptions(CallOptions.acceptIpcCompression(ipcCompressionCodecs)); + } + return builder.build(); } catch (final SQLException e) { try { allocator.close(); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/FlightSqlIpcCompressionExample.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/FlightSqlIpcCompressionExample.java new file mode 100644 index 0000000000..7dca87f37d --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/FlightSqlIpcCompressionExample.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.example; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.util.Properties; +import java.util.concurrent.TimeUnit; + +/** Runs a Flight SQL query with optional negotiated IPC compression. */ +public final class FlightSqlIpcCompressionExample { + private FlightSqlIpcCompressionExample() {} + + public static void main(String[] args) throws Exception { + if (args.length != 3) { + System.err.println( + "Usage: FlightSqlIpcCompressionExample "); + System.exit(2); + } + + final Properties properties = new Properties(); + if (!"none".equalsIgnoreCase(args[2])) { + properties.setProperty("ipcCompression", args[2]); + } + + final long started = System.nanoTime(); + long rows = 0; + int columns; + try (Connection connection = DriverManager.getConnection(args[0], properties); + Statement statement = connection.createStatement(); + ResultSet results = statement.executeQuery(args[1])) { + final ResultSetMetaData metadata = results.getMetaData(); + columns = metadata.getColumnCount(); + while (results.next()) { + rows++; + for (int column = 1; column <= columns; column++) { + results.getObject(column); + } + } + } + final long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started); + System.out.printf( + "compression=%s rows=%d columns=%d elapsed_ms=%d%n", args[2], rows, columns, elapsedMillis); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java index d0ba74dbcc..2efb15d0d0 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java @@ -16,10 +16,13 @@ */ package org.apache.arrow.driver.jdbc.utils; +import com.google.common.base.Splitter; import java.sql.SQLException; import java.time.Duration; import java.util.Arrays; import java.util.HashMap; +import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Properties; @@ -30,6 +33,7 @@ import org.apache.arrow.flight.FlightCallHeaders; import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.compression.CompressionUtil; import org.apache.calcite.avatica.ConnectionConfig; import org.apache.calcite.avatica.ConnectionConfigImpl; import org.apache.calcite.avatica.ConnectionProperty; @@ -182,6 +186,31 @@ public boolean useClientCache() { return ArrowFlightConnectionProperty.USE_CLIENT_CACHE.getBoolean(properties); } + /** IPC body compression codecs to advertise, in preference order. */ + public CompressionUtil.CodecType[] getIpcCompressionCodecs() throws SQLException { + final String value = ArrowFlightConnectionProperty.IPC_COMPRESSION.getString(properties); + if (value == null || value.trim().isEmpty() || "none".equalsIgnoreCase(value.trim())) { + return new CompressionUtil.CodecType[0]; + } + final List names = Splitter.on(',').splitToList(value); + final CompressionUtil.CodecType[] codecs = new CompressionUtil.CodecType[names.size()]; + for (int i = 0; i < names.size(); i++) { + final String name = names.get(i).trim().toLowerCase(Locale.ROOT); + try { + codecs[i] = + "lz4".equals(name) + ? CompressionUtil.CodecType.LZ4_FRAME + : CompressionUtil.CodecType.valueOf(name.toUpperCase(Locale.ROOT)); + if (codecs[i] == CompressionUtil.CodecType.NO_COMPRESSION) { + throw new IllegalArgumentException(); + } + } catch (IllegalArgumentException e) { + throw new SQLException("Unsupported IPC compression codec: " + names.get(i).trim(), e); + } + } + return codecs; + } + /** * Gets the {@link CallOption}s from this {@link ConnectionConfig}. * @@ -267,6 +296,7 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty { CATALOG("catalog", null, Type.STRING, false), CONNECT_TIMEOUT_MILLIS("connectTimeoutMs", 10000, Type.NUMBER, false), USE_CLIENT_CACHE("useClientCache", true, Type.BOOLEAN, false), + IPC_COMPRESSION("ipcCompression", null, Type.STRING, false), // OAuth configuration properties OAUTH_FLOW("oauth.flow", null, Type.STRING, false), diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java index f71114e1b5..53fc14510d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java @@ -45,6 +45,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.compression.CompressionUtil; import org.junit.jupiter.api.extension.AfterAllCallback; import org.junit.jupiter.api.extension.BeforeAllCallback; import org.junit.jupiter.api.extension.ExtensionContext; @@ -69,6 +70,7 @@ public class FlightServerTestExtension private final Authentication authentication; private final CertKeyPair certKeyPair; private final File mTlsCACert; + private final CompressionUtil.CodecType[] ipcCompressionCodecs; private final InterceptorMiddleware.Factory interceptorFactory = new InterceptorMiddleware.Factory(); @@ -80,7 +82,8 @@ private FlightServerTestExtension( final FlightSqlProducer producer, final Authentication authentication, final CertKeyPair certKeyPair, - final File mTlsCACert) { + final File mTlsCACert, + final CompressionUtil.CodecType[] ipcCompressionCodecs) { this.properties = Preconditions.checkNotNull(properties); this.config = Preconditions.checkNotNull(config); this.allocator = Preconditions.checkNotNull(allocator); @@ -88,6 +91,7 @@ private FlightServerTestExtension( this.authentication = authentication; this.certKeyPair = certKeyPair; this.mTlsCACert = mTlsCACert; + this.ipcCompressionCodecs = ipcCompressionCodecs; } /** @@ -136,6 +140,10 @@ public Connection getConnection(String timezone) throws SQLException { return this.createDataSource().getConnection(); } + public void setIpcCompression(String codecs) { + properties.put("ipcCompression", codecs); + } + private void setUseEncryption(boolean useEncryption) { properties.put("useEncryption", useEncryption); } @@ -154,6 +162,9 @@ private FlightServer initiateServer(Location location) throws IOException { FlightServer.builder(allocator, location, producer) .headerAuthenticator(authentication.authenticate()) .middleware(FlightServerMiddleware.Key.of("KEY"), interceptorFactory); + if (ipcCompressionCodecs.length > 0) { + builder.ipcCompression(ipcCompressionCodecs); + } if (certKeyPair != null) { builder.useTls(certKeyPair.cert, certKeyPair.key); } @@ -237,6 +248,7 @@ public static final class Builder { private Authentication authentication; private CertKeyPair certKeyPair; private File mTlsCACert; + private CompressionUtil.CodecType[] ipcCompressionCodecs = new CompressionUtil.CodecType[0]; public Builder() { this.properties = new Properties(); @@ -266,6 +278,11 @@ public Builder authentication(final Authentication authentication) { return this; } + public Builder ipcCompression(CompressionUtil.CodecType... codecs) { + ipcCompressionCodecs = codecs.clone(); + return this; + } + /** * Enable TLS on the server. * @@ -303,7 +320,8 @@ public FlightServerTestExtension build() { producer, authentication, certKeyPair, - mTlsCACert); + mTlsCACert, + ipcCompressionCodecs); } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/IpcCompressionTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/IpcCompressionTest.java new file mode 100644 index 0000000000..e64a3d2daf --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/IpcCompressionTest.java @@ -0,0 +1,67 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import static org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; +import static org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers.assertLegacyRegularSqlResultSet; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; +import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; +import org.apache.arrow.flight.FlightConstants; +import org.apache.arrow.flight.FlightMethod; +import org.apache.arrow.vector.compression.CompressionUtil; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +public class IpcCompressionTest { + @RegisterExtension + public static final FlightServerTestExtension FLIGHT_SERVER = + new FlightServerTestExtension.Builder() + .authentication( + new UserPasswordAuthentication.Builder() + .user( + FlightServerTestExtension.DEFAULT_USER, + FlightServerTestExtension.DEFAULT_PASSWORD) + .build()) + .producer(CoreMockedSqlProducers.getLegacyProducer()) + .ipcCompression(CompressionUtil.CodecType.LZ4_FRAME, CompressionUtil.CodecType.ZSTD) + .build(); + + @BeforeAll + public static void configureCompression() { + FLIGHT_SERVER.setIpcCompression("zstd,lz4_frame"); + } + + @Test + public void readsCompressedResults() throws Exception { + try (Connection connection = FLIGHT_SERVER.getConnection(false); + Statement statement = connection.createStatement(); + ResultSet results = statement.executeQuery(LEGACY_REGULAR_SQL_CMD)) { + assertLegacyRegularSqlResultSet(results); + } + assertEquals( + "zstd,lz4_frame", + FLIGHT_SERVER + .getInterceptorFactory() + .getHeader(FlightMethod.DO_GET, FlightConstants.IPC_ACCEPT_COMPRESSION_HEADER)); + } +} From 9072fbc85dfbb017181cdc25b7183329a1f222e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Wed, 25 Mar 2026 10:56:45 +0000 Subject: [PATCH 2/2] fix(flight-sql-jdbc): allow zstd natives in shaded driver --- .../apache/arrow/driver/jdbc/ITDriverJarValidation.java | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java index 145744ad38..ef8ea67577 100644 --- a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java +++ b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java @@ -56,6 +56,13 @@ public class ITDriverJarValidation { public static final Set ALLOWED_PREFIXES = ImmutableSet.of( "org/apache/arrow/driver/jdbc/", // Driver code + // zstd-jni classes and native libraries must keep their original names + "com/github/luben/zstd/", + "aix/", + "darwin/", + "freebsd/", + "linux/", + "win/", "META-INF/maven/", // Maven metadata (useful for security scanner "META-INF/services/", // ServiceLoader implementations "META-INF/license/",