xiangfu0 commented on code in PR #19800:
URL: https://github.com/apache/pinot/pull/19800#discussion_r4225334254
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -186,15 +192,74 @@ public void processBlock(MseBlock.Data block) {
}
}
- /// Get aggregation result limited to first `maxRows` rows, ordered with
`comparator`.
- public List<Object[]> getResult(Comparator<Object[]> comparator, int
maxRows) {
+ /// Select up to `maxRows` groups using `collations`, preserving all
intermediate groups when trim is unsafe.
+ public List<Object[]> getResult(List<RelFieldCollation> collations, int
maxRows) {
+ int numKeys = _groupKeyIds.length;
+ int numCollations = collations.size();
+ AggregationFunction[] sortFunctions = new
AggregationFunction[numCollations];
+ boolean finalizeSortValues = false;
+ if (_aggType.isOutputIntermediateFormat() && !_leafReturnFinalResult) {
+ if (_groupIdGenerator.getNumGroups() <= maxRows) {
Review Comment:
Minor (test gap): this early-return branch — the common production case,
since the default trim capacity is max(5*limit, 5000) and servers routinely
hold fewer groups than that — is not exercised by any of the new tests: every
ordered-trim test uses 6 groups against a trim capacity of 5, so `numGroups <=
maxRows` is always false. A future edit that breaks this branch (e.g. moving it
below the collation scan so funnel/sketch intermediates get finalized when
groups are few, or returning sorted-but-finalized rows) would pass the whole
suite. One parameterized case with trim capacity >= group count (e.g.
`MSE_MIN_GROUP_TRIM_SIZE` large enough that `getTableCapacity` exceeds 6)
asserting all groups survive with intact intermediates would pin it.
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java:
##########
@@ -186,15 +192,74 @@ public void processBlock(MseBlock.Data block) {
}
}
- /// Get aggregation result limited to first `maxRows` rows, ordered with
`comparator`.
- public List<Object[]> getResult(Comparator<Object[]> comparator, int
maxRows) {
+ /// Select up to `maxRows` groups using `collations`, preserving all
intermediate groups when trim is unsafe.
+ public List<Object[]> getResult(List<RelFieldCollation> collations, int
maxRows) {
+ int numKeys = _groupKeyIds.length;
+ int numCollations = collations.size();
+ AggregationFunction[] sortFunctions = new
AggregationFunction[numCollations];
+ boolean finalizeSortValues = false;
+ if (_aggType.isOutputIntermediateFormat() && !_leafReturnFinalResult) {
+ if (_groupIdGenerator.getNumGroups() <= maxRows) {
+ return getResult(maxRows);
+ }
+ for (int i = 0; i < numCollations; i++) {
+ int index = collations.get(i).getFieldIndex();
+ if (index < numKeys) {
+ continue;
+ }
+ AggregationFunction aggFunction = _aggFunctions[index - numKeys];
+ if (aggFunction.getIntermediateResultColumnType().getStoredType() !=
ColumnDataType.OBJECT) {
+ continue;
+ }
+ // These intermediate objects already compare by the aggregate's
numeric value.
+ switch (aggFunction.getType()) {
+ case AVG, AVGMV, MINMAXRANGE, MINMAXRANGEMV, SUMPRECISION:
+ continue;
+ default:
+ break;
+ }
+ // Serialized sketches and array results lack a safe scalar trim
order. Funnel finalization consumes events.
+ // Preserve every group for the downstream aggregate and sort instead
of finalizing these intermediates.
+ if (!aggFunction.getFinalResultColumnType().getStoredType().isNumber()
+ || aggFunction.getType().name().startsWith("FUNNEL")) {
+ return getResult(Integer.MAX_VALUE);
+ }
+ sortFunctions[i] = aggFunction;
+ finalizeSortValues = true;
+ }
+ }
+ if (!finalizeSortValues) {
+ return getResult(new SortUtils.SortComparator(collations, true),
maxRows, Function.identity(),
+ Function.identity());
+ }
+ SortUtils.SortComparator[] comparators = new
SortUtils.SortComparator[numCollations];
+ for (int i = 0; i < numCollations; i++) {
+ comparators[i] = new
SortUtils.SortComparator(List.of(collations.get(i).withFieldIndex(i)), true);
+ }
+ Comparator<SortedRow> comparator = (r1, r2) -> {
+ for (int i = 0; i < numCollations; i++) {
+ int index = collations.get(i).getFieldIndex();
Review Comment:
Nit (hot path): the lambda re-evaluates `collations.get(i).getFieldIndex()`
on every comparator invocation, for every collation level — O(numGroups log
trimSize) compares against up to `numGroupsLimit` (100K) groups. The indexes
are fixed when the comparator is built; precomputing them next to
`sortFunctions` removes two virtual calls per level per compare:
```suggestion
int[] sortColumnIndexes = new int[numCollations];
for (int i = 0; i < numCollations; i++) {
sortColumnIndexes[i] = collations.get(i).getFieldIndex();
}
Comparator<SortedRow> comparator = (r1, r2) -> {
for (int i = 0; i < numCollations; i++) {
int index = sortColumnIndexes[i];
```
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
