Disable inter-segment concurrency for composite/terms agg

We currently disable concurrency for terms aggregation due to precision errors. We should do the same for composite agg when one of the sources is terms, as it is subject to the same precision errors.
javanna · Nov 2, 2023 · f9b8998 · f9b8998
1 parent 86c2718
commit f9b8998
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 0 deletions.
diff --git a/...a/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java b/...a/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java
@@ -302,4 +302,14 @@ public boolean equals(Object obj) {
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.ZERO;
     }
+
+    @Override
+    public boolean supportsParallelCollection() {
+        for (CompositeValuesSourceBuilder<?> source : sources) {
+            if (source.supportsParallelCollection() == false) {
+                return false;
+            }
+        }
+        return true;
+    }
 }
diff --git a/.../org/elasticsearch/search/aggregations/bucket/composite/CompositeValuesSourceBuilder.java b/.../org/elasticsearch/search/aggregations/bucket/composite/CompositeValuesSourceBuilder.java
@@ -325,4 +325,13 @@ public final CompositeValuesSourceConfig build(AggregationContext context) throw
     protected ZoneId timeZone() {
         return null;
     }
+
+    /**
+     * Return false if this composite values source does not support parallel collection.
+     * As a result, a request including a composite aggregation with such source is always executed sequentially despite concurrency is
+     * enabled for the query phase.
+     */
+    public boolean supportsParallelCollection() {
+        return true;
+    }
 }
diff --git a/...java/org/elasticsearch/search/aggregations/bucket/composite/TermsValuesSourceBuilder.java b/...java/org/elasticsearch/search/aggregations/bucket/composite/TermsValuesSourceBuilder.java
@@ -215,4 +215,9 @@ protected CompositeValuesSourceConfig innerBuild(ValuesSourceRegistry registry,
         return registry.getAggregator(REGISTRY_KEY, config)
             .apply(config, name, script() != null, format(), missingBucket(), missingOrder(), order());
     }
+
+    @Override
+    public boolean supportsParallelCollection() {
+        return false;
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java b/server/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java
@@ -34,6 +34,11 @@
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.AbstractSearchTestCase;
 import org.elasticsearch.search.SearchExtBuilder;
+import org.elasticsearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.composite.DateHistogramValuesSourceBuilder;
+import org.elasticsearch.search.aggregations.bucket.composite.GeoTileGridValuesSourceBuilder;
+import org.elasticsearch.search.aggregations.bucket.composite.HistogramValuesSourceBuilder;
+import org.elasticsearch.search.aggregations.bucket.composite.TermsValuesSourceBuilder;
 import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.MaxAggregationBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
@@ -921,6 +926,34 @@ public void testSupportsParallelCollection() {
             searchSourceBuilder.aggregation(new TermsAggregationBuilder("terms"));
             assertFalse(searchSourceBuilder.supportsParallelCollection());
         }
+        {
+            SearchSourceBuilder searchSourceBuilder = newSearchSourceBuilder.get();
+            searchSourceBuilder.aggregation(
+                new CompositeAggregationBuilder("name", Collections.singletonList(new DateHistogramValuesSourceBuilder("name")))
+            );
+            assertTrue(searchSourceBuilder.supportsParallelCollection());
+        }
+        {
+            SearchSourceBuilder searchSourceBuilder = newSearchSourceBuilder.get();
+            searchSourceBuilder.aggregation(
+                new CompositeAggregationBuilder("name", Collections.singletonList(new HistogramValuesSourceBuilder("name")))
+            );
+            assertTrue(searchSourceBuilder.supportsParallelCollection());
+        }
+        {
+            SearchSourceBuilder searchSourceBuilder = newSearchSourceBuilder.get();
+            searchSourceBuilder.aggregation(
+                new CompositeAggregationBuilder("name", Collections.singletonList(new GeoTileGridValuesSourceBuilder("name")))
+            );
+            assertTrue(searchSourceBuilder.supportsParallelCollection());
+        }
+        {
+            SearchSourceBuilder searchSourceBuilder = newSearchSourceBuilder.get();
+            searchSourceBuilder.aggregation(
+                new CompositeAggregationBuilder("name", Collections.singletonList(new TermsValuesSourceBuilder("name")))
+            );
+            assertFalse(searchSourceBuilder.supportsParallelCollection());
+        }
         {
             SearchSourceBuilder searchSourceBuilder = newSearchSourceBuilder.get();
             searchSourceBuilder.collapse(CollapseBuilderTests.randomCollapseBuilder());