From e157fb089a8f1863b49ea14841d7653582a60117 Mon Sep 17 00:00:00 2001
From: Jihoon Son <jihoonson@apache.org>
Date: Fri, 10 Apr 2020 09:05:38 -0700
Subject: [PATCH] Fix wrong cardinality computation in BufferArrayGrouper
 (#9655)

* Fix wrong cardinality computation in BufferArrayGrouper

* fix javadoc
---
 .../epinephelinae/BufferArrayGrouper.java        | 16 ++++++++++++----
 .../epinephelinae/GroupByQueryEngineV2.java      |  6 +++++-
 .../epinephelinae/BufferArrayGrouperTest.java    | 11 ++++-------
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouper.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouper.java
index 9b912e40560..864b6a3b16e 100644
--- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouper.java
+++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouper.java
@@ -69,12 +69,20 @@ public class BufferArrayGrouper implements VectorGrouper, IntGrouper
   @Nullable
   private int[] vAggregationRows = null;
 
-  static long requiredBufferCapacity(
-      int cardinality,
-      AggregatorFactory[] aggregatorFactories
-  )
+  /**
+   * Computes required buffer capacity for a grouping key of the given cardinaltiy and aggregatorFactories.
+   * This method assumes that the given cardinality doesn't count nulls.
+   *
+   * Returns -1 if cardinality + 1 (for null) > Integer.MAX_VALUE. Returns computed required buffer capacity
+   * otherwise.
+   */
+  static long requiredBufferCapacity(int cardinality, AggregatorFactory[] aggregatorFactories)
   {
     final long cardinalityWithMissingValue = computeCardinalityWithMissingValue(cardinality);
+    // Cardinality should be in the integer range. See DimensionDictionarySelector.
+    if (cardinalityWithMissingValue > Integer.MAX_VALUE) {
+      return -1;
+    }
     final long recordSize = Arrays.stream(aggregatorFactories)
                                   .mapToLong(AggregatorFactory::getMaxIntermediateSizeWithNulls)
                                   .sum();
diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java
index fd3213fe9a8..8e3fe05ae7e 100644
--- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java
+++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java
@@ -300,7 +300,11 @@ public class GroupByQueryEngineV2
       );
 
       // Check that all keys and aggregated values can be contained in the buffer
-      return requiredBufferCapacity <= buffer.capacity() ? cardinality : -1;
+      if (requiredBufferCapacity < 0 || requiredBufferCapacity > buffer.capacity()) {
+        return -1;
+      } else {
+        return cardinality;
+      }
     } else {
       return -1;
     }
diff --git a/processing/src/test/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouperTest.java b/processing/src/test/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouperTest.java
index c05d97632af..d61700a5c26 100644
--- a/processing/src/test/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouperTest.java
+++ b/processing/src/test/java/org/apache/druid/query/groupby/epinephelinae/BufferArrayGrouperTest.java
@@ -110,20 +110,17 @@ public class BufferArrayGrouperTest
     final long[] requiredSizes;
 
     if (NullHandling.sqlCompatible()) {
-      // We need additional size to store nullability information.
-      requiredSizes = new long[]{19, 101, 19595788279L, 19595788288L};
+      // We need additional space to store nulls.
+      requiredSizes = new long[]{19, 101, 19595788279L, -1};
     } else {
-      requiredSizes = new long[]{17, 90, 17448304632L, 17448304640L};
+      requiredSizes = new long[]{17, 90, 17448304632L, -1};
     }
 
     for (int i = 0; i < cardinalityArray.length; i++) {
       Assert.assertEquals(
           StringUtils.format("cardinality[%d]", cardinalityArray[i]),
           requiredSizes[i],
-          BufferArrayGrouper.requiredBufferCapacity(
-              cardinalityArray[i],
-              aggregatorFactories
-          )
+          BufferArrayGrouper.requiredBufferCapacity(cardinalityArray[i], aggregatorFactories)
       );
     }
   }