Support Dynamic Pruning in Cardinality Aggregation (#13821)

* Cardinality aggregation dynamic pruning changes Signed-off-by: bowenlan-amzn <[email protected]> * Reading Signed-off-by: bowenlan-amzn <[email protected]> * remaining disjunction scorer full understand Signed-off-by: bowenlan-amzn <[email protected]> * utilize competitive iterator api to perform pruning Signed-off-by: bowenlan-amzn <[email protected]> * handle missing input Signed-off-by: bowenlan-amzn <[email protected]> * add change log Signed-off-by: bowenlan-amzn <[email protected]> * clean up Signed-off-by: bowenlan-amzn <[email protected]> * Clean up Signed-off-by: bowenlan-amzn <[email protected]> * Test fix Signed-off-by: bowenlan-amzn <[email protected]> * Do all the scoring within Cardinality Signed-off-by: bowenlan-amzn <[email protected]> * clean unnecessary Signed-off-by: bowenlan-amzn <[email protected]> * fix Signed-off-by: bowenlan-amzn <[email protected]> * Add dynamic flag for this feature Signed-off-by: bowenlan-amzn <[email protected]> * Add random test, small bug fix Signed-off-by: bowenlan-amzn <[email protected]> * address comment Signed-off-by: bowenlan-amzn <[email protected]> * Address comments Signed-off-by: bowenlan-amzn <[email protected]> * address comments Signed-off-by: bowenlan-amzn <[email protected]> --------- Signed-off-by: bowenlan-amzn <[email protected]> Co-authored-by: Rishabh Maurya <[email protected]>
opensearch-project · Jun 11, 2024 · 6d67099 · 6d67099
1 parent c38dfef
commit 6d67099
Show file tree

Hide file tree

Showing 10 changed files with 624 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add support for query level resource usage tracking ([#13172](https://github.com/opensearch-project/OpenSearch/pull/13172))
 - [Query Insights] Add cpu and memory metrics to top n queries ([#13739](https://github.com/opensearch-project/OpenSearch/pull/13739))
 - Derived field object type support ([#13720](https://github.com/opensearch-project/OpenSearch/pull/13720))
+- Support Dynamic Pruning in Cardinality Aggregation ([#13821](https://github.com/opensearch-project/OpenSearch/pull/13821))
 
 ### Dependencies
 - Bump `com.github.spullara.mustache.java:compiler` from 0.9.10 to 0.9.13 ([#13329](https://github.com/opensearch-project/OpenSearch/pull/13329), [#13559](https://github.com/opensearch-project/OpenSearch/pull/13559))

diff --git a/...rc/internalClusterTest/java/org/opensearch/search/aggregations/metrics/CardinalityIT.java b/...rc/internalClusterTest/java/org/opensearch/search/aggregations/metrics/CardinalityIT.java
@@ -34,6 +34,7 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse;
 import org.opensearch.action.index.IndexRequestBuilder;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.common.settings.Settings;
@@ -59,6 +60,7 @@
 import static java.util.Collections.emptyMap;
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
+import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD;
 import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.search.aggregations.AggregationBuilders.cardinality;
 import static org.opensearch.search.aggregations.AggregationBuilders.global;
@@ -255,6 +257,36 @@ public void testSingleValuedString() throws Exception {
         assertCount(count, numDocs);
     }
 
+    public void testDisableDynamicPruning() throws Exception {
+        SearchResponse response = client().prepareSearch("idx")
+            .addAggregation(cardinality("cardinality").precisionThreshold(precisionThreshold).field("str_value"))
+            .get();
+        assertSearchResponse(response);
+
+        Cardinality count1 = response.getAggregations().get("cardinality");
+
+        final ClusterUpdateSettingsResponse updateSettingResponse = client().admin()
+            .cluster()
+            .prepareUpdateSettings()
+            .setTransientSettings(Settings.builder().put(CARDINALITY_AGGREGATION_PRUNING_THRESHOLD.getKey(), 0))
+            .get();
+        assertEquals(updateSettingResponse.getTransientSettings().get(CARDINALITY_AGGREGATION_PRUNING_THRESHOLD.getKey()), "0");
+
+        response = client().prepareSearch("idx")
+            .addAggregation(cardinality("cardinality").precisionThreshold(precisionThreshold).field("str_value"))
+            .get();
+        assertSearchResponse(response);
+        Cardinality count2 = response.getAggregations().get("cardinality");
+
+        assertEquals(count1, count2);
+
+        client().admin()
+            .cluster()
+            .prepareUpdateSettings()
+            .setTransientSettings(Settings.builder().putNull(CARDINALITY_AGGREGATION_PRUNING_THRESHOLD.getKey()))
+            .get();
+    }
+
     public void testSingleValuedNumeric() throws Exception {
         SearchResponse response = client().prepareSearch("idx")
             .addAggregation(cardinality("cardinality").precisionThreshold(precisionThreshold).field(singleNumericField()))

diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java
@@ -541,6 +541,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.MAX_OPEN_PIT_CONTEXT,
                 SearchService.MAX_PIT_KEEPALIVE_SETTING,
                 SearchService.MAX_AGGREGATION_REWRITE_FILTERS,
+                SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD,
                 CreatePitController.PIT_INIT_KEEP_ALIVE,
                 Node.WRITE_PORTS_FILE_SETTING,
                 Node.NODE_NAME_SETTING,

diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -106,6 +106,7 @@
 import java.util.function.Function;
 import java.util.function.LongSupplier;
 
+import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD;
 import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.search.SearchService.MAX_AGGREGATION_REWRITE_FILTERS;
 
@@ -189,6 +190,7 @@ final class DefaultSearchContext extends SearchContext {
     private final boolean concurrentSearchSettingsEnabled;
     private final SetOnce<Boolean> requestShouldUseConcurrentSearch = new SetOnce<>();
     private final int maxAggRewriteFilters;
+    private final int cardinalityAggregationPruningThreshold;
 
     DefaultSearchContext(
         ReaderContext readerContext,
@@ -244,6 +246,7 @@ final class DefaultSearchContext extends SearchContext {
         this.requestToAggReduceContextBuilder = requestToAggReduceContextBuilder;
 
         this.maxAggRewriteFilters = evaluateFilterRewriteSetting();
+        this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
     }
 
     @Override
@@ -1010,4 +1013,16 @@ private int evaluateFilterRewriteSetting() {
         }
         return 0;
     }
+
+    @Override
+    public int cardinalityAggregationPruningThreshold() {
+        return cardinalityAggregationPruningThreshold;
+    }
+
+    private int evaluateCardinalityAggregationPruningThreshold() {
+        if (clusterService != null) {
+            return clusterService.getClusterSettings().get(CARDINALITY_AGGREGATION_PRUNING_THRESHOLD);
+        }
+        return 0;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -289,6 +289,15 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
+    // value 0 can disable dynamic pruning optimization in cardinality aggregation
+    public static final Setting<Integer> CARDINALITY_AGGREGATION_PRUNING_THRESHOLD = Setting.intSetting(
+        "search.dynamic_pruning.cardinality_aggregation.max_allowed_cardinality",
+        100,
+        0,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final int DEFAULT_SIZE = 10;
     public static final int DEFAULT_FROM = 0;