public static DataSchema canonicalizeDataSchemaForAggregation(QueryContext queryContext, DataSchema dataSchema) {
    List<Pair<AggregationFunction, FilterContext>> filteredAggregationFunctions =
        queryContext.getFilteredAggregationFunctions();
    assert filteredAggregationFunctions != null;
    int numAggregations = filteredAggregationFunctions.size();
    Preconditions.checkState(dataSchema.size() == numAggregations,
        "BUG: Expect same number of aggregations and columns in data schema, got %s aggregations, %s columns in data "
            + "schema", numAggregations, dataSchema.size());
    String[] columnNames = new String[numAggregations];
    for (int i = 0; i < numAggregations; i++) {
      Pair<AggregationFunction, FilterContext> pair = filteredAggregationFunctions.get(i);
      AggregationFunction aggregationFunction = pair.getLeft();
      columnNames[i] = AggregationFunctionUtils.getResultColumnName(aggregationFunction, pair.getRight());
    }
    return new DataSchema(columnNames, dataSchema.getColumnDataTypes());
  }
@Test
  public void testCanonicalizeDataSchemaForAggregation() {
    QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT SUM(col1 + col2) FROM testTable");
    // Intentionally make data schema not matching the string representation of the expression
    DataSchema dataSchema = new DataSchema(new String[]{"sum(col1+col2)"}, new ColumnDataType[]{ColumnDataType.DOUBLE});
    DataSchema canonicalDataSchema =
        ReducerDataSchemaUtils.canonicalizeDataSchemaForAggregation(queryContext, dataSchema);
    assertEquals(canonicalDataSchema,
        new DataSchema(new String[]{"sum(plus(col1,col2))"}, new ColumnDataType[]{ColumnDataType.DOUBLE}));

    queryContext = QueryContextConverterUtils.getQueryContext("SELECT SUM(col1 + 1), MIN(col2 + 2) FROM testTable");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"sum(col1+1)", "min(col2+2)"},
        new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE});
    canonicalDataSchema = ReducerDataSchemaUtils.canonicalizeDataSchemaForAggregation(queryContext, dataSchema);
    assertEquals(canonicalDataSchema, new DataSchema(new String[]{"sum(plus(col1,'1'))", "min(plus(col2,'2'))"},
        new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE}));

    queryContext = QueryContextConverterUtils.getQueryContext(
        "SELECT MAX(col1 + 1) FILTER(WHERE col3 > 0) - MIN(col2 + 2) FILTER(WHERE col4 > 0) FROM testTable");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"max(col1+1)", "min(col2+2)"},
        new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE});
    canonicalDataSchema = ReducerDataSchemaUtils.canonicalizeDataSchemaForAggregation(queryContext, dataSchema);
    assertEquals(canonicalDataSchema, new DataSchema(
        new String[]{"max(plus(col1,'1')) FILTER(WHERE col3 > '0')", "min(plus(col2,'2')) FILTER(WHERE col4 > '0')"},
        new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE}));
  }