add numServersQueried and numServersResponded to MSQE response (#14806)

apache · Jan 22, 2025 · f845a97 · f845a97
1 parent 7356304
commit f845a97
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 2 deletions.
diff --git a/.../src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/.../src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -64,6 +65,7 @@
 import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
 import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
 import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.routing.QueryServerInstance;
 import org.apache.pinot.query.routing.WorkerManager;
 import org.apache.pinot.query.runtime.MultiStageStatsTreeBuilder;
 import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
@@ -199,8 +201,13 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
     }
 
     DispatchableSubPlan dispatchableSubPlan = queryPlanResult.getQueryPlan();
-    Set<String> tableNames = queryPlanResult.getTableNames();
 
+    Set<QueryServerInstance> servers = new HashSet<>();
+    for (DispatchablePlanFragment planFragment: dispatchableSubPlan.getQueryStageList()) {
+      servers.addAll(planFragment.getServerInstances());
+    }
+
+    Set<String> tableNames = queryPlanResult.getTableNames();
     _brokerMetrics.addMeteredGlobalValue(BrokerMeter.MULTI_STAGE_QUERIES_GLOBAL, 1);
     for (String tableName : tableNames) {
       _brokerMetrics.addMeteredTableValue(tableName, BrokerMeter.MULTI_STAGE_QUERIES, 1);
@@ -277,8 +284,12 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       BrokerResponseNativeV2 brokerResponse = new BrokerResponseNativeV2();
       brokerResponse.setResultTable(queryResults.getResultTable());
       brokerResponse.setTablesQueried(tableNames);
-      // TODO: Add servers queried/responded stats
       brokerResponse.setBrokerReduceTimeMs(queryResults.getBrokerReduceTimeMs());
+      // MSE cannot finish if a single queried server did not respond, so we can use the same count for
+      // both the queried and responded stats. Minus one prevents the broker to be included in the count
+      // (it will always be included because of the root of the query plan)
+      brokerResponse.setNumServersQueried(servers.size() - 1);
+      brokerResponse.setNumServersResponded(servers.size() - 1);
 
       // Attach unavailable segments
       int numUnavailableSegments = 0;

diff --git a/...sts/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java b/...sts/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java
@@ -1347,6 +1347,16 @@ public void testConcurrentQueries() {
     executorService.shutdownNow();
   }
 
+  @Test
+  public void testNumServersQueried() throws Exception {
+    String query = "select * from mytable limit 10";
+    JsonNode jsonNode = postQuery(query);
+    JsonNode numServersQueried = jsonNode.get("numServersQueried");
+    assertNotNull(numServersQueried);
+    assertTrue(numServersQueried.isInt());
+    assertTrue(numServersQueried.asInt() > 0);
+  }
+
   private void checkQueryResultForDBTest(String column, String tableName)
       throws Exception {
     checkQueryResultForDBTest(column, tableName, null, null);

diff --git a/...anner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java b/...anner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java
@@ -22,6 +22,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.apache.pinot.core.routing.TimeBoundaryInfo;
 import org.apache.pinot.query.planner.PlanFragment;
 import org.apache.pinot.query.routing.QueryServerInstance;
@@ -112,4 +113,8 @@ public void setServerInstanceToWorkerIdMap(Map<QueryServerInstance, List<Integer
     _serverInstanceToWorkerIdMap.clear();
     _serverInstanceToWorkerIdMap.putAll(serverInstanceToWorkerIdMap);
   }
+
+  public Set<QueryServerInstance> getServerInstances() {
+    return _serverInstanceToWorkerIdMap.keySet();
+  }
 }