This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 3381db70831 Support is_leaf_return_final_result hint for aggregation
without GROUP BY (#20063)
3381db70831 is described below
commit 3381db70831a8fe3e5e04c4439dcb8997912a501
Author: Xiaotian (Jackie) Jiang <[email protected]>
AuthorDate: Tue Oct 27 18:05:37 2026 -0700
Support is_leaf_return_final_result hint for aggregation without GROUP BY
(#20063)
---
.../pinot/calcite/rel/hint/PinotHintOptions.java | 8 ++++
.../query/runtime/operator/AggregateOperator.java | 8 ++--
.../operator/MultistageAggregationExecutor.java | 38 ++++++++++++---
.../plan/server/ServerPlanRequestVisitor.java | 40 +++++++++-------
.../runtime/operator/AggregateOperatorTest.java | 56 ++++++++++++++++++++++
.../SerializedIntermediateResultMergeTest.java | 2 +-
.../src/test/resources/queries/QueryHints.json | 8 ++++
7 files changed, 130 insertions(+), 30 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
index 36b43b0f734..58facc68314 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
@@ -44,7 +44,15 @@ public class PinotHintOptions {
public static class AggregateOptions {
public static final String IS_PARTITIONED_BY_GROUP_BY_KEYS =
"is_partitioned_by_group_by_keys";
+
+ /// Asks the leaf stage to return final results instead of intermediate
results for aggregations, and the final
+ /// stage to merge the final results. Can be applied to aggregation and
group-by queries where the aggregated values
+ /// are partitioned across the leaf stage workers, but the group keys (if
any) are not. A leaf stage worker
+ /// finalizes the OFFLINE and REALTIME parts of a hybrid table separately,
so the aggregated values must be
+ /// partitioned between them as well. Similar to the
`serverReturnFinalResultKeyUnpartitioned` (group-by) and
+ /// `serverReturnFinalResult` (aggregation) query options of the
single-stage engine.
public static final String IS_LEAF_RETURN_FINAL_RESULT =
"is_leaf_return_final_result";
+
public static final String IS_SKIP_LEAF_STAGE_GROUP_BY =
"is_skip_leaf_stage_group_by";
/// Enables trimming of aggregation intermediate results by pushing down
order by and limit,
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
index 68d80696468..55cc6819032 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java
@@ -136,12 +136,12 @@ public class AggregateOperator extends MultiStageOperator
{
// Initialize the appropriate executor.
AggregateNode.AggType aggType = node.getAggType();
- // TODO: Allow leaf return final result for non-group-by queries
boolean leafReturnFinalResult = node.isLeafReturnFinalResult();
_isGroupBy = !groupKeys.isEmpty();
if (!_isGroupBy) {
_aggregationExecutor =
- new MultistageAggregationExecutor(_aggFunctions, filterArgIds,
maxFilterArgId, aggType, _resultSchema);
+ new MultistageAggregationExecutor(_aggFunctions, filterArgIds,
maxFilterArgId, aggType, leafReturnFinalResult,
+ _resultSchema);
_groupByExecutor = null;
} else {
_groupByExecutor =
@@ -272,8 +272,8 @@ public class AggregateOperator extends MultiStageOperator {
}
}
if (_groupByExecutor.getNumGroups() >=
_groupByExecutor.getNumGroupsWarningLimit()) {
- LOGGER.warn("numGroups reached warning limit: {} (actual: {})",
- _groupByExecutor.getNumGroupsWarningLimit(),
_groupByExecutor.getNumGroups());
+ LOGGER.warn("numGroups reached warning limit: {} (actual: {})",
_groupByExecutor.getNumGroupsWarningLimit(),
+ _groupByExecutor.getNumGroups());
_statMap.merge(StatKey.NUM_GROUPS_WARNING_LIMIT_REACHED, true);
}
return dataBlock;
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
index 4f059bf573d..47d76627ce3 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java
@@ -41,6 +41,7 @@ public class MultistageAggregationExecutor {
private final int[] _filterArgIds;
private final int _maxFilterArgId;
private final AggType _aggType;
+ private final boolean _leafReturnFinalResult;
private final DataSchema _resultSchema;
// Result holders for each mode.
@@ -48,11 +49,12 @@ public class MultistageAggregationExecutor {
private final Object[] _mergeResultHolder;
public MultistageAggregationExecutor(AggregationFunction[] aggFunctions,
int[] filterArgIds, int maxFilterArgId,
- AggType aggType, DataSchema resultSchema) {
+ AggType aggType, boolean leafReturnFinalResult, DataSchema resultSchema)
{
_aggFunctions = aggFunctions;
_filterArgIds = filterArgIds;
_maxFilterArgId = maxFilterArgId;
_aggType = aggType;
+ _leafReturnFinalResult = leafReturnFinalResult;
_resultSchema = resultSchema;
int numFunctions = aggFunctions.length;
@@ -72,6 +74,8 @@ public class MultistageAggregationExecutor {
public void processBlock(MseBlock.Data block) {
if (!_aggType.isInputIntermediateFormat()) {
processAggregate(block);
+ } else if (_leafReturnFinalResult) {
+ processMergeFinalResults(block);
} else {
processMerge(block);
}
@@ -84,19 +88,28 @@ public class MultistageAggregationExecutor {
AggregationFunction aggFunction = _aggFunctions[i];
Object value;
switch (_aggType) {
- case LEAF:
- value =
aggFunction.extractAggregationResult(_aggregateResultHolder[i]);
+ case LEAF: {
+ Object intermediateResult =
aggFunction.extractAggregationResult(_aggregateResultHolder[i]);
+ value = _leafReturnFinalResult ?
aggFunction.extractFinalResult(intermediateResult) : intermediateResult;
break;
+ }
case INTERMEDIATE:
value = _mergeResultHolder[i];
break;
- case FINAL:
- value = aggFunction.extractFinalResult(_mergeResultHolder[i]);
+ case FINAL: {
+ Object mergedResult = _mergeResultHolder[i];
+ // The merged result stays null when there is no input row, even if
the leaf returns final results. In that
+ // case, extract the final result of the empty input, e.g. 0 for
COUNT.
+ value = _leafReturnFinalResult && mergedResult != null
+ ? mergedResult
+ : aggFunction.extractFinalResult(mergedResult);
break;
- case DIRECT:
+ }
+ case DIRECT: {
Object intermediate =
aggFunction.extractAggregationResult(_aggregateResultHolder[i]);
value = aggFunction.extractFinalResult(intermediate);
break;
+ }
default:
throw new UnsupportedOperationException("Unsupported aggTyp: " +
_aggType);
}
@@ -158,8 +171,19 @@ public class MultistageAggregationExecutor {
}
Object[] intermediateResults =
AggregateOperator.getIntermediateResults(aggFunction, block);
for (Object intermediateResult : intermediateResults) {
+ _mergeResultHolder[i] = AggregationFunctionUtils.merge(aggFunction,
_mergeResultHolder[i], intermediateResult);
+ }
+ }
+ }
+
+ private void processMergeFinalResults(MseBlock.Data block) {
+ for (int i = 0; i < _aggFunctions.length; i++) {
+ AggregationFunction aggFunction = _aggFunctions[i];
+ Object[] finalResults =
AggregateOperator.getIntermediateResults(aggFunction, block);
+ for (Object finalResult : finalResults) {
_mergeResultHolder[i] =
- AggregationFunctionUtils.merge(aggFunction, _mergeResultHolder[i],
intermediateResult);
+ AggregationFunctionUtils.mergeFinalResult(aggFunction,
(Comparable) _mergeResultHolder[i],
+ (Comparable) finalResult);
}
}
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
index 2a3c35429be..1281be54cf3 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
@@ -49,7 +49,7 @@ import org.apache.pinot.query.planner.plannode.ValueNode;
import org.apache.pinot.query.planner.plannode.WindowNode;
import org.apache.pinot.query.runtime.blocks.MseBlock;
import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerResult;
-import org.apache.pinot.spi.utils.CommonConstants;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
@@ -75,24 +75,29 @@ public class ServerPlanRequestVisitor implements
PlanNodeVisitor<Void, ServerPla
if (!groupByList.isEmpty()) {
pinotQuery.setGroupByList(groupByList);
}
- /// GROUP BY GROUPING SETS / ROLLUP / CUBE: push the per-set row
expansion down to the single-stage engine. The
- /// per-set column-index lists are over groupByList (==
node.getGroupKeys()), so they transfer directly. The
- /// single-stage leaf expands each row across the sets and appends the
synthetic $groupingId ordinal column; the
- /// multi-stage final stage then groups on it (it is one of the group
keys), keeping the grouping sets distinct.
+ // GROUP BY GROUPING SETS / ROLLUP / CUBE: push the per-set row
expansion down to the single-stage engine. The
+ // per-set column-index lists are over groupByList (==
node.getGroupKeys()), so they transfer directly. The
+ // single-stage leaf expands each row across the sets and appends the
synthetic $groupingId ordinal column; the
+ // multi-stage final stage then groups on it (it is one of the group
keys), keeping the grouping sets distinct.
if (node.isGroupingSets()) {
pinotQuery.setGroupingSets(node.getGroupingSets());
}
- List<Expression> selectList =
CalciteRexExpressionParser.convertAggregateList(groupByList, node.getAggCalls(),
- node.getFilterArgs(), pinotQuery.getSelectList());
+ List<Expression> selectList =
+ CalciteRexExpressionParser.convertAggregateList(groupByList,
node.getAggCalls(), node.getFilterArgs(),
+ pinotQuery.getSelectList());
for (Expression expression : selectList) {
applyTimestampIndex(expression, pinotQuery);
}
pinotQuery.setSelectList(selectList);
if (node.getAggType() == AggregateNode.AggType.DIRECT) {
-
pinotQuery.putToQueryOptions(CommonConstants.Broker.Request.QueryOptionKey.SERVER_RETURN_FINAL_RESULT,
"true");
+
pinotQuery.putToQueryOptions(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true");
} else if (node.isLeafReturnFinalResult()) {
- pinotQuery.putToQueryOptions(
-
CommonConstants.Broker.Request.QueryOptionKey.SERVER_RETURN_FINAL_RESULT_KEY_UNPARTITIONED,
"true");
+ // SERVER_RETURN_FINAL_RESULT_KEY_UNPARTITIONED only applies to
group-by queries.
+ if (groupByList.isEmpty()) {
+
pinotQuery.putToQueryOptions(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true");
+ } else {
+
pinotQuery.putToQueryOptions(QueryOptionKey.SERVER_RETURN_FINAL_RESULT_KEY_UNPARTITIONED,
"true");
+ }
}
int limit = node.getLimit();
if (limit > 0) {
@@ -136,8 +141,8 @@ public class ServerPlanRequestVisitor implements
PlanNodeVisitor<Void, ServerPla
if (visit(node.getInputs().get(0), context)) {
PinotQuery pinotQuery = context.getPinotQuery();
if (pinotQuery.getFilterExpression() == null) {
- Expression expression =
CalciteRexExpressionParser.toExpression(node.getCondition(),
- pinotQuery.getSelectList());
+ Expression expression =
+ CalciteRexExpressionParser.toExpression(node.getCondition(),
pinotQuery.getSelectList());
applyTimestampIndex(expression, pinotQuery);
pinotQuery.setFilterExpression(expression);
} else {
@@ -219,8 +224,8 @@ public class ServerPlanRequestVisitor implements
PlanNodeVisitor<Void, ServerPla
if (rex.getType() == EnrichedJoinNode.FilterProjectRexType.FILTER) {
// filter logic here
if (pinotQuery.getFilterExpression() == null) {
- Expression expression =
CalciteRexExpressionParser.toExpression(rex.getFilter(),
- pinotQuery.getSelectList());
+ Expression expression =
+ CalciteRexExpressionParser.toExpression(rex.getFilter(),
pinotQuery.getSelectList());
applyTimestampIndex(expression, pinotQuery);
pinotQuery.setFilterExpression(expression);
} else {
@@ -230,8 +235,7 @@ public class ServerPlanRequestVisitor implements
PlanNodeVisitor<Void, ServerPla
} else {
// project logic here
List<Expression> selectList =
- CalciteRexExpressionParser.convertRexNodes(
- rex.getProjectAndResultSchema().getProject(),
+
CalciteRexExpressionParser.convertRexNodes(rex.getProjectAndResultSchema().getProject(),
pinotQuery.getSelectList());
for (Expression expression : selectList) {
applyTimestampIndex(expression, pinotQuery);
@@ -269,8 +273,8 @@ public class ServerPlanRequestVisitor implements
PlanNodeVisitor<Void, ServerPla
public Void visitProject(ProjectNode node, ServerPlanRequestContext context)
{
if (visit(node.getInputs().get(0), context)) {
PinotQuery pinotQuery = context.getPinotQuery();
- List<Expression> selectList =
CalciteRexExpressionParser.convertRexNodes(node.getProjects(),
- pinotQuery.getSelectList());
+ List<Expression> selectList =
+ CalciteRexExpressionParser.convertRexNodes(node.getProjects(),
pinotQuery.getSelectList());
for (Expression expression : selectList) {
applyTimestampIndex(expression, pinotQuery);
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
index cc23e88ffbb..11a6a86840e 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java
@@ -426,6 +426,62 @@ public class AggregateOperatorTest {
assertTrue(result.nextBlock().isSuccess());
}
+ /// Each leaf returns the final results over its own partition of values,
and the final stage merges them.
+ @Test
+ public void testLeafReturnFinalResultWithoutGroupBy() {
+ OpChainExecutionContext context = OperatorTestUtil.getContext(Map.of());
+ List<RexExpression.FunctionCall> aggCalls =
+ List.of(new RexExpression.FunctionCall(INT, "DISTINCTCOUNT",
List.of(new RexExpression.InputRef(0))),
+ new RexExpression.FunctionCall(DOUBLE, "SUM", List.of(new
RexExpression.InputRef(1))));
+ DataSchema inputSchema = new DataSchema(new String[]{"value", "num"}, new
ColumnDataType[]{INT, INT});
+ DataSchema resultSchema = new DataSchema(new String[]{"distinctCount",
"sum"}, new ColumnDataType[]{INT, DOUBLE});
+ BlockListMultiStageOperator.Builder partials = new
BlockListMultiStageOperator.Builder(context, resultSchema);
+ // Each leaf has 5 values of its own, each appearing twice.
+ for (int partition = 0; partition < 2; partition++) {
+ BlockListMultiStageOperator.Builder input = new
BlockListMultiStageOperator.Builder(context, inputSchema);
+ for (int value = 0; value < 5; value++) {
+ input.addRow(partition * 10 + value, 1);
+ input.addRow(partition * 10 + value, 1);
+ }
+ AggregateOperator leaf = new AggregateOperator(context,
input.buildWithEos(),
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1, -1), List.of(),
+ AggType.LEAF, true, null, 0));
+ MseBlock leafBlock = leaf.nextBlock();
+ List<Object[]> leafRows = ((MseBlock.Data)
leafBlock).asRowHeap().getRows();
+ assertEquals(leafRows.size(), 1);
+ assertEquals(leafRows.get(0), new Object[]{5, 10.0});
+ partials.addBlock(leafBlock);
+ assertTrue(leaf.nextBlock().isSuccess());
+ }
+ AggregateOperator result = new AggregateOperator(context,
partials.buildWithEos(),
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1, -1), List.of(),
+ AggType.FINAL, true, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+ assertEquals(rows.get(0), new Object[]{10, 20.0});
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
+ /// Without any input row, the final stage returns the results of the empty
input.
+ @Test
+ public void testLeafReturnFinalResultWithoutGroupByEmptyInput() {
+ OpChainExecutionContext context = OperatorTestUtil.getContext(Map.of());
+ List<RexExpression.FunctionCall> aggCalls =
+ List.of(new RexExpression.FunctionCall(INT, "DISTINCTCOUNT",
List.of(new RexExpression.InputRef(0))),
+ new RexExpression.FunctionCall(DOUBLE, "SUM", List.of(new
RexExpression.InputRef(1))));
+ DataSchema resultSchema = new DataSchema(new String[]{"distinctCount",
"sum"}, new ColumnDataType[]{INT, DOUBLE});
+ AggregateOperator result = new AggregateOperator(context,
+ new BlockListMultiStageOperator.Builder(context,
resultSchema).buildWithEos(),
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1, -1), List.of(),
+ AggType.FINAL, true, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+ assertEquals(rows.get(0), new Object[]{0, null});
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
@DataProvider
public Object[][] approximatedResultCases() {
return new Object[][]{{true, 1, true}, {true, 10, false}, {false, 1,
true}, {false, 10, false}};
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
index 00932c93219..9dfd70ce1aa 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SerializedIntermediateResultMergeTest.java
@@ -127,7 +127,7 @@ public class SerializedIntermediateResultMergeTest {
}
private static MultistageAggregationExecutor newAggregationExecutor() {
- return new MultistageAggregationExecutor(newFunctions(), new int[]{-1,
-1}, -1, AggType.FINAL,
+ return new MultistageAggregationExecutor(newFunctions(), new int[]{-1,
-1}, -1, AggType.FINAL, false,
AGGREGATION_RESULT_SCHEMA);
}
diff --git a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
index 1bbf8754e02..6323dacd803 100644
--- a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
+++ b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
@@ -95,6 +95,14 @@
"description": "Group by non-partition column but aggregate on
partition column with partition parallelism and hint",
"sql": "SELECT /*+ aggOptions(is_leaf_return_final_result='true') */
{tbl1}.name, COUNT(*), SUM({tbl1}.num), COUNT(DISTINCT {tbl1}.num) FROM {tbl1}
/*+ tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4', partition_parallelism='2') */ GROUP BY {tbl1}.name"
},
+ {
+ "description": "Aggregate on partition column without group by with
hint",
+ "sql": "SELECT /*+ aggOptions(is_leaf_return_final_result='true') */
COUNT(*), SUM({tbl1}.num), COUNT(DISTINCT {tbl1}.num) FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */"
+ },
+ {
+ "description": "Aggregate on partition column without group by with
partition parallelism and hint",
+ "sql": "SELECT /*+ aggOptions(is_leaf_return_final_result='true') */
COUNT(*), SUM({tbl1}.num), COUNT(DISTINCT {tbl1}.num) FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4', partition_parallelism='2') */"
+ },
{
"description": "Skip leaf stage aggregation with GROUP BY hint",
"sql": "SELECT /*+ aggOptions(is_skip_leaf_stage_group_by='true') */
{tbl1}.name, COUNT(*), SUM({tbl1}.num), MIN({tbl1}.num), COUNT(DISTINCT
{tbl1}.num) FROM {tbl1} WHERE {tbl1}.num >= 0 GROUP BY {tbl1}.name"
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
