diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java index a7cb9ced7..4c1fbf761 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java +++ b/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 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); } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java index c121d9489..bd3113f8b 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/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);