DX-103340: Allow using 1GB+ data buffers in variable width vectors

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

@@ -17,8 +17,6 @@
 
 package org.apache.arrow.vector;
 
-import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt;
-
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -30,6 +28,7 @@
 import org.apache.arrow.memory.util.ArrowBufPointer;
 import org.apache.arrow.memory.util.ByteFunctionHelpers;
 import org.apache.arrow.memory.util.CommonUtil;
+import static org.apache.arrow.memory.util.LargeMemoryUtil.capAtMaxInt;
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.compare.VectorVisitor;
@@ -568,10 +567,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

@@ -17,17 +17,6 @@
 
 package org.apache.arrow.vector;
 
-import static org.apache.arrow.vector.TestUtils.newVarBinaryVector;
-import static org.apache.arrow.vector.TestUtils.newVarCharVector;
-import static org.apache.arrow.vector.TestUtils.newVector;
-import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
-import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertTrue;
-
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
@@ -42,6 +31,9 @@
 import org.apache.arrow.memory.rounding.DefaultRoundingPolicy;
 import org.apache.arrow.memory.util.ArrowBufPointer;
 import org.apache.arrow.memory.util.CommonUtil;
+import static org.apache.arrow.vector.TestUtils.newVarBinaryVector;
+import static org.apache.arrow.vector.TestUtils.newVarCharVector;
+import static org.apache.arrow.vector.TestUtils.newVector;
 import org.apache.arrow.vector.compare.Range;
 import org.apache.arrow.vector.compare.RangeEqualsVisitor;
 import org.apache.arrow.vector.compare.VectorEqualsVisitor;
@@ -58,6 +50,7 @@
 import org.apache.arrow.vector.holders.NullableVarCharHolder;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.testing.ValueVectorDataPopulator;
+import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.UnionMode;
@@ -70,6 +63,12 @@
 import org.apache.arrow.vector.util.Text;
 import org.apache.arrow.vector.util.TransferPair;
 import org.junit.After;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -92,7 +91,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

@@ -17,12 +17,11 @@
 
 package org.apache.arrow.vector;
 
-import static org.junit.Assert.*;
-
 import java.nio.charset.StandardCharsets;
 
 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;
@@ -39,6 +38,9 @@
 import org.apache.arrow.vector.util.DataSizeRoundingUtil;
 import org.junit.After;
 import org.junit.Assert;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -223,6 +225,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

@@ -17,17 +17,14 @@
 
 package org.apache.arrow.vector.util;
 
-import static junit.framework.TestCase.assertEquals;
-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.Collections;
 import java.util.List;
 
 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;
@@ -56,6 +53,10 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import static junit.framework.TestCase.assertEquals;
+import static junit.framework.TestCase.assertTrue;
 
 /**
  * Test cases for {@link VectorAppender}.
@@ -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);