handle overflow for MutableOffHeapByteArrayStore buffer starting si…

…ze (apache#13215)
gortiz · Jun 14, 2024 · b4a6e83 · b4a6e83
1 parent 9670cd0
commit b4a6e83
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/...main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java b/...main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
@@ -170,15 +170,20 @@ public void close()
   private final int _startSize;
 
   @VisibleForTesting
-  public int getStartSize() {
-    return _startSize;
+  public static int getStartSize(int numArrays, int avgArrayLen) {
+    // For each array, we store the array and its startoffset (4 bytes)
+    long estimatedSize = numArrays * ((long) avgArrayLen + 4);
+    if (estimatedSize > 0 && estimatedSize <= Integer.MAX_VALUE) {
+      return (int) estimatedSize;
+    }
+    return Integer.MAX_VALUE;
   }
 
   public MutableOffHeapByteArrayStore(PinotDataBufferMemoryManager memoryManager, String allocationContext,
       int numArrays, int avgArrayLen) {
     _memoryManager = memoryManager;
     _allocationContext = allocationContext;
-    _startSize = numArrays * (avgArrayLen + 4); // For each array, we store the array and its startoffset (4 bytes)
+    _startSize = getStartSize(numArrays, avgArrayLen);
     expand(_startSize);
   }
 

diff --git a/.../java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java b/.../java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
@@ -29,6 +29,7 @@
 public class MutableOffHeapByteArrayStoreTest {
 
   private PinotDataBufferMemoryManager _memoryManager;
+  private static final int ONE_GB = 1024 * 1024 * 1024;
 
   @BeforeClass
   public void setUp() {
@@ -44,8 +45,11 @@ public void tearDown()
   @Test
   public void maxValueTest()
       throws Exception {
-    MutableOffHeapByteArrayStore store = new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", 1024, 32);
-    final int arrSize = store.getStartSize();
+    int numArrays = 1024;
+    int avgArrayLen = 32;
+    MutableOffHeapByteArrayStore store =
+        new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen);
+    final int arrSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen);
     byte[] dataIn = new byte[arrSize - 4];
     for (int i = 0; i < dataIn.length; i++) {
       dataIn[i] = (byte) (i % Byte.MAX_VALUE);
@@ -56,11 +60,21 @@ public void maxValueTest()
     store.close();
   }
 
+  @Test
+  public void startSizeTest() {
+    Assert.assertEquals(MutableOffHeapByteArrayStore.getStartSize(1, ONE_GB), ONE_GB + 4);
+    Assert.assertEquals(MutableOffHeapByteArrayStore.getStartSize(3, ONE_GB), Integer.MAX_VALUE);
+    Assert.assertEquals(MutableOffHeapByteArrayStore.getStartSize(5, ONE_GB), Integer.MAX_VALUE);
+  }
+
   @Test
   public void overflowTest()
       throws Exception {
-    MutableOffHeapByteArrayStore store = new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", 1024, 32);
-    final int maxSize = store.getStartSize() - 4;
+    int numArrays = 1024;
+    int avgArrayLen = 32;
+    MutableOffHeapByteArrayStore store =
+        new MutableOffHeapByteArrayStore(_memoryManager, "stringColumn", numArrays, avgArrayLen);
+    final int maxSize = MutableOffHeapByteArrayStore.getStartSize(numArrays, avgArrayLen) - 4;
 
     byte[] b1 = new byte[3];
     for (int i = 0; i < b1.length; i++) {