Check bounds on index in setPosition()

Author: martin-traverse

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeListProducer.java

@@ -60,6 +60,9 @@ public void skipNull() {
 
   @Override
   public void setPosition(int index) {
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds");
+    }
     int delegateOffset = vector.getOffsetBuffer().getInt(index * (long) Integer.BYTES);
     delegate.setPosition(delegateOffset);
     super.setPosition(index);

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroListProducer.java

@@ -57,6 +57,9 @@ public void produce(Encoder encoder) throws IOException {
 
   @Override
   public void setPosition(int index) {
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds");
+    }
     int delegateOffset = vector.getOffsetBuffer().getInt(index * (long) Integer.BYTES);
     delegate.setPosition(delegateOffset);
     super.setPosition(index);

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroMapProducer.java

@@ -55,6 +55,9 @@ public void produce(Encoder encoder) throws IOException {
 
   @Override
   public void setPosition(int index) {
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds");
+    }
     int delegateOffset = vector.getOffsetBuffer().getInt(index * (long) Integer.BYTES);
     delegate.setPosition(delegateOffset);
     super.setPosition(index);

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullableProducer.java

@@ -58,6 +58,9 @@ public void skipNull() {
 
   @Override
   public void setPosition(int index) {
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds");
+    }
     delegate.setPosition(index);
     super.setPosition(index);
   }

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStructProducer.java

@@ -55,6 +55,9 @@ public void skipNull() {
 
   @Override
   public void setPosition(int index) {
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds: " + index);
+    }
     for (Producer<?> delegate : delegates) {
       delegate.setPosition(index);
     }

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/BaseAvroProducer.java

@@ -44,6 +44,10 @@ public void skipNull() {
 
   @Override
   public void setPosition(int index) {
+    // currentIndex == value is a valid state, no more values will be produced
+    if (index < 0 || index > vector.getValueCount()) {
+      throw new IllegalArgumentException("Index out of bounds");
+    }
     currentIndex = index;
   }
 

adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java

@@ -2294,6 +2294,9 @@ public void testWriteMap() throws Exception {
       root.setRowCount(rowCount);
       root.allocateNew();
 
+      // Total number of entries that will be writen to each vector
+      int entryCount = 5 * 4 * 3;
+
       // Set test data for intList
       BaseWriter.MapWriter writer = intMapVector.getWriter();
       for (int i = 0; i < rowCount; i++) {
@@ -2307,6 +2310,9 @@ public void testWriteMap() throws Exception {
         writer.endMap();
       }
 
+      // Update count for data vector (map writer does not do this)
+      intMapVector.getDataVector().setValueCount(entryCount);
+
       // Set test data for stringList
       BaseWriter.MapWriter stringWriter = stringMapVector.getWriter();
       for (int i = 0; i < rowCount; i++) {
@@ -2320,6 +2326,9 @@ public void testWriteMap() throws Exception {
         stringWriter.endMap();
       }
 
+      // Update count for data vector (map writer does not do this)
+      stringMapVector.getDataVector().setValueCount(entryCount);
+
       // Set test data for dateList
       BaseWriter.MapWriter dateWriter = dateMapVector.getWriter();
       for (int i = 0; i < rowCount; i++) {
@@ -2333,6 +2342,9 @@ public void testWriteMap() throws Exception {
         dateWriter.endMap();
       }
 
+      // Update count for data vector (map writer does not do this)
+      dateMapVector.getDataVector().setValueCount(entryCount);
+
       File dataFile = new File(TMP, "testWriteMap.avro");
 
       // Write an AVRO block using the producer classes
@@ -2431,6 +2443,9 @@ public void testWriteNullableMap() throws Exception {
       writer.endEntry();
       writer.endMap();
 
+      // Update count for data vector (map writer does not do this)
+      nullEntriesVector.getDataVector().setValueCount(3);
+
       // Set test data for stringList
       BaseWriter.MapWriter nullMapWriter = nullMapVector.getWriter();
       nullMapWriter.writeNull();
@@ -2448,6 +2463,9 @@ public void testWriteNullableMap() throws Exception {
       nullMapWriter.endEntry();
       nullMapWriter.endMap();
 
+      // Update count for data vector (map writer does not do this)
+      nullMapVector.getDataVector().setValueCount(2);
+
       // Set test data for dateList
       BaseWriter.MapWriter nullBothWriter = nullBothVector.getWriter();
       nullBothWriter.writeNull();
@@ -2465,6 +2483,9 @@ public void testWriteNullableMap() throws Exception {
       nullBothWriter.endEntry();
       nullBothWriter.endMap();
 
+      // Update count for data vector (map writer does not do this)
+      nullBothVector.getDataVector().setValueCount(2);
+
       File dataFile = new File(TMP, "testWriteNullableMap.avro");
 
       // Write an AVRO block using the producer classes