apacheGH-37829: [Java] Avoid resizing data buffer twice when appendin…

…g variable length vectors
hrishisd · Sep 24, 2023 · 785c228 · 785c228
1 parent e83c23b
commit 785c228
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 4 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java b/java/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java
@@ -116,7 +116,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) {
 
     // make sure there is enough capacity
     while (targetVector.getValueCapacity() < newValueCount) {
-      targetVector.reAlloc();
+      ((BaseVariableWidthVector) targetVector).reallocValidityAndOffsetBuffers();
     }
     while (targetVector.getDataBuffer().capacity() < newValueCapacity) {
       ((BaseVariableWidthVector) targetVector).reallocDataBuffer();
@@ -170,7 +170,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) {
 
     // make sure there is enough capacity
     while (targetVector.getValueCapacity() < newValueCount) {
-      targetVector.reAlloc();
+      ((BaseLargeVariableWidthVector) targetVector).reallocValidityAndOffsetBuffers();
     }
     while (targetVector.getDataBuffer().capacity() < newValueCapacity) {
       ((BaseLargeVariableWidthVector) targetVector).reallocDataBuffer();

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java b/java/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java
@@ -21,11 +21,13 @@
 import static junit.framework.TestCase.assertTrue;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BaseValueVector;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.BitVector;
 import org.apache.arrow.vector.Float4Vector;
@@ -63,7 +65,8 @@ public class TestVectorAppender {
 
   @Before
   public void prepare() {
-    allocator = new RootAllocator(1024 * 1024);
+    // Permit allocating 4 vectors of max size.
+    allocator = new RootAllocator(4 * BaseValueVector.MAX_ALLOCATION_SIZE);
   }
 
   @After
@@ -185,6 +188,27 @@ public void testAppendEmptyVariableWidthVector() {
     }
   }
 
+  @Test
+  public void testAppendLargeAndSmallVariableVectorsWithinLimit() {
+    int sixteenthOfMaxAllocation = Math.toIntExact(BaseValueVector.MAX_ALLOCATION_SIZE / 16);
+    try (VarCharVector target = makeVarCharVec(1, sixteenthOfMaxAllocation);
+        VarCharVector delta = makeVarCharVec(sixteenthOfMaxAllocation, 1)
+    ) {
+      new VectorAppender(delta).visit(target, null);
+      new VectorAppender(target).visit(delta, null);
+    }
+  }
+
+  private VarCharVector makeVarCharVec(int numElements, int bytesPerElement) {
+    var v = new VarCharVector("text", allocator);
+    v.allocateNew((long) numElements * bytesPerElement, numElements);
+    for (int i = 0; i < numElements; i++) {
+      v.setSafe(i, "a".repeat(bytesPerElement).getBytes(StandardCharsets.UTF_8));
+    }
+    v.setValueCount(numElements);
+    return v;
+  }
+
   @Test
   public void testAppendLargeVariableWidthVector() {
     final int length1 = 5;

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/util/TestVectorSchemaRootAppender.java b/java/vector/src/test/java/org/apache/arrow/vector/util/TestVectorSchemaRootAppender.java
@@ -50,7 +50,7 @@ public void shutdown() {
   }
 
   @Test
-  public void testVectorScehmaRootAppend() {
+  public void testVectorSchemaRootAppend() {
     final int length1 = 5;
     final int length2 = 3;
     final int length3 = 2;