Add IT for num_reduced_phases with batched query execution (elastic#134312)

Author: benchaplin

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search/120_batch_reduce_size.yml

@@ -1,7 +1,4 @@
 setup:
-  - skip:
-      awaits_fix: "TODO fix this test, the response with batched execution is not deterministic enough for the available matchers"
-
   - do:
       indices.create:
           index: test_1

server/src/internalClusterTest/java/org/elasticsearch/action/search/BatchedQueryPhaseIT.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.action.search;
+
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
+import org.elasticsearch.cluster.routing.IndexRoutingTable;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.elasticsearch.action.search.SearchType.QUERY_THEN_FETCH;
+import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
+import static org.hamcrest.Matchers.equalTo;
+
+public class BatchedQueryPhaseIT extends ESIntegTestCase {
+
+    public void testNumReducePhases() {
+        String indexName = "test-idx";
+        assertAcked(
+            prepareCreate(indexName).setMapping("title", "type=keyword")
+                .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0))
+        );
+        for (int i = 0; i < 100; i++) {
+            prepareIndex(indexName).setId(Integer.toString(i)).setSource("title", "testing" + i).get();
+        }
+        refresh();
+
+        final String coordinatorNode = internalCluster().getRandomNodeName();
+        final String coordinatorNodeId = getNodeId(coordinatorNode);
+        assertNoFailuresAndResponse(
+            client(coordinatorNode).prepareSearch(indexName)
+                .setBatchedReduceSize(2)
+                .addAggregation(terms("terms").field("title"))
+                .setSearchType(QUERY_THEN_FETCH),
+            response -> {
+                Map<String, Integer> shardsPerNode = getNodeToShardCountMap(indexName);
+                // Shards are not batched if they are already on the coordinating node or if there is only one per data node.
+                final int coordinatorShards = shardsPerNode.getOrDefault(coordinatorNodeId, 0);
+                final long otherSingleShardNodes = shardsPerNode.entrySet()
+                    .stream()
+                    .filter(entry -> entry.getKey().equals(coordinatorNodeId) == false)
+                    .filter(entry -> entry.getValue() == 1)
+                    .count();
+                final int numNotBatchedShards = coordinatorShards + (int) otherSingleShardNodes;
+
+                // Because batched_reduce_size = 2, whenever two or more shard results exist on the coordinating node, they will be
+                // partially reduced (batched queries do not count towards num_reduce_phases).
+                // Hence, the formula: (# of NOT batched shards) - 1.
+                final int expectedNumReducePhases = Math.max(1, numNotBatchedShards - 1);
+                assertThat(response.getNumReducePhases(), equalTo(expectedNumReducePhases));
+            }
+        );
+    }
+
+    private Map<String, Integer> getNodeToShardCountMap(String indexName) {
+        ClusterState clusterState = clusterAdmin().prepareState(TEST_REQUEST_TIMEOUT).get().getState();
+        IndexRoutingTable indexRoutingTable = clusterState.routingTable(ProjectId.DEFAULT).index(indexName);
+        if (indexRoutingTable == null) {
+            return Collections.emptyMap();
+        }
+
+        Map<String, Integer> nodeToShardCount = new HashMap<>();
+        for (int shardId = 0; shardId < indexRoutingTable.size(); shardId++) {
+            IndexShardRoutingTable shardRoutingTable = indexRoutingTable.shard(shardId);
+            for (int copy = 0; copy < shardRoutingTable.size(); copy++) {
+                ShardRouting shardRouting = shardRoutingTable.shard(copy);
+                String nodeId = shardRouting.currentNodeId();
+                if (nodeId != null) {
+                    nodeToShardCount.merge(nodeId, 1, Integer::sum);
+                }
+            }
+        }
+
+        return nodeToShardCount;
+    }
+}

server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java

@@ -95,6 +95,7 @@ public class QueryPhaseResultConsumer extends ArraySearchPhaseResults<SearchPhas
     final TopDocsStats topDocsStats;
     private volatile MergeResult mergeResult;
     private volatile boolean hasPartialReduce;
+    // Note: at this time, numReducePhases does not count reductions that occur on the data node as part of batched query execution.
     private volatile int numReducePhases;
 
     /**