(pinot) branch master updated: handle overflow for `MutableOffHeapByteArrayStore` buffer starting size (#13215)

[jackie]

This is an automated email from the ASF dual-hosted git repository.

jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new 9b75bffcb2 handle overflow for `MutableOffHeapByteArrayStore` buffer 
starting size (#13215)
9b75bffcb2 is described below

commit 9b75bffcb2cd9373022b2ff0ac63327803fc79b4
Author: Christopher Peck <27231838+itschrisp...@users.noreply.github.com>
AuthorDate: Fri Jun 7 17:41:35 2024 -0700

    handle overflow for `MutableOffHeapByteArrayStore` buffer starting size 
(#13215)
---
 .../writer/impl/MutableOffHeapByteArrayStore.java  | 11 ++++++++---
 .../impl/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 54323e8fc3..8ddd6d6a87 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 class MutableOffHeapByteArrayStore implements 
Closeable {
   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 049bad01ef..f747579431 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 @@ import org.testng.annotations.Test;
 public class MutableOffHeapByteArrayStoreTest {
 
   private PinotDataBufferMemoryManager _memoryManager;
+  private static final int ONE_GB = 1024 * 1024 * 1024;
 
   @BeforeClass
   public void setUp() {
@@ -44,8 +45,11 @@ public class MutableOffHeapByteArrayStoreTest {
   @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 class MutableOffHeapByteArrayStoreTest {
     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++) {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org