From 9b75bffcb2cd9373022b2ff0ac63327803fc79b4 Mon Sep 17 00:00:00 2001
From: Christopher Peck <27231838+itschrispeck@users.noreply.github.com>
Date: Fri, 7 Jun 2024 17:41:35 -0700
Subject: [PATCH] handle overflow for `MutableOffHeapByteArrayStore` buffer
 starting size (#13215)

---
 .../impl/MutableOffHeapByteArrayStore.java    | 11 +++++++---
 .../MutableOffHeapByteArrayStoreTest.java     | 22 +++++++++++++++----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
index 54323e8fc3c6..8ddd6d6a87ef 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
+++ b/pinot-segment-local/src/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/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
index 049bad01ef1f..f74757943144 100644
--- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
+++ b/pinot-segment-local/src/test/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++) {