This is an automated email from the ASF dual-hosted git repository. gortiz pushed a commit to branch feature/mse-push-stats in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 144de6d8a01e068525a577c7e95b0ad36b4dff5d Author: Gonzalo Ortiz <[email protected]> AuthorDate: Mon Jun 8 19:59:04 2026 +0200 fix(mse): copy folded pipeline-breaker stats so a leaf opchain's operators aren't duplicated LeafOperator.calculateUpstreamStats() handed back the shared mutable _pipelineBreakerStats instance. The base MultiStageOperator.calculateStats() mutates the returned object (appending the LEAF entry to its flat operator list), and calculateStats() runs more than once per opchain -- once when the MailboxSendOperator serializes its EOS stats, and again from the scheduler's completion callback that feeds the stream-stats listener. Each call appended the leaf's own operators again, so a pipeline-breaker leaf reported a doubled flat list (e.g. [RECEIVE, PB, LEAF, SEND, LEAF, SEND]). In stream mode this broke the stats-tree encoder: the live operator tree (MAILBOX_SEND -> LEAF, grafted with PIPELINE_BREAKER -> MAILBOX_RECEIVE = 4 nodes) no longer matched the inflated flat list (6), so encode() threw treeSize != flatSize, the leaf's stats were dropped, and the semi-join leaf stage rendered as an empty stage on the broker. A non-pipeline-breaker leaf was immune because it returns a fresh emptyStats() each call. Return a deep copy instead, leaving the shared instance pristine. Adds a LeafOperatorTest regression asserting calculateStats() is idempotent for a pipeline-breaker leaf. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> --- .../pinot/query/runtime/operator/LeafOperator.java | 8 +++- .../query/runtime/operator/LeafOperatorTest.java | 48 ++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java index cb4a8c3b424..289ccc53c75 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LeafOperator.java @@ -169,8 +169,14 @@ public class LeafOperator extends MultiStageOperator { @Override protected MultiStageQueryStats calculateUpstreamStats() { + // Return a COPY, not the shared instance. calculateStats() runs more than once per opchain (e.g. when the + // MailboxSendOperator serializes its EOS stats and again from the scheduler's completion callback that feeds the + // stream-stats listener), and the base calculateStats() appends this operator's entry to the returned object. + // Handing back the shared _pipelineBreakerStats would append LEAF (+ the downstream MAILBOX_SEND) once per call, + // duplicating the leaf opchain's own operators in the flat stats list. A non-pipeline-breaker leaf is immune + // because it returns a fresh emptyStats() each call. return _pipelineBreakerStats != null - ? _pipelineBreakerStats + ? MultiStageQueryStats.copy(_pipelineBreakerStats) : MultiStageQueryStats.emptyStats(_context.getStageId()); } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java index 685055773a2..b56840d05c7 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java @@ -54,7 +54,9 @@ import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUt import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.SuccessMseBlock; +import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerOperator; import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.utils.JsonUtils; import org.mockito.Mock; @@ -150,6 +152,52 @@ public class LeafOperatorTest { operator.close(); } + @Test + public void calculateStatsIsIdempotentWhenFoldingPipelineBreaker() { + // Regression for a pipeline-breaker leaf whose own operators were duplicated in the flat stats. calculateStats() + // runs more than once per opchain (the MailboxSendOperator serializing its EOS stats, then again from the + // scheduler completion callback that feeds the stream-stats listener). The folded pipeline-breaker stats are a + // shared mutable instance, so LeafOperator.calculateUpstreamStats() must hand back a COPY -- otherwise each call + // appends this leaf's own LEAF entry again, inflating the flat operator list and breaking the stats-tree encoder + // with a treeSize != flatSize mismatch. + DataSchema schema = new DataSchema(new String[]{"intCol"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT}); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT intCol FROM tbl"); + List<BaseResultsBlock> dataBlocks = Collections.singletonList( + new SelectionResultsBlock(schema, Collections.singletonList(new Object[]{1}), queryContext)); + InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock()); + QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock); + + // The leaf folds two operators (MAILBOX_RECEIVE + PIPELINE_BREAKER) ahead of its own LEAF entry. + OpChainExecutionContext context = OperatorTestUtil.getTracingContext(); + MultiStageQueryStats pipelineBreakerStats = new MultiStageQueryStats.Builder(context.getStageId()) + .customizeOpen(open -> open + .addLastOperator(MultiStageOperator.Type.MAILBOX_RECEIVE, + new StatMap<>(BaseMailboxReceiveOperator.StatKey.class)) + .addLastOperator(MultiStageOperator.Type.PIPELINE_BREAKER, + new StatMap<>(PipelineBreakerOperator.StatKey.class))) + .build(); + + LeafOperator operator = new LeafOperator(context, mockQueryRequests(1), schema, queryExecutor, _executorService, + pipelineBreakerStats); + _operatorRef.set(operator); + + // Drain so the leaf produces its stats. + while (!operator.nextBlock().isEos()) { + // consume blocks + } + + int firstCount = operator.calculateStats().getCurrentStats().getLastOperatorIndex() + 1; + int secondCount = operator.calculateStats().getCurrentStats().getLastOperatorIndex() + 1; + + // MAILBOX_RECEIVE + PIPELINE_BREAKER (folded) + LEAF (this operator) = 3. A shared upstream instance would + // append LEAF a second time on the second call, yielding 4. + assertEquals(firstCount, 3, "Expected MAILBOX_RECEIVE, PIPELINE_BREAKER, LEAF"); + assertEquals(secondCount, firstCount, "calculateStats() must be idempotent for a pipeline-breaker leaf"); + + operator.close(); + } + @Test public void shouldHandleDesiredDataSchemaConversionCorrectly() { // Given: --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
