yashmayya commented on code in PR #19066:
URL: https://github.com/apache/pinot/pull/19066#discussion_r3641764127
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java:
##########
@@ -158,28 +207,7 @@ private void mergeBlock(GroupByResultsBlock resultsBlock) {
_numGroupsWarningLimitReached = true;
}
List<IntermediateRecord> intermediateRecords =
resultsBlock.getIntermediateRecords();
- if (intermediateRecords == null) {
- AggregationGroupByResult aggregationGroupByResult =
resultsBlock.getAggregationGroupByResult();
- if (aggregationGroupByResult != null) {
- try {
- Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator =
aggregationGroupByResult.getGroupKeyIterator();
- int mergedKeys = 0;
- while (groupKeyIterator.hasNext()) {
-
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++,
EXPLAIN_NAME);
- GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
- Object[] keys = groupKey._keys;
- Object[] values = Arrays.copyOf(keys, _numColumns);
- int groupId = groupKey._groupId;
- for (int i = 0; i < _numAggregationFunctions; i++) {
- values[_numGroupByExpressions + i] =
aggregationGroupByResult.getResultForGroupId(i, groupId);
- }
- _indexedTable.upsert(new Key(keys), new Record(values));
- }
- } finally {
- aggregationGroupByResult.closeGroupKeyGenerator();
- }
- }
- } else {
+ if (intermediateRecords != null) {
Review Comment:
Since raw blocks should never reach the consumer anymore, silently ignoring
them here makes a future regression hard to spot — if a new enqueue path ever
skips `detachFromWorkerThreadState`, we'd quietly drop an entire segment's
results. Can we fail loudly instead?
```java
} else if (resultsBlock.getAggregationGroupByResult() != null) {
throw new IllegalStateException("Raw group-by result reached the consumer
thread; it must be detached on the worker thread");
}
```
##########
pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java:
##########
@@ -238,6 +240,111 @@ public void testCountAggregation() {
}
}
+ /// Regression test for a data race in the streaming group-by leaf (see
+ /// {@link StreamingGroupByCombineOperator#detachFromWorkerThreadState}).
+ ///
+ /// <p>The per-segment {@link
org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator}
+ /// switches from a per-instance array holder to a REUSED thread-local
{@code IntGroupIdMap} once a segment's
+ /// group cardinality exceeds {@code maxInitialResultHolderCapacity}
(default 10000). If the streaming operator
+ /// hands a raw per-segment result to the consumer thread, that consumer
iterates the thread-local map while
+ /// the producing worker clears/expands it for its next segment — corrupting
group ids and blowing past the
+ /// result holder bounds. This test forces the thread-local holder (>10000
groups/segment) across many
+ /// segments with a low flush threshold, and asserts the merged sums are
exact over several iterations.
+ @Test
+ public void testHighCardinalityConcurrentMergeIsCorrect()
Review Comment:
All tests in this class go through the raw-result path, so the pass-through
branch of `detachFromWorkerThreadState` (blocks that already carry intermediate
records) is uncovered. Might be worth one ORDER BY test with a small segment
trim size to lock in that trimmed blocks flow through unchanged.
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java:
##########
@@ -143,6 +144,54 @@ protected BaseResultsBlock getNextBlock() {
return attachExecutionStats(new MetadataResultsBlock());
}
+ /// Detaches a per-segment group-by result from the worker thread's reused
thread-local group-key state.
+ ///
+ /// <p>A raw {@link AggregationGroupByResult} is backed by the group-key map
returned by
+ /// {@link
org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator}'s
thread-locals,
+ /// which the same worker thread clears and repopulates for its next
segment. If such a block were handed to
+ /// the consumer thread (as {@link BaseStreamingCombineOperator} does), the
consumer would iterate that map
+ /// while the worker mutates it — a data race that yields corrupted group
ids and out-of-bounds result-holder
+ /// access. We therefore materialize the result into self-contained {@link
IntermediateRecord}s here, on the
+ /// producing worker thread, before hand-off — mirroring how the
non-streaming
+ /// {@link org.apache.pinot.core.operator.combine.GroupByCombineOperator}
merges inline on the worker thread.
+ ///
+ /// <p>Blocks that already carry intermediate records (order-by /
segment-trim path) are self-contained and
+ /// returned unchanged.
+ @Override
+ protected GroupByResultsBlock
detachFromWorkerThreadState(GroupByResultsBlock resultsBlock) {
+ AggregationGroupByResult aggregationGroupByResult =
resultsBlock.getAggregationGroupByResult();
+ if (aggregationGroupByResult == null ||
resultsBlock.getIntermediateRecords() != null) {
+ return resultsBlock;
+ }
+ List<IntermediateRecord> records = new ArrayList<>();
Review Comment:
nit: `aggregationGroupByResult.getNumGroups()` is O(1), so this could be
presized — saves the grow-copies on exactly the >10k-group segments this path
targets.
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java:
##########
@@ -143,6 +144,54 @@ protected BaseResultsBlock getNextBlock() {
return attachExecutionStats(new MetadataResultsBlock());
}
+ /// Detaches a per-segment group-by result from the worker thread's reused
thread-local group-key state.
+ ///
+ /// <p>A raw {@link AggregationGroupByResult} is backed by the group-key map
returned by
+ /// {@link
org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator}'s
thread-locals,
+ /// which the same worker thread clears and repopulates for its next
segment. If such a block were handed to
+ /// the consumer thread (as {@link BaseStreamingCombineOperator} does), the
consumer would iterate that map
+ /// while the worker mutates it — a data race that yields corrupted group
ids and out-of-bounds result-holder
+ /// access. We therefore materialize the result into self-contained {@link
IntermediateRecord}s here, on the
+ /// producing worker thread, before hand-off — mirroring how the
non-streaming
+ /// {@link org.apache.pinot.core.operator.combine.GroupByCombineOperator}
merges inline on the worker thread.
+ ///
+ /// <p>Blocks that already carry intermediate records (order-by /
segment-trim path) are self-contained and
+ /// returned unchanged.
+ @Override
+ protected GroupByResultsBlock
detachFromWorkerThreadState(GroupByResultsBlock resultsBlock) {
+ AggregationGroupByResult aggregationGroupByResult =
resultsBlock.getAggregationGroupByResult();
+ if (aggregationGroupByResult == null ||
resultsBlock.getIntermediateRecords() != null) {
+ return resultsBlock;
+ }
+ List<IntermediateRecord> records = new ArrayList<>();
+ try {
+ Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator =
aggregationGroupByResult.getGroupKeyIterator();
+ int extractedKeys = 0;
+ while (groupKeyIterator.hasNext()) {
+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(extractedKeys++,
EXPLAIN_NAME);
+ GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
+ Object[] keys = groupKey._keys;
+ Object[] values = Arrays.copyOf(keys, _numColumns);
Review Comment:
Pre-existing rather than introduced here (the old consumer-side loop had the
same math), but since this loop now mirrors `GroupByCombineOperator`: that one
sizes with `_numKeyColumns = _queryContext.getNumGroupByKeyColumns()`, which
includes the synthetic `$groupingId` key column for grouping sets, whereas this
class uses `_numGroupByExpressions`.
`GroupByUtils.buildGroupingSetsResultsBlock` returns a raw block when per-set
trim doesn't trigger, with the discriminator as the trailing key, and the
`CombinePlanNode` gate doesn't exclude grouping sets from the streaming path —
so the first aggregation value would clobber the `$groupingId` slot here. Not
reachable via MSE today (grouping sets aren't pushed to leaf stages), so fine
as a follow-up too, but aligning this with the non-streaming operator is a
one-liner while the loop is being touched.
--
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]
