This is an automated email from the ASF dual-hosted git repository.
xiangfu0 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 de2ae09d7fe Fix MSE group trimming for intermediate aggregation
results (#19800)
de2ae09d7fe is described below
commit de2ae09d7fe368ec806affdca24c5d50516d1ea1
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Oct 12 20:14:24 2026 -0700
Fix MSE group trimming for intermediate aggregation results (#19800)
* Fix MSE group trimming for intermediate aggregation results
Extract final stored sort values once per group while preserving
intermediate
aggregation payloads for downstream merging. Copy consumptive funnel queues
before extracting their sort values.
How to reproduce:
Enable MSE group trimming on a grouped query following a join, order by
DISTINCTCOUNTHLL or DISTINCTCOUNTSMARTHLL, and set a small
mse_min_group_trim_size. LEAF or INTERMEDIATE aggregation compares a sketch
or set as Comparable and fails with ClassCastException.
* Avoid extra allocations in scalar MSE group trimming
Keep the original raw-row heap when sort columns already have their final
stored representation. Only decorate intermediate object values or values
whose final stored type differs. Cover scalar COUNT trimming through leaf,
intermediate, and final aggregation in the existing regression fixture.
How to reproduce:
Trim 100,000 intermediate SUM groups to 5,000 rows. The first fix allocated
a SortedRow and a cloned result array for every group, increasing allocation
from about 4.0 MB to 8.8 MB per trim. The fast path restores baseline
allocation.
* Preserve safe ordering and intermediate state during MSE group trim
Keep SUMPRECISION, AVG and range intermediates on their natural
Comparable heap path. Preserve all groups for raw-sketch, array and
funnel sort keys that cannot safely be finalized for intermediate trim.
Use compact lazy sort-key caches and a shared heap loop, and avoid
sorting when intermediate groups already fit within the trim capacity.
How to reproduce:
Enable MSE ordered group trimming with LIMIT 1 and minimum trim size 1,
then ORDER BY SUMPRECISION(value) DESC over grouped sums 5,6,7,8,9,100.
The prior trim converts BigDecimal to STRING and drops the 100 group.
HISTOGRAM ordered trimming also casts double[] to Comparable.
Validation: 40 focused runtime tests and affected-module formatting,
checkstyle and license checks pass. Four regression controls fail on
the previous PR head.
* Cache MSE trim sort indexes and cover untrimmed intermediate results
Precompute object-comparator field indexes outside heap comparisons.
Extend the existing HLL parameterized regression with a single case
where all six groups fit below the trim capacity and preserve their
intermediate sketches through merging.
How to reproduce:
Run ordered LEAF aggregation with six groups, LIMIT 1 and
MSE_MIN_GROUP_TRIM_SIZE=10. Verify all groups survive and their
two disjoint partials merge to cardinalities 2 through 12.
* Simplify MSE group trimming with native final sort values
Use cached native final aggregate keys for intermediate group trimming,
following SSE, without function-name exclusions or numeric-only guards.
Preserve funnel intermediate queues, bitmaps and sets during extraction
so selected rows remain mergeable. Keep the scalar heap fast path.
How to reproduce:
Enable MSE group trimming after a join, order by DISTINCTCOUNTHLL or
DISTINCTCOUNTSMARTHLL, and exceed the trim capacity. The original heap
compares an intermediate sketch/set as Comparable. Extracting funnel
sort values also consumes queue events or intersects partial step sets,
losing data required by later merging.
Validation: 58 focused tests pass; seven new funnel cases fail on the
previous implementations. Both affected modules pass spotless, license
formatting, checkstyle and license checks.
* Avoid eager copies when finalizing set funnels
Read the first-step count directly, stop at empty intersections, and copy
only the smaller operand while preserving primitive numeric containers.
Reuse the private intersection on later steps when it is already smaller.
Keep intermediate sets intact for later extraction and merging.
How to reproduce: finalize a single-step set funnel, or a funnel with a
large first-step population and an empty second step. The prior finalizer
copies every first-step member despite needing only its size. A boxed copy
also collapses distinct NaN payloads in float and double sets.
Validation: 64 focused funnel/MSE tests and all four pinot-core quality
checks pass. Three regressions fail with the preceding SetMergeStrategy
and pass with this change.
* Avoid eager copies when finalizing bitmap funnels
Read first-step cardinality directly. Create a private bitmap intersection
only when two nonempty steps need intersecting, reuse it for later steps,
and stop at empty populations or intersections without mutating inputs.
How to reproduce: finalize a single-step bitmap funnel, or one with a
large first population and an empty second step. The previous finalizer
clones the first population although only its cardinality is needed.
Validation: 27 focused funnel tests and all four pinot-core formatting,
checkstyle and license checks pass, with no warnings on added lines.
Both new regressions fail with the preceding bitmap finalizer.
* Reuse object sort entries during group trimming
Bound query-local sort wrappers and key arrays to the trim capacity plus
one by recycling rejected or evicted entries. Reset cached keys only after an
entry leaves the heap, preserving native comparison and intermediate payloads.
How to reproduce: trim 100,000 intermediate OBJECT groups to 500 or 5,000
rows and profile allocations; the previous implementation creates a wrapper and
key array for every candidate.
Validation: 67 focused MSE/funnel tests and all four affected-module
formatting, checkstyle and license checks pass. The reuse regression rejects a
missing-cache-reset mutation.
---
.../aggregation/function/AggregationFunction.java | 2 +
.../function/funnel/BitmapMergeStrategy.java | 16 +-
.../function/funnel/SetMergeStrategy.java | 54 +++++-
.../FunnelCompleteCountAggregationFunction.java | 1 +
...unnelEventsFunctionEvalAggregationFunction.java | 1 +
.../window/FunnelMatchStepAggregationFunction.java | 1 +
.../window/FunnelMaxStepAggregationFunction.java | 1 +
...FunnelStepDurationStatsAggregationFunction.java | 1 +
.../function/funnel/FunnelFinalResultTest.java | 193 ++++++++++++++++++++
.../query/runtime/operator/AggregateOperator.java | 16 +-
.../operator/MultistageGroupByExecutor.java | 108 ++++++++++-
.../runtime/operator/AggregateOperatorTest.java | 198 +++++++++++++++++++++
.../operator/MultistageGroupByExecutorTest.java | 47 +++++
13 files changed, 613 insertions(+), 26 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
index 79ca2736e8a..ac42303dee6 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java
@@ -241,6 +241,8 @@ public interface AggregationFunction<IntermediateResult,
FinalResult extends Com
ColumnDataType getFinalResultColumnType();
/// Extracts the final result used in the broker response from the given
intermediate result.
+ /// When final results are extracted for group-by sorting, preserve the
intermediate information needed for
+ /// subsequent merging and final extraction.
///
/// A `null` intermediate result means nothing was aggregated, and this
method decides what that means for this
/// particular aggregation. It is the only place where that per-function
answer is expressed, so it must never
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/BitmapMergeStrategy.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/BitmapMergeStrategy.java
index 0be8bbadbe0..490c1f404ce 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/BitmapMergeStrategy.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/BitmapMergeStrategy.java
@@ -41,11 +41,19 @@ class BitmapMergeStrategy implements
MergeStrategy<List<RoaringBitmap>> {
@Override
public LongArrayList extractFinalResult(List<RoaringBitmap> stepsBitmaps) {
long[] result = new long[_numSteps];
- result[0] = stepsBitmaps.get(0).getCardinality();
+ RoaringBitmap intersection = stepsBitmaps.get(0);
+ result[0] = intersection.getCardinality();
for (int i = 1; i < _numSteps; i++) {
- // intersect this step with previous step
- stepsBitmaps.get(i).and(stepsBitmaps.get(i - 1));
- result[i] = stepsBitmaps.get(i).getCardinality();
+ RoaringBitmap step = stepsBitmaps.get(i);
+ if (intersection.isEmpty() || step.isEmpty()) {
+ break;
+ }
+ if (i == 1) {
+ intersection = RoaringBitmap.and(intersection, step);
+ } else {
+ intersection.and(step);
+ }
+ result[i] = intersection.getCardinality();
}
return LongArrayList.wrap(result);
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/SetMergeStrategy.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/SetMergeStrategy.java
index d5ae4bc987f..ba2ea449362 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/SetMergeStrategy.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/SetMergeStrategy.java
@@ -18,7 +18,16 @@
*/
package org.apache.pinot.core.query.aggregation.function.funnel;
+import it.unimi.dsi.fastutil.doubles.DoubleOpenHashSet;
+import it.unimi.dsi.fastutil.doubles.DoubleSet;
+import it.unimi.dsi.fastutil.floats.FloatOpenHashSet;
+import it.unimi.dsi.fastutil.floats.FloatSet;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSet;
import it.unimi.dsi.fastutil.longs.LongArrayList;
+import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
+import it.unimi.dsi.fastutil.longs.LongSet;
+import java.util.HashSet;
import java.util.List;
import java.util.Set;
@@ -41,12 +50,49 @@ class SetMergeStrategy implements MergeStrategy<List<Set>> {
@Override
public LongArrayList extractFinalResult(List<Set> stepsSets) {
long[] result = new long[_numSteps];
- result[0] = stepsSets.get(0).size();
+ Set<?> intersection = stepsSets.get(0);
+ result[0] = intersection.size();
for (int i = 1; i < _numSteps; i++) {
- // intersect this step with previous step
- stepsSets.get(i).retainAll(stepsSets.get(i - 1));
- result[i] = stepsSets.get(i).size();
+ Set<?> step = stepsSets.get(i);
+ if (intersection.isEmpty() || step.isEmpty()) {
+ break;
+ }
+ // Preserve the intermediate sets and copy only the smaller population.
+ if (i == 1 || intersection.size() > step.size()) {
+ intersection = intersection.size() <= step.size()
+ ? intersect(intersection, step)
+ : intersect(step, intersection);
+ } else {
+ intersection.retainAll(step);
+ }
+ result[i] = intersection.size();
}
return LongArrayList.wrap(result);
}
+
+ private static Set<?> intersect(Set<?> smaller, Set<?> larger) {
+ if (smaller instanceof IntSet intSet && larger instanceof IntSet
otherIntSet) {
+ IntSet intersection = new IntOpenHashSet(intSet);
+ intersection.retainAll(otherIntSet);
+ return intersection;
+ }
+ if (smaller instanceof LongSet longSet && larger instanceof LongSet
otherLongSet) {
+ LongSet intersection = new LongOpenHashSet(longSet);
+ intersection.retainAll(otherLongSet);
+ return intersection;
+ }
+ if (smaller instanceof FloatSet floatSet && larger instanceof FloatSet
otherFloatSet) {
+ FloatSet intersection = new FloatOpenHashSet(floatSet);
+ intersection.retainAll(otherFloatSet);
+ return intersection;
+ }
+ if (smaller instanceof DoubleSet doubleSet && larger instanceof DoubleSet
otherDoubleSet) {
+ DoubleSet intersection = new DoubleOpenHashSet(doubleSet);
+ intersection.retainAll(otherDoubleSet);
+ return intersection;
+ }
+ Set<?> intersection = new HashSet<>(smaller);
+ intersection.retainAll(larger);
+ return intersection;
+ }
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelCompleteCountAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelCompleteCountAggregationFunction.java
index 02a0392d850..186144d2143 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelCompleteCountAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelCompleteCountAggregationFunction.java
@@ -51,6 +51,7 @@ public class FunnelCompleteCountAggregationFunction extends
FunnelBaseAggregatio
if (stepEvents == null || stepEvents.isEmpty()) {
return totalCompletedRounds;
}
+ stepEvents = new PriorityQueue<>(stepEvents);
ArrayDeque<FunnelStepEvent> slidingWindow = new ArrayDeque<>();
while (!stepEvents.isEmpty()) {
fillWindow(stepEvents, slidingWindow);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelEventsFunctionEvalAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelEventsFunctionEvalAggregationFunction.java
index 62dd1e7d0f1..c2261bff918 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelEventsFunctionEvalAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelEventsFunctionEvalAggregationFunction.java
@@ -490,6 +490,7 @@ public class FunnelEventsFunctionEvalAggregationFunction
if (stepEvents == null || stepEvents.isEmpty()) {
return finalResults;
}
+ stepEvents = new PriorityQueue<>(stepEvents);
ArrayDeque<FunnelStepEventWithExtraFields> slidingWindow = new
ArrayDeque<>();
while (!stepEvents.isEmpty()) {
fillWindow(stepEvents, slidingWindow);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMatchStepAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMatchStepAggregationFunction.java
index 0c08dee146d..7de02133b19 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMatchStepAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMatchStepAggregationFunction.java
@@ -56,6 +56,7 @@ public class FunnelMatchStepAggregationFunction extends
FunnelBaseAggregationFun
if (stepEvents == null || stepEvents.isEmpty()) {
return result;
}
+ stepEvents = new PriorityQueue<>(stepEvents);
ArrayDeque<FunnelStepEvent> slidingWindow = new ArrayDeque<>();
while (!stepEvents.isEmpty()) {
fillWindow(stepEvents, slidingWindow);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMaxStepAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMaxStepAggregationFunction.java
index 7cac01529c9..01ea8750a87 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMaxStepAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelMaxStepAggregationFunction.java
@@ -51,6 +51,7 @@ public class FunnelMaxStepAggregationFunction extends
FunnelBaseAggregationFunct
if (stepEvents == null || stepEvents.isEmpty()) {
return finalMaxStep;
}
+ stepEvents = new PriorityQueue<>(stepEvents);
ArrayDeque<FunnelStepEvent> slidingWindow = new ArrayDeque<>();
while (!stepEvents.isEmpty()) {
fillWindow(stepEvents, slidingWindow);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelStepDurationStatsAggregationFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelStepDurationStatsAggregationFunction.java
index b162aabba43..f957464cafe 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelStepDurationStatsAggregationFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/funnel/window/FunnelStepDurationStatsAggregationFunction.java
@@ -97,6 +97,7 @@ public class FunnelStepDurationStatsAggregationFunction
extends FunnelBaseAggreg
if (stepEvents == null || stepEvents.isEmpty()) {
return new DoubleArrayList();
}
+ stepEvents = new PriorityQueue<>(stepEvents);
Map<Integer, List<Object>> stepValueAggregators = initValueAggregator();
boolean hasMatchedFunnel = false;
ArrayDeque<FunnelStepEvent> slidingWindow = new ArrayDeque<>();
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/funnel/FunnelFinalResultTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/funnel/FunnelFinalResultTest.java
new file mode 100644
index 00000000000..4e615237a4d
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/funnel/FunnelFinalResultTest.java
@@ -0,0 +1,193 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.function.funnel;
+
+import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
+import it.unimi.dsi.fastutil.doubles.DoubleOpenHashSet;
+import it.unimi.dsi.fastutil.floats.FloatOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.longs.LongArrayList;
+import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
+import it.unimi.dsi.fastutil.objects.ObjectArrayList;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.PriorityQueue;
+import java.util.Set;
+import org.apache.pinot.common.request.Literal;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.FunctionContext;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.funnel.window.FunnelCompleteCountAggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.funnel.window.FunnelEventsFunctionEvalAggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.funnel.window.FunnelMatchStepAggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.funnel.window.FunnelMaxStepAggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.funnel.window.FunnelStepDurationStatsAggregationFunction;
+import org.roaringbitmap.RoaringBitmap;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
+
+/// Final extraction must preserve partial funnel state for subsequent merging.
+public class FunnelFinalResultTest {
+ @DataProvider
+ public Object[][] funnels() {
+ ExpressionContext step0 = ExpressionContext.forIdentifier("step0");
+ ExpressionContext step1 = ExpressionContext.forIdentifier("step1");
+ List<ExpressionContext> arguments =
List.of(ExpressionContext.forIdentifier("ts"),
+ ExpressionContext.forLiteral(Literal.longValue(1000)),
ExpressionContext.forLiteral(Literal.intValue(2)),
+ step0, step1);
+ List<ExpressionContext> durationArguments = new ArrayList<>(arguments);
+
durationArguments.add(ExpressionContext.forLiteral(Literal.stringValue("DURATIONFUNCTIONS=COUNT")));
+ List<ExpressionContext> eventsArguments = new ArrayList<>(arguments);
+ eventsArguments.add(ExpressionContext.forLiteral(Literal.intValue(0)));
+ List<ExpressionContext> countArguments = List.of(
+ ExpressionContext.forFunction(
+ new FunctionContext(FunctionContext.Type.TRANSFORM, "steps",
List.of(step0, step1))),
+ ExpressionContext.forFunction(new
FunctionContext(FunctionContext.Type.TRANSFORM, "correlateby",
+ List.of(ExpressionContext.forIdentifier("userId")))));
+ List<ExpressionContext> setArguments = new ArrayList<>(countArguments);
+ setArguments.add(ExpressionContext.forFunction(new
FunctionContext(FunctionContext.Type.TRANSFORM, "settings",
+ List.of(ExpressionContext.forLiteral(Literal.stringValue("set"))))));
+ return new Object[][]{
+ {new FunnelMaxStepAggregationFunction(arguments, true), events(0),
events(1), 1, 2},
+ {new FunnelCompleteCountAggregationFunction(arguments, true),
events(0), events(1), 0, 1},
+ {new FunnelMatchStepAggregationFunction(arguments, true), events(0),
events(1),
+ IntArrayList.wrap(new int[]{1, 0}), IntArrayList.wrap(new int[]{1,
1})},
+ {new FunnelStepDurationStatsAggregationFunction(durationArguments,
true), events(0), events(1),
+ DoubleArrayList.wrap(new double[]{1, 0}), DoubleArrayList.wrap(new
double[]{1, 1})},
+ {new FunnelEventsFunctionEvalAggregationFunction(eventsArguments,
true), eventsWithExtraFields(0),
+ eventsWithExtraFields(1), new ObjectArrayList<>(List.of("0")), new
ObjectArrayList<>(List.of("1, 0"))},
+ {new FunnelCountAggregationFunctionFactory(countArguments, true).get(),
+ List.of(RoaringBitmap.bitmapOf(1), RoaringBitmap.bitmapOf(2)),
+ List.of(RoaringBitmap.bitmapOf(2), RoaringBitmap.bitmapOf(1)),
+ LongArrayList.wrap(new long[]{1, 0}), LongArrayList.wrap(new
long[]{2, 2})},
+ {new FunnelCountAggregationFunctionFactory(setArguments, true).get(),
+ List.of(new HashSet<>(List.of(1)), new HashSet<>(List.of(2))),
+ List.of(new HashSet<>(List.of(2)), new HashSet<>(List.of(1))),
+ LongArrayList.wrap(new long[]{1, 0}), LongArrayList.wrap(new
long[]{2, 2})}
+ };
+ }
+
+ @Test(dataProvider = "funnels")
+ public <I, F extends Comparable<?>> void
testFinalExtractionPreservesIntermediateResult(
+ AggregationFunction<I, F> function,
+ I intermediateResult, I otherIntermediateResult, F
expectedPartialResult, F expectedMergedResult) {
+ assertEquals(function.extractFinalResult(intermediateResult),
expectedPartialResult);
+ assertEquals(function.extractFinalResult(intermediateResult),
expectedPartialResult);
+
assertEquals(function.extractFinalResult(function.merge(intermediateResult,
otherIntermediateResult)),
+ expectedMergedResult);
+ }
+
+ @Test
+ public void testBitmapFinalExtractionAvoidsUnnecessaryCopies() {
+ RoaringBitmap firstStep = mock(RoaringBitmap.class);
+ when(firstStep.getCardinality()).thenReturn(1_000_000);
+ when(firstStep.clone()).thenThrow(new AssertionError("First-step
population must not be copied"));
+ assertEquals(new
BitmapMergeStrategy(1).extractFinalResult(List.of(firstStep)),
+ LongArrayList.wrap(new long[]{1_000_000}));
+ assertEquals(new BitmapMergeStrategy(3).extractFinalResult(
+ List.of(firstStep, new RoaringBitmap(), RoaringBitmap.bitmapOf(2))),
+ LongArrayList.wrap(new long[]{1_000_000, 0, 0}));
+ assertEquals(new BitmapMergeStrategy(2).extractFinalResult(List.of(new
RoaringBitmap(), firstStep)),
+ LongArrayList.wrap(new long[]{0, 0}));
+ verify(firstStep, never()).clone();
+ }
+
+ @Test
+ public void testBitmapFinalExtractionPreservesStepPopulations() {
+ RoaringBitmap laterStep = mock(RoaringBitmap.class);
+ when(laterStep.isEmpty()).thenThrow(new AssertionError("Steps after an
empty intersection must not be inspected"));
+ List<RoaringBitmap> steps = List.of(RoaringBitmap.bitmapOf(1, 2, 3),
RoaringBitmap.bitmapOf(2, 3),
+ RoaringBitmap.bitmapOf(3, 4), RoaringBitmap.bitmapOf(5, 6), laterStep);
+ int[][] originalSteps = steps.subList(0,
4).stream().map(RoaringBitmap::toArray).toArray(int[][]::new);
+ BitmapMergeStrategy strategy = new BitmapMergeStrategy(5);
+ LongArrayList expectedCounts = LongArrayList.wrap(new long[]{3, 2, 1, 0,
0});
+ assertEquals(strategy.extractFinalResult(steps), expectedCounts);
+ assertEquals(strategy.extractFinalResult(steps), expectedCounts);
+ for (int i = 0; i < originalSteps.length; i++) {
+ assertEquals(steps.get(i).toArray(), originalSteps[i]);
+ }
+ verify(laterStep, never()).isEmpty();
+ }
+
+ @Test
+ public void testSetFinalExtractionAvoidsTraversingLargeFirstStep() {
+ Set<?> firstStep = mock(Set.class);
+ when(firstStep.size()).thenReturn(1_000_000);
+ when(firstStep.contains(2)).thenReturn(true);
+ when(firstStep.iterator()).thenThrow(new AssertionError("Large first-step
population must not be traversed"));
+ assertEquals(new
SetMergeStrategy(1).extractFinalResult(List.of(firstStep)),
+ LongArrayList.wrap(new long[]{1_000_000}));
+ assertEquals(new SetMergeStrategy(3).extractFinalResult(List.of(firstStep,
Set.of(), Set.of(2))),
+ LongArrayList.wrap(new long[]{1_000_000, 0, 0}));
+ assertEquals(new SetMergeStrategy(3).extractFinalResult(List.of(firstStep,
Set.of(2, 3), Set.of(2, 4, 5))),
+ LongArrayList.wrap(new long[]{1_000_000, 1, 1}));
+ verify(firstStep, never()).iterator();
+ verify(firstStep, never()).toArray();
+ }
+
+ @DataProvider
+ public Object[][] stepSets() {
+ float floatNaN1 = Float.intBitsToFloat(0x7fc00001);
+ float floatNaN2 = Float.intBitsToFloat(0x7fc00002);
+ double doubleNaN1 = Double.longBitsToDouble(0x7ff8000000000001L);
+ double doubleNaN2 = Double.longBitsToDouble(0x7ff8000000000002L);
+ return new Object[][]{
+ {List.of(new IntOpenHashSet(new int[]{1, 2, 3}), new
IntOpenHashSet(new int[]{2, 3}),
+ new IntOpenHashSet(new int[]{3, 4, 5, 6})), new long[]{3, 2, 1}},
+ {List.of(new LongOpenHashSet(new long[]{1, 2, 3}), new
LongOpenHashSet(new long[]{2, 3}),
+ new LongOpenHashSet(new long[]{3, 4, 5, 6})), new long[]{3, 2, 1}},
+ {List.of(new FloatOpenHashSet(new float[]{floatNaN1, floatNaN2, -0.0f,
0.0f}),
+ new FloatOpenHashSet(new float[]{floatNaN1, floatNaN2, 0.0f}),
+ new FloatOpenHashSet(new float[]{floatNaN2, 0.0f, 1, 2})), new
long[]{4, 3, 2}},
+ {List.of(new DoubleOpenHashSet(new double[]{doubleNaN1, doubleNaN2,
-0.0, 0.0}),
+ new DoubleOpenHashSet(new double[]{doubleNaN1, doubleNaN2, 0.0}),
+ new DoubleOpenHashSet(new double[]{doubleNaN2, 0.0, 1, 2})), new
long[]{4, 3, 2}},
+ {List.of(Set.of("1", "2", "3"), Set.of("2", "3"), Set.of("3", "4",
"5", "6")), new long[]{3, 2, 1}}
+ };
+ }
+
+ @Test(dataProvider = "stepSets")
+ public void testSetFinalExtractionPreservesStepPopulations(List<Set> steps,
long[] expectedCounts) {
+ Object[][] originalSteps =
steps.stream().map(Set::toArray).toArray(Object[][]::new);
+ SetMergeStrategy strategy = new SetMergeStrategy(3);
+ assertEquals(strategy.extractFinalResult(steps),
LongArrayList.wrap(expectedCounts));
+ assertEquals(strategy.extractFinalResult(steps),
LongArrayList.wrap(expectedCounts));
+ for (int i = 0; i < steps.size(); i++) {
+ assertEquals(steps.get(i).toArray(), originalSteps[i]);
+ }
+ }
+
+ private static PriorityQueue<FunnelStepEvent> events(int step) {
+ return new PriorityQueue<>(List.of(new FunnelStepEvent(100L * (step + 1),
step)));
+ }
+
+ private static PriorityQueue<FunnelStepEventWithExtraFields>
eventsWithExtraFields(int step) {
+ return new PriorityQueue<>(
+ List.of(new FunnelStepEventWithExtraFields(new FunnelStepEvent(100L *
(step + 1), step), List.of())));
+ }
+}
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 4717d5a2b96..c13f7b09890 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
@@ -21,7 +21,6 @@ package org.apache.pinot.query.runtime.operator;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
-import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -51,7 +50,6 @@ import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.runtime.blocks.MseBlock;
import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
-import org.apache.pinot.query.runtime.operator.utils.SortUtils;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.utils.CommonConstants.Server;
@@ -88,10 +86,8 @@ public class AggregateOperator extends MultiStageOperator {
// trimming - related members
private final int _groupTrimSize;
- /// Comparator is used in priority queue, and the order is reversed so that
peek() returns the smallest row to be
- /// compared with other rows.
@Nullable
- private final Comparator<Object[]> _comparator;
+ private final List<RelFieldCollation> _collations;
public AggregateOperator(OpChainExecutionContext context, MultiStageOperator
input, AggregateNode node) {
super(context);
@@ -114,7 +110,7 @@ public class AggregateOperator extends MultiStageOperator {
_input = input;
int groupTrimSize = Integer.MAX_VALUE;
- Comparator<Object[]> comparator = null;
+ List<RelFieldCollation> trimCollations = null;
int limit = node.getLimit();
int minGroupTrimSize = getMinGroupTrimSize(node.getNodeHint(),
context.getOpChainMetadata());
if (limit > 0 && minGroupTrimSize > 0) {
@@ -127,12 +123,12 @@ public class AggregateOperator extends MultiStageOperator
{
} else {
groupTrimSize = GroupByUtils.getTableCapacity(limit, minGroupTrimSize);
if (groupTrimSize < Integer.MAX_VALUE) {
- comparator = new SortUtils.SortComparator(collations, true);
+ trimCollations = collations;
}
}
}
_groupTrimSize = groupTrimSize;
- _comparator = comparator;
+ _collations = trimCollations;
_errorOnNumGroupsLimit = getErrorOnNumGroupsLimit(node.getNodeHint(),
context.getOpChainMetadata());
@@ -244,8 +240,8 @@ public class AggregateOperator extends MultiStageOperator {
} else {
assert _groupByExecutor != null;
List<Object[]> rows;
- if (_comparator != null) {
- rows = _groupByExecutor.getResult(_comparator, _groupTrimSize);
+ if (_collations != null) {
+ rows = _groupByExecutor.getResult(_collations, _groupTrimSize);
} else {
rows = _groupByExecutor.getResult(_groupTrimSize);
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
index 5409aa7c463..af296abf93e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
@@ -19,13 +19,17 @@
package org.apache.pinot.query.runtime.operator;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
+import java.util.function.BiFunction;
+import java.util.function.Function;
import javax.annotation.Nullable;
+import org.apache.calcite.rel.RelFieldCollation;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.common.datablock.DataBlock;
import org.apache.pinot.common.request.context.ExpressionContext;
@@ -43,6 +47,7 @@ import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.runtime.blocks.MseBlock;
import org.apache.pinot.query.runtime.operator.groupby.GroupIdGenerator;
import org.apache.pinot.query.runtime.operator.groupby.GroupIdGeneratorFactory;
+import org.apache.pinot.query.runtime.operator.utils.SortUtils;
import org.apache.pinot.query.runtime.operator.utils.TypeUtils;
import org.apache.pinot.spi.utils.CommonConstants.Server;
import org.roaringbitmap.PeekableIntIterator;
@@ -52,6 +57,8 @@ import org.roaringbitmap.RoaringBitmap;
/// Class that executes the keyed group by aggregations for the multistage
AggregateOperator.
@SuppressWarnings({"rawtypes", "unchecked"})
public class MultistageGroupByExecutor {
+ private static final Object UNINITIALIZED_SORT_VALUE = new Object();
+
private final int[] _groupKeyIds;
private final AggregationFunction[] _aggFunctions;
private final int[] _filterArgIds;
@@ -186,15 +193,65 @@ public class MultistageGroupByExecutor {
}
}
- /// 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` and native final
aggregate values for intermediate rows.
+ 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;
+ }
+ sortFunctions[i] = aggFunction;
+ finalizeSortValues = true;
+ }
+ }
+ if (!finalizeSortValues) {
+ return getResult(new SortUtils.SortComparator(collations, true),
maxRows, (row, reusable) -> row,
+ Function.identity());
+ }
+ SortUtils.SortComparator[] comparators = new
SortUtils.SortComparator[numCollations];
+ int[] sortColumnIndexes = new int[numCollations];
+ for (int i = 0; i < numCollations; i++) {
+ RelFieldCollation collation = collations.get(i);
+ sortColumnIndexes[i] = collation.getFieldIndex();
+ comparators[i] = new
SortUtils.SortComparator(List.of(collation.withFieldIndex(i)), true);
+ }
+ Comparator<SortedRow> comparator = (r1, r2) -> {
+ for (int i = 0; i < numCollations; i++) {
+ int index = sortColumnIndexes[i];
+ cacheSortValue(r1, i, index, sortFunctions[i]);
+ cacheSortValue(r2, i, index, sortFunctions[i]);
+ int result = comparators[i].compare(r1._sortKeys, r2._sortKeys);
+ if (result != 0) {
+ return result;
+ }
+ }
+ return 0;
+ };
+ return getResult(comparator, maxRows, (row, reusable) -> getSortedRow(row,
numCollations, reusable),
+ row -> row._row);
+ }
+
+ /// Select groups with a shared heap loop; scalar rows use identity
functions without allocating heap wrappers.
+ private <T> List<Object[]> getResult(Comparator<T> comparator, int maxRows,
BiFunction<Object[], T, T> decorate,
+ Function<T, Object[]> unwrap) {
int numGroups = Math.min(_groupIdGenerator.getNumGroups(), maxRows);
if (numGroups == 0) {
return List.of();
}
- // TODO: Change it to use top-K algorithm
- PriorityQueue<Object[]> sortedRows = new PriorityQueue<>(numGroups,
comparator);
+ PriorityQueue<T> sortedRows = new PriorityQueue<>(numGroups, comparator);
int numKeys = _groupKeyIds.length;
int numFunctions = _aggFunctions.length;
ColumnDataType[] resultStoredTypes =
_resultSchema.getStoredColumnDataTypes();
@@ -203,18 +260,21 @@ public class MultistageGroupByExecutor {
int idx = 0;
while (idx < numGroups && groupKeyIterator.hasNext()) {
- Object[] row = getRow(groupKeyIterator, numKeys, numFunctions,
resultStoredTypes);
+ T row = decorate.apply(getRow(groupKeyIterator, numKeys, numFunctions,
resultStoredTypes), null);
sortedRows.add(row);
idx++;
}
+ T reusable = null;
while (groupKeyIterator.hasNext()) {
idx++;
// TODO: allocate new array row only if row enters set
- Object[] row = getRow(groupKeyIterator, numKeys, numFunctions,
resultStoredTypes);
+ T row = decorate.apply(getRow(groupKeyIterator, numKeys, numFunctions,
resultStoredTypes), reusable);
if (comparator.compare(sortedRows.peek(), row) < 0) {
- sortedRows.poll();
+ reusable = sortedRows.poll();
sortedRows.offer(row);
+ } else {
+ reusable = row;
}
}
@@ -223,13 +283,45 @@ public class MultistageGroupByExecutor {
int resultSize = sortedRows.size();
ArrayList<Object[]> result = new ArrayList<>(sortedRows.size());
for (int i = resultSize - 1; i >= 0; i--) {
- result.add(sortedRows.poll());
+ result.add(unwrap.apply(sortedRows.poll()));
}
// reverse priority queue order because comparators are reversed
Collections.reverse(result);
return result;
}
+ private SortedRow getSortedRow(Object[] row, int numCollations, @Nullable
SortedRow reusable) {
+ if (reusable == null) {
+ reusable = new SortedRow(row, new Object[numCollations]);
+ } else {
+ reusable._row = row;
+ }
+ // Only rejected or evicted entries are reused; never change a row or
cached key while it is in the heap.
+ Arrays.fill(reusable._sortKeys, UNINITIALIZED_SORT_VALUE);
+ return reusable;
+ }
+
+ private void cacheSortValue(SortedRow row, int sortIndex, int columnIndex,
+ @Nullable AggregationFunction aggFunction) {
+ if (row._sortKeys[sortIndex] == UNINITIALIZED_SORT_VALUE) {
+ Object value = row._row[columnIndex];
+ // Like SSE, compare native final objects before output conversion (e.g.
BigDecimal, SerializedHLL, lists).
+ // Extract lazily and retain the intermediate row for merging. Cache
null results as well as non-null results.
+ row._sortKeys[sortIndex] = aggFunction != null ?
aggFunction.extractFinalResult(value) : value;
+ }
+ }
+
+ /// Query-local heap entry caching ORDER BY values. Its row and cache can be
reset only outside the heap.
+ private static final class SortedRow {
+ private Object[] _row;
+ private final Object[] _sortKeys;
+
+ private SortedRow(Object[] row, Object[] sortKeys) {
+ _row = row;
+ _sortKeys = sortKeys;
+ }
+ }
+
/// Get aggregation result limited to `maxRows` rows.
public List<Object[]> getResult(int maxRows) {
int numGroups = Math.min(_groupIdGenerator.getNumGroups(), maxRows);
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 60ab3e50356..0e68a59e350 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
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.query.runtime.operator;
+import java.math.BigDecimal;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -44,6 +45,7 @@ import org.apache.pinot.spi.utils.CommonConstants.Server;
import org.mockito.Mock;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.BOOLEAN;
@@ -55,6 +57,7 @@ import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.MockitoAnnotations.openMocks;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertEqualsNoOrder;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
@@ -351,6 +354,201 @@ public class AggregateOperatorTest {
assertEquals(getAggregateOperator(context, nodeHint, 5, List.of(new
RelFieldCollation(1))).getGroupTrimSize(), 30);
}
+ @DataProvider
+ public Object[][] orderedGroupTrimCases() {
+ return new Object[][]{
+ {"COUNT", null, false, 1},
+ {"COUNT", null, true, 1},
+ {"DISTINCTCOUNT", null, false, 1},
+ {"DISTINCTCOUNT", null, true, 1},
+ {"DISTINCTCOUNTHLL", null, false, 1},
+ {"DISTINCTCOUNTHLL", null, true, 1},
+ {"DISTINCTCOUNTSMARTHLL", null, false, 1},
+ {"DISTINCTCOUNTSMARTHLL", null, true, 1},
+ {"DISTINCTCOUNTSMARTHLL", "threshold=1", false, 1},
+ {"DISTINCTCOUNTSMARTHLL", "threshold=1", true, 1},
+ {"DISTINCTCOUNTHLL", null, false, 10}
+ };
+ }
+
+ @Test(dataProvider = "orderedGroupTrimCases")
+ public void testOrderedGroupTrimPreservesIntermediateResults(String
functionName, @Nullable String parameters,
+ boolean leafReturnFinalResult, int minGroupTrimSize) {
+ OpChainExecutionContext context = OperatorTestUtil.getContext(
+ Map.of(QueryOptionKey.MSE_MIN_GROUP_TRIM_SIZE,
Integer.toString(minGroupTrimSize)));
+ ColumnDataType finalType = functionName.equals("DISTINCTCOUNTHLL") ||
functionName.equals("COUNT")
+ ? ColumnDataType.LONG
+ : INT;
+ List<RexExpression> operands = parameters == null
+ ? List.of(new RexExpression.InputRef(1))
+ : List.of(new RexExpression.InputRef(1), new
RexExpression.Literal(STRING, parameters));
+ List<RexExpression.FunctionCall> aggCalls =
+ List.of(new RexExpression.FunctionCall(finalType, functionName,
operands));
+ List<RelFieldCollation> collations =
+ List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.DESCENDING), new RelFieldCollation(0));
+ DataSchema inputSchema = new DataSchema(new String[]{"group", "value"},
new ColumnDataType[]{INT, INT});
+ DataSchema resultSchema =
+ new DataSchema(new String[]{"group", "distinctCount"}, new
ColumnDataType[]{INT, finalType});
+ ColumnDataType intermediateType = leafReturnFinalResult ||
functionName.equals("COUNT")
+ ? finalType
+ : ColumnDataType.OBJECT;
+ DataSchema intermediateSchema =
+ new DataSchema(new String[]{"group", "distinctCount"}, new
ColumnDataType[]{INT, intermediateType});
+ BlockListMultiStageOperator.Builder partials = new
BlockListMultiStageOperator.Builder(context, intermediateSchema);
+ // Each leaf has six groups with cardinalities 1..6. Disjoint partials
must merge to cardinalities 2..12.
+ for (int partition = 0; partition < 2; partition++) {
+ BlockListMultiStageOperator.Builder input = new
BlockListMultiStageOperator.Builder(context, inputSchema);
+ for (int group = 1; group <= 6; group++) {
+ for (int value = 0; value < group; value++) {
+ input.addRow(group, partition * 10 + value);
+ }
+ }
+ AggregateOperator leaf = new AggregateOperator(context,
input.buildWithEos(),
+ new AggregateNode(-1, intermediateSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1),
+ List.of(0), AggType.LEAF, leafReturnFinalResult, collations, 1));
+ partials.addBlock(leaf.nextBlock());
+ assertTrue(leaf.nextBlock().isSuccess());
+ }
+ AggregateOperator result = new AggregateOperator(context,
partials.buildWithEos(),
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.FINAL, leafReturnFinalResult, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), minGroupTrimSize == 1 ? 5 : 6);
+ Map<Integer, Long> counts = new HashMap<>();
+ for (Object[] row : rows) {
+ counts.put((Integer) row[0], ((Number) row[1]).longValue());
+ }
+ assertEquals(counts, minGroupTrimSize == 1
+ ? Map.of(2, 4L, 3, 6L, 4, 8L, 5, 10L, 6, 12L)
+ : Map.of(1, 2L, 2, 4L, 3, 6L, 4, 8L, 5, 10L, 6, 12L));
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
+ @Test
+ public void testOrderedSumPrecisionGroupTrimUsesNumericResults() {
+ OpChainExecutionContext context =
+
OperatorTestUtil.getContext(Map.of(QueryOptionKey.MSE_MIN_GROUP_TRIM_SIZE,
"1"));
+ DataSchema inputSchema = new DataSchema(new String[]{"group", "value"},
new ColumnDataType[]{INT, INT});
+ List<RexExpression.FunctionCall> aggCalls = List.of(new
RexExpression.FunctionCall(ColumnDataType.BIG_DECIMAL,
+ "SUMPRECISION", List.of(new RexExpression.InputRef(1))));
+ DataSchema intermediateSchema =
+ new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT,
ColumnDataType.OBJECT});
+ AggregateOperator leaf = new AggregateOperator(context,
+ new BlockListMultiStageOperator(context,
OperatorTestUtil.block(inputSchema, new Object[]{1, 5},
+ new Object[]{2, 6}, new Object[]{3, 7}, new Object[]{4, 8}, new
Object[]{5, 9}, new Object[]{6, 100})),
+ new AggregateNode(-1, intermediateSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.LEAF, false, List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.DESCENDING)), 1));
+ DataSchema resultSchema =
+ new DataSchema(new String[]{"group", "sum"}, new ColumnDataType[]{INT,
ColumnDataType.BIG_DECIMAL});
+ AggregateOperator result = new AggregateOperator(context, leaf,
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.FINAL, false, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 5);
+ assertEquals(rows.stream().map(row -> row[0]).sorted().toList(),
List.of(2, 3, 4, 5, 6));
+ assertEquals(rows.stream().filter(row ->
row[0].equals(6)).findFirst().orElseThrow()[1], BigDecimal.valueOf(100));
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
+ @Test
+ public void testOrderedGroupTrimPreservesFunnelEvents() {
+ OpChainExecutionContext context =
+
OperatorTestUtil.getContext(Map.of(QueryOptionKey.MSE_MIN_GROUP_TRIM_SIZE,
"1"));
+ DataSchema inputSchema = new DataSchema(new String[]{"group", "timestamp",
"step"},
+ new ColumnDataType[]{INT, ColumnDataType.LONG, BOOLEAN});
+ List<RexExpression.FunctionCall> aggCalls = List.of(new
RexExpression.FunctionCall(INT, "FUNNELMAXSTEP",
+ List.of(new RexExpression.InputRef(1), new
RexExpression.Literal(ColumnDataType.LONG, 1000L),
+ new RexExpression.Literal(INT, 1), new
RexExpression.InputRef(2))));
+ DataSchema intermediateSchema =
+ new DataSchema(new String[]{"group", "funnel"}, new
ColumnDataType[]{INT, ColumnDataType.OBJECT});
+ AggregateOperator leaf = new AggregateOperator(context,
+ new BlockListMultiStageOperator(context,
+ OperatorTestUtil.block(inputSchema, new Object[]{1, 1L, 1}, new
Object[]{2, 2L, 1}, new Object[]{3, 3L, 1},
+ new Object[]{4, 4L, 1}, new Object[]{5, 5L, 1}, new
Object[]{6, 6L, 1})),
+ new AggregateNode(-1, intermediateSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.LEAF, false, List.of(new RelFieldCollation(1), new
RelFieldCollation(0)), 1));
+ DataSchema resultSchema = new DataSchema(new String[]{"group", "funnel"},
new ColumnDataType[]{INT, INT});
+ AggregateOperator result = new AggregateOperator(context, leaf,
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.FINAL, false, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.stream().map(row -> row[0]).sorted().toList(),
List.of(1, 2, 3, 4, 5));
+ rows.forEach(row -> assertEquals(row[1], 1));
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
+ @Test
+ public void testOrderedRawHllGroupTrimUsesNativeCardinalities() {
+ OpChainExecutionContext context =
+
OperatorTestUtil.getContext(Map.of(QueryOptionKey.MSE_MIN_GROUP_TRIM_SIZE,
"1"));
+ DataSchema inputSchema = new DataSchema(new String[]{"group", "value"},
new ColumnDataType[]{INT, INT});
+ BlockListMultiStageOperator.Builder inputBuilder = new
BlockListMultiStageOperator.Builder(context, inputSchema);
+ for (int group = 1; group <= 6; group++) {
+ for (int value = 0; value < group; value++) {
+ inputBuilder.addRow(group, group * 10 + value);
+ }
+ }
+ MseBlock.Data input = (MseBlock.Data)
inputBuilder.buildWithEos().nextBlock();
+ List<RexExpression.FunctionCall> aggCalls = List.of(new
RexExpression.FunctionCall(STRING, "DISTINCTCOUNTRAWHLL",
+ List.of(new RexExpression.InputRef(1), new RexExpression.Literal(INT,
12))));
+ List<RelFieldCollation> collations =
+ List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.DESCENDING), new RelFieldCollation(0));
+ DataSchema resultSchema = new DataSchema(new String[]{"group", "hll"}, new
ColumnDataType[]{INT, STRING});
+ AggregateOperator direct = new AggregateOperator(context, new
BlockListMultiStageOperator(context, input),
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.DIRECT, false, null, 0));
+ DataSchema intermediateSchema =
+ new DataSchema(new String[]{"group", "hll"}, new ColumnDataType[]{INT,
ColumnDataType.OBJECT});
+ AggregateOperator leaf = new AggregateOperator(context, new
BlockListMultiStageOperator(context, input),
+ new AggregateNode(-1, intermediateSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.LEAF, false, collations, 1));
+ AggregateOperator result = new AggregateOperator(context, leaf,
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.FINAL, false, null, 0));
+
+ List<Object[]> expected = ((MseBlock.Data)
direct.nextBlock()).asRowHeap().getRows().stream()
+ .filter(row -> !row[0].equals(1)).toList();
+ List<Object[]> actual = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(actual.stream().map(row -> row[0]).sorted().toList(),
List.of(2, 3, 4, 5, 6));
+ assertEqualsNoOrder(actual.stream().map(row -> List.of(row)).toArray(),
+ expected.stream().map(row -> List.of(row)).toArray());
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
+ @Test
+ public void testOrderedHistogramGroupTrimUsesNativeLists() {
+ OpChainExecutionContext context =
+
OperatorTestUtil.getContext(Map.of(QueryOptionKey.MSE_MIN_GROUP_TRIM_SIZE,
"1"));
+ DataSchema inputSchema = new DataSchema(new String[]{"group", "value"},
new ColumnDataType[]{INT, INT});
+ List<RexExpression.FunctionCall> aggCalls = List.of(new
RexExpression.FunctionCall(ColumnDataType.DOUBLE_ARRAY,
+ "HISTOGRAM", List.of(new RexExpression.InputRef(1), new
RexExpression.Literal(DOUBLE, 0.0),
+ new RexExpression.Literal(DOUBLE, 10.0), new
RexExpression.Literal(INT, 10))));
+ DataSchema intermediateSchema =
+ new DataSchema(new String[]{"group", "histogram"}, new
ColumnDataType[]{INT, ColumnDataType.OBJECT});
+ BlockListMultiStageOperator.Builder input = new
BlockListMultiStageOperator.Builder(context, inputSchema);
+ for (int group = 1; group <= 6; group++) {
+ for (int value = 0; value < group; value++) {
+ input.addRow(group, 1);
+ }
+ }
+ AggregateOperator leaf = new AggregateOperator(context,
input.buildWithEos(),
+ new AggregateNode(-1, intermediateSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.LEAF, false, List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.DESCENDING)), 1));
+ DataSchema resultSchema =
+ new DataSchema(new String[]{"group", "histogram"}, new
ColumnDataType[]{INT, ColumnDataType.DOUBLE_ARRAY});
+ AggregateOperator result = new AggregateOperator(context, leaf,
+ new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY,
List.of(), aggCalls, List.of(-1), List.of(0),
+ AggType.FINAL, false, null, 0));
+
+ List<Object[]> rows = ((MseBlock.Data)
result.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.stream().map(row -> row[0]).sorted().toList(),
List.of(2, 3, 4, 5, 6));
+ rows.forEach(row -> assertEquals(((double[]) row[1])[1], ((Integer)
row[0]).doubleValue()));
+ assertTrue(result.nextBlock().isSuccess());
+ }
+
private AggregateOperator getAggregateOperator(OpChainExecutionContext
context, PlanNode.NodeHint nodeHint, int limit,
@Nullable List<RelFieldCollation> collations) {
List<RexExpression.FunctionCall> aggCalls = List.of(getSum(new
RexExpression.InputRef(1)));
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
index 6a9956e6be1..505a92af6f5 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
@@ -22,19 +22,23 @@ import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import org.apache.calcite.rel.RelFieldCollation;
import org.apache.pinot.common.request.context.ExpressionContext;
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.AvgAggregationFunction;
import
org.apache.pinot.core.query.aggregation.function.CountAggregationFunction;
import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.segment.local.customobject.AvgPair;
import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
@@ -134,6 +138,49 @@ public class MultistageGroupByExecutorTest {
assertEquals(executor.getNumGroups(), expected.size());
}
+ @Test
+ public void testOrderedIntermediateGroupsReuseSortEntries() {
+ DataSchema schema = new DataSchema(new String[]{"group", "avg"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.OBJECT});
+ AvgAggregationFunction function =
+ new
AvgAggregationFunction(List.of(ExpressionContext.forIdentifier("$1")), true);
+ MultistageGroupByExecutor executor = new MultistageGroupByExecutor(new
int[]{0},
+ new AggregationFunction[]{function}, new int[]{-1}, -1,
AggType.INTERMEDIATE, false, schema, Map.of(), null);
+ Double[] averages = {null, 10.0, 0.0, 30.0, 5.0, 30.0, null, 20.0, 40.0,
1.0, 40.0, 0.0};
+ AvgPair[] pairs = new AvgPair[averages.length];
+ Object[][] rows = new Object[averages.length][];
+ for (int i = 0; i < averages.length; i++) {
+ pairs[i] = averages[i] == null ? null : new AvgPair(averages[i] * (i +
1), i + 1);
+ rows[i] = new Object[]{i, pairs[i]};
+ }
+ executor.processBlock(OperatorTestUtil.block(schema, rows));
+
+ // More than top-K + 1 candidates exercise reuse of both rejected and
evicted entries, with ties and null keys.
+ List<List<RelFieldCollation>> orders = List.of(
+ List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.LAST),
+ new RelFieldCollation(0)),
+ List.of(new RelFieldCollation(1,
RelFieldCollation.Direction.ASCENDING, RelFieldCollation.NullDirection.FIRST),
+ new RelFieldCollation(0)));
+ int[][] expectedGroups = {{8, 10, 3}, {0, 6, 2}};
+ for (int repeat = 0; repeat < 2; repeat++) {
+ for (int order = 0; order < orders.size(); order++) {
+ List<Object[]> result = executor.getResult(orders.get(order), 3);
+ assertEquals(result.size(), 3);
+ for (int i = 0; i < result.size(); i++) {
+ int group = expectedGroups[order][i];
+ assertEquals(result.get(i)[0], group);
+ assertSame(result.get(i)[1], pairs[group]);
+ }
+ }
+ }
+ for (int i = 0; i < pairs.length; i++) {
+ if (pairs[i] != null) {
+ assertEquals(pairs[i].getSum(), averages[i] * (i + 1));
+ assertEquals(pairs[i].getCount(), i + 1);
+ }
+ }
+ }
+
private static MultistageGroupByExecutor newExecutor(int numKeys, boolean
leafReturnFinalResult, int groupLimit) {
return newExecutor(numKeys, leafReturnFinalResult, groupLimit, -1, false);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
