MINOR: Backport Flight IPC compression negotiation

.github/workflows/jarbuild.yml

@@ -267,6 +267,10 @@ jobs:
           # tries to replace /usr/local/bin/2to3 and so on and causes
           # a conflict error.
           brew update
+          # The runner has stale OpenSSL 1.1 links not tracked by brew unlink.
+          # Relink its installed OpenSSL 3 before upgrading Python dependencies.
+          brew unlink openssl@3
+          brew link --overwrite openssl@3
           for python_package in $(brew list | grep python@); do
             brew install --overwrite ${python_package}
           done

compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java

@@ -44,7 +44,11 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu
     uncompressedBuffer.getBytes(/* index= */ 0, inBytes);
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     try (InputStream in = new ByteArrayInputStream(inBytes);
-        OutputStream out = new FramedLZ4CompressorOutputStream(baos)) {
+        // The default reserves 4 MiB for every Arrow buffer, including tiny validity buffers.
+        OutputStream out =
+            new FramedLZ4CompressorOutputStream(
+                baos,
+                new FramedLZ4CompressorOutputStream.Parameters(frameBlockSize(inBytes.length)))) {
       IOUtils.copy(in, out);
     } catch (IOException e) {
       throw new RuntimeException(e);
@@ -59,6 +63,19 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu
     return compressedBuffer;
   }
 
+  private static FramedLZ4CompressorOutputStream.BlockSize frameBlockSize(int length) {
+    if (length <= 65536) {
+      return FramedLZ4CompressorOutputStream.BlockSize.K64;
+    }
+    if (length <= 262144) {
+      return FramedLZ4CompressorOutputStream.BlockSize.K256;
+    }
+    if (length <= 1048576) {
+      return FramedLZ4CompressorOutputStream.BlockSize.M1;
+    }
+    return FramedLZ4CompressorOutputStream.BlockSize.M4;
+  }
+
   @Override
   protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) {
     Preconditions.checkArgument(

compression/src/test/java/org/apache/arrow/compression/TestLz4FrameBlockSize.java

@@ -0,0 +1,131 @@
+/*
+ * 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.compression;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Random;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.compression.CompressionUtil;
+import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorOutputStream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+/** Regression coverage for bounded frame buffers and multi-block interoperability. */
+class TestLz4FrameBlockSize {
+  @Test
+  void preservesCompressionAcrossSmallBlockBoundaries() {
+    byte[] pattern = new byte[50000];
+    new Random(1323).nextBytes(pattern);
+    byte[] expected = new byte[1000000];
+    for (int i = 0; i < expected.length; i++) {
+      expected[i] = pattern[i % pattern.length];
+    }
+    try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      ArrowBuf input = allocator.buffer(expected.length);
+      input.setBytes(0, expected);
+      input.writerIndex(expected.length);
+      Lz4CompressionCodec codec = new Lz4CompressionCodec();
+      ArrowBuf compressed = codec.compress(allocator, input);
+      long size = compressed.writerIndex();
+      try (ArrowBuf restored = codec.decompress(allocator, compressed)) {
+        byte[] actual = new byte[expected.length];
+        restored.getBytes(0, actual);
+        assertArrayEquals(expected, actual);
+        assertTrue(size < expected.length / 4, "Large buffers must retain cross-block matches");
+      }
+    }
+  }
+
+  @Test
+  void stillReadsFourMiBFrames() throws IOException {
+    byte[] expected = new byte[131072];
+    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+    try (FramedLZ4CompressorOutputStream out = new FramedLZ4CompressorOutputStream(bytes)) {
+      out.write(expected);
+    }
+    try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      ArrowBuf framed =
+          allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + bytes.size());
+      framed.setBytes(
+          0,
+          ByteBuffer.allocate((int) CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH)
+              .order(ByteOrder.LITTLE_ENDIAN)
+              .putLong(expected.length)
+              .array());
+      framed.setBytes(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, bytes.toByteArray());
+      framed.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + bytes.size());
+      assertEquals(7, (framed.getByte(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + 5) >>> 4) & 7);
+      try (ArrowBuf restored = new Lz4CompressionCodec().decompress(allocator, framed)) {
+        byte[] actual = new byte[expected.length];
+        restored.getBytes(0, actual);
+        assertArrayEquals(expected, actual);
+      }
+    }
+  }
+
+  @ParameterizedTest
+  @CsvSource({
+    "1024, 4",
+    "65536, 4",
+    "65537, 5",
+    "262144, 5",
+    "262145, 6",
+    "1048576, 6",
+    "1048577, 7",
+    "5242880, 7"
+  })
+  void usesAppropriateFrameBlocksAndRoundTrips(int length, int expectedBlockSizeId) {
+    byte[] expected = new byte[length];
+    for (int i = 0; i < length; i++) {
+      expected[i] = (byte) (i % 17);
+    }
+    Lz4CompressionCodec codec = new Lz4CompressionCodec();
+    try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      ArrowBuf input = allocator.buffer(length);
+      input.setBytes(0, expected);
+      input.writerIndex(length);
+      // compress and decompress each consume their input buffer.
+      ArrowBuf compressed = codec.compress(allocator, input);
+      try {
+        // The LZ4 frame BD byte follows its four-byte magic and FLG byte.
+        int blockSizeId =
+            (compressed.getByte(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + 5) >>> 4) & 7;
+        assertEquals(
+            expectedBlockSizeId, blockSizeId, "LZ4 frame block allocation must fit the input");
+      } catch (Throwable failure) {
+        compressed.close();
+        throw failure;
+      }
+      try (ArrowBuf restored = codec.decompress(allocator, compressed)) {
+        assertEquals(length, restored.writerIndex());
+        byte[] actual = new byte[length];
+        restored.getBytes(0, actual);
+        assertArrayEquals(expected, actual);
+      }
+    }
+  }
+}

flight/flight-core/pom.xml

@@ -108,6 +108,11 @@ under the License.
       <artifactId>grpc-services</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-compression</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-inprocess</artifactId>

flight/flight-core/src/main/java/org/apache/arrow/flight/CallOptions.java

@@ -17,14 +17,37 @@
 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 {
   public static CallOption timeout(long duration, TimeUnit unit) {
     return new Timeout(duration, unit);
   }
 
+  /**
+   * Advertise support for IPC body compression codecs in preference order.
+   *
+   * <p>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 extends AbstractStub<T>> T wrapStub(T stub, CallOption[] options) {
     for (CallOption option : options) {
       if (option instanceof GrpcCallOption) {

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<ArrowMessage> 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<ArrowMessage> messageCallback)
+      throws Exception {
     final Set<Long> 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)) {

flight/flight-core/src/main/java/org/apache/arrow/flight/FlightBindingService.java

@@ -26,13 +26,15 @@
 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;
 import org.apache.arrow.flight.impl.Flight;
 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<CompressionUtil.CodecType> ipcCompressionCodecs) {
     this.allocator = allocator;
-    this.delegate = new FlightService(allocator, producer, authHandler, executor);
+    this.delegate =
+        new FlightService(allocator, producer, authHandler, executor, ipcCompressionCodecs);
   }
 
   public static MethodDescriptor<Flight.Ticket, ArrowMessage> getDoGetDescriptor(

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<ServerHeaderMiddleware> HEADER_KEY =
       FlightServerMiddleware.Key.of("org.apache.arrow.flight.ServerHeaderMiddleware");
 

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<KeyFactory<?>> interceptors;
     // Keep track of inserted interceptors
     private final Set<String> interceptorKeys;
+    private List<CompressionUtil.CodecType> 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.
+     *
+     * <p>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