OA-1063: Add metadata-based aggregation support for SUM function

pinot-core/src/main/java/org/apache/pinot/core/operator/query/NonScanBasedAggregationOperator.java

@@ -144,6 +144,10 @@ protected AggregationResultsBlock getNextBlock() {
           result = getDistinctCountSmartHLLResult(Objects.requireNonNull(dataSource.getDictionary()),
               (DistinctCountSmartHLLAggregationFunction) aggregationFunction);
           break;
+        case SUM:
+          // dataSource.getDataSourceMetadata().getSumValue() is always present for sum aggregation function.
+          result = dataSource.getDataSourceMetadata().getSumValue().get();
+          break;
         default:
           throw new IllegalStateException(
               "Non-scan based aggregation operator does not support function type: " + aggregationFunction.getType());

pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationPlanNode.java

@@ -37,6 +37,7 @@
 import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.segment.spi.SegmentContext;
 import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
 import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader;
 
 import static org.apache.pinot.segment.spi.AggregationFunctionType.COUNT;
@@ -63,6 +64,7 @@
 import static org.apache.pinot.segment.spi.AggregationFunctionType.MINMAXRANGEMV;
 import static org.apache.pinot.segment.spi.AggregationFunctionType.MINMV;
 import static org.apache.pinot.segment.spi.AggregationFunctionType.SEGMENTPARTITIONEDDISTINCTCOUNT;
+import static org.apache.pinot.segment.spi.AggregationFunctionType.SUM;
 
 
 /**
@@ -79,7 +81,7 @@ public class AggregationPlanNode implements PlanNode {
 
   // DISTINCTCOUNT excluded because consuming segment metadata contains unknown cardinality when there is no dictionary
   private static final EnumSet<AggregationFunctionType> METADATA_BASED_FUNCTIONS =
-      EnumSet.of(COUNT, MIN, MINMV, MAX, MAXMV, MINMAXRANGE, MINMAXRANGEMV);
+      EnumSet.of(COUNT, MIN, MINMV, MAX, MAXMV, MINMAXRANGE, MINMAXRANGEMV, SUM);
 
   private final IndexSegment _indexSegment;
   private final SegmentContext _segmentContext;
@@ -213,8 +215,12 @@ private boolean isFitForNonScanBasedPlan() {
         }
       }
       if (METADATA_BASED_FUNCTIONS.contains(aggregationFunction.getType())) {
-        if (dataSource.getDataSourceMetadata().getMaxValue() != null
-            && dataSource.getDataSourceMetadata().getMinValue() != null) {
+        DataSourceMetadata metadata = dataSource.getDataSourceMetadata();
+        if (aggregationFunction.getType() == SUM) {
+          if (metadata.getSumValue().isPresent()) {
+            continue;
+          }
+        } else if (metadata.getMaxValue() != null && metadata.getMinValue() != null) {
           continue;
         }
       }

pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java

@@ -116,6 +116,7 @@ public void buildSegment()
     segmentGeneratorConfig.setTableName("testTable");
     segmentGeneratorConfig.setSegmentName(SEGMENT_NAME);
     segmentGeneratorConfig.setOutDir(INDEX_DIR.getAbsolutePath());
+    segmentGeneratorConfig.setEnableColumnSumValue(true);
 
     // Build the index segment.
     SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl();
@@ -158,82 +159,89 @@ public void deleteSegment() {
 
   @Test(dataProvider = "testPlanMakerDataProvider")
   public void testPlanMaker(String query, Class<? extends Operator<?>> operatorClass,
-      Class<? extends Operator<?>> upsertOperatorClass) {
+      Class<? extends Operator<?>> upsertOperatorClass, String name) {
     QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
     Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_indexSegment), queryContext).run();
-    assertTrue(operatorClass.isInstance(operator));
+    assertTrue(operatorClass.isInstance(operator), name);
 
     SegmentContext segmentContext = new SegmentContext(_upsertIndexSegment);
     segmentContext.setQueryableDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(_upsertIndexSegment));
     Operator<?> upsertOperator = PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext).run();
-    assertTrue(upsertOperatorClass.isInstance(upsertOperator));
+    assertTrue(upsertOperatorClass.isInstance(upsertOperator), name);
   }
 
   @DataProvider(name = "testPlanMakerDataProvider")
   public Object[][] testPlanMakerDataProvider() {
     List<Object[]> entries = new ArrayList<>();
     // Selection
     entries.add(new Object[]{
-        "select * from testTable", SelectionOnlyOperator.class, SelectionOnlyOperator.class
+        "select * from testTable", SelectionOnlyOperator.class, SelectionOnlyOperator.class, "select-all"
     });
     // Selection
     entries.add(new Object[]{
-        "select column1,column5 from testTable", SelectionOnlyOperator.class, SelectionOnlyOperator.class
+        "select column1,column5 from testTable", SelectionOnlyOperator.class, SelectionOnlyOperator.class,
+        "select-columns"
     });
     // Selection with filter
     entries.add(new Object[]{
-        "select * from testTable where daysSinceEpoch > 100", SelectionOnlyOperator.class, SelectionOnlyOperator.class
+        "select * from testTable where daysSinceEpoch > 100", SelectionOnlyOperator.class, SelectionOnlyOperator.class,
+        "select-all-with-filter"
     });
     // COUNT from metadata (via non-scan-based aggregation because it is an all-match)
     entries.add(new Object[]{
-        "select count(*) from testTable", NonScanBasedAggregationOperator.class, FastFilteredCountOperator.class
+        "select count(*) from testTable", NonScanBasedAggregationOperator.class, FastFilteredCountOperator.class,
+        "count-all"
     });
     // COUNT from metadata with a non-match-all filter, fast filtered count is to be used
     entries.add(new Object[]{
         "select count(*) from testTable where column1 <= 10", FastFilteredCountOperator.class,
-        FastFilteredCountOperator.class
+        FastFilteredCountOperator.class, "count-all-with-non-match-all-filter"
     });
     // COUNT from metadata with match all filter
     entries.add(new Object[]{
         "select count(*) from testTable where column1 > 10", NonScanBasedAggregationOperator.class,
-        FastFilteredCountOperator.class
+        FastFilteredCountOperator.class, "count-all-with-match-all-filter"
     });
     // MIN/MAX from dictionary
     entries.add(new Object[]{
         "select max(daysSinceEpoch),min(daysSinceEpoch) from testTable", NonScanBasedAggregationOperator.class,
-        AggregationOperator.class
+        AggregationOperator.class, "min-max-from-dictionary"
     });
     // MIN/MAX from dictionary with match all filter
     entries.add(new Object[]{
         "select max(daysSinceEpoch),min(daysSinceEpoch) from testTable where column1 > 10",
-        NonScanBasedAggregationOperator.class, AggregationOperator.class
+        NonScanBasedAggregationOperator.class, AggregationOperator.class,
+        "min-max-from-dictionary-with-match-all-filter"
     });
     // MINMAXRANGE from dictionary
     entries.add(new Object[]{
         "select minmaxrange(daysSinceEpoch) from testTable", NonScanBasedAggregationOperator.class,
-        AggregationOperator.class
+        AggregationOperator.class, "minmaxrange-from-dictionary"
     });
     // MINMAXRANGE from dictionary with match all filter
     entries.add(new Object[]{
         "select minmaxrange(daysSinceEpoch) from testTable where column1 > 10", NonScanBasedAggregationOperator.class,
-        AggregationOperator.class
+        AggregationOperator.class, "minmaxrange-from-dictionary-with-match-all-filter"
     });
-    // Aggregation
+    // SUM from metadata (default segment version is v3, which stores sum metadata for numeric SV columns)
     entries.add(new Object[]{
-        "select sum(column1) from testTable", AggregationOperator.class, AggregationOperator.class
+        "select sum(column1) from testTable", NonScanBasedAggregationOperator.class,
+        AggregationOperator.class, "sum-from-metadata"
     });
     // Aggregation group-by
     entries.add(new Object[]{
-        "select sum(column1) from testTable group by daysSinceEpoch", GroupByOperator.class, GroupByOperator.class
+        "select sum(column1) from testTable group by daysSinceEpoch", GroupByOperator.class,
+        GroupByOperator.class, "sum-from-metadata-group-by"
     });
     // COUNT from metadata, MIN from dictionary
     entries.add(new Object[]{
-        "select count(*),min(column17) from testTable", NonScanBasedAggregationOperator.class, AggregationOperator.class
+        "select count(*),min(column17) from testTable", NonScanBasedAggregationOperator.class,
+        AggregationOperator.class, "count-from-metadata-min-from-dictionary"
     });
     // Aggregation group-by
     entries.add(new Object[]{
         "select count(*),min(daysSinceEpoch) from testTable group by daysSinceEpoch", GroupByOperator.class,
-        GroupByOperator.class
+        GroupByOperator.class, "count-from-metadata-min-from-dictionary-group-by"
     });
 
     return entries.toArray(new Object[entries.size()][]);

pinot-core/src/test/java/org/apache/pinot/core/plan/maker/NonScanBasedSumAggregationPlanMakerTest.java

@@ -0,0 +1,181 @@
+/**
+ * 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.plan.maker;
+
+import java.io.File;
+import java.net.URL;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.query.AggregationOperator;
+import org.apache.pinot.core.operator.query.NonScanBasedAggregationOperator;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver;
+import org.apache.pinot.segment.spi.creator.SegmentVersion;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.ingestion.IngestionConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.TimeGranularitySpec;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.AfterTest;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeTest;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Tests that SUM aggregation uses the non-scan-based path when segment metadata contains sum values (v3 segments),
+ * and falls back to scan-based aggregation when sum metadata is absent (non-v3 segments).
+ */
+public class NonScanBasedSumAggregationPlanMakerTest {
+  private static final String AVRO_DATA = "data" + File.separator + "test_data-sv.avro";
+  private static final File INDEX_DIR =
+      new File(FileUtils.getTempDirectory(), "NonScanBasedSumAggregationPlanMakerTest");
+  private static final String V3_SEGMENT_NAME = "testTable_v3_segment";
+  private static final String V1_SEGMENT_NAME = "testTable_v1_segment";
+  private static final InstancePlanMakerImplV2 PLAN_MAKER = new InstancePlanMakerImplV2();
+
+  private IndexSegment _v3Segment;
+  private IndexSegment _v1Segment;
+
+  @BeforeTest
+  public void buildSegments()
+      throws Exception {
+    FileUtils.deleteQuietly(INDEX_DIR);
+
+    URL resource = getClass().getClassLoader().getResource(AVRO_DATA);
+    assertNotNull(resource);
+    String filePath = resource.getFile();
+
+    Schema schema = new Schema.SchemaBuilder().setSchemaName("testTable")
+        .addMetric("column1", FieldSpec.DataType.INT)
+        .addMetric("column3", FieldSpec.DataType.INT)
+        .addSingleValueDimension("column5", FieldSpec.DataType.STRING)
+        .addSingleValueDimension("column6", FieldSpec.DataType.INT)
+        .addSingleValueDimension("column7", FieldSpec.DataType.INT)
+        .addSingleValueDimension("column9", FieldSpec.DataType.INT)
+        .addSingleValueDimension("column11", FieldSpec.DataType.STRING)
+        .addSingleValueDimension("column12", FieldSpec.DataType.STRING)
+        .addMetric("column17", FieldSpec.DataType.INT)
+        .addMetric("column18", FieldSpec.DataType.INT)
+        .addTime(new TimeGranularitySpec(DataType.INT, TimeUnit.DAYS, "daysSinceEpoch"), null).build();
+
+    IngestionConfig ingestionConfig = new IngestionConfig();
+    ingestionConfig.setSegmentTimeValueCheck(false);
+    ingestionConfig.setRowTimeValueCheck(false);
+    TableConfig tableConfig =
+        new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable").setTimeColumnName("daysSinceEpoch")
+            .setIngestionConfig(ingestionConfig).build();
+
+    // Build v3 segment (has sum metadata)
+    SegmentGeneratorConfig v3Config = new SegmentGeneratorConfig(tableConfig, schema);
+    v3Config.setInputFilePath(filePath);
+    v3Config.setTableName("testTable");
+    v3Config.setSegmentName(V3_SEGMENT_NAME);
+    v3Config.setOutDir(INDEX_DIR.getAbsolutePath());
+    v3Config.setSegmentVersion(SegmentVersion.v3);
+    v3Config.setEnableColumnSumValue(true);
+    SegmentIndexCreationDriver v3Driver = new SegmentIndexCreationDriverImpl();
+    v3Driver.init(v3Config);
+    v3Driver.build();
+
+    // Build v1 segment (no sum metadata)
+    SegmentGeneratorConfig v1Config = new SegmentGeneratorConfig(tableConfig, schema);
+    v1Config.setInputFilePath(filePath);
+    v1Config.setTableName("testTable");
+    v1Config.setSegmentName(V1_SEGMENT_NAME);
+    v1Config.setOutDir(INDEX_DIR.getAbsolutePath());
+    v1Config.setSegmentVersion(SegmentVersion.v1);
+    SegmentIndexCreationDriver v1Driver = new SegmentIndexCreationDriverImpl();
+    v1Driver.init(v1Config);
+    v1Driver.build();
+  }
+
+  @BeforeClass
+  public void loadSegments()
+      throws Exception {
+    ServerMetrics.register(mock(ServerMetrics.class));
+    _v3Segment = ImmutableSegmentLoader.load(new File(INDEX_DIR, V3_SEGMENT_NAME), ReadMode.heap);
+    _v1Segment = ImmutableSegmentLoader.load(new File(INDEX_DIR, V1_SEGMENT_NAME), ReadMode.heap);
+  }
+
+  @AfterClass
+  public void destroySegments() {
+    _v3Segment.destroy();
+    _v1Segment.destroy();
+  }
+
+  @AfterTest
+  public void deleteSegments() {
+    FileUtils.deleteQuietly(INDEX_DIR);
+  }
+
+  @Test
+  public void testSumUsesNonScanPathForV3Segment() {
+    String query = "select sum(column1) from testTable";
+    QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+    Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_v3Segment), queryContext).run();
+    assertTrue(operator instanceof NonScanBasedAggregationOperator,
+        "SUM on v3 segment should use NonScanBasedAggregationOperator");
+  }
+
+  @Test
+  public void testSumFallsBackToScanForV1Segment() {
+    String query = "select sum(column1) from testTable";
+    QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+    Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_v1Segment), queryContext).run();
+    assertTrue(operator instanceof AggregationOperator,
+        "SUM on v1 segment should fall back to AggregationOperator");
+  }
+
+  @Test
+  public void testSumWithFilterFallsBackToScan() {
+    // SUM with non-match-all filter cannot use non-scan path
+    String query = "select sum(column1) from testTable where column1 <= 10";
+    QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+    Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_v3Segment), queryContext).run();
+    assertTrue(operator instanceof AggregationOperator,
+        "SUM with filter should fall back to AggregationOperator");
+  }
+
+  @Test
+  public void testCountAndSumCombinedUsesNonScanForV3() {
+    String query = "select count(*), sum(column1) from testTable";
+    QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+    Operator<?> operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_v3Segment), queryContext).run();
+    assertTrue(operator instanceof NonScanBasedAggregationOperator,
+        "COUNT(*) + SUM on v3 segment should use NonScanBasedAggregationOperator");
+  }
+}

pinot-core/src/test/java/org/apache/pinot/queries/BaseSingleValueQueriesTest.java

@@ -125,6 +125,7 @@ public void buildSegment()
     segmentGeneratorConfig.setInputFilePath(filePath);
     segmentGeneratorConfig.setTableName(RAW_TABLE_NAME);
     segmentGeneratorConfig.setOutDir(INDEX_DIR.getAbsolutePath());
+    segmentGeneratorConfig.setEnableColumnSumValue(true);
 
     // Build the index segment.
     SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl();

pinot-core/src/test/java/org/apache/pinot/queries/CastQueriesTest.java

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Random;
 import org.apache.commons.io.FileUtils;
+import org.apache.pinot.core.operator.BaseOperator;
 import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
 import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
 import org.apache.pinot.core.operator.query.AggregationOperator;
@@ -141,8 +142,9 @@ public void testCastSum() {
     String query = "select cast(sum(" + X_COL + ") as int), "
         + "cast(sum(" + Y_COL + ") as int) "
         + "from " + RAW_TABLE_NAME;
-    AggregationOperator aggregationOperator = getOperator(query);
-    List<Object> aggregationResult = aggregationOperator.nextBlock().getResults();
+    // SUM on v3 numeric SV columns uses NonScanBasedAggregationOperator (metadata-based)
+    BaseOperator<AggregationResultsBlock> operator = getOperator(query);
+    List<Object> aggregationResult = operator.nextBlock().getResults();
     assertNotNull(aggregationResult);
     assertEquals(aggregationResult.size(), 2);
     assertEquals(((Number) aggregationResult.get(0)).intValue(), NUM_RECORDS / 2);

pinot-core/src/test/java/org/apache/pinot/queries/InterSegmentAggregationSingleValueQueriesTest.java

@@ -151,12 +151,14 @@ public void testMin() {
   public void testSum() {
     String query = "SELECT SUM(column1) AS v1, SUM(column3) AS v2 FROM testTable";
 
+    // Without filter, query should be answered by NonScanBasedAggregationOperator (numEntriesScannedPostFilter = 0)
+    // using sum metadata stored in the v3 segment.
     BrokerResponseNative brokerResponse = getBrokerResponse(query);
     DataSchema expectedDataSchema =
         new DataSchema(new String[]{"v1", "v2"}, new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE});
     ResultTable expectedResultTable = new ResultTable(expectedDataSchema,
         Collections.singletonList(new Object[]{129268741751388.0, 129156636756600.0}));
-    QueriesTestUtils.testInterSegmentsResult(brokerResponse, 120000L, 0L, 240000L, 120000L, expectedResultTable);
+    QueriesTestUtils.testInterSegmentsResult(brokerResponse, 120000L, 0L, 0L, 120000L, expectedResultTable);
 
     brokerResponse = getBrokerResponse(query + FILTER);
     expectedResultTable = new ResultTable(expectedDataSchema,