This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 18d0fbe186d Fix alias handling for empty response on SSE queries with
post aggregation (#17815)
18d0fbe186d is described below
commit 18d0fbe186d56136a215227672e23838fe39aa61
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Mar 5 09:54:30 2026 -0800
Fix alias handling for empty response on SSE queries with post aggregation
(#17815)
---
.../broker/requesthandler/EmptyResponseUtils.java | 92 ++++++++++++++--------
.../requesthandler/EmptyResponseUtilsTest.java | 34 ++++++++
2 files changed, 93 insertions(+), 33 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtils.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtils.java
index a6c7278ff72..bc940b15498 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtils.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtils.java
@@ -33,6 +33,7 @@ import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
import
org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
+import org.apache.pinot.core.query.reduce.PostAggregationHandler;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.query.QueryEnvironment;
import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter;
@@ -80,59 +81,84 @@ public class EmptyResponseUtils {
private static ResultTable buildEmptyAggregationResultTable(QueryContext
queryContext) {
List<Pair<AggregationFunction, FilterContext>>
filteredAggregationFunctions =
queryContext.getFilteredAggregationFunctions();
- List<String> aliases = queryContext.getAliasList();
assert filteredAggregationFunctions != null;
int numAggregations = filteredAggregationFunctions.size();
- String[] columnNames = new String[numAggregations];
- ColumnDataType[] columnDataTypes = new ColumnDataType[numAggregations];
- Object[] row = new Object[numAggregations];
+
+ // Build pre-post-aggregation DataSchema from individual aggregation
functions
+ String[] preAggColumnNames = new String[numAggregations];
+ ColumnDataType[] preAggColumnDataTypes = new
ColumnDataType[numAggregations];
+ Object[] rawRow = new Object[numAggregations];
for (int i = 0; i < numAggregations; i++) {
Pair<AggregationFunction, FilterContext> pair =
filteredAggregationFunctions.get(i);
AggregationFunction aggregationFunction = pair.getLeft();
- if (aliases.size() == numAggregations && aliases.get(i) != null) {
- columnNames[i] = aliases.get(i);
- } else {
- columnNames[i] =
AggregationFunctionUtils.getResultColumnName(aggregationFunction,
pair.getRight());
- }
- columnDataTypes[i] = aggregationFunction.getFinalResultColumnType();
+ preAggColumnNames[i] =
AggregationFunctionUtils.getResultColumnName(aggregationFunction,
pair.getRight());
+ preAggColumnDataTypes[i] =
aggregationFunction.getFinalResultColumnType();
Object finalResult = aggregationFunction.extractFinalResult(
aggregationFunction.extractAggregationResult(aggregationFunction.createAggregationResultHolder()));
- row[i] = finalResult != null ? columnDataTypes[i].convert(finalResult) :
null;
+ rawRow[i] = finalResult != null ?
preAggColumnDataTypes[i].convert(finalResult) : null;
}
- return new ResultTable(new DataSchema(columnNames, columnDataTypes),
List.<Object[]>of(row));
+ DataSchema prePostAggDataSchema = new DataSchema(preAggColumnNames,
preAggColumnDataTypes);
+
+ // Use PostAggregationHandler to evaluate post-aggregation expressions and
reorder columns
+ PostAggregationHandler postAggregationHandler = new
PostAggregationHandler(queryContext, prePostAggDataSchema);
+ DataSchema resultDataSchema = postAggregationHandler.getResultDataSchema();
+ Object[] resultRow = postAggregationHandler.getResult(rawRow);
+
+ applyAliases(queryContext, resultDataSchema);
+ return new ResultTable(resultDataSchema, List.<Object[]>of(resultRow));
}
private static ResultTable buildEmptyGroupByResultTable(QueryContext
queryContext) {
List<Pair<AggregationFunction, FilterContext>>
filteredAggregationFunctions =
queryContext.getFilteredAggregationFunctions();
- List<String> aliases = queryContext.getAliasList();
List<ExpressionContext> groupByExpressions =
queryContext.getGroupByExpressions();
assert filteredAggregationFunctions != null && groupByExpressions != null;
- int numColumns = groupByExpressions.size() +
filteredAggregationFunctions.size();
- String[] columnNames = new String[numColumns];
- ColumnDataType[] columnDataTypes = new ColumnDataType[numColumns];
- int index = 0;
- for (ExpressionContext groupByExpression : groupByExpressions) {
- if (aliases.size() == numColumns && aliases.get(index) != null) {
- columnNames[index] = aliases.get(index);
- } else {
- columnNames[index] = groupByExpression.toString();
- }
+ int numGroupByExpressions = groupByExpressions.size();
+ int numAggregations = filteredAggregationFunctions.size();
+ int numColumns = numGroupByExpressions + numAggregations;
+
+ // Build pre-post-aggregation DataSchema with group-by columns followed by
aggregation columns
+ String[] preAggColumnNames = new String[numColumns];
+ ColumnDataType[] preAggColumnDataTypes = new ColumnDataType[numColumns];
+ for (int i = 0; i < numGroupByExpressions; i++) {
+ preAggColumnNames[i] = groupByExpressions.get(i).toString();
// Use STRING column data type as default for group-by expressions
- columnDataTypes[index] = ColumnDataType.STRING;
- index++;
+ preAggColumnDataTypes[i] = ColumnDataType.STRING;
}
- for (Pair<AggregationFunction, FilterContext> pair :
filteredAggregationFunctions) {
+ for (int i = 0; i < numAggregations; i++) {
+ Pair<AggregationFunction, FilterContext> pair =
filteredAggregationFunctions.get(i);
AggregationFunction aggregationFunction = pair.getLeft();
- if (aliases.size() == numColumns && aliases.get(index) != null) {
- columnNames[index] = aliases.get(index);
- } else {
- columnNames[index] =
AggregationFunctionUtils.getResultColumnName(aggregationFunction,
pair.getRight());
+ preAggColumnNames[numGroupByExpressions + i] =
+ AggregationFunctionUtils.getResultColumnName(aggregationFunction,
pair.getRight());
+ preAggColumnDataTypes[numGroupByExpressions + i] =
aggregationFunction.getFinalResultColumnType();
+ }
+ DataSchema prePostAggDataSchema = new DataSchema(preAggColumnNames,
preAggColumnDataTypes);
+
+ // Use PostAggregationHandler to evaluate post-aggregation expressions and
reorder columns
+ PostAggregationHandler postAggregationHandler = new
PostAggregationHandler(queryContext, prePostAggDataSchema);
+ DataSchema resultDataSchema = postAggregationHandler.getResultDataSchema();
+
+ applyAliases(queryContext, resultDataSchema);
+ return new ResultTable(resultDataSchema, List.of());
+ }
+
+ private static void applyAliases(QueryContext queryContext, DataSchema
dataSchema) {
+ List<String> aliasList = queryContext.getAliasList();
+ if (aliasList == null || aliasList.isEmpty()) {
+ return;
+ }
+ String[] columnNames = dataSchema.getColumnNames();
+ List<ExpressionContext> selectExpressions =
queryContext.getSelectExpressions();
+ int numSelectExpressions = selectExpressions.size();
+ if (columnNames.length != numSelectExpressions || aliasList.size() !=
numSelectExpressions) {
+ return;
+ }
+ for (int i = 0; i < numSelectExpressions; i++) {
+ String alias = aliasList.get(i);
+ if (alias != null) {
+ columnNames[i] = alias;
}
- columnDataTypes[index] = aggregationFunction.getFinalResultColumnType();
- index++;
}
- return new ResultTable(new DataSchema(columnNames, columnDataTypes),
List.of());
}
/// Tries to fill an [DataSchema] when no row has been returned.
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtilsTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtilsTest.java
index c9ceafc8e0a..ca0295b1974 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtilsTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/EmptyResponseUtilsTest.java
@@ -140,6 +140,40 @@ public class EmptyResponseUtilsTest {
assertTrue(resultTable.getRows().isEmpty());
}
+ @Test
+ public void testBuildEmptyResultTableWithPostAggregation() {
+ // Aggregation with post-aggregation expression and aliases
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT SUM(a) * 1.0 / COUNT(*) AS rate, SUM(a) AS total FROM
testTable WHERE foo = 'bar'");
+ ResultTable resultTable =
EmptyResponseUtils.buildEmptyResultTable(queryContext);
+ DataSchema dataSchema = resultTable.getDataSchema();
+ assertEquals(dataSchema.getColumnNames(), new String[]{"rate", "total"});
+ assertEquals(dataSchema.getColumnDataTypes(),
+ new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.DOUBLE});
+ List<Object[]> rows = resultTable.getRows();
+ assertEquals(rows.size(), 1);
+
+ // Aggregation with post-aggregation expression without alias
+ queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT SUM(a) + MAX(b) FROM testTable WHERE foo = 'bar'");
+ resultTable = EmptyResponseUtils.buildEmptyResultTable(queryContext);
+ dataSchema = resultTable.getDataSchema();
+ assertEquals(dataSchema.size(), 1);
+ assertEquals(dataSchema.getColumnDataTypes(), new
ColumnDataType[]{ColumnDataType.DOUBLE});
+ rows = resultTable.getRows();
+ assertEquals(rows.size(), 1);
+
+ // Group-by with post-aggregation expression
+ queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT c, SUM(a) * 1.0 / COUNT(*) AS rate FROM testTable WHERE foo =
'bar' GROUP BY c");
+ resultTable = EmptyResponseUtils.buildEmptyResultTable(queryContext);
+ dataSchema = resultTable.getDataSchema();
+ assertEquals(dataSchema.getColumnNames(), new String[]{"c", "rate"});
+ assertEquals(dataSchema.getColumnDataTypes(),
+ new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.DOUBLE});
+ assertTrue(resultTable.getRows().isEmpty());
+ }
+
@Test
public void testBuildEmptyResultTableWithDistinctCountRawHLL() {
// Test DISTINCTCOUNTRAWHLL aggregation with empty results
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
