Startree aggregate on mv columns

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AvgMVAggregationFunction.java

@@ -22,10 +22,12 @@
 import java.util.Map;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
 import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
 import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
 import org.apache.pinot.segment.local.customobject.AvgPair;
 import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
 
 
 public class AvgMVAggregationFunction extends AvgAggregationFunction {
@@ -43,8 +45,24 @@ public AggregationFunctionType getType() {
   public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
-    double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
+    if (blockValSet.getValueType() == DataType.BYTES) {
+      // Serialized AvgPair
+      byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      AvgPair avgPair = new AvgPair();
+      forEachNotNull(length, blockValSet, (from, to) -> {
+        for (int i = from; i < to; i++) {
+          AvgPair value = ObjectSerDeUtils.AVG_PAIR_SER_DE.deserialize(bytesValues[i]);
+          avgPair.apply(value);
+        }
+      });
+      if (avgPair.getCount() != 0) {
+        updateAggregationResult(aggregationResultHolder, avgPair.getSum(), avgPair.getCount());
+      }
+      return;
+    }
+
+    double[][] valuesArray = blockValSet.getDoubleValuesMV();
     AvgPair avgPair = new AvgPair();
     forEachNotNull(length, blockValSet, (from, to) -> {
       for (int i = from; i < to; i++) {
@@ -53,7 +71,6 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
         }
       }
     });
-
     if (avgPair.getCount() != 0) {
       updateAggregationResult(aggregationResultHolder, avgPair.getSum(), avgPair.getCount());
     }
@@ -63,21 +80,32 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
   public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
-    double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
-    forEachNotNull(length, blockValSet, (from, to) -> {
-      for (int i = from; i < to; i++) {
-        aggregateOnGroupKey(groupKeyArray[i], groupByResultHolder, valuesArray[i]);
-      }
-    });
+    if (blockValSet.getValueType() != DataType.BYTES) {
+      double[][] valuesArray = blockValSet.getDoubleValuesMV();
+      forEachNotNull(length, blockValSet, (from, to) -> {
+        for (int i = from; i < to; i++) {
+          aggregateOnGroupKey(groupKeyArray[i], groupByResultHolder, valuesArray[i]);
+        }
+      });
+    } else {
+      // Serialized AvgPair.
+      byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      forEachNotNull(length, blockValSet, (from, to) -> {
+        for (int i = from; i < to; i++) {
+          AvgPair avgPair = ObjectSerDeUtils.AVG_PAIR_SER_DE.deserialize(bytesValues[i]);
+          updateGroupByResult(groupKeyArray[i], groupByResultHolder, avgPair.getSum(), avgPair.getCount());
+        }
+      });
+    }
   }
 
   @Override
   public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
-    double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
+    double[][] valuesArray = blockValSet.getDoubleValuesMV();
     forEachNotNull(length, blockValSet, (from, to) -> {
       for (int i = from; i < to; i++) {
         double[] values = valuesArray[i];

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CountMVAggregationFunction.java

@@ -53,6 +53,17 @@ public List<ExpressionContext> getInputExpressions() {
   public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
+    if (blockValSet.isSingleValue()) {
+      long[] valueArray = blockValSet.getLongValuesSV();
+      long count = 0;
+      for (int i = 0; i < length; i++) {
+        count += valueArray[i];
+      }
+      aggregationResultHolder.setValue(aggregationResultHolder.getDoubleResult() + count);
+      return;
+    }
+
     int[] valueArray = blockValSet.getNumMVEntries();
     // Hack to make count effectively final for use in the lambda (we know that there aren't concurrent access issues
     // with forEachNotNull)
@@ -69,6 +80,16 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
   public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
+    if (blockValSet.isSingleValue()) {
+      long[] valueArray = blockValSet.getLongValuesSV();
+      for (int i = 0; i < length; i++) {
+        int groupKey = groupKeyArray[i];
+        groupByResultHolder.setValueForKey(groupKey, groupByResultHolder.getDoubleResult(groupKey) + valueArray[i]);
+      }
+      return;
+    }
+
     int[] valueArray = blockValSet.getNumMVEntries();
     forEachNotNull(length, blockValSet, (from, to) -> {
       for (int i = from; i < to; i++) {
@@ -82,6 +103,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol
   public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
     int[] valueArray = blockValSet.getNumMVEntries();
     forEachNotNull(length, blockValSet, (from, to) -> {
       for (int i = from; i < to; i++) {

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/SumMVAggregationFunction.java

@@ -42,6 +42,18 @@ public AggregationFunctionType getType() {
   public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
+    if (blockValSet.isSingleValue()) {
+      // StarTree pre-aggregated values
+      double[] valueArray = blockValSet.getDoubleValuesSV();
+      double sum = 0.0;
+      for (int i = 0; i < length; i++) {
+        sum += valueArray[i];
+      }
+      updateAggregationResultHolder(aggregationResultHolder, sum);
+      return;
+    }
+
     double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
     Double sum;
@@ -62,6 +74,17 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
   public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
+    if (blockValSet.isSingleValue()) {
+      // StarTree pre-aggregated values
+      double[] valueArray = blockValSet.getDoubleValuesSV();
+      for (int i = 0; i < length; i++) {
+        int groupKey = groupKeyArray[i];
+        groupByResultHolder.setValueForKey(groupKey, groupByResultHolder.getDoubleResult(groupKey) + valueArray[i]);
+      }
+      return;
+    }
+
     double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
     if (_nullHandlingEnabled) {
@@ -95,6 +118,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol
   public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
       Map<ExpressionContext, BlockValSet> blockValSetMap) {
     BlockValSet blockValSet = blockValSetMap.get(_expression);
+
     double[][] valuesArray = blockValSet.getDoubleValuesMV();
 
     if (_nullHandlingEnabled) {

pinot-core/src/test/java/org/apache/pinot/core/startree/v2/AvgMVStarTreeV2Test.java

@@ -0,0 +1,57 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.startree.v2;
+
+import java.util.Random;
+import org.apache.pinot.segment.local.aggregator.AvgMVValueAggregator;
+import org.apache.pinot.segment.local.aggregator.ValueAggregator;
+import org.apache.pinot.segment.local.customobject.AvgPair;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+
+import static org.testng.Assert.assertEquals;
+
+
+public class AvgMVStarTreeV2Test extends BaseStarTreeV2Test<Object, AvgPair> {
+
+  @Override
+  ValueAggregator<Object, AvgPair> getValueAggregator() {
+    return new AvgMVValueAggregator();
+  }
+
+  @Override
+  DataType getRawValueType() {
+    return DataType.INT;
+  }
+
+  @Override
+  Object getRandomRawValue(Random random) {
+    return random.nextInt();
+  }
+
+  @Override
+  protected void assertAggregatedValue(AvgPair starTreeResult, AvgPair nonStarTreeResult) {
+    assertEquals(starTreeResult.getSum(), nonStarTreeResult.getSum(), 1e-5);
+    assertEquals(starTreeResult.getCount(), nonStarTreeResult.getCount());
+  }
+
+  @Override
+  protected boolean isAggColSingleValueField() {
+    return false;
+  }
+}

pinot-core/src/test/java/org/apache/pinot/core/startree/v2/BaseStarTreeV2Test.java

@@ -98,7 +98,7 @@ abstract class BaseStarTreeV2Test<R, A> {
   private static final String DIMENSION1 = "d1__COLUMN_NAME";
   private static final String DIMENSION2 = "DISTINCTCOUNTRAWHLL__d2";
   private static final int DIMENSION_CARDINALITY = 100;
-  private static final String METRIC = "m";
+  private static final String AGG_COL = "m";
 
   // Supported filters
   private static final String QUERY_FILTER_AND = String.format(" WHERE %1$s = 0 AND %2$s < 10", DIMENSION1, DIMENSION2);
@@ -151,7 +151,8 @@ public void setUp()
     DataType rawValueType = getRawValueType();
     // Raw value type will be null for COUNT aggregation function
     if (rawValueType != null) {
-      schemaBuilder.addMetric(METRIC, rawValueType);
+      schemaBuilder.addDimensionField(AGG_COL, rawValueType,
+              dimensionFieldSpec -> dimensionFieldSpec.setSingleValueField(isAggColSingleValueField()));
     }
     Schema schema = schemaBuilder.build();
     TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build();
@@ -162,7 +163,7 @@ public void setUp()
       segmentRecord.putValue(DIMENSION1, RANDOM.nextInt(DIMENSION_CARDINALITY));
       segmentRecord.putValue(DIMENSION2, RANDOM.nextInt(DIMENSION_CARDINALITY));
       if (rawValueType != null) {
-        segmentRecord.putValue(METRIC, getRandomRawValue(RANDOM));
+        segmentRecord.putValue(AGG_COL, getRandomRawValue(RANDOM));
       }
       segmentRecords.add(segmentRecord);
     }
@@ -175,8 +176,9 @@ public void setUp()
     driver.build();
 
     StarTreeIndexConfig starTreeIndexConfig = new StarTreeIndexConfig(Arrays.asList(DIMENSION1, DIMENSION2), null, null,
-        Collections.singletonList(new StarTreeAggregationConfig(METRIC, _valueAggregator.getAggregationType().getName(),
-            null, getCompressionCodec(), true, getIndexVersion(), null, null)), MAX_LEAF_RECORDS);
+        Collections.singletonList(new StarTreeAggregationConfig(AGG_COL,
+                _valueAggregator.getAggregationType().getName(), null, getCompressionCodec(),
+                true, getIndexVersion(), null, null)), MAX_LEAF_RECORDS);
     File indexDir = new File(TEMP_DIR, SEGMENT_NAME);
     // Randomly build star-tree using on-heap or off-heap mode
     MultipleTreesBuilder.BuildMode buildMode =
@@ -196,9 +198,9 @@ String getAggregation(AggregationFunctionType aggregationType) {
     } else if (aggregationType == AggregationFunctionType.PERCENTILEEST
         || aggregationType == AggregationFunctionType.PERCENTILETDIGEST) {
       // Append a percentile number for percentile functions
-      return String.format("%s(%s, 50)", aggregationType.getName(), METRIC);
+      return String.format("%s(%s, 50)", aggregationType.getName(), AGG_COL);
     } else {
-      return String.format("%s(%s)", aggregationType.getName(), METRIC);
+      return String.format("%s(%s)", aggregationType.getName(), AGG_COL);
     }
   }
 
@@ -476,7 +478,16 @@ private Map<List<Integer>, List<Object>> computeNonStarTreeResult(FilterPlanNode
 
   private Object getNextRawValue(int docId, ForwardIndexReader reader, ForwardIndexReaderContext readerContext,
       Dictionary dictionary) {
-    return dictionary.get(reader.getDictId(docId, readerContext));
+    if (isAggColSingleValueField()) {
+      return dictionary.get(reader.getDictId(docId, readerContext));
+    } else {
+      int[] dictIds = reader.getDictIdMV(docId, readerContext);
+        Object[] rawValue = new Object[dictIds.length];
+        for (int i = 0; i < dictIds.length; i++) {
+          rawValue[i] = dictionary.get(dictIds[i]);
+        }
+        return rawValue;
+    }
   }
 
   /**
@@ -499,6 +510,10 @@ Integer getIndexVersion() {
     return version > 1 ? version : null;
   }
 
+  protected boolean isAggColSingleValueField() {
+    return true;
+  }
+
   abstract ValueAggregator<R, A> getValueAggregator();
 
   abstract DataType getRawValueType();

pinot-core/src/test/java/org/apache/pinot/core/startree/v2/CountMVStarTreeV2Test.java

@@ -0,0 +1,55 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.startree.v2;
+
+import java.util.Random;
+import org.apache.pinot.segment.local.aggregator.CountMVValueAggregator;
+import org.apache.pinot.segment.local.aggregator.ValueAggregator;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+
+import static org.testng.Assert.assertEquals;
+
+
+public class CountMVStarTreeV2Test extends BaseStarTreeV2Test<Object, Long> {
+
+  @Override
+  ValueAggregator<Object, Long> getValueAggregator() {
+    return new CountMVValueAggregator();
+  }
+
+  @Override
+  DataType getRawValueType() {
+    return DataType.INT;
+  }
+
+  @Override
+  Object getRandomRawValue(Random random) {
+    return random.nextInt();
+  }
+
+  @Override
+  protected void assertAggregatedValue(Long starTreeResult, Long nonStarTreeResult) {
+    assertEquals(starTreeResult, nonStarTreeResult);
+  }
+
+  @Override
+  protected boolean isAggColSingleValueField() {
+    return false;
+  }
+}