GH-1464: Keep vectors usable when loading their validity buffer fails

vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java

@@ -503,8 +503,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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();

vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java

@@ -304,8 +304,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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();

vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java

@@ -320,8 +320,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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();

vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java

@@ -354,10 +354,13 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> ownBuffers
     ArrowBuf viewBuf = ownBuffers.get(1);
     List<ArrowBuf> 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));

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;
       }

vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java

@@ -170,8 +170,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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());

vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java

@@ -255,8 +255,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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);
 

vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java

@@ -202,8 +202,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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();

vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java

@@ -213,8 +213,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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);
 

vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java

@@ -201,8 +201,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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();

vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java

@@ -166,8 +166,11 @@ public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> 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());
   }

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);

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<Arguments> 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<BufferAllocator, FieldVector> create, int bufferCount) {
+    return Arguments.of(create, bufferCount);
+  }
+
+  @ParameterizedTest
+  @MethodSource("vectors")
+  public void testFailedValidityLoadKeepsVector(
+      Function<BufferAllocator, FieldVector> create, int bufferCount) {
+    try (BufferAllocator allocator = new RootAllocator(64 * 1024)) {
+      FieldVector vector = create.apply(allocator);
+      vector.setInitialCapacity(16);
+      vector.allocateNew();
+      List<ArrowBuf> 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());
+    }
+  }
+}