Disable concurrent aggs for Diversified Sampler and Sampler aggs (opensearch-project#11087)

Signed-off-by: Jay Deng <jayd0104@gmail.com>

Author: jed326

CHANGELOG.md

@@ -125,6 +125,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add instrumentation for indexing in transport bulk action and transport shard bulk action. ([#10273](https://github.com/opensearch-project/OpenSearch/pull/10273))
 - [BUG] Disable sort optimization for HALF_FLOAT ([#10999](https://github.com/opensearch-project/OpenSearch/pull/10999))
 - Performance improvement for MultiTerm Queries on Keyword fields ([#7057](https://github.com/opensearch-project/OpenSearch/issues/7057))
+- Disable concurrent aggs for Diversified Sampler and Sampler aggs ([#11087](https://github.com/opensearch-project/OpenSearch/issues/11087))
 
 ### Deprecated
 

server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/DiversifiedSamplerIT.java

@@ -33,9 +33,9 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import org.opensearch.action.admin.indices.refresh.RefreshRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.SearchType;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.query.TermQueryBuilder;
@@ -132,13 +132,14 @@ public void setupSuiteScopeCluster() throws Exception {
             client().prepareIndex("test")
                 .setId("" + i)
                 .setSource("author", parts[5], "name", parts[2], "genre", parts[8], "price", Float.parseFloat(parts[3]))
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                 .get();
             client().prepareIndex("idx_unmapped_author")
                 .setId("" + i)
                 .setSource("name", parts[2], "genre", parts[8], "price", Float.parseFloat(parts[3]))
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                 .get();
         }
-        client().admin().indices().refresh(new RefreshRequest("test")).get();
     }
 
     public void testIssue10719() throws Exception {
@@ -221,10 +222,6 @@ public void testNestedDiversity() throws Exception {
     }
 
     public void testNestedSamples() throws Exception {
-        assumeFalse(
-            "Concurrent search case muted pending fix: https://github.com/opensearch-project/OpenSearch/issues/10046",
-            internalCluster().clusterService().getClusterSettings().get(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING)
-        );
         // Test samples nested under samples
         int MAX_DOCS_PER_AUTHOR = 1;
         int MAX_DOCS_PER_GENRE = 2;

server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/SamplerIT.java

@@ -34,9 +34,9 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import org.opensearch.action.admin.indices.refresh.RefreshRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.SearchType;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.query.TermQueryBuilder;
@@ -132,13 +132,14 @@ public void setupSuiteScopeCluster() throws Exception {
             client().prepareIndex("test")
                 .setId("" + i)
                 .setSource("author", parts[5], "name", parts[2], "genre", parts[8], "price", Float.parseFloat(parts[3]))
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                 .get();
             client().prepareIndex("idx_unmapped_author")
                 .setId("" + i)
                 .setSource("name", parts[2], "genre", parts[8], "price", Float.parseFloat(parts[3]))
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
                 .get();
         }
-        client().admin().indices().refresh(new RefreshRequest("test")).get();
     }
 
     public void testIssue10719() throws Exception {
@@ -195,6 +196,23 @@ public void testSimpleSampler() throws Exception {
         assertThat(maxBooksPerAuthor, equalTo(3L));
     }
 
+    public void testSimpleSamplerShardSize() throws Exception {
+        final int SHARD_SIZE = randomIntBetween(1, 3);
+        SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(SHARD_SIZE);
+        sampleAgg.subAggregation(terms("authors").field("author"));
+        SearchResponse response = client().prepareSearch("test")
+            .setSearchType(SearchType.QUERY_THEN_FETCH)
+            .setQuery(new TermQueryBuilder("genre", "fantasy"))
+            .setFrom(0)
+            .setSize(60)
+            .addAggregation(sampleAgg)
+            .get();
+        assertSearchResponse(response);
+        Sampler sample = response.getAggregations().get("sample");
+        Terms authors = sample.getAggregations().get("authors");
+        assertEquals(SHARD_SIZE * NUM_SHARDS, sample.getDocCount());
+    }
+
     public void testUnmappedChildAggNoDiversity() throws Exception {
         SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(100);
         sampleAgg.subAggregation(terms("authors").field("author"));

server/src/main/java/org/opensearch/search/aggregations/bucket/sampler/DiversifiedAggregatorFactory.java

@@ -162,6 +162,6 @@ public InternalAggregation buildEmptyAggregation() {
 
     @Override
     protected boolean supportsConcurrentSegmentSearch() {
-        return true;
+        return false;
     }
 }

server/src/main/java/org/opensearch/search/aggregations/bucket/sampler/SamplerAggregatorFactory.java

@@ -75,6 +75,6 @@ public Aggregator createInternal(
 
     @Override
     protected boolean supportsConcurrentSegmentSearch() {
-        return true;
+        return false;
     }
 }