npawar 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_r371436880
##########
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:
When will this `if` block be needed? These will always be in the group by
clause, and hence be handled by 269 right?
----------------------------------------------------------------
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:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
