From f2f8e38f9424bcacf3946197c9afcd50ef1d58fa Mon Sep 17 00:00:00 2001
From: Richard Startin <richard@startree.ai>
Date: Tue, 4 Jan 2022 17:14:05 +0000
Subject: [PATCH] unpeel group by 3 ways to enable vectorization (#7949)

---
 .../DictionaryBasedGroupKeyGenerator.java     | 49 ++++++++++++++++---
 1 file changed, 41 insertions(+), 8 deletions(-)

diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
index 1163759f9bbd..f3e868bda691 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
@@ -257,19 +257,52 @@ private class ArrayBasedHolder implements RawKeyHolder {
 
     @Override
     public void processSingleValue(int numDocs, int[] outGroupIds) {
-      if (_numGroupByExpressions == 1) {
-        processSingleValue(numDocs, _singleValueDictIds[0], outGroupIds);
-      } else {
-        processSingleValueGeneric(numDocs, outGroupIds);
+      switch (_numGroupByExpressions) {
+        case 1:
+          processSingleValue(numDocs, _singleValueDictIds[0], outGroupIds);
+          return;
+        case 2:
+          processSingleValue(numDocs, _singleValueDictIds[0], _singleValueDictIds[1], outGroupIds);
+          return;
+        case 3:
+          processSingleValue(numDocs, _singleValueDictIds[0], _singleValueDictIds[1], _singleValueDictIds[2],
+              outGroupIds);
+          return;
+        default:
       }
+      processSingleValueGeneric(numDocs, outGroupIds);
     }
 
     private void processSingleValue(int numDocs, int[] dictIds, int[] outGroupIds) {
       System.arraycopy(dictIds, 0, outGroupIds, 0, numDocs);
-      for (int i = 0; i < numDocs && _numKeys < _globalGroupIdUpperBound; i++) {
-        if (!_flags[outGroupIds[i]]) {
-          _numKeys++;
-          _flags[outGroupIds[i]] = true;
+      markGroups(numDocs, outGroupIds);
+    }
+
+    private void processSingleValue(int numDocs, int[] dictIds0, int[] dictIds1, int[] outGroupIds) {
+      for (int i = 0; i < numDocs; i++) {
+        outGroupIds[i] = dictIds1[i] * _cardinalities[0] + dictIds0[i];
+      }
+      markGroups(numDocs, outGroupIds);
+    }
+
+    private void processSingleValue(int numDocs, int[] dictIds0, int[] dictIds1, int[] dictIds2, int[] outGroupIds) {
+      int cardinality = _cardinalities[0] * _cardinalities[1];
+      for (int i = 0; i < numDocs; i++) {
+        outGroupIds[i] = dictIds2[i] * cardinality + dictIds1[i] * _cardinalities[0] + dictIds0[i];
+      }
+      markGroups(numDocs, outGroupIds);
+    }
+
+    private void markGroups(int numDocs, int[] groupIds) {
+      if (_numKeys < _globalGroupIdUpperBound) {
+        for (int i = 0; i < numDocs; i++) {
+          if (!_flags[groupIds[i]]) {
+            _numKeys++;
+            _flags[groupIds[i]] = true;
+            if (_numKeys == _globalGroupIdUpperBound) {
+              return;
+            }
+          }
         }
       }
     }