Skip to content

Commit 41e1c31

Browse files
pj-workspacepj-workspace
committed
GH-1234: Release serialization buffers when compression fails
Generated-by: OpenAI Codex
1 parent 19a3ff5 commit 41e1c31

4 files changed

Lines changed: 215 additions & 36 deletions

File tree

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.util.stream.Stream;
3838
import org.apache.arrow.memory.ArrowBuf;
3939
import org.apache.arrow.memory.BufferAllocator;
40+
import org.apache.arrow.memory.OutOfMemoryException;
4041
import org.apache.arrow.memory.RootAllocator;
4142
import org.apache.arrow.util.AutoCloseables;
4243
import org.apache.arrow.vector.IntVector;
@@ -64,6 +65,7 @@
6465
import org.junit.jupiter.params.ParameterizedTest;
6566
import org.junit.jupiter.params.provider.Arguments;
6667
import org.junit.jupiter.params.provider.MethodSource;
68+
import org.junit.jupiter.params.provider.ValueSource;
6769

6870
/** Test cases for {@link CompressionCodec}s. */
6971
class TestCompressionCodec {
@@ -79,6 +81,27 @@ void terminate() {
7981
allocator.close();
8082
}
8183

84+
@ParameterizedTest
85+
@ValueSource(booleans = {false, true})
86+
void testCompressionAllocationFailureReleasesSource(boolean useZstd) {
87+
CompressionCodec codec = useZstd ? new ZstdCompressionCodec() : new Lz4CompressionCodec();
88+
try (IntVector vector = new IntVector("values", allocator);
89+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
90+
vector.allocateNew(1);
91+
vector.set(0, 42);
92+
root.setRowCount(1);
93+
long allocatedBefore = allocator.getAllocatedMemory();
94+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
95+
allocator.setLimit(allocatedBefore);
96+
VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
97+
98+
assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
99+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
100+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
101+
assertEquals(42, vector.get(0));
102+
}
103+
}
104+
82105
static Collection<Arguments> codecs() {
83106
List<Arguments> params = new ArrayList<>();
84107

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

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.ArrayList;
2020
import java.util.List;
2121
import org.apache.arrow.memory.ArrowBuf;
22+
import org.apache.arrow.util.AutoCloseables;
2223
import org.apache.arrow.vector.compression.CompressionCodec;
2324
import org.apache.arrow.vector.compression.CompressionUtil;
2425
import org.apache.arrow.vector.compression.NoCompressionCodec;
@@ -78,18 +79,23 @@ public ArrowRecordBatch getRecordBatch() {
7879
List<ArrowFieldNode> nodes = new ArrayList<>();
7980
List<ArrowBuf> buffers = new ArrayList<>();
8081
List<Long> variadicBufferCounts = new ArrayList<>();
81-
for (FieldVector vector : root.getFieldVectors()) {
82-
appendNodes(vector, nodes, buffers, variadicBufferCounts);
82+
try {
83+
for (FieldVector vector : root.getFieldVectors()) {
84+
appendNodes(vector, nodes, buffers, variadicBufferCounts);
85+
}
86+
// Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
87+
return new ArrowRecordBatch(
88+
root.getRowCount(),
89+
nodes,
90+
buffers,
91+
CompressionUtil.createBodyCompression(codec),
92+
variadicBufferCounts,
93+
alignBuffers, /*retainBuffers*/
94+
false);
95+
} catch (RuntimeException | Error e) {
96+
AutoCloseables.close(e, buffers);
97+
throw e;
8398
}
84-
// Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
85-
return new ArrowRecordBatch(
86-
root.getRowCount(),
87-
nodes,
88-
buffers,
89-
CompressionUtil.createBodyCompression(codec),
90-
variadicBufferCounts,
91-
alignBuffers, /*retainBuffers*/
92-
false);
9399
}
94100

95101
private long getVariadicBufferCount(FieldVector vector) {

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

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -29,34 +29,34 @@ public abstract class AbstractCompressionCodec implements CompressionCodec {
2929

3030
@Override
3131
public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
32-
// GH-1116: capture writerIndex() once so the empty-buffer check, size
33-
// comparison, and uncompressed-length prefix all see the same value.
34-
long uncompressedLength = uncompressedBuffer.writerIndex();
32+
try (uncompressedBuffer) {
33+
// GH-1116: capture writerIndex() once so the empty-buffer check, size
34+
// comparison, and uncompressed-length prefix all see the same value.
35+
long uncompressedLength = uncompressedBuffer.writerIndex();
36+
37+
if (uncompressedLength == 0L) {
38+
// shortcut for empty buffer
39+
ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
40+
compressedBuffer.setLong(0, 0);
41+
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
42+
return compressedBuffer;
43+
}
44+
45+
ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
46+
long compressedLength =
47+
compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
48+
49+
if (compressedLength > uncompressedLength) {
50+
// compressed buffer is larger, send the raw buffer
51+
compressedBuffer.close();
52+
// XXX: this makes a copy of uncompressedBuffer
53+
compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
54+
} else {
55+
writeUncompressedLength(compressedBuffer, uncompressedLength);
56+
}
3557

36-
if (uncompressedLength == 0L) {
37-
// shortcut for empty buffer
38-
ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
39-
compressedBuffer.setLong(0, 0);
40-
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
41-
uncompressedBuffer.close();
4258
return compressedBuffer;
4359
}
44-
45-
ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
46-
long compressedLength =
47-
compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
48-
49-
if (compressedLength > uncompressedLength) {
50-
// compressed buffer is larger, send the raw buffer
51-
compressedBuffer.close();
52-
// XXX: this makes a copy of uncompressedBuffer
53-
compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
54-
} else {
55-
writeUncompressedLength(compressedBuffer, uncompressedLength);
56-
}
57-
58-
uncompressedBuffer.close();
59-
return compressedBuffer;
6060
}
6161

6262
@Override

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

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.vector;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertSame;
21+
import static org.junit.jupiter.api.Assertions.assertThrows;
22+
23+
import java.util.Collections;
24+
import java.util.List;
25+
import org.apache.arrow.memory.ArrowBuf;
26+
import org.apache.arrow.memory.BufferAllocator;
27+
import org.apache.arrow.memory.OutOfMemoryException;
28+
import org.apache.arrow.memory.RootAllocator;
29+
import org.apache.arrow.vector.compression.AbstractCompressionCodec;
30+
import org.apache.arrow.vector.compression.CompressionUtil;
31+
import org.junit.jupiter.api.Test;
32+
import org.junit.jupiter.params.ParameterizedTest;
33+
import org.junit.jupiter.params.provider.ValueSource;
34+
35+
/** Exception-safety tests for record batch serialization. */
36+
class TestVectorUnloaderFailure {
37+
38+
@ParameterizedTest
39+
@ValueSource(longs = {0, 16})
40+
void compressionAllocationFailureReleasesBuffers(long availableBytes) {
41+
try (BufferAllocator allocator = new RootAllocator();
42+
IntVector vector = new IntVector("values", allocator);
43+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
44+
vector.allocateNew(1);
45+
vector.set(0, 42);
46+
root.setRowCount(1);
47+
long allocatedBefore = allocator.getAllocatedMemory();
48+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
49+
allocator.setLimit(allocatedBefore + availableBytes);
50+
VectorUnloader unloader = new VectorUnloader(root, true, new CopyCodec(), true);
51+
52+
assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
53+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
54+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
55+
assertEquals(42, vector.get(0));
56+
}
57+
}
58+
59+
@ParameterizedTest
60+
@ValueSource(booleans = {false, true})
61+
void compressionFailurePreservesExceptionAndSource(boolean throwError) {
62+
Throwable failure =
63+
throwError
64+
? new OutOfMemoryError("compression failed")
65+
: new IllegalStateException("compression failed");
66+
try (BufferAllocator allocator = new RootAllocator();
67+
IntVector vector = new IntVector("values", allocator);
68+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
69+
vector.allocateNew(1);
70+
vector.set(0, 42);
71+
root.setRowCount(1);
72+
long allocatedBefore = allocator.getAllocatedMemory();
73+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
74+
CopyCodec codec =
75+
new CopyCodec() {
76+
private int calls;
77+
78+
@Override
79+
protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
80+
if (++calls == 2) {
81+
if (failure instanceof Error) {
82+
throw (Error) failure;
83+
}
84+
throw (RuntimeException) failure;
85+
}
86+
return super.doCompress(allocator, input);
87+
}
88+
};
89+
VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
90+
91+
assertSame(failure, assertThrows(failure.getClass(), unloader::getRecordBatch));
92+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
93+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
94+
assertEquals(42, vector.get(0));
95+
}
96+
}
97+
98+
@Test
99+
void emptyBufferAllocationFailureReleasesInput() {
100+
try (BufferAllocator allocator = new RootAllocator()) {
101+
ArrowBuf input = allocator.buffer(8);
102+
allocator.setLimit(allocator.getAllocatedMemory());
103+
104+
assertThrows(OutOfMemoryException.class, () -> new CopyCodec().compress(allocator, input));
105+
assertEquals(0, input.getReferenceManager().getRefCount());
106+
assertEquals(0, allocator.getAllocatedMemory());
107+
}
108+
}
109+
110+
@Test
111+
void invalidLaterVectorReleasesPreviouslyRetainedBuffers() {
112+
try (BufferAllocator allocator = new RootAllocator();
113+
IntVector first = new IntVector("first", allocator)) {
114+
FieldVector invalid =
115+
new NullVector("invalid") {
116+
@Override
117+
public List<ArrowBuf> getFieldBuffers() {
118+
return Collections.singletonList(allocator.getEmpty());
119+
}
120+
};
121+
try (VectorSchemaRoot root = VectorSchemaRoot.of(first, invalid)) {
122+
first.allocateNew(1);
123+
first.set(0, 42);
124+
root.setRowCount(1);
125+
int referencesBefore = first.getDataBuffer().getReferenceManager().getRefCount();
126+
127+
assertThrows(IllegalArgumentException.class, new VectorUnloader(root)::getRecordBatch);
128+
assertEquals(referencesBefore, first.getDataBuffer().getReferenceManager().getRefCount());
129+
assertEquals(42, first.get(0));
130+
}
131+
}
132+
}
133+
134+
private static class CopyCodec extends AbstractCompressionCodec {
135+
@Override
136+
protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
137+
return CompressionUtil.packageRawBuffer(allocator, input);
138+
}
139+
140+
@Override
141+
protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf input) {
142+
throw new UnsupportedOperationException();
143+
}
144+
145+
@Override
146+
public CompressionUtil.CodecType getCodecType() {
147+
return CompressionUtil.CodecType.LZ4_FRAME;
148+
}
149+
}
150+
}

0 commit comments

Comments
 (0)