[GitHub] [incubator-pinot] fx19880617 commented on a change in pull request #5019: Make output schema to match selection list for aggregation groupbys

[GitBox <gi...@apache.org>]

fx19880617 commented on a change in pull request #5019: Make output schema to match selection list for aggregation groupbys
URL: https://github.com/apache/incubator-pinot/pull/5019#discussion_r371456688
 
 

 ##########
 File path: pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java
 ##########
 @@ -194,14 +210,104 @@ private void setSQLGroupByInResultTable(BrokerResponseNative brokerResponseNativ
         values[index] = _aggregationFunctions[aggNum++].extractFinalResult(values[index]);
         index++;
       }
-      rows.add(values);
+      if (_sqlSelectionList != null) {
+        Object[] finalValues = new Object[_sqlSelectionList.size()];
+        for (int i = 0; i < finalSchemaMapIdx.length; i++) {
+          if (finalSchemaMapIdx[i] == -1) {
+            finalValues[i] = null;
+          } else {
+            finalValues[i] = values[finalSchemaMapIdx[i]];
+          }
+        }
+        rows.add(finalValues);
+      } else {
+        rows.add(values);
+      }
       numRows++;
     }
 
     DataSchema finalDataSchema = getSQLResultTableSchema(dataSchema);
+    if (_sqlSelectionList != null) {
+      int columnSize = _sqlSelectionList.size();
+      String[] columns = new String[columnSize];
+      DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[columnSize];
+      for (int i = 0; i < columnSize; i++) {
+        if (finalSchemaMapIdx[i] == -1) {
+          columns[i] = RequestUtils.prettyPrint(_sqlSelectionList.get(i));
+          columnDataTypes[i] = DataSchema.ColumnDataType.STRING;
+        } else {
+          columns[i] = finalDataSchema.getColumnName(finalSchemaMapIdx[i]);
+          columnDataTypes[i] = finalDataSchema.getColumnDataType(finalSchemaMapIdx[i]);
+        }
+      }
+      finalDataSchema = new DataSchema(columns, columnDataTypes);
+    }
     brokerResponseNative.setResultTable(new ResultTable(finalDataSchema, rows));
   }
 
+  /**
+   * Generate index mapping based on selection expression to DataTable schema, which is groupBy columns,
+   * then aggregation functions.
+   * @param dataSchema
+   *
+   * @return a mapping from final schema idx to corresponding idx in data table schema.
+   */
+  private int[] getFinalSchemaMapIdx(DataSchema dataSchema) {
+    int[] finalSchemaMapIdx = new int[_sqlSelectionList.size()];
+    int nextAggregationIdx = _numGroupBy;
+    for (int i = 0; i < _sqlSelectionList.size(); i++) {
+      finalSchemaMapIdx[i] = getExpressionMapIdx(dataSchema, _sqlSelectionList.get(i), nextAggregationIdx);
+      if (finalSchemaMapIdx[i] == nextAggregationIdx) {
+        nextAggregationIdx++;
+      }
+    }
+    return finalSchemaMapIdx;
+  }
+
+  private int getExpressionMapIdx(DataSchema dataSchema, Expression expression, int nextAggregationIdx) {
+    // Check if expression matches groupBy list.
+    int idxFromGroupByList = getGroupByIdx(_groupByList, expression);
+    if (idxFromGroupByList != -1) {
+      return idxFromGroupByList;
+    }
+    // Handle all functions
+    if (expression.getFunctionCall() != null) {
+      // handle AS
+      if (expression.getFunctionCall().getOperator().equalsIgnoreCase("AS")) {
+        return getExpressionMapIdx(dataSchema, expression.getFunctionCall().getOperands().get(0), nextAggregationIdx);
+      }
+      // Return next aggregation idx.
+      return nextAggregationIdx;
+    }
+    // Handle identifier, which is a column.
+    if (expression.getIdentifier() != null) {
 
 Review comment:
   Right, will remove this block

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org