GH-1234: Release serialization buffers when compression fails

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

@@ -37,6 +37,7 @@
 import java.util.stream.Stream;
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.IntVector;
@@ -64,6 +65,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /** Test cases for {@link CompressionCodec}s. */
 class TestCompressionCodec {
@@ -79,6 +81,27 @@ void terminate() {
     allocator.close();
   }
 
+  @ParameterizedTest
+  @ValueSource(booleans = {false, true})
+  void testCompressionAllocationFailureReleasesSource(boolean useZstd) {
+    CompressionCodec codec = useZstd ? new ZstdCompressionCodec() : new Lz4CompressionCodec();
+    try (IntVector vector = new IntVector("values", allocator);
+        VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+      vector.allocateNew(1);
+      vector.set(0, 42);
+      root.setRowCount(1);
+      long allocatedBefore = allocator.getAllocatedMemory();
+      int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
+      allocator.setLimit(allocatedBefore);
+      VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
+
+      assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
+      assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+      assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
+      assertEquals(42, vector.get(0));
+    }
+  }
+
   static Collection<Arguments> codecs() {
     List<Arguments> params = new ArrayList<>();
 

vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java

@@ -19,6 +19,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.compression.CompressionCodec;
 import org.apache.arrow.vector.compression.CompressionUtil;
 import org.apache.arrow.vector.compression.NoCompressionCodec;
@@ -78,18 +79,23 @@ public ArrowRecordBatch getRecordBatch() {
     List<ArrowFieldNode> nodes = new ArrayList<>();
     List<ArrowBuf> buffers = new ArrayList<>();
     List<Long> variadicBufferCounts = new ArrayList<>();
-    for (FieldVector vector : root.getFieldVectors()) {
-      appendNodes(vector, nodes, buffers, variadicBufferCounts);
+    try {
+      for (FieldVector vector : root.getFieldVectors()) {
+        appendNodes(vector, nodes, buffers, variadicBufferCounts);
+      }
+      // Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
+      return new ArrowRecordBatch(
+          root.getRowCount(),
+          nodes,
+          buffers,
+          CompressionUtil.createBodyCompression(codec),
+          variadicBufferCounts,
+          alignBuffers, /*retainBuffers*/
+          false);
+    } catch (RuntimeException | Error e) {
+      AutoCloseables.close(e, buffers);
+      throw e;
     }
-    // Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
-    return new ArrowRecordBatch(
-        root.getRowCount(),
-        nodes,
-        buffers,
-        CompressionUtil.createBodyCompression(codec),
-        variadicBufferCounts,
-        alignBuffers, /*retainBuffers*/
-        false);
   }
 
   private long getVariadicBufferCount(FieldVector vector) {

vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java

@@ -29,34 +29,34 @@ public abstract class AbstractCompressionCodec implements CompressionCodec {
 
   @Override
   public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
-    // GH-1116: capture writerIndex() once so the empty-buffer check, size
-    // comparison, and uncompressed-length prefix all see the same value.
-    long uncompressedLength = uncompressedBuffer.writerIndex();
+    try (uncompressedBuffer) {
+      // GH-1116: capture writerIndex() once so the empty-buffer check, size
+      // comparison, and uncompressed-length prefix all see the same value.
+      long uncompressedLength = uncompressedBuffer.writerIndex();
+
+      if (uncompressedLength == 0L) {
+        // shortcut for empty buffer
+        ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
+        compressedBuffer.setLong(0, 0);
+        compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
+        return compressedBuffer;
+      }
+
+      ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
+      long compressedLength =
+          compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
+
+      if (compressedLength > uncompressedLength) {
+        // compressed buffer is larger, send the raw buffer
+        compressedBuffer.close();
+        // XXX: this makes a copy of uncompressedBuffer
+        compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
+      } else {
+        writeUncompressedLength(compressedBuffer, uncompressedLength);
+      }
 
-    if (uncompressedLength == 0L) {
-      // shortcut for empty buffer
-      ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
-      compressedBuffer.setLong(0, 0);
-      compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
-      uncompressedBuffer.close();
       return compressedBuffer;
     }
-
-    ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
-    long compressedLength =
-        compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
-
-    if (compressedLength > uncompressedLength) {
-      // compressed buffer is larger, send the raw buffer
-      compressedBuffer.close();
-      // XXX: this makes a copy of uncompressedBuffer
-      compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
-    } else {
-      writeUncompressedLength(compressedBuffer, uncompressedLength);
-    }
-
-    uncompressedBuffer.close();
-    return compressedBuffer;
   }
 
   @Override

vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java

@@ -0,0 +1,152 @@
+/*
+ * 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.vector;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.Collections;
+import java.util.List;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.compression.AbstractCompressionCodec;
+import org.apache.arrow.vector.compression.CompressionUtil;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+/** Exception-safety tests for record batch serialization. */
+class TestVectorUnloaderFailure {
+
+  @Test
+  void compressionAllocationFailureReleasesBuffers() {
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector vector = new IntVector("values", allocator);
+        VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+      vector.allocateNew(1);
+      vector.set(0, 42);
+      root.setRowCount(1);
+      long allocatedBefore = allocator.getAllocatedMemory();
+      int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
+      // Any compression allocation must fail, regardless of allocator rounding.
+      allocator.setLimit(allocatedBefore);
+      VectorUnloader unloader = new VectorUnloader(root, true, new SimulatedLz4Codec(), true);
+
+      assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
+      assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+      assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
+      assertEquals(42, vector.get(0));
+    }
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {false, true})
+  void compressionFailurePreservesExceptionAndSource(boolean throwError) {
+    Throwable failure =
+        throwError
+            ? new OutOfMemoryError("compression failed")
+            : new IllegalStateException("compression failed");
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector vector = new IntVector("values", allocator);
+        VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+      vector.allocateNew(1);
+      vector.set(0, 42);
+      root.setRowCount(1);
+      long allocatedBefore = allocator.getAllocatedMemory();
+      int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
+      SimulatedLz4Codec codec =
+          new SimulatedLz4Codec() {
+            private int calls;
+
+            @Override
+            protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
+              if (++calls == 2) {
+                if (failure instanceof Error) {
+                  throw (Error) failure;
+                }
+                throw (RuntimeException) failure;
+              }
+              return super.doCompress(allocator, input);
+            }
+          };
+      VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
+
+      assertSame(failure, assertThrows(failure.getClass(), unloader::getRecordBatch));
+      assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+      assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
+      assertEquals(42, vector.get(0));
+    }
+  }
+
+  @Test
+  void emptyBufferAllocationFailureReleasesInput() {
+    try (BufferAllocator allocator = new RootAllocator()) {
+      ArrowBuf input = allocator.buffer(8);
+      allocator.setLimit(allocator.getAllocatedMemory());
+
+      assertThrows(
+          OutOfMemoryException.class, () -> new SimulatedLz4Codec().compress(allocator, input));
+      assertEquals(0, input.getReferenceManager().getRefCount());
+      assertEquals(0, allocator.getAllocatedMemory());
+    }
+  }
+
+  @Test
+  void invalidLaterVectorReleasesPreviouslyRetainedBuffers() {
+    try (BufferAllocator allocator = new RootAllocator();
+        IntVector first = new IntVector("first", allocator)) {
+      FieldVector invalid =
+          new NullVector("invalid") {
+            @Override
+            public List<ArrowBuf> getFieldBuffers() {
+              return Collections.singletonList(allocator.getEmpty());
+            }
+          };
+      try (VectorSchemaRoot root = VectorSchemaRoot.of(first, invalid)) {
+        first.allocateNew(1);
+        first.set(0, 42);
+        root.setRowCount(1);
+        int referencesBefore = first.getDataBuffer().getReferenceManager().getRefCount();
+
+        assertThrows(IllegalArgumentException.class, new VectorUnloader(root)::getRecordBatch);
+        assertEquals(referencesBefore, first.getDataBuffer().getReferenceManager().getRefCount());
+        assertEquals(42, first.get(0));
+      }
+    }
+  }
+
+  // Compression-only stub for allocation and ownership tests, not LZ4 round trips.
+  private static class SimulatedLz4Codec extends AbstractCompressionCodec {
+    @Override
+    protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
+      return CompressionUtil.packageRawBuffer(allocator, input);
+    }
+
+    @Override
+    protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf input) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public CompressionUtil.CodecType getCodecType() {
+      return CompressionUtil.CodecType.LZ4_FRAME;
+    }
+  }
+}