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 1b552624f52 Add a broker hook to prepare each MSE stage before
dispatch (#20104)
1b552624f52 is described below
commit 1b552624f52a96926c67d340d9c27a8de390a4ef
Author: Yash Mayya <[email protected]>
AuthorDate: Fri Oct 30 12:59:39 2026 -0700
Add a broker hook to prepare each MSE stage before dispatch (#20104)
---
.../MultiStageBrokerRequestHandler.java | 15 +++++
.../MultiStageBrokerRequestHandlerTest.java | 71 ++++++++++++++++++++++
2 files changed, 86 insertions(+)
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 c68506ca14e..1d1f9fd9809 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
@@ -660,6 +660,16 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
RequestContext requestContext, HttpHeaders httpHeaders) {
}
+ /// Extension hook for preparing a stage right before the broker sends it to
its workers, to run the query or to
+ /// explain it on the servers. To run the query, the broker calls it for
every stage, including stage 0 (the broker
+ /// reduce). To explain the query on the servers, it calls it only for the
leaf stages that have segments.
+ ///
+ /// Subclasses can add entries to the custom properties of the stage
+ /// ([DispatchablePlanFragment#getCustomProperties()]), which its workers
get in their stage metadata. They must not
+ /// replace the existing entries, such as the table name and the time
boundary that servers use for leaf stages.
+ protected void prepareStageForDispatch(DispatchablePlanFragment stage,
Map<String, String> queryOptions) {
+ }
+
/// Estimates the total number of server query threads the given plan will
consume; used to acquire permits from the
/// multi-stage query throttler before dispatch. Subclasses can override to
accurately
/// estimate the real per-server work.
@@ -902,6 +912,10 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
return new BrokerResponseNative(QueryErrorCode.TOO_MANY_REQUESTS,
errorMessage);
}
+ for (DispatchablePlanFragment stage :
dispatchableSubPlan.getQueryStageMap().values()) {
+ prepareStageForDispatch(stage, query.getOptions());
+ }
+
// Short-circuit: if all leaf stages are empty (all segments pruned or
table has no data),
// run only the broker reduce stage locally. No server dispatch is
attempted.
boolean allLeafStagesEmpty = dispatchableSubPlan.isAllLeafStagesEmpty();
@@ -1151,6 +1165,7 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
private Collection<PlanNode> requestPhysicalPlan(DispatchablePlanFragment
fragment,
RequestContext requestContext, long queryTimeoutMs, Map<String, String>
queryOptions) {
List<PlanNode> stagePlans;
+ prepareStageForDispatch(fragment, queryOptions);
try {
stagePlans = _queryDispatcher.explain(requestContext, fragment,
queryTimeoutMs, queryOptions);
} catch (Exception e) {
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 149228f41d5..b89ae2fc01b 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
@@ -37,6 +37,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.response.BrokerResponse;
import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.core.routing.MockRoutingManagerFactory;
import org.apache.pinot.core.routing.MultiClusterRoutingContext;
import org.apache.pinot.core.routing.RoutingManager;
import org.apache.pinot.core.transport.ServerInstance;
@@ -58,6 +59,7 @@ import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.auth.broker.RequesterIdentity;
import org.apache.pinot.spi.env.PinotConfiguration;
import
org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory;
+import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.query.QueryThreadContext;
import org.apache.pinot.spi.trace.QueryFingerprint;
import org.apache.pinot.spi.trace.RequestContext;
@@ -72,13 +74,16 @@ import org.apache.pinot.spi.utils.NetUtils;
import org.apache.pinot.sql.parsers.CalciteSqlParser;
import org.apache.pinot.sql.parsers.SqlNodeAndOptions;
import org.apache.pinot.util.TestUtils;
+import org.mockito.ArgumentCaptor;
import org.mockito.MockedConstruction;
import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
@@ -153,6 +158,72 @@ public class MultiStageBrokerRequestHandlerTest extends
QueryEnvironmentTestBase
"onQueryCompletion hook must be called with the BrokerResponse from
handleRequest for MSE");
}
+ /// The broker calls
[MultiStageBrokerRequestHandler#prepareStageForDispatch] for each stage before
it sends the stage
+ /// to its workers, both to run the query and to explain it on the servers.
+ @Test
+ public void testPrepareStageForDispatchHook()
+ throws Exception {
+ int port1 = NetUtils.findOpenPort();
+ int port2 = NetUtils.findOpenPort();
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1,
port2);
+ TABLE_SCHEMAS.forEach((tableName, schema) -> factory.registerTable(schema,
tableName));
+ SERVER1_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port1, tableName, segment)));
+ SERVER2_SEGMENTS.forEach(
+ (tableName, segments) -> segments.forEach(segment ->
factory.registerSegment(port2, tableName, segment)));
+ RoutingManager routingManager = factory.buildRoutingManager(null);
+ WorkerManager workerManager =
+ new WorkerManager("Broker_localhost", "localhost",
NetUtils.findOpenPort(), routingManager);
+ PinotConfiguration config = new PinotConfiguration();
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
"localhost");
+ config.setProperty(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
Integer.toString(NetUtils.findOpenPort()));
+ BrokerQueryEventListenerFactory.init(config);
+ QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class);
+ when(queryQuotaManager.acquire(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireDatabase(anyString())).thenReturn(true);
+ when(queryQuotaManager.acquireApplication(anyString())).thenReturn(true);
+ AtomicReference<DispatchableSubPlan> throttledPlan = new
AtomicReference<>();
+ try (MockedConstruction<QueryDispatcher> dispatchers =
mockConstruction(QueryDispatcher.class)) {
+ MultiStageBrokerRequestHandler handler = new
MultiStageBrokerRequestHandler(config, "testBrokerId",
+ new BrokerRequestIdGenerator(), routingManager, new
AllowAllAccessControlFactory(), queryQuotaManager,
+ factory.buildTableCache(), mock(MultiStageQueryThrottler.class),
mock(FailureDetector.class),
+ ThreadAccountantUtils.getNoOpAccountant(), null, workerManager,
workerManager) {
+ @Override
+ protected void prepareStageForDispatch(DispatchablePlanFragment stage,
Map<String, String> queryOptions) {
+ stage.getCustomProperties().put("testKey", "stage" +
stage.getPlanFragment().getFragmentId());
+ }
+
+ @Override
+ protected int getEstimatedNumQueryThreads(DispatchableSubPlan
dispatchableSubPlan) {
+ throttledPlan.set(dispatchableSubPlan);
+ return super.getEstimatedNumQueryThreads(dispatchableSubPlan);
+ }
+ };
+ try {
+ String query = "SELECT col1, COUNT(*) FROM a GROUP BY col1";
+
+ // The mocked throttler stops the query before it runs
+ BrokerResponse response = handler.handleRequest(query);
+ assertEquals(response.getExceptions().get(0).getErrorCode(),
QueryErrorCode.EXECUTION_TIMEOUT.getId());
+ Map<Integer, DispatchablePlanFragment> stages =
throttledPlan.get().getQueryStageMap();
+ assertEquals(stages.size(), 3);
+ stages.forEach(
+ (stageId, stage) ->
assertEquals(stage.getCustomProperties().get("testKey"), "stage" + stageId));
+
+ handler.handleRequest("SET " + QueryOptionKey.EXPLAIN_ASKING_SERVERS +
" = true; EXPLAIN PLAN FOR " + query);
+ ArgumentCaptor<DispatchablePlanFragment> explainedStages =
+ ArgumentCaptor.forClass(DispatchablePlanFragment.class);
+ verify(dispatchers.constructed().get(0), atLeastOnce())
+ .explain(any(), explainedStages.capture(), anyLong(), any());
+ for (DispatchablePlanFragment stage : explainedStages.getAllValues()) {
+ assertEquals(stage.getCustomProperties().get("testKey"), "stage" +
stage.getPlanFragment().getFragmentId());
+ }
+ } finally {
+ handler.shutDown();
+ }
+ }
+ }
+
@Test
public void
testApplyBrokerDefaultQueryOptionsInjectsStreamingGroupByFlushThreshold()
throws Exception {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
