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.

Author: jbonofre

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