From 12e2228510fa9bccb9321f7ac987b0408e0ac29c Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Wed, 24 Nov 2021 13:14:58 -0800
Subject: [PATCH] RowBasedGrouperHelper: Set hasMultipleValues = false in
 capabilities. (#11954)

Useful because it enables anything that consumes groupBy results to
potentially operate more efficiently.
---
 .../epinephelinae/RowBasedGrouperHelper.java  | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/RowBasedGrouperHelper.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/RowBasedGrouperHelper.java
index 1b50df77d9d..b4f583668fe 100644
--- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/RowBasedGrouperHelper.java
+++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/RowBasedGrouperHelper.java
@@ -60,6 +60,7 @@ import org.apache.druid.query.ordering.StringComparators;
 import org.apache.druid.segment.BaseDoubleColumnValueSelector;
 import org.apache.druid.segment.BaseFloatColumnValueSelector;
 import org.apache.druid.segment.BaseLongColumnValueSelector;
+import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionHandlerUtils;
@@ -67,6 +68,7 @@ import org.apache.druid.segment.DimensionSelector;
 import org.apache.druid.segment.RowAdapter;
 import org.apache.druid.segment.RowBasedColumnSelectorFactory;
 import org.apache.druid.segment.column.ColumnCapabilities;
+import org.apache.druid.segment.column.ColumnCapabilitiesImpl;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.column.ValueType;
@@ -381,10 +383,27 @@ public class RowBasedGrouperHelper
           }
         };
 
+    // Decorate "signature" so that it returns hasMultipleValues = false. (groupBy does not return multiple values.)
+    final ColumnInspector decoratedSignature = new ColumnInspector()
+    {
+      @Nullable
+      @Override
+      public ColumnCapabilities getColumnCapabilities(String column)
+      {
+        final ColumnCapabilities baseCapabilities = signature.getColumnCapabilities(column);
+
+        if (baseCapabilities == null || baseCapabilities.hasMultipleValues().isFalse()) {
+          return baseCapabilities;
+        } else {
+          return ColumnCapabilitiesImpl.copyOf(baseCapabilities).setHasMultipleValues(false);
+        }
+      }
+    };
+
     return RowBasedColumnSelectorFactory.create(
         adapter,
         supplier::get,
-        () -> signature,
+        () -> decoratedSignature,
         false
     );
   }