diff --git a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java index d2d2921649..4378be96ac 100644 --- a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java +++ b/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 codecs() { List params = new ArrayList<>(); diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java index 342f210b82..9b9db18906 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java +++ b/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 nodes = new ArrayList<>(); List buffers = new ArrayList<>(); List 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) { diff --git a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java b/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java index b108173c82..39b5ceae3d 100644 --- a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java +++ b/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 diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java new file mode 100644 index 0000000000..828b81fb95 --- /dev/null +++ b/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 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; + } + } +}