From 21e72e628a532c513f28a87bc5eb8edccaf29e60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Thu, 12 Nov 2026 15:07:25 +0100 Subject: [PATCH 1/2] Keep vectors usable when loading their validity buffer fails loadFieldBuffers released the vector's validity buffer before BitVectorHelper.loadValidityBuffer created the new one. When that failed, for example on an allocation failure, the vector still pointed at the released buffer, and closing it later released the buffer again. Create the new validity buffer first in every vector that loads one. loadValidityBuffer also treated a missing validity buffer with an unknown null count (-1) as the mixed case, and kept the null or empty source buffer: a NullPointerException for a NULL C Data buffer, or a vector whose validity buffer is too small. Treat it as all valid, and reject a missing validity buffer when the null count says that only some values are null. --- .../arrow/vector/BaseFixedWidthVector.java | 5 +- .../vector/BaseLargeVariableWidthVector.java | 5 +- .../arrow/vector/BaseVariableWidthVector.java | 5 +- .../vector/BaseVariableWidthViewVector.java | 5 +- .../apache/arrow/vector/BitVectorHelper.java | 23 +++-- .../vector/complex/FixedSizeListVector.java | 5 +- .../arrow/vector/complex/LargeListVector.java | 5 +- .../vector/complex/LargeListViewVector.java | 5 +- .../arrow/vector/complex/ListVector.java | 5 +- .../arrow/vector/complex/ListViewVector.java | 5 +- .../arrow/vector/complex/StructVector.java | 5 +- .../arrow/vector/TestBitVectorHelper.java | 33 +++++++ .../arrow/vector/TestLoadFieldBuffers.java | 87 +++++++++++++++++++ 13 files changed, 176 insertions(+), 17 deletions(-) create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java index df1ac74f9b..b00f2413ab 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -503,8 +503,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf dataBuffer = ownBuffers.get(1); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); refreshValueCapacity(); diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java index 3fac195786..e50a24a573 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java @@ -304,8 +304,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); valueBuffer.getReferenceManager().release(); diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java index d5bd167256..ed03f4d24a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java @@ -320,8 +320,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); valueBuffer.getReferenceManager().release(); diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java index d9c48e6a85..6026a4929a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java @@ -354,10 +354,13 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf viewBuf = ownBuffers.get(1); List dataBufs = ownBuffers.subList(2, ownBuffers.size()); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuf, allocator); + this.clear(); this.viewBuffer = viewBuf.getReferenceManager().retain(viewBuf, allocator); - this.validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuf, allocator); + this.validityBuffer = newValidityBuffer; for (ArrowBuf dataBuf : dataBufs) { this.dataBuffers.add(dataBuf.getReferenceManager().retain(dataBuf, allocator)); diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java index bc2c3da98f..d85e686f42 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java +++ b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java @@ -309,6 +309,8 @@ public static byte getBitsFromNextByte(ArrowBuf data, int index, int offset) { * @param allocator The allocator to use for creating a new buffer if necessary. * @return A new buffer that is either allocated or points to the same memory as * sourceValidityBuffer. + * @throws IllegalArgumentException if the source validity buffer is missing, but the null count + * says that only some values are null */ public static ArrowBuf loadValidityBuffer( final ArrowFieldNode fieldNode, @@ -317,17 +319,24 @@ public static ArrowBuf loadValidityBuffer( final int valueCount = fieldNode.getLength(); ArrowBuf newBuffer = null; - // Create a new validity buffer iff both of the following are true: - // - validity buffer is not present, that is, it is either null or empty (in the case of - // IPC for instance). - // - values are either all NULLs or all non-NULLs + // Create a new validity buffer iff the validity buffer is not present, that is, it is either + // null or empty (in the case of IPC for instance). The values are then either all NULLs, or + // all non-NULLs when the null count is 0 or unknown (-1). boolean isValidityBufferNull = sourceValidityBuffer == null || sourceValidityBuffer.capacity() == 0; - if (isValidityBufferNull - && (fieldNode.getNullCount() == 0 || fieldNode.getNullCount() == valueCount)) { + if (isValidityBufferNull) { + final int nullCount = fieldNode.getNullCount(); + final boolean allNull = valueCount > 0 && nullCount == valueCount; + if (!allNull && nullCount != 0 && nullCount != -1) { + throw new IllegalArgumentException( + "Validity buffer is missing for " + + valueCount + + " values with a null count of " + + nullCount); + } newBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount)); newBuffer.setZero(0, newBuffer.capacity()); - if (fieldNode.getNullCount() != 0) { + if (allNull) { /* all NULLs */ return newBuffer; } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java index e3b4ab477f..59b9c30382 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java @@ -170,8 +170,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java index f3fc6faafb..ed066b9174 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java @@ -255,8 +255,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java index 2da7eb057e..8f0c2e97b4 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java @@ -202,8 +202,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); sizeBuffer.getReferenceManager().release(); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java index f2dee09154..42deb85fad 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java @@ -213,8 +213,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java index d41f61e291..c9171cf76f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java @@ -201,8 +201,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; offsetBuffer.getReferenceManager().release(); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); sizeBuffer.getReferenceManager().release(); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java index 5e5bb7fc21..d5ca2ad01d 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java @@ -166,8 +166,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); + // Load the new validity buffer first, so that a failure leaves the vector unchanged + ArrowBuf newValidityBuffer = + BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); validityBuffer.getReferenceManager().release(); - validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); + validityBuffer = newValidityBuffer; valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java b/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java index 68aa61962b..65b4bcc7f5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java @@ -18,6 +18,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import org.apache.arrow.memory.ArrowBuf; @@ -276,6 +277,38 @@ public void testLoadValidityBuffer() { } } + @Test + public void testLoadMissingValidityBufferWithUnknownNullCount() { + try (RootAllocator allocator = new RootAllocator(1024); + ArrowBuf empty = allocator.buffer(0)) { + // Without a validity buffer, an unknown null count (-1) means that no value is null + ArrowFieldNode fieldNode = new ArrowFieldNode(1024, -1); + for (ArrowBuf src : new ArrowBuf[] {null, empty}) { + try (ArrowBuf buf = BitVectorHelper.loadValidityBuffer(fieldNode, src, allocator)) { + assertEquals(128, allocator.getAllocatedMemory()); + for (int i = 0; i < 128; i++) { + assertEquals((byte) 0xff, buf.getByte(i)); + } + } + } + } + } + + @Test + public void testLoadMissingValidityBufferWithSomeNullsRejected() { + try (RootAllocator allocator = new RootAllocator(1024); + ArrowBuf empty = allocator.buffer(0)) { + ArrowFieldNode fieldNode = new ArrowFieldNode(1024, 100); + assertThrows( + IllegalArgumentException.class, + () -> BitVectorHelper.loadValidityBuffer(fieldNode, null, allocator)); + assertThrows( + IllegalArgumentException.class, + () -> BitVectorHelper.loadValidityBuffer(fieldNode, empty, allocator)); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + private void concatAndVerify( ArrowBuf buf1, int count1, ArrowBuf buf2, int count2, ArrowBuf output) { BitVectorHelper.concatBits(buf1, count1, buf2, count2, output); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java b/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java new file mode 100644 index 0000000000..d9c785ac12 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java @@ -0,0 +1,87 @@ +/* + * 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.assertThrows; + +import java.util.Collections; +import java.util.List; +import java.util.function.Function; +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.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.complex.LargeListVector; +import org.apache.arrow.vector.complex.LargeListViewVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.ListViewVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** Checks that a vector whose buffers fail to load can still be used and closed. */ +public class TestLoadFieldBuffers { + + static Stream vectors() { + return Stream.of( + vector(a -> new IntVector("v", a), 2), + vector(a -> new VarCharVector("v", a), 3), + vector(a -> new LargeVarCharVector("v", a), 3), + vector(a -> new ViewVarCharVector("v", a), 2), + vector(a -> ListVector.empty("v", a), 2), + vector(a -> LargeListVector.empty("v", a), 2), + vector(a -> FixedSizeListVector.empty("v", 2, a), 1), + vector(a -> ListViewVector.empty("v", a), 3), + vector(a -> LargeListViewVector.empty("v", a), 3), + vector(a -> StructVector.empty("v", a), 1)); + } + + private static Arguments vector(Function create, int bufferCount) { + return Arguments.of(create, bufferCount); + } + + @ParameterizedTest + @MethodSource("vectors") + public void testFailedValidityLoadKeepsVector( + Function create, int bufferCount) { + try (BufferAllocator allocator = new RootAllocator(64 * 1024)) { + FieldVector vector = create.apply(allocator); + vector.setInitialCapacity(16); + vector.allocateNew(); + List buffers = Collections.nCopies(bufferCount, allocator.getEmpty()); + + // Without a validity buffer, 1,000,000 values need a 125 KB validity buffer, which is more + // than the allocator allows + assertThrows( + OutOfMemoryException.class, + () -> vector.loadFieldBuffers(new ArrowFieldNode(1_000_000, 0), buffers)); + // A missing validity buffer is rejected when only some values are null + assertThrows( + IllegalArgumentException.class, + () -> vector.loadFieldBuffers(new ArrowFieldNode(16, 3), buffers)); + + // The vector still owns its own buffers, so closing it releases each of them once + vector.close(); + assertEquals(0, allocator.getAllocatedMemory()); + } + } +} From 16aadfe4bf69bfdc459acf70947608383d727973 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Fri, 13 Nov 2026 06:18:36 +0100 Subject: [PATCH 2/2] Clear vectors before loading their buffers Instead of loading the new validity buffer before releasing the old one, release the vector's own buffers first and reset its value count, then load. A failure while loading the validity buffer, or while retaining any other buffer, now leaves the vector empty but usable, without the extra memory of keeping the old validity buffer. List, list view, fixed size list and struct vectors keep their child vectors. BaseVariableWidthViewVector already cleared itself first: restore it, since loading its validity buffer before clear() leaked that buffer when retaining the view buffer failed. --- .../arrow/vector/BaseFixedWidthVector.java | 9 ++-- .../vector/BaseLargeVariableWidthVector.java | 10 ++--- .../arrow/vector/BaseVariableWidthVector.java | 10 ++--- .../vector/BaseVariableWidthViewVector.java | 5 +-- .../vector/complex/FixedSizeListVector.java | 11 ++--- .../arrow/vector/complex/LargeListVector.java | 14 ++++--- .../vector/complex/LargeListViewVector.java | 15 +++---- .../arrow/vector/complex/ListVector.java | 14 ++++--- .../arrow/vector/complex/ListViewVector.java | 15 +++---- .../arrow/vector/complex/StructVector.java | 11 ++--- .../arrow/vector/TestLoadFieldBuffers.java | 42 +++++++++++++++++-- 11 files changed, 93 insertions(+), 63 deletions(-) diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java index b00f2413ab..ff0d5a586e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -503,12 +503,9 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf dataBuffer = ownBuffers.get(1); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - valueBuffer.getReferenceManager().release(); + // Clear the vector first, so that a failure while loading leaves it empty but usable + clear(); + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); refreshValueCapacity(); diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java index e50a24a573..d40dc8ab06 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java @@ -304,14 +304,10 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Clear the vector first, so that a failure while loading leaves it empty but usable + clear(); + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); - valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); lastSet = fieldNode.getLength() - 1; diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java index ed03f4d24a..94c4a14eb8 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java @@ -320,14 +320,10 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Clear the vector first, so that a failure while loading leaves it empty but usable + clear(); + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); - valueBuffer.getReferenceManager().release(); valueBuffer = dataBuffer.getReferenceManager().retain(dataBuffer, allocator); lastSet = fieldNode.getLength() - 1; diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java index 6026a4929a..d9c48e6a85 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java @@ -354,13 +354,10 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf viewBuf = ownBuffers.get(1); List dataBufs = ownBuffers.subList(2, ownBuffers.size()); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuf, allocator); - this.clear(); this.viewBuffer = viewBuf.getReferenceManager().retain(viewBuf, allocator); - this.validityBuffer = newValidityBuffer; + this.validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuf, allocator); for (ArrowBuf dataBuf : dataBufs) { this.dataBuffers.add(dataBuf.getReferenceManager().retain(dataBuf, allocator)); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java index 59b9c30382..62ba453dd8 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java @@ -170,11 +170,12 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; + // Release the validity buffer first, so that a failure while loading leaves the vector empty + // but usable. Unlike clear(), this keeps the child vectors. + validityBuffer = releaseBuffer(validityBuffer); + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java index ed066b9174..b24ab8b192 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java @@ -255,12 +255,14 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Release this vector's buffers first, so that a failure while loading leaves it empty but + // usable. Unlike clear(), this keeps the child vectors. + validityBuffer = releaseBuffer(validityBuffer); + offsetBuffer = releaseBuffer(offsetBuffer); + lastSet = -1; + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java index 8f0c2e97b4..4d53137cc9 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java @@ -202,14 +202,15 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Release this vector's buffers first, so that a failure while loading leaves it empty but + // usable. Unlike clear(), this keeps the child vectors. + validityBuffer = releaseBuffer(validityBuffer); + offsetBuffer = releaseBuffer(offsetBuffer); + sizeBuffer = releaseBuffer(sizeBuffer); + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); - sizeBuffer.getReferenceManager().release(); sizeBuffer = szBuffer.getReferenceManager().retain(szBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java index 42deb85fad..7130fa0b57 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java @@ -213,12 +213,14 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Release this vector's buffers first, so that a failure while loading leaves it empty but + // usable. Unlike clear(), this keeps the child vectors. + validityBuffer = releaseBuffer(validityBuffer); + offsetBuffer = releaseBuffer(offsetBuffer); + lastSet = -1; + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java index c9171cf76f..2f9e608645 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java @@ -201,14 +201,15 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; - offsetBuffer.getReferenceManager().release(); + // Release this vector's buffers first, so that a failure while loading leaves it empty but + // usable. Unlike clear(), this keeps the child vectors. + validityBuffer = releaseBuffer(validityBuffer); + offsetBuffer = releaseBuffer(offsetBuffer); + sizeBuffer = releaseBuffer(sizeBuffer); + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); offsetBuffer = offBuffer.getReferenceManager().retain(offBuffer, allocator); - sizeBuffer.getReferenceManager().release(); sizeBuffer = szBuffer.getReferenceManager().retain(szBuffer, allocator); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java index d5ca2ad01d..7a71796b23 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java @@ -166,11 +166,12 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); - // Load the new validity buffer first, so that a failure leaves the vector unchanged - ArrowBuf newValidityBuffer = - BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); - validityBuffer.getReferenceManager().release(); - validityBuffer = newValidityBuffer; + // Release the validity buffer first, so that a failure while loading leaves the vector empty + // but usable. Unlike clear(), this keeps the child vectors. + clearValidityBuffer(); + valueCount = 0; + + validityBuffer = BitVectorHelper.loadValidityBuffer(fieldNode, bitBuffer, allocator); valueCount = fieldNode.getLength(); validityAllocationSizeInBytes = checkedCastToInt(validityBuffer.capacity()); } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java b/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java index d9c785ac12..bc7f5a0c93 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.function.Function; @@ -38,7 +39,7 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -/** Checks that a vector whose buffers fail to load can still be used and closed. */ +/** Checks that a vector whose buffers fail to load is left empty, and can still be closed. */ public class TestLoadFieldBuffers { static Stream vectors() { @@ -55,18 +56,24 @@ static Stream vectors() { vector(a -> StructVector.empty("v", a), 1)); } + static Stream vectorsWithRetainedBuffers() { + // The validity buffer is the only buffer of a fixed size list or a struct vector + return vectors().filter(arguments -> (int) arguments.get()[1] > 1); + } + private static Arguments vector(Function create, int bufferCount) { return Arguments.of(create, bufferCount); } @ParameterizedTest @MethodSource("vectors") - public void testFailedValidityLoadKeepsVector( + public void testFailedValidityLoadClearsVector( Function create, int bufferCount) { try (BufferAllocator allocator = new RootAllocator(64 * 1024)) { FieldVector vector = create.apply(allocator); vector.setInitialCapacity(16); vector.allocateNew(); + vector.setValueCount(16); List buffers = Collections.nCopies(bufferCount, allocator.getEmpty()); // Without a validity buffer, 1,000,000 values need a 125 KB validity buffer, which is more @@ -74,12 +81,41 @@ public void testFailedValidityLoadKeepsVector( assertThrows( OutOfMemoryException.class, () -> vector.loadFieldBuffers(new ArrowFieldNode(1_000_000, 0), buffers)); + assertEquals(0, vector.getValueCount()); // A missing validity buffer is rejected when only some values are null assertThrows( IllegalArgumentException.class, () -> vector.loadFieldBuffers(new ArrowFieldNode(16, 3), buffers)); + assertEquals(0, vector.getValueCount()); + + // Closing the vector releases each of its buffers once + vector.close(); + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @ParameterizedTest + @MethodSource("vectorsWithRetainedBuffers") + public void testFailedBufferRetainClearsVector( + Function create, int bufferCount) { + try (BufferAllocator allocator = new RootAllocator(); + BufferAllocator otherAllocator = new RootAllocator(); + ArrowBuf otherBuffer = otherAllocator.buffer(64)) { + FieldVector vector = create.apply(allocator); + vector.setInitialCapacity(16); + vector.allocateNew(); + vector.setValueCount(16); + // The validity buffer is created by the vector, but the next buffer belongs to another root + // allocator, so the vector's allocator can't retain it + List buffers = new ArrayList<>(Collections.nCopies(bufferCount, otherBuffer)); + buffers.set(0, allocator.getEmpty()); + + assertThrows( + IllegalStateException.class, + () -> vector.loadFieldBuffers(new ArrowFieldNode(16, 0), buffers)); + assertEquals(0, vector.getValueCount()); - // The vector still owns its own buffers, so closing it releases each of them once + // Closing the vector releases each of its buffers once vector.close(); assertEquals(0, allocator.getAllocatedMemory()); }