GH-1466: Add an option to fully validate IPC batches when reading them

vector/src/main/java/org/apache/arrow/vector/ipc/ArrowReader.java

@@ -35,6 +35,7 @@
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.DictionaryUtility;
+import org.apache.arrow.vector.util.ValueVectorUtility;
 import org.apache.arrow.vector.util.VectorBatchAppender;
 
 /** Abstract class to read Schema and ArrowRecordBatches. */
@@ -48,6 +49,8 @@ public abstract class ArrowReader implements DictionaryProvider, AutoCloseable {
 
   private final CompressionCodec.Factory compressionFactory;
 
+  private boolean validateBatches = false;
+
   protected ArrowReader(BufferAllocator allocator) {
     this(allocator, CompressionCodec.Factory.INSTANCE);
   }
@@ -57,6 +60,34 @@ protected ArrowReader(BufferAllocator allocator, CompressionCodec.Factory compre
     this.compressionFactory = compressionFactory;
   }
 
+  /**
+   * Sets whether every record batch and dictionary batch is checked with {@link
+   * ValueVectorUtility#validateFull} once it is loaded. This is off by default, because the check
+   * reads every value of the batch.
+   *
+   * <p>Turn it on when the data comes from an untrusted source. The IPC format does not guarantee
+   * that offsets, views or run ends stay within the buffers they point into, and reading such a
+   * batch can read past the end of a buffer, without bounds checks when {@code
+   * arrow.enable_unsafe_memory_access} is set. With this option, a batch that fails the check is
+   * cleared and the load throws {@link
+   * org.apache.arrow.vector.validate.ValidateUtil.ValidateException}.
+   *
+   * @param validateBatches whether to check each loaded batch
+   */
+  public void setValidateBatches(boolean validateBatches) {
+    this.validateBatches = validateBatches;
+  }
+
+  /**
+   * Returns whether every loaded batch is checked with {@link ValueVectorUtility#validateFull}.
+   *
+   * @return true if loaded batches are checked
+   * @see #setValidateBatches(boolean)
+   */
+  public boolean isValidateBatches() {
+    return validateBatches;
+  }
+
   /**
    * Returns the vector schema root. This will be loaded with new values on every call to
    * loadNextBatch.
@@ -214,6 +245,14 @@ protected void loadRecordBatch(ArrowRecordBatch batch) {
     } finally {
       batch.close();
     }
+    if (validateBatches) {
+      try {
+        ValueVectorUtility.validateFull(root);
+      } catch (RuntimeException e) {
+        root.clear();
+        throw e;
+      }
+    }
   }
 
   /**
@@ -250,5 +289,15 @@ private void load(ArrowDictionaryBatch dictionaryBatch, FieldVector vector) {
     } finally {
       dictionaryBatch.close();
     }
+    // A delta is checked before it is appended, so that an invalid delta leaves the dictionary
+    // unchanged
+    if (validateBatches) {
+      try {
+        ValueVectorUtility.validateFull(vector);
+      } catch (RuntimeException e) {
+        vector.clear();
+        throw e;
+      }
+    }
   }
 }

vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java

@@ -34,7 +34,9 @@
 import org.apache.arrow.vector.complex.DenseUnionVector;
 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.NonNullableStructVector;
 import org.apache.arrow.vector.complex.RunEndEncodedVector;
 import org.apache.arrow.vector.complex.UnionVector;
@@ -111,6 +113,15 @@ private void validateDataBuffer(ValueVector vector, long minCapacity) {
         dataBuffer.capacity());
   }
 
+  private void validateSizeBuffer(ArrowBuf sizeBuffer, long minCapacity) {
+    validateOrThrow(sizeBuffer != null, "The size buffer is null.");
+    validateOrThrow(
+        sizeBuffer.capacity() >= minCapacity,
+        "Not enough capacity for the size buffer. Minimum capacity %s, actual capacity %s.",
+        minCapacity,
+        sizeBuffer.capacity());
+  }
+
   private void validateTypeBuffer(ArrowBuf typeBuf, long minCapacity) {
     validateOrThrow(typeBuf != null, "The type buffer is null.");
     validateOrThrow(
@@ -248,6 +259,34 @@ public Void visit(LargeListVector vector, Void value) {
     return null;
   }
 
+  @Override
+  public Void visit(ListViewVector vector, Void value) {
+    int valueCount = vector.getValueCount();
+    validateVectorCommon(vector);
+    validateValidityBuffer(vector, valueCount);
+    validateOffsetBuffer(vector, (long) valueCount * ListViewVector.OFFSET_WIDTH);
+    validateSizeBuffer(vector.getSizeBuffer(), (long) valueCount * ListViewVector.SIZE_WIDTH);
+    FieldVector dataVector = vector.getDataVector();
+    if (dataVector != null) {
+      dataVector.accept(this, null);
+    }
+    return null;
+  }
+
+  @Override
+  public Void visit(LargeListViewVector vector, Void value) {
+    int valueCount = vector.getValueCount();
+    validateVectorCommon(vector);
+    validateValidityBuffer(vector, valueCount);
+    validateOffsetBuffer(vector, (long) valueCount * LargeListViewVector.OFFSET_WIDTH);
+    validateSizeBuffer(vector.getSizeBuffer(), (long) valueCount * LargeListViewVector.SIZE_WIDTH);
+    FieldVector dataVector = vector.getDataVector();
+    if (dataVector != null) {
+      dataVector.accept(this, null);
+    }
+    return null;
+  }
+
   @Override
   public Void visit(NonNullableStructVector vector, Void value) {
     int valueCount = vector.getValueCount();

vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java

@@ -18,20 +18,26 @@
 
 import static org.apache.arrow.vector.validate.ValidateUtil.validateOrThrow;
 
+import java.util.List;
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.vector.BaseFixedWidthVector;
+import org.apache.arrow.vector.BaseIntVector;
 import org.apache.arrow.vector.BaseLargeVariableWidthVector;
 import org.apache.arrow.vector.BaseVariableWidthVector;
 import org.apache.arrow.vector.BaseVariableWidthViewVector;
 import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.NullVector;
 import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.compare.VectorVisitor;
 import org.apache.arrow.vector.complex.DenseUnionVector;
 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.NonNullableStructVector;
+import org.apache.arrow.vector.complex.RunEndEncodedVector;
 import org.apache.arrow.vector.complex.UnionVector;
 
 /** Utility for validating vector data. */
@@ -121,10 +127,48 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) {
 
   @Override
   public Void visit(BaseVariableWidthViewVector vector, Void value) {
+    validateViews(vector);
     vector.validateScalars();
     return null;
   }
 
+  private void validateViews(BaseVariableWidthViewVector vector) {
+    ArrowBuf viewBuffer = vector.getDataBuffer();
+    List<ArrowBuf> dataBuffers = vector.getDataBuffers();
+    for (int i = 0; i < vector.getValueCount(); i++) {
+      if (vector.isNull(i)) {
+        continue;
+      }
+      long start = (long) i * BaseVariableWidthViewVector.ELEMENT_SIZE;
+      int length = viewBuffer.getInt(start);
+      validateOrThrow(length >= 0, "The view at position %s has a negative length: %s.", i, length);
+      if (length > BaseVariableWidthViewVector.INLINE_SIZE) {
+        long bufferIndexStart =
+            start
+                + BaseVariableWidthViewVector.LENGTH_WIDTH
+                + BaseVariableWidthViewVector.PREFIX_WIDTH;
+        int bufferIndex = viewBuffer.getInt(bufferIndexStart);
+        int offset =
+            viewBuffer.getInt(bufferIndexStart + BaseVariableWidthViewVector.BUF_INDEX_WIDTH);
+        validateOrThrow(
+            bufferIndex >= 0 && bufferIndex < dataBuffers.size(),
+            "The view at position %s refers to data buffer %s, but there are %s data buffers.",
+            i,
+            bufferIndex,
+            dataBuffers.size());
+        long capacity = dataBuffers.get(bufferIndex).capacity();
+        validateOrThrow(
+            offset >= 0 && (long) offset + length <= capacity,
+            "The view at position %s (offset %s, length %s) is outside data buffer %s of %s bytes.",
+            i,
+            offset,
+            length,
+            bufferIndex,
+            capacity);
+      }
+    }
+  }
+
   @Override
   public Void visit(ListVector vector, Void value) {
     validateOffsetBuffer(vector, vector.getValueCount());
@@ -137,7 +181,7 @@ public Void visit(ListVector vector, Void value) {
 
   @Override
   public Void visit(FixedSizeListVector vector, Void value) {
-    validateOffsetBuffer(vector, vector.getValueCount());
+    // A fixed size list has no offset buffer, and the buffer validation checks the inner length
     ValueVector innerVector = vector.getDataVector();
     if (innerVector != null) {
       innerVector.accept(this, null);
@@ -155,6 +199,91 @@ public Void visit(LargeListVector vector, Void value) {
     return null;
   }
 
+  @Override
+  public Void visit(ListViewVector vector, Void value) {
+    ValueVector innerVector = vector.getDataVector();
+    int innerValueCount = innerVector == null ? 0 : innerVector.getValueCount();
+    for (int i = 0; i < vector.getValueCount(); i++) {
+      if (vector.isNull(i)) {
+        continue;
+      }
+      int offset = vector.getOffsetBuffer().getInt((long) i * ListViewVector.OFFSET_WIDTH);
+      int size = vector.getSizeBuffer().getInt((long) i * ListViewVector.SIZE_WIDTH);
+      validateListView(i, offset, size, innerValueCount);
+    }
+    if (innerVector != null) {
+      innerVector.accept(this, null);
+    }
+    return null;
+  }
+
+  @Override
+  public Void visit(LargeListViewVector vector, Void value) {
+    ValueVector innerVector = vector.getDataVector();
+    int innerValueCount = innerVector == null ? 0 : innerVector.getValueCount();
+    for (int i = 0; i < vector.getValueCount(); i++) {
+      if (vector.isNull(i)) {
+        continue;
+      }
+      long offset = vector.getOffsetBuffer().getLong((long) i * LargeListViewVector.OFFSET_WIDTH);
+      long size = vector.getSizeBuffer().getLong((long) i * LargeListViewVector.SIZE_WIDTH);
+      validateListView(i, offset, size, innerValueCount);
+    }
+    if (innerVector != null) {
+      innerVector.accept(this, null);
+    }
+    return null;
+  }
+
+  private void validateListView(int index, long offset, long size, int innerValueCount) {
+    validateOrThrow(
+        size >= 0, "The size at position %s of the list view is negative: %s.", index, size);
+    validateOrThrow(
+        size == 0 || (offset >= 0 && offset <= innerValueCount - size),
+        "The list view at position %s (offset %s, size %s) is outside the inner vector of %s"
+            + " values.",
+        index,
+        offset,
+        size,
+        innerValueCount);
+  }
+
+  @Override
+  public Void visit(RunEndEncodedVector vector, Void value) {
+    FieldVector runEndsVector = vector.getRunEndsVector();
+    FieldVector valuesVector = vector.getValuesVector();
+    if (runEndsVector != null) {
+      int runCount = runEndsVector.getValueCount();
+      // Run ends are positive and strictly increasing, and each run has a value. Before its
+      // children are initialized, the run ends vector is empty and not an integer vector.
+      if (runCount > 0) {
+        BaseIntVector runEnds = (BaseIntVector) runEndsVector;
+        long previousEnd = 0;
+        for (int i = 0; i < runCount; i++) {
+          long end = runEnds.getValueAsLong(i);
+          validateOrThrow(
+              end > previousEnd,
+              "The run end at position %s (%s) is not greater than the previous run end (%s).",
+              i,
+              end,
+              previousEnd);
+          previousEnd = end;
+        }
+        int valueCount = valuesVector == null ? 0 : valuesVector.getValueCount();
+        validateOrThrow(
+            valueCount >= runCount,
+            "The values vector has %s values, fewer than the %s runs.",
+            valueCount,
+            runCount);
+      }
+      runEndsVector.accept(this, null);
+    }
+    if (valuesVector != null) {
+      valuesVector.accept(this, null);
+    }
+    return null;
+  }
+
   @Override
   public Void visit(NonNullableStructVector vector, Void value) {
     for (ValueVector subVector : vector.getChildrenFromFields()) {
@@ -182,7 +311,12 @@ public Void visit(DenseUnionVector vector, Void value) {
       byte typeId = vector.getTypeId(i);
       ValueVector subVector = vector.getVectorByType(typeId);
       validateOrThrow(
-          offset < subVector.getValueCount(),
+          subVector != null,
+          "Dense union vector type id %s at position %s has no sub-vector.",
+          typeId,
+          i);
+      validateOrThrow(
+          offset >= 0 && offset < subVector.getValueCount(),
           "Dense union vector offset exceeds sub-vector boundary. Vector offset %s, sub vector size %s",
           offset,
           subVector.getValueCount());

vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java

@@ -67,7 +67,10 @@
 import org.apache.arrow.vector.complex.DenseUnionVector;
 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.MapVector;
 import org.apache.arrow.vector.complex.NonNullableStructVector;
 import org.apache.arrow.vector.complex.RunEndEncodedVector;
 import org.apache.arrow.vector.complex.UnionVector;
@@ -392,7 +395,9 @@ public Void visit(BaseVariableWidthViewVector vector, Void value) {
 
   @Override
   public Void visit(ListVector vector, Void value) {
-    validateVectorCommon(vector, ArrowType.List.class);
+    // A MapVector is a ListVector of type Map
+    validateVectorCommon(
+        vector, vector instanceof MapVector ? ArrowType.Map.class : ArrowType.List.class);
     ValueVector innerVector = vector.getDataVector();
     if (innerVector != null) {
       innerVector.accept(this, null);
@@ -428,6 +433,26 @@ public Void visit(LargeListVector vector, Void value) {
     return null;
   }
 
+  @Override
+  public Void visit(ListViewVector vector, Void value) {
+    validateVectorCommon(vector, ArrowType.ListView.class);
+    ValueVector innerVector = vector.getDataVector();
+    if (innerVector != null) {
+      innerVector.accept(this, null);
+    }
+    return null;
+  }
+
+  @Override
+  public Void visit(LargeListViewVector vector, Void value) {
+    validateVectorCommon(vector, ArrowType.LargeListView.class);
+    ValueVector innerVector = vector.getDataVector();
+    if (innerVector != null) {
+      innerVector.accept(this, null);
+    }
+    return null;
+  }
+
   @Override
   public Void visit(NonNullableStructVector vector, Void value) {
     validateVectorCommon(vector, ArrowType.Struct.class);

vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java

@@ -2934,8 +2934,9 @@ public void testValidateInvalidOffsets() {
       offsets.setInt(12, 1024);
       vector.setValueCount(1);
       vector.setIndexDefined(0);
-      var e = assertThrows(IndexOutOfBoundsException.class, vector::validateFull);
-      assertTrue(e.getMessage().contains("index: 1024"));
+      // The view is checked against the data buffer before the value is read
+      var e = assertThrows(ValidateUtil.ValidateException.class, vector::validateFull);
+      assertTrue(e.getMessage().contains("(offset 1024, length 64) is outside data buffer 0"));
     }
   }
 }