GH-301: Allow adding a vector at the end of VectorSchemaRoot

Author: axreldable

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

@@ -199,13 +199,18 @@ public FieldVector getVector(int index) {
    */
   public VectorSchemaRoot addVector(int index, FieldVector vector) {
     Preconditions.checkNotNull(vector);
-    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
+    Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size());
     List<FieldVector> newVectors = new ArrayList<>();
-    for (int i = 0; i < fieldVectors.size(); i++) {
-      if (i == index) {
-        newVectors.add(vector);
+    if (index == fieldVectors.size()) {
+      newVectors.addAll(fieldVectors);
+      newVectors.add(vector);
+    } else {
+      for (int i = 0; i < fieldVectors.size(); i++) {
+        if (i == index) {
+          newVectors.add(vector);
+        }
+        newVectors.add(fieldVectors.get(i));
       }
-      newVectors.add(fieldVectors.get(i));
     }
     return new VectorSchemaRoot(newVectors);
   }

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

@@ -171,6 +171,26 @@ public void testAddVector() {
     }
   }
 
+  @Test
+  public void testAddVectorAtEnd() {
+    try (final IntVector intVector1 = new IntVector("intVector1", allocator);
+        final IntVector intVector2 = new IntVector("intVector2", allocator);
+        final IntVector intVector3 = new IntVector("intVector3", allocator); ) {
+
+      VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2));
+      assertEquals(2, original.getFieldVectors().size());
+
+      VectorSchemaRoot newRecordBatch = original.addVector(2, intVector3);
+      assertEquals(3, newRecordBatch.getFieldVectors().size());
+      assertEquals(intVector1, newRecordBatch.getFieldVectors().get(0));
+      assertEquals(intVector2, newRecordBatch.getFieldVectors().get(1));
+      assertEquals(intVector3, newRecordBatch.getFieldVectors().get(2));
+
+      original.close();
+      newRecordBatch.close();
+    }
+  }
+
   @Test
   public void testRemoveVector() {
     try (final IntVector intVector1 = new IntVector("intVector1", allocator);