This is an automated email from the ASF dual-hosted git repository.
yashmayya 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 236401b2dc0 Add totalMemAllocatedBytes to multi-stage query responses
(#19816)
236401b2dc0 is described below
commit 236401b2dc0f64cf5ce700920769751479691511
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Oct 15 13:09:11 2026 -0700
Add totalMemAllocatedBytes to multi-stage query responses (#19816)
---
.../MultiStageBrokerRequestHandler.java | 23 +++++++++++---
.../MultiStageBrokerRequestHandlerTest.java | 37 ++++++++++++++++++++++
.../response/broker/BrokerResponseNativeV2.java | 19 +++++++++--
.../MultiNodesOfflineClusterIntegrationTest.java | 30 ++++++++++++++++++
.../query/runtime/operator/MultiStageOperator.java | 6 ++++
5 files changed, 107 insertions(+), 8 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index adf76c39c2e..007d6c82a81 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -1181,6 +1181,23 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
return stagePlans;
}
+ @VisibleForTesting
+ static void mergeStageStats(BrokerResponseNativeV2 brokerResponse,
+ List<MultiStageQueryStats.StageStats.Closed> queryStats) {
+ for (MultiStageQueryStats.StageStats.Closed stageStats : queryStats) {
+ if (stageStats != null) { // for example pipeline breaker may not have
stats
+ stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse,
stats));
+ if (!stageStats.isEmpty()) {
+ // The stage root's allocation covers all the operators of its op
chain. It is read by name because stage
+ // roots have different stat key types. Leaf threads and pipeline
breakers are added in
+ // MultiStageOperator.Type#mergeInto.
+ brokerResponse.addTotalMemAllocatedBytes(
+
stageStats.getLastOperatorStats().getUnsafe("ALLOCATED_MEMORY_BYTES", 0L));
+ }
+ }
+ }
+ }
+
private void fillOldBrokerResponseStats(BrokerResponseNativeV2
brokerResponse,
List<MultiStageQueryStats.StageStats.Closed> queryStats,
DispatchableSubPlan dispatchableSubPlan,
@Nullable List<QueryDispatcher.QueryResult.StageCoverage> stageCoverage,
@@ -1191,11 +1208,7 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
MultiStageStatsTreeBuilder treeBuilder = new
MultiStageStatsTreeBuilder(queryStageMap, queryStats,
stageStatsTrees);
brokerResponse.setStageStats(treeBuilder.jsonStatsByStage(0));
- for (MultiStageQueryStats.StageStats.Closed stageStats : queryStats) {
- if (stageStats != null) { // for example pipeline breaker may not have
stats
- stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse,
stats));
- }
- }
+ mergeStageStats(brokerResponse, queryStats);
// Broker-pruned segments are computed during routing in WorkerManager,
not reported by servers.
// Inject the count directly since it doesn't flow through the
LeafOperator stat pipeline.
long numSegmentsPrunedByBroker =
dispatchableSubPlan.getNumSegmentsPrunedByBroker();
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
index d2d99c2ebbd..d13a209398f 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
@@ -30,6 +30,7 @@ import org.apache.pinot.broker.api.AccessControl;
import org.apache.pinot.broker.broker.AllowAllAccessControlFactory;
import org.apache.pinot.broker.queryquota.QueryQuotaManager;
import org.apache.pinot.common.config.provider.TableCache;
+import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.failuredetector.FailureDetector;
import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.response.BrokerResponse;
@@ -42,6 +43,12 @@ import
org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
import org.apache.pinot.query.planner.physical.PinotDispatchPlanner;
import org.apache.pinot.query.routing.WorkerManager;
+import org.apache.pinot.query.runtime.operator.BaseMailboxReceiveOperator;
+import org.apache.pinot.query.runtime.operator.LeafOperator;
+import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
+import org.apache.pinot.query.runtime.operator.MultiStageOperator;
+import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
+import org.apache.pinot.query.runtime.plan.pipeline.PipelineBreakerOperator;
import org.apache.pinot.query.service.dispatch.QueryDispatcher;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.auth.broker.RequesterIdentity;
@@ -199,6 +206,36 @@ public class MultiStageBrokerRequestHandlerTest extends
QueryEnvironmentTestBase
"Per-query SET must take precedence over the broker default");
}
+ @Test
+ public void testMergeStageStatsTotalMemAllocatedBytes() {
+ StatMap<BaseMailboxReceiveOperator.StatKey> rootReceive = new
StatMap<>(BaseMailboxReceiveOperator.StatKey.class)
+ .merge(BaseMailboxReceiveOperator.StatKey.ALLOCATED_MEMORY_BYTES,
100_000L);
+ StatMap<BaseMailboxReceiveOperator.StatKey> breakerReceive =
+ new StatMap<>(BaseMailboxReceiveOperator.StatKey.class)
+ .merge(BaseMailboxReceiveOperator.StatKey.ALLOCATED_MEMORY_BYTES,
1L);
+ StatMap<PipelineBreakerOperator.StatKey> breaker = new
StatMap<>(PipelineBreakerOperator.StatKey.class)
+ .merge(PipelineBreakerOperator.StatKey.ALLOCATED_MEMORY_BYTES, 10L);
+ StatMap<LeafOperator.StatKey> leaf = new
StatMap<>(LeafOperator.StatKey.class)
+ .merge(LeafOperator.StatKey.ALLOCATED_MEMORY_BYTES, 100L)
+ .merge(LeafOperator.StatKey.THREAD_MEM_ALLOCATED_BYTES, 1_000L);
+ StatMap<MailboxSendOperator.StatKey> send = new
StatMap<>(MailboxSendOperator.StatKey.class)
+ .merge(MailboxSendOperator.StatKey.ALLOCATED_MEMORY_BYTES, 10_000L);
+
+ // Stage 1 is a leaf stage with a pipeline breaker, flattened the same way
in legacy and stream stats modes
+ BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2();
+ MultiStageBrokerRequestHandler.mergeStageStats(brokerResponse, List.of(
+ new
MultiStageQueryStats.StageStats.Closed(List.of(MultiStageOperator.Type.MAILBOX_RECEIVE),
+ List.of(rootReceive)),
+ new MultiStageQueryStats.StageStats.Closed(
+ List.of(MultiStageOperator.Type.MAILBOX_RECEIVE,
MultiStageOperator.Type.PIPELINE_BREAKER,
+ MultiStageOperator.Type.LEAF,
MultiStageOperator.Type.MAILBOX_SEND),
+ List.of(breakerReceive, breaker, leaf, send))));
+
+ // Op chain roots (stage roots and the breaker) plus the leaf threads, but
not the breaker's receive or the leaf's
+ // own op chain allocation, which their roots already include
+ assertEquals(brokerResponse.getTotalMemAllocatedBytes(), 100_000L +
10_000L + 1_000L + 10L);
+ }
+
@Test
public void
testApplyBrokerDefaultQueryOptionsNoStreamingDistinctInjectionWhenConfigUnset()
throws Exception {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
index 827c800d037..e5738163357 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java
@@ -53,9 +53,10 @@ import org.apache.pinot.common.response.ProcessingException;
"realtimeResponseSerializationCpuTimeNs", "offlineTotalCpuTimeNs",
"realtimeTotalCpuTimeNs",
"explainPlanNumEmptyFilterSegments",
"explainPlanNumMatchAllFilterSegments", "traceInfo", "tablesQueried",
"offlineThreadMemAllocatedBytes", "realtimeThreadMemAllocatedBytes",
"offlineResponseSerMemAllocatedBytes",
- "realtimeResponseSerMemAllocatedBytes", "offlineTotalMemAllocatedBytes",
"realtimeTotalMemAllocatedBytes", "pools",
- "rlsFiltersApplied", "approximateFunctionApplied", "groupsTrimmed",
"mseLiteLeafStageLimitReached",
- "mseLiteLeafStageEffectiveLimit", "mseLiteFanOutAdjustedLimitApplied",
"responseMetadata"
+ "realtimeResponseSerMemAllocatedBytes", "offlineTotalMemAllocatedBytes",
"realtimeTotalMemAllocatedBytes",
+ "totalMemAllocatedBytes", "pools", "rlsFiltersApplied",
"approximateFunctionApplied", "groupsTrimmed",
+ "mseLiteLeafStageLimitReached", "mseLiteLeafStageEffectiveLimit",
"mseLiteFanOutAdjustedLimitApplied",
+ "responseMetadata"
})
public class BrokerResponseNativeV2 implements BrokerResponse {
private final StatMap<StatKey> _brokerStats = new StatMap<>(StatKey.class);
@@ -81,6 +82,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse
{
/// virtual rows can be generated. For example, in a join query, the number
of rows can be more than the number of
/// rows in the table.
private long _maxRowsInOperator;
+ private long _totalMemAllocatedBytes;
private String _requestId;
private String _clientRequestId;
private String _brokerId;
@@ -272,6 +274,17 @@ public class BrokerResponseNativeV2 implements
BrokerResponse {
_maxRowsInOperator = Math.max(_maxRowsInOperator, maxRows);
}
+ /// Returns the heap bytes allocated by the query across all its stages and
workers, as reported in the stage stats.
+ /// This is the total allocated (including short-lived objects), not the
peak heap used. It is 0 unless thread
+ /// allocated bytes measurement is enabled on the brokers and servers.
+ public long getTotalMemAllocatedBytes() {
+ return _totalMemAllocatedBytes;
+ }
+
+ public void addTotalMemAllocatedBytes(long bytes) {
+ _totalMemAllocatedBytes += bytes;
+ }
+
@Override
public long getTimeUsedMs() {
return _timeUsedMs;
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiNodesOfflineClusterIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiNodesOfflineClusterIntegrationTest.java
index c5849bba8cc..04774cf461c 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiNodesOfflineClusterIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiNodesOfflineClusterIntegrationTest.java
@@ -326,6 +326,36 @@ public class MultiNodesOfflineClusterIntegrationTest
extends OfflineClusterInteg
assertEquals(result.get("numServersQueried").intValue(), 1);
}
+ @Test
+ public void testMultiStageTotalMemAllocatedBytes()
+ throws Exception {
+ setUseMultiStageQueryEngine(true);
+
+ // The semi-join query uses a pipeline breaker
+ for (String query : List.of("SELECT Carrier, COUNT(*) FROM mytable GROUP
BY Carrier",
+ "SELECT Carrier FROM mytable WHERE Origin IN (SELECT Origin FROM
mytable WHERE ArrDelay > 0) LIMIT 10")) {
+ JsonNode result = postQuery("SET streamStats = false; " + query);
+ long totalMemAllocatedBytes =
result.path("totalMemAllocatedBytes").asLong();
+ assertTrue(totalMemAllocatedBytes > 0);
+ assertEquals(totalMemAllocatedBytes,
sumAllocatedBytes(result.get("stageStats"), true));
+ }
+ }
+
+ /// Sums the allocation of each op chain root (the top node, each
MAILBOX_SEND and each pipeline breaker) and of the
+ /// single-stage threads of each LEAF. The legacy (non-stream) stats tree
renders a pipeline breaker as the
+ /// MAILBOX_RECEIVE child of its LEAF, but that node carries the stats of
the PIPELINE_BREAKER operator itself.
+ private static long sumAllocatedBytes(JsonNode node, boolean isOpChainRoot) {
+ long bytes = node.path("threadMemAllocatedBytes").asLong();
+ if (isOpChainRoot) {
+ bytes += node.path("allocatedMemoryBytes").asLong();
+ }
+ boolean isLeaf = node.path("type").asText().equals("LEAF");
+ for (JsonNode child : node.path("children")) {
+ bytes += sumAllocatedBytes(child, isLeaf ||
child.path("type").asText().equals("MAILBOX_SEND"));
+ }
+ return bytes;
+ }
+
// Disabled because segments might not be server partitioned with multiple
servers
@Test(enabled = false)
@Override
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
index 88dafd2e012..c702d42a3c4 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java
@@ -422,6 +422,9 @@ public abstract class MultiStageOperator implements
Operator<MseBlock>, AutoClos
@SuppressWarnings("unchecked")
StatMap<LeafOperator.StatKey> stats = (StatMap<LeafOperator.StatKey>)
map;
response.mergeMaxRowsInOperator(stats.getLong(LeafOperator.StatKey.EMITTED_ROWS));
+ // The single-stage threads of the leaf are not covered by the stage
root's allocation, which
+ // MultiStageBrokerRequestHandler#mergeStageStats adds
+
response.addTotalMemAllocatedBytes(stats.getLong(LeafOperator.StatKey.THREAD_MEM_ALLOCATED_BYTES));
StatMap<BrokerResponseNativeV2.StatKey> brokerStats = new
StatMap<>(BrokerResponseNativeV2.StatKey.class);
for (LeafOperator.StatKey statKey : stats.keySet()) {
@@ -495,6 +498,9 @@ public abstract class MultiStageOperator implements
Operator<MseBlock>, AutoClos
@SuppressWarnings("unchecked")
StatMap<PipelineBreakerOperator.StatKey> stats =
(StatMap<PipelineBreakerOperator.StatKey>) map;
response.mergeMaxRowsInOperator(stats.getLong(PipelineBreakerOperator.StatKey.EMITTED_ROWS));
+ // Pipeline breakers run in their own op chain, so they are not
covered by the stage root's allocation, which
+ // MultiStageBrokerRequestHandler#mergeStageStats adds
+
response.addTotalMemAllocatedBytes(stats.getLong(PipelineBreakerOperator.StatKey.ALLOCATED_MEMORY_BYTES));
}
},
SORT_OR_LIMIT(10, SortOperator.StatKey.class) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
