ArrayAggregation: Use long to avoid overflow (#5544)

* ArrayAggregation: Use long to avoid overflow

* Add Tests

processing/src/main/java/io/druid/query/groupby/epinephelinae/BufferArrayGrouper.java

@@ -61,7 +61,7 @@ public class BufferArrayGrouper implements IntGrouper
   private ByteBuffer usedFlagBuffer;
   private ByteBuffer valBuffer;
 
-  static int requiredBufferCapacity(
+  static long requiredBufferCapacity(
       int cardinality,
       AggregatorFactory[] aggregatorFactories
   )
@@ -72,7 +72,7 @@ public class BufferArrayGrouper implements IntGrouper
                                  .sum();
 
     return getUsedFlagBufferCapacity(cardinalityWithMissingValue) +  // total used flags size
-           cardinalityWithMissingValue * recordSize;                 // total values size
+        (long) cardinalityWithMissingValue * recordSize;                 // total values size
   }
 
   /**

processing/src/main/java/io/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java

@@ -219,7 +219,7 @@ public class GroupByQueryEngineV2
       final AggregatorFactory[] aggregatorFactories = query
           .getAggregatorSpecs()
           .toArray(new AggregatorFactory[query.getAggregatorSpecs().size()]);
-      final int requiredBufferCapacity = BufferArrayGrouper.requiredBufferCapacity(
+      final long requiredBufferCapacity = BufferArrayGrouper.requiredBufferCapacity(
           cardinality,
           aggregatorFactories
       );

processing/src/test/java/io/druid/query/groupby/epinephelinae/BufferArrayGrouperTest.java

@@ -86,4 +86,21 @@ public class BufferArrayGrouperTest
     grouper.init();
     return grouper;
   }
+
+  @Test
+  public void testRequiredBufferCapacity()
+  {
+    int[] cardinalityArray = new int[] {1, 10, Integer.MAX_VALUE - 1};
+    AggregatorFactory[] aggregatorFactories = new AggregatorFactory[] {
+        new LongSumAggregatorFactory("sum", "sum")
+    };
+
+    long[] requiredSizes = new long[] {17, 90, 16911433721L};
+
+    for (int i = 0; i < cardinalityArray.length; i++) {
+      Assert.assertEquals(requiredSizes[i], BufferArrayGrouper.requiredBufferCapacity(
+          cardinalityArray[i],
+          aggregatorFactories));
+    }
+  }
 }