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

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

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