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