public static Pair<DataSchema, int[]> getResultTableDataSchemaAndColumnIndices(QueryContext queryContext,
      DataSchema dataSchema) {
    List<ExpressionContext> selectExpressions = queryContext.getSelectExpressions();
    int numSelectExpressions = selectExpressions.size();
    ColumnDataType[] columnDataTypesInDataSchema = dataSchema.getColumnDataTypes();
    int numColumnsInDataSchema = columnDataTypesInDataSchema.length;

    // No order-by expression
    // NOTE: Order-by expressions are ignored for queries with LIMIT 0.
    List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions();
    if (orderByExpressions == null || queryContext.getLimit() == 0) {
      // For 'SELECT *', use the server response data schema as the final results data schema.
      if ((numSelectExpressions == 1 && selectExpressions.get(0).equals(IDENTIFIER_STAR))) {
        int[] columnIndices = new int[numColumnsInDataSchema];
        for (int i = 0; i < numColumnsInDataSchema; i++) {
          columnIndices[i] = i;
        }
        return Pair.of(dataSchema, columnIndices);
      }

      // For select without duplicate columns, the order of the final selection columns is the same as the order of the
      // columns in the data schema.
      if (numSelectExpressions == numColumnsInDataSchema) {
        String[] columnNames = new String[numSelectExpressions];
        int[] columnIndices = new int[numSelectExpressions];
        for (int i = 0; i < numSelectExpressions; i++) {
          columnNames[i] = selectExpressions.get(i).toString();
          columnIndices[i] = i;
        }
        return Pair.of(new DataSchema(columnNames, columnDataTypesInDataSchema), columnIndices);
      }

      // For select with duplicate columns, construct a map from expression to index with the same order as the data
      // schema, then look up the selection expressions.
      Object2IntOpenHashMap<ExpressionContext> expressionIndexMap = new Object2IntOpenHashMap<>(numColumnsInDataSchema);
      for (ExpressionContext selectExpression : selectExpressions) {
        expressionIndexMap.putIfAbsent(selectExpression, expressionIndexMap.size());
      }
      Preconditions.checkState(expressionIndexMap.size() == numColumnsInDataSchema,
          "BUG: Expect same number of deduped columns in SELECT clause and in data schema, got %s before dedup and %s"
              + " after dedup in SELECT clause, %s in data schema", numSelectExpressions, expressionIndexMap.size(),
          numColumnsInDataSchema);
      String[] columnNames = new String[numSelectExpressions];
      ColumnDataType[] columnDataTypes = new ColumnDataType[numSelectExpressions];
      int[] columnIndices = new int[numSelectExpressions];
      for (int i = 0; i < numSelectExpressions; i++) {
        ExpressionContext selectExpression = selectExpressions.get(i);
        int columnIndex = expressionIndexMap.getInt(selectExpression);
        columnNames[i] = selectExpression.toString();
        columnDataTypes[i] = columnDataTypesInDataSchema[columnIndex];
        columnIndices[i] = columnIndex;
      }
      return Pair.of(new DataSchema(columnNames, columnDataTypes), columnIndices);
    }

    // For 'SELECT *' with order-by, exclude transform functions from the returned columns and sort.
    if (numSelectExpressions == 1 && selectExpressions.get(0).equals(IDENTIFIER_STAR)) {
      String[] columnNamesInDataSchema = dataSchema.getColumnNames();
      List<Integer> columnIndexList = new ArrayList<>(columnNamesInDataSchema.length);
      for (int i = 0; i < columnNamesInDataSchema.length; i++) {
        if (columnNamesInDataSchema[i].indexOf('(') == -1) {
          columnIndexList.add(i);
        }
      }
      columnIndexList.sort(Comparator.comparing(o -> columnNamesInDataSchema[o]));
      int numColumns = columnIndexList.size();
      String[] columnNames = new String[numColumns];
      ColumnDataType[] columnDataTypes = new ColumnDataType[numColumns];
      int[] columnIndices = new int[numColumns];
      for (int i = 0; i < numColumns; i++) {
        int columnIndex = columnIndexList.get(i);
        columnNames[i] = columnNamesInDataSchema[columnIndex];
        columnDataTypes[i] = columnDataTypesInDataSchema[columnIndex];
        columnIndices[i] = columnIndex;
      }
      return Pair.of(new DataSchema(columnNames, columnDataTypes), columnIndices);
    }

    // For other order-by queries, construct a map from expression to index with the same order as the data schema,
    // then look up the selection expressions.
    Object2IntOpenHashMap<ExpressionContext> expressionIndexMap = new Object2IntOpenHashMap<>(numColumnsInDataSchema);
    // NOTE: Order-by expressions are already deduped in QueryContext.
    for (OrderByExpressionContext orderByExpression : orderByExpressions) {
      expressionIndexMap.put(orderByExpression.getExpression(), expressionIndexMap.size());
    }
    for (ExpressionContext selectExpression : selectExpressions) {
      expressionIndexMap.putIfAbsent(selectExpression, expressionIndexMap.size());
    }
    String[] columnNames = new String[numSelectExpressions];
    ColumnDataType[] columnDataTypes = new ColumnDataType[numSelectExpressions];
    int[] columnIndices = new int[numSelectExpressions];
    if (expressionIndexMap.size() == numColumnsInDataSchema) {
      for (int i = 0; i < numSelectExpressions; i++) {
        ExpressionContext selectExpression = selectExpressions.get(i);
        int columnIndex = expressionIndexMap.getInt(selectExpression);
        columnNames[i] = selectExpression.toString();
        columnDataTypes[i] = columnDataTypesInDataSchema[columnIndex];
        columnIndices[i] = columnIndex;
      }
    } else {
      // When all segments are pruned on the server side, the data schema will only contain the columns in the SELECT
      // clause, and data type for all columns are set to STRING. See ResultBlocksUtils for details.
      for (int i = 0; i < numSelectExpressions; i++) {
        columnNames[i] = selectExpressions.get(i).toString();
        columnDataTypes[i] = ColumnDataType.STRING;
        columnIndices[i] = i;
      }
    }
    return Pair.of(new DataSchema(columnNames, columnDataTypes), columnIndices);
  }
@Test
  public void testGetResultTableColumnIndices() {
    // Select * without order-by
    QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable");
    DataSchema dataSchema = new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    });
    Pair<DataSchema, int[]> pair =
        SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{0, 1, 2});

    // Select * without order-by, all the segments are pruned on the server side
    dataSchema = new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING});
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING}));

    // Select * with order-by but LIMIT 0
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY col1 LIMIT 0");
    dataSchema = new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{0, 1, 2});

    // Select * with order-by but LIMIT 0, all the segments are pruned on the server side
    dataSchema = new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING});
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING}));

    // Select columns without order-by
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT col1 + 1, col2 + 2 FROM testTable");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "add(col2+2)"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{0, 1});

    // Select columns without order-by, all the segments are pruned on the server side
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "add(col2+2)"}, new ColumnDataType[]{
        ColumnDataType.STRING, ColumnDataType.STRING
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')"}, new ColumnDataType[]{
        ColumnDataType.STRING, ColumnDataType.STRING
    }));

    // Select duplicate columns without order-by
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT col1 + 1, col2 + 2, col1 + 1 FROM testTable");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "add(col2+2)"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(),
        new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')", "plus(col1,'1')"}, new ColumnDataType[]{
            ColumnDataType.DOUBLE, ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
        }));
    assertEquals(pair.getRight(), new int[]{0, 1, 0});

    // Select duplicate columns without order-by, all the segments are pruned on the server side
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "add(col2+2)", "add(col1+1)"}, new ColumnDataType[]{
        ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(),
        new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')", "plus(col1,'1')"}, new ColumnDataType[]{
            ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING
        }));

    // Select * with order-by
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY col3");
    dataSchema = new DataSchema(new String[]{"col3", "col1", "col2"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.INT, ColumnDataType.LONG
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{1, 2, 0});

    // Select * with order-by, all the segments are pruned on the server side
    dataSchema = new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING});
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING}));

    // Select * ordering on function
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY col1 + col2");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+col2)", "col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{1, 2, 3});

    // Select * ordering on function, all the segments are pruned on the server side
    dataSchema = new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING});
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING}));

    // Select * ordering on both column and function
    queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY col1 + col2, col2");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+col2)", "col2", "col1", "col3"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.LONG, ColumnDataType.INT, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"col1", "col2", "col3"}, new ColumnDataType[]{
        ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE
    }));
    assertEquals(pair.getRight(), new int[]{2, 1, 3});

    // Select * ordering on both column and function, all the segments are pruned on the server side
    dataSchema = new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING});
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"*"}, new ColumnDataType[]{ColumnDataType.STRING}));

    // Select columns with order-by
    queryContext = QueryContextConverterUtils.getQueryContext(
        "SELECT col1 + 1, col3, col2 + 2 FROM testTable ORDER BY col2 + 2, col4");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col2+2)", "col4", "add(col1+1)", "col3"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.STRING, ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"plus(col1,'1')", "col3", "plus(col2,'2')"},
        new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE, ColumnDataType.DOUBLE}));
    assertEquals(pair.getRight(), new int[]{2, 3, 0});

    // Select columns with order-by, all the segments are pruned on the server side
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "col3", "add(col2+2)"}, new ColumnDataType[]{
        ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(), new DataSchema(new String[]{"plus(col1,'1')", "col3", "plus(col2,'2')"},
        new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING}));

    // Select duplicate columns with order-by
    queryContext = QueryContextConverterUtils.getQueryContext(
        "SELECT col1 + 1, col2 + 2, col1 + 1 FROM testTable ORDER BY col2 + 2, col4");
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col2+2)", "col4", "add(col1+1)"}, new ColumnDataType[]{
        ColumnDataType.DOUBLE, ColumnDataType.STRING, ColumnDataType.DOUBLE
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(),
        new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')", "plus(col1,'1')"}, new ColumnDataType[]{
            ColumnDataType.DOUBLE, ColumnDataType.DOUBLE, ColumnDataType.DOUBLE
        }));
    assertEquals(pair.getRight(), new int[]{2, 0, 2});

    // Select duplicate columns with order-by, all the segments are pruned on the server side
    // Intentionally make data schema not matching the string representation of the expression
    dataSchema = new DataSchema(new String[]{"add(col1+1)", "add(col2+2)", "add(col1+1)"}, new ColumnDataType[]{
        ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING
    });
    pair = SelectionOperatorUtils.getResultTableDataSchemaAndColumnIndices(queryContext, dataSchema);
    assertEquals(pair.getLeft(),
        new DataSchema(new String[]{"plus(col1,'1')", "plus(col2,'2')", "plus(col1,'1')"}, new ColumnDataType[]{
            ColumnDataType.STRING, ColumnDataType.STRING, ColumnDataType.STRING
        }));
  }