Add size check in DictionaryBuildingSingleValuedRowBasedKeySerdeHelper in putToKeyBuffer

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/RowBasedGrouperHelper.java

@@ -158,13 +158,13 @@
   /**
    * Create a {@link Grouper} that groups according to the dimensions and aggregators in "query", along with
    * an {@link Accumulator} that accepts ResultRows and forwards them to the grouper.
-   *
+   * <p>
    * The pair will operate in one of two modes:
-   *
+   * <p>
    * 1) Combining mode (used if "subquery" is null). In this mode, filters from the "query" are ignored, and
    * its aggregators are converted into combining form. The input ResultRows are assumed to be partially-grouped
    * results originating from the provided "query".
-   *
+   * <p>
    * 2) Subquery mode (used if "subquery" is nonnull). In this mode, filters from the "query" (both intervals
    * and dim filters) are respected, and its aggregators are used in standard (not combining) form. The input
    * ResultRows are assumed to be results originating from the provided "subquery".
@@ -758,7 +758,10 @@
         case COMPLEX:
           return (InputRawSupplierColumnSelectorStrategy<ColumnValueSelector>)
               columnSelector ->
-                  () -> DimensionHandlerUtils.convertObjectToType(columnSelector.getObject(), capabilities.toColumnType());
+                  () -> DimensionHandlerUtils.convertObjectToType(
+                      columnSelector.getObject(),
+                      capabilities.toColumnType()
+                  );
         default:
           throw new IAE("Cannot create query type helper from invalid type [%s]", capabilities.asTypeString());
       }
@@ -1619,6 +1622,11 @@
         final Object obj = key.getKey()[idx];
         int id = getReverseDictionary().getInt(obj);
         if (id == DimensionDictionary.ABSENT_VALUE_ID) {
+          int size = estimatedKeySize(key);
+          if (currentEstimatedSize + size > maxDictionarySize) {
+            return false;
+          }
+          currentEstimatedSize += size;
           id = getDictionary().size();
           getReverseDictionary().put(obj, id);
           getDictionary().add(obj);
@@ -1627,6 +1635,8 @@
         return true;
       }
 
+      abstract int estimatedKeySize(RowBasedKey key);
+
       @Override
       public void getFromByteBuffer(ByteBuffer buffer, int initialOffset, int dimValIdx, Object[] dimValues)
       {
@@ -1692,6 +1702,23 @@
         }
       }
 
+      @Override
+      int estimatedKeySize(RowBasedKey key)
+      {
+        Object[] obj = key.getKey();
+        int size = 0;
+        for (Object o : obj) {
+          if (o instanceof String) {
+            size += estimateStringKeySize((String) o);
+          } else if (o != null) {
+            size += estimateStringKeySize(String.valueOf(o));
+          } else {
+            size += DictionaryBuildingUtils.estimateEntryFootprint(0);
+          }
+        }
+        return size;
+      }
+
       @Override
       public BufferComparator getBufferComparator()
       {
@@ -1720,6 +1747,7 @@
 
     private class ArrayNumericRowBasedKeySerdeHelper extends DictionaryBuildingSingleValuedRowBasedKeySerdeHelper
     {
+      private final int elementSize;
       private final BufferComparator bufferComparator;
       private final List<Object[]> dictionary;
       private final Object2IntMap<Object[]> reverseDictionary;
@@ -1732,6 +1760,19 @@
       {
         super(keyBufferPosition);
         final TypeSignature<ValueType> elementType = arrayType.getElementType();
+        switch (elementType.getType()) {
+          case LONG:
+            elementSize = Long.BYTES;
+            break;
+          case FLOAT:
+            elementSize = Float.BYTES;
+            break;
+          case DOUBLE:
+            elementSize = Double.BYTES;
+            break;
+          default:
+            throw DruidException.defensive("Expecting primitive numeric types");
+        }
         this.dictionary = getDictionaryForType(elementType);
         this.reverseDictionary = getReverseDictionaryForType(elementType);
         this.bufferComparator = (lhsBuffer, rhsBuffer, lhsPosition, rhsPosition) -> {
@@ -1779,6 +1820,13 @@
         }
       }
 
+      @Override
+      int estimatedKeySize(RowBasedKey key)
+      {
+        Object[] numericKey = key.getKey();
+        return numericKey.length * elementSize;
+      }
+
       @Override
       public BufferComparator getBufferComparator()
       {
@@ -1829,6 +1877,17 @@
             );
       }
 
+      @Override
+      int estimatedKeySize(RowBasedKey key)
+      {
+        Object[] stringKey = key.getKey();
+        int size = 0;
+        for (Object obj : stringKey) {
+          size += (int) estimateStringKeySize((String) obj);
+        }
+        return size;
+      }
+
       @Override
       public BufferComparator getBufferComparator()
       {
@@ -1936,7 +1995,6 @@
        * this returns -1.
        *
        * @param s a string
-       *
        * @return id for this string, or -1
        */
       private int addToDictionary(final String s)