(pinot) 01/01: fix(mse): copy folded pipeline-breaker stats so a leaf opchain's operators aren't duplicated

[go...@apache.org]

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 <19...@users.noreply.github.com>
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) <no...@anthropic.com>
---
 .../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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org