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] 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()); + } + } +}