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..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,9 +503,9 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf dataBuffer = ownBuffers.get(1); - validityBuffer.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.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..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,11 +304,10 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); - validityBuffer.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.getReferenceManager().release(); 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 d5bd167256..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,11 +320,10 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf dataBuffer = ownBuffers.get(2); - validityBuffer.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.getReferenceManager().release(); 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/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..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,7 +170,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); - validityBuffer.getReferenceManager().release(); + // 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(); 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..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,9 +255,14 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); - validityBuffer.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.getReferenceManager().release(); 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 2da7eb057e..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,11 +202,15 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); - validityBuffer.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.getReferenceManager().release(); 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 f2dee09154..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,9 +213,14 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); ArrowBuf offBuffer = ownBuffers.get(1); - validityBuffer.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.getReferenceManager().release(); 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 d41f61e291..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,11 +201,15 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf offBuffer = ownBuffers.get(1); ArrowBuf szBuffer = ownBuffers.get(2); - validityBuffer.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.getReferenceManager().release(); 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 5e5bb7fc21..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,7 +166,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List ownBuffers ArrowBuf bitBuffer = ownBuffers.get(0); - validityBuffer.getReferenceManager().release(); + // 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/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..bc7f5a0c93 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestLoadFieldBuffers.java @@ -0,0 +1,123 @@ +/* + * 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.ArrayList; +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 is left empty, and can still be 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)); + } + + 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 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 + // than the allocator allows + 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()); + + // Closing the vector releases each of its buffers once + vector.close(); + assertEquals(0, allocator.getAllocatedMemory()); + } + } +}