Add Hybrid Cardinality collector to prioritize Ordinals Collector

anandpatel9998 · anandpatel9998 · commit 3d9432ee6f8b · 2025-10-07T09:07:38.000-07:00
Current cardinality aggregator logic selects DirectCollector over OrdinalsCollector when relative memory overhead due to OrdinalsCollector (compared to DirectCollector) is higher. Because of this relative memory consumption logic, DirectCollector is selected for high cardinality aggregation queries. DirectCollector is slower compared to OrdinalsCollector. This default selection leads to higher search latency even when Opensearch process have available memory to use ordinals collector for faster query performance.

There is no way to figure out memory requirement for nested aggregation because number of buckets are dynamically created as we traverse through all the matching document ids. To overcome this limitation, this change create a hybrid collector which will first use Ordinals Collector and will switch to DirectCollector if memory usage for Ordinals Collector Increase beyond certain threshold. When Hybrid collector switch from Ordinals Collector to Direct Collector, it will utilize already computed aggregation data from Ordinals Collector so that we do not have to rebuild aggregation result using Direct Collector.

Signed-off-by: Anand Pravinbhai Patel &lt;anapat@amazon.com&gt;
diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/170_cardinality_metric.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/170_cardinality_metric.yml
@@ -287,6 +287,7 @@ setup:
   - gt: { profile.shards.0.aggregations.0.breakdown.post_collection: 0 }
   - match: { profile.shards.0.aggregations.0.debug.empty_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.numeric_collectors_used: 0 }
-  - gt: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
+  - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_overhead_too_high: 0 }
   - match: { profile.shards.0.aggregations.0.debug.string_hashing_collectors_used: 0 }
+  - gt: { profile.shards.0.aggregations.0.debug.hybrid_collectors_used: 0 }
diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/170_cardinality_metric_unsigned.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/170_cardinality_metric_unsigned.yml
@@ -282,6 +282,7 @@ setup:
   - gt: { profile.shards.0.aggregations.0.breakdown.post_collection: 0 }
   - match: { profile.shards.0.aggregations.0.debug.empty_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.numeric_collectors_used: 0 }
-  - gt: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
+  - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_overhead_too_high: 0 }
   - match: { profile.shards.0.aggregations.0.debug.string_hashing_collectors_used: 0 }
+  - gt: { profile.shards.0.aggregations.0.debug.hybrid_collectors_used: 0 }
diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java
@@ -159,6 +159,7 @@
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.SearchService;
 import org.opensearch.search.aggregations.MultiBucketConsumerService;
+import org.opensearch.search.aggregations.metrics.CardinalityAggregator;
 import org.opensearch.search.backpressure.settings.NodeDuressSettings;
 import org.opensearch.search.backpressure.settings.SearchBackpressureSettings;
 import org.opensearch.search.backpressure.settings.SearchShardTaskSettings;
@@ -574,6 +575,8 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.INDICES_MAX_CLAUSE_COUNT_SETTING,
                 SearchService.SEARCH_MAX_QUERY_STRING_LENGTH,
                 SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD,
+                CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_ENABLED,
+                CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_MEMORY_THRESHOLD,
                 SearchService.KEYWORD_INDEX_OR_DOC_VALUES_ENABLED,
                 CreatePitController.PIT_INIT_KEEP_ALIVE,
                 Node.WRITE_PORTS_FILE_SETTING,
diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -55,6 +55,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.BigArrays;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
@@ -75,6 +76,7 @@
 import org.opensearch.search.aggregations.BucketCollectorProcessor;
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.SearchContextAggregations;
+import org.opensearch.search.aggregations.metrics.CardinalityAggregator;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.collapse.CollapseContext;
 import org.opensearch.search.deciders.ConcurrentSearchDecision;
@@ -220,6 +222,8 @@ final class DefaultSearchContext extends SearchContext {
     private final int maxAggRewriteFilters;
     private final int filterRewriteSegmentThreshold;
     private final int cardinalityAggregationPruningThreshold;
+    private final boolean cardinalityAggregationHybridCollectorEnabled;
+    private final long cardinalityAggregationHybridCollectorMemoryThreshold;
     private final int bucketSelectionStrategyFactor;
     private final boolean keywordIndexOrDocValuesEnabled;
 
@@ -287,6 +291,8 @@ final class DefaultSearchContext extends SearchContext {
         this.maxAggRewriteFilters = evaluateFilterRewriteSetting();
         this.filterRewriteSegmentThreshold = evaluateAggRewriteFilterSegThreshold();
         this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
+        this.cardinalityAggregationHybridCollectorEnabled = evaluateCardinalityAggregationHybridCollectorEnabled();
+        this.cardinalityAggregationHybridCollectorMemoryThreshold = evaluateCardinalityAggregationHybridCollectorMemoryThreshold();
         this.bucketSelectionStrategyFactor = evaluateBucketSelectionStrategyFactor();
         this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
         this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
@@ -1238,6 +1244,16 @@ public int cardinalityAggregationPruningThreshold() {
         return cardinalityAggregationPruningThreshold;
     }
 
+    @Override
+    public boolean cardinalityAggregationHybridCollectorEnabled() {
+        return cardinalityAggregationHybridCollectorEnabled;
+    }
+
+    @Override
+    public long cardinalityAggregationHybridCollectorMemoryThreshold() {
+        return cardinalityAggregationHybridCollectorMemoryThreshold;
+    }
+
     @Override
     public int bucketSelectionStrategyFactor() {
         return bucketSelectionStrategyFactor;
@@ -1255,6 +1271,22 @@ private int evaluateCardinalityAggregationPruningThreshold() {
         return 0;
     }
 
+    private boolean evaluateCardinalityAggregationHybridCollectorEnabled() {
+        if (clusterService != null) {
+            return clusterService.getClusterSettings().get(CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_ENABLED);
+        }
+        return false;
+    }
+
+    private long evaluateCardinalityAggregationHybridCollectorMemoryThreshold() {
+        if (clusterService != null) {
+            ByteSizeValue threshold = clusterService.getClusterSettings()
+                .get(CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_MEMORY_THRESHOLD);
+            return threshold.getBytes();
+        }
+        return Runtime.getRuntime().maxMemory() / 100; // 1% default
+    }
+
     private int evaluateBucketSelectionStrategyFactor() {
         if (clusterService != null) {
             return clusterService.getClusterSettings().get(BUCKET_SELECTION_STRATEGY_FACTOR_SETTING);
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java
@@ -57,11 +57,13 @@
 import org.opensearch.common.hash.MurmurHash3;
 import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.lease.Releasables;
+import org.opensearch.common.settings.Setting;
 import org.opensearch.common.util.BigArrays;
 import org.opensearch.common.util.BitArray;
 import org.opensearch.common.util.BitMixer;
 import org.opensearch.common.util.LongArray;
 import org.opensearch.common.util.ObjectArray;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.index.fielddata.SortedBinaryDocValues;
 import org.opensearch.index.fielddata.SortedNumericDoubleValues;
@@ -88,6 +90,27 @@ public class CardinalityAggregator extends NumericMetricsAggregator.SingleValue
 
     private static final Logger logger = LogManager.getLogger(CardinalityAggregator.class);
 
+    /**
+     * Setting to enable/disable hybrid collector for cardinality aggregation.
+     */
+    public static final Setting<Boolean> CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_ENABLED = Setting.boolSetting(
+        "search.aggregations.cardinality.hybrid_collector.enabled",
+        true,
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
+    /**
+     * Setting for hybrid collector memory threshold. Supports both percentage (e.g., "1%")
+     * and absolute values (e.g., "10mb", "1gb").
+     */
+    public static final Setting<ByteSizeValue> CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_MEMORY_THRESHOLD = Setting.memorySizeSetting(
+        "search.aggregations.cardinality.hybrid_collector.memory_threshold",
+        "1%",
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
     final CardinalityAggregatorFactory.ExecutionMode executionMode;
     final int precision;
     final ValuesSource valuesSource;
@@ -103,6 +126,7 @@ public class CardinalityAggregator extends NumericMetricsAggregator.SingleValue
     int emptyCollectorsUsed;
     int numericCollectorsUsed;
     int ordinalsCollectorsUsed;
+    int hybridCollectorsUsed;
     int ordinalsCollectorsOverheadTooHigh;
     int stringHashingCollectorsUsed;
     int dynamicPrunedSegments;
@@ -157,14 +181,24 @@ private Collector pickCollector(LeafReaderContext ctx) throws IOException {
                 collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays());
             } else if (executionMode == null) {
                 // no hint provided, fall back to heuristics
-                final long ordinalsMemoryUsage = OrdinalsCollector.memoryOverhead(maxOrd);
-                final long countsMemoryUsage = HyperLogLogPlusPlus.memoryUsage(precision);
-                // only use ordinals if they don't increase memory usage by more than 25%
-                if (ordinalsMemoryUsage < countsMemoryUsage / 4) {
-                    ordinalsCollectorsUsed++;
-                    collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays());
+                // Check if hybrid collector is enabled
+                boolean hybridCollectorEnabled = context.cardinalityAggregationHybridCollectorEnabled();
+
+                if (hybridCollectorEnabled) {
+                    // Use HybridCollector with configurable memory threshold
+                    long memoryThreshold = context.cardinalityAggregationHybridCollectorMemoryThreshold();
+                    MurmurHash3Values hashValues = MurmurHash3Values.hash(source.bytesValues(ctx));
+                    hybridCollectorsUsed++;
+                    collector = new HybridCollector(counts, ordinalValues, hashValues, context.bigArrays(), memoryThreshold);
                 } else {
-                    ordinalsCollectorsOverheadTooHigh++;
+                    final long ordinalsMemoryUsage = OrdinalsCollector.memoryOverhead(maxOrd);
+                    final long countsMemoryUsage = HyperLogLogPlusPlus.memoryUsage(precision);
+                    if (ordinalsMemoryUsage < countsMemoryUsage / 4) {
+                        ordinalsCollectorsUsed++;
+                        collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays());
+                    } else {
+                        ordinalsCollectorsOverheadTooHigh++;
+                    }
                 }
             }
         }
@@ -339,6 +373,7 @@ public void collectDebugInfo(BiConsumer<String, Object> add) {
         add.accept("empty_collectors_used", emptyCollectorsUsed);
         add.accept("numeric_collectors_used", numericCollectorsUsed);
         add.accept("ordinals_collectors_used", ordinalsCollectorsUsed);
+        add.accept("hybrid_collectors_used", hybridCollectorsUsed);
         add.accept("ordinals_collectors_overhead_too_high", ordinalsCollectorsOverheadTooHigh);
         add.accept("string_hashing_collectors_used", stringHashingCollectorsUsed);
         add.accept("dynamic_pruned_segments", dynamicPrunedSegments);
@@ -564,6 +599,7 @@ public static long memoryOverhead(long maxOrd) {
         private final int maxOrd;
         private final HyperLogLogPlusPlus counts;
         private ObjectArray<BitArray> visitedOrds;
+        private long currentMemoryUsage = 0;
 
         OrdinalsCollector(HyperLogLogPlusPlus counts, SortedSetDocValues values, BigArrays bigArrays) {
             if (values.getValueCount() > Integer.MAX_VALUE) {
@@ -583,6 +619,8 @@ public void collect(int doc, long bucketOrd) throws IOException {
             if (bits == null) {
                 bits = new BitArray(maxOrd, bigArrays);
                 visitedOrds.set(bucketOrd, bits);
+                // Update memory usage when new BitArray is created
+                currentMemoryUsage += memoryOverhead(maxOrd);
             }
             if (values.advanceExact(doc)) {
                 int count = values.docValueCount();
@@ -627,6 +665,10 @@ public void postCollect() throws IOException {
             }
         }
 
+        public long getCurrentMemoryUsage() {
+            return currentMemoryUsage;
+        }
+
         @Override
         public void close() {
             for (int i = 0; i < visitedOrds.size(); i++) {
@@ -761,4 +803,71 @@ public long nextValue() throws IOException {
             }
         }
     }
+
+    /**
+     * Hybrid Collector that starts with OrdinalsCollector and switches to DirectCollector
+     * when memory consumption exceeds a threshold.
+     */
+    static class HybridCollector extends Collector {
+        private final HyperLogLogPlusPlus counts;
+        private final MurmurHash3Values hashValues;
+        private final long memoryThreshold;
+
+        private Collector activeCollector;
+        private final OrdinalsCollector ordinalsCollector;
+        private boolean switchedToDirectCollector = false;
+
+        HybridCollector(
+            HyperLogLogPlusPlus counts,
+            SortedSetDocValues ordinalValues,
+            MurmurHash3Values hashValues,
+            BigArrays bigArrays,
+            long memoryThreshold
+        ) {
+            this.counts = counts;
+            this.hashValues = hashValues;
+            this.memoryThreshold = memoryThreshold;
+
+            // Start with OrdinalsCollector
+            this.ordinalsCollector = new OrdinalsCollector(counts, ordinalValues, bigArrays);
+            this.activeCollector = ordinalsCollector;
+        }
+
+        @Override
+        public void collect(int doc, long bucketOrd) throws IOException {
+            activeCollector.collect(doc, bucketOrd);
+
+            // Check memory usage after collection
+            if (!switchedToDirectCollector && ordinalsCollector.getCurrentMemoryUsage() > memoryThreshold) {
+                switchToDirectCollector();
+            }
+        }
+
+        private void switchToDirectCollector() throws IOException {
+            // Switching to Direct Collector because memory consumption from Ordinals Collector is high.
+            // Post collect all the already computed data.
+            ordinalsCollector.postCollect();
+            ordinalsCollector.close();
+            logger.debug("switching to direct collector");
+
+            // Pass computed data to Direct Collector.
+            activeCollector = new DirectCollector(counts, hashValues);
+            switchedToDirectCollector = true;
+        }
+
+        @Override
+        public void postCollect() throws IOException {
+            activeCollector.postCollect();
+        }
+
+        @Override
+        public void close() {
+            Releasables.close(activeCollector);
+        }
+
+        // Visible for testing
+        Collector getActiveCollector() {
+            return activeCollector;
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -555,6 +555,14 @@ public int cardinalityAggregationPruningThreshold() {
         return 0;
     }
 
+    public boolean cardinalityAggregationHybridCollectorEnabled() {
+        return false;
+    }
+
+    public long cardinalityAggregationHybridCollectorMemoryThreshold() {
+        return Runtime.getRuntime().maxMemory() / 100; // 1% of available memory default
+    }
+
     public int bucketSelectionStrategyFactor() {
         return SearchService.DEFAULT_BUCKET_SELECTION_STRATEGY_FACTOR;
     }
diff --git a/server/src/test/java/org/opensearch/search/aggregations/metrics/CardinalityAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/metrics/CardinalityAggregatorTests.java

Original file line number	Diff line number	Diff line change
`@@ -555,6 +555,14 @@ public int cardinalityAggregationPruningThreshold() {`
`555`	`555`	`return 0;`
`556`	`556`	`}`
`557`	`557`
	`558`	`+ public boolean cardinalityAggregationHybridCollectorEnabled() {`
	`559`	`+ return false;`
	`560`	`+ }`
	`561`	`+`
	`562`	`+ public long cardinalityAggregationHybridCollectorMemoryThreshold() {`
	`563`	`+ return Runtime.getRuntime().maxMemory() / 100; // 1% of available memory default`
	`564`	`+ }`
	`565`	`+`
`558`	`566`	`public int bucketSelectionStrategyFactor() {`
`559`	`567`	`return SearchService.DEFAULT_BUCKET_SELECTION_STRATEGY_FACTOR;`
`560`	`568`	`}`