DX-103340: Allow using 1GB+ data buffers in variable width vectors (#98)

Backport of apache/arrow-java#722

Author: timhurskidremio

java/pom.xml

@@ -478,8 +478,8 @@
               <io.netty.tryReflectionSetAccessible>true</io.netty.tryReflectionSetAccessible>
               <user.timezone>UTC</user.timezone>
               <!-- Note: changing the below configuration might increase the max allocation size for a vector
-              which in turn can cause OOM. -->
-              <arrow.vector.max_allocation_bytes>1048576</arrow.vector.max_allocation_bytes>
+              which in turn can cause OOM. Using 2MB - 1byte to simulate the defaul limit of 2^31 - 1 bytes. -->
+              <arrow.vector.max_allocation_bytes>2097151</arrow.vector.max_allocation_bytes>
             </systemPropertyVariables>
           </configuration>
         </plugin>

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

@@ -568,10 +568,13 @@ public void reallocDataBuffer(long desiredAllocSize) {
       return;
     }
 
-    final long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize);
+    final long newAllocationSize =
+        Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE);
     assert newAllocationSize >= 1;
 
-    checkDataBufferSize(newAllocationSize);
+    if (newAllocationSize < desiredAllocSize) {
+      checkDataBufferSize(desiredAllocSize);
+    }
 
     final ArrowBuf newBuf = allocator.buffer(newAllocationSize);
     newBuf.setBytes(0, valueBuffer, 0, valueBuffer.capacity());

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

@@ -92,7 +92,7 @@ public void init() {
   private static final byte[] STR5 = "EEE5".getBytes(utf8Charset);
   private static final byte[] STR6 = "FFFFF6".getBytes(utf8Charset);
   private static final int MAX_VALUE_COUNT =
-      (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 7);
+      (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 9);
   private static final int MAX_VALUE_COUNT_8BYTE = (int) (MAX_VALUE_COUNT / 2);
 
   @After

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

@@ -23,6 +23,7 @@
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.memory.util.CommonUtil;
 import org.apache.arrow.vector.complex.DenseUnionVector;
 import org.apache.arrow.vector.complex.FixedSizeListVector;
 import org.apache.arrow.vector.complex.ListVector;
@@ -223,6 +224,17 @@ public void testVariableAllocateAfterReAlloc() throws Exception {
     }
   }
 
+  @Test
+  public void testVariableReAllocAbove1GB() throws Exception {
+    try (final VarCharVector vector = new VarCharVector("", allocator)) {
+      long desiredSizeAboveLastPowerOf2 =
+          CommonUtil.nextPowerOfTwo(BaseVariableWidthVector.MAX_ALLOCATION_SIZE) / 2 + 1;
+      vector.reallocDataBuffer(desiredSizeAboveLastPowerOf2);
+
+      assertTrue(vector.getDataBuffer().capacity() >= desiredSizeAboveLastPowerOf2);
+    }
+  }
+
   @Test
   public void testLargeVariableAllocateAfterReAlloc() throws Exception {
     try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) {

java/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java

@@ -28,6 +28,7 @@
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.memory.util.CommonUtil;
 import org.apache.arrow.vector.BaseValueVector;
 import org.apache.arrow.vector.BigIntVector;
 import org.apache.arrow.vector.BitVector;
@@ -191,7 +192,15 @@ public void testAppendEmptyVariableWidthVector() {
 
   @Test
   public void testAppendLargeAndSmallVariableVectorsWithinLimit() {
-    int sixteenthOfMaxAllocation = Math.toIntExact(BaseValueVector.MAX_ALLOCATION_SIZE / 16);
+    // Using the max power of 2 allocation size to avoid hitting the max limit at round ups
+    long maxPowerOfTwoAllocationSize =
+        CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE);
+    if (maxPowerOfTwoAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) {
+      maxPowerOfTwoAllocationSize =
+          CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE / 2);
+    }
+
+    int sixteenthOfMaxAllocation = Math.toIntExact(maxPowerOfTwoAllocationSize / 16);
     try (VarCharVector target = makeVarCharVec(1, sixteenthOfMaxAllocation);
         VarCharVector delta = makeVarCharVec(sixteenthOfMaxAllocation, 1)) {
       new VectorAppender(delta).visit(target, null);