Add support for dynamic pruning to cardinality aggregations on low-cardinality keyword fields.

docs/changelog/92060.yaml

@@ -0,0 +1,6 @@
+pr: 92060
+summary: Add support for dynamic pruning to cardinality aggregations on low-cardinality
+  keyword fields
+area: Aggregations
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/common/util/BitArray.java

@@ -39,6 +39,18 @@ public void set(long index) {
         bits.set(wordNum, bits.get(wordNum) | bitmask(index));
     }
 
+    /**
+     * Set the {@code index}th bit and return {@code true} if the bit was set already.
+     */
+    public boolean getAndSet(long index) {
+        long wordNum = wordNum(index);
+        bits = bigArrays.grow(bits, wordNum + 1);
+        long word = bits.get(wordNum);
+        long bitMask = bitmask(index);
+        bits.set(wordNum, word | bitMask);
+        return (word & bitMask) != 0;
+    }
+
     /** this = this OR other */
     public void or(BitArray other) {
         or(other.bits);

server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorFactory.java

@@ -157,10 +157,17 @@ public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
                 if (valuesSourceConfig.hasValues()) {
                     if (valuesSourceConfig.getValuesSource() instanceof final ValuesSource.Bytes.WithOrdinals source) {
                         if (executionMode.useGlobalOrdinals(context, source, precision)) {
+                            final String field;
+                            if (valuesSourceConfig.alignesWithSearchIndex()) {
+                                field = valuesSourceConfig.fieldType().name();
+                            } else {
+                                field = null;
+                            }
                             final long maxOrd = source.globalMaxOrd(context.searcher().getIndexReader());
                             return new GlobalOrdCardinalityAggregator(
                                 name,
                                 source,
+                                field,
                                 precision,
                                 Math.toIntExact(maxOrd),
                                 context,

server/src/main/java/org/elasticsearch/search/aggregations/metrics/GlobalOrdCardinalityAggregator.java

@@ -8,9 +8,16 @@
 
 package org.elasticsearch.search.aggregations.metrics;
 
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.SortedSetDocValues;
+import org.apache.lucene.index.Terms;
+import org.apache.lucene.index.TermsEnum;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.PriorityQueue;
 import org.elasticsearch.common.hash.MurmurHash3;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BitArray;
@@ -26,18 +33,35 @@
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
+import java.util.function.BiConsumer;
 
 /**
  * An aggregator that computes approximate counts of unique values
  * using global ords.
  */
 public class GlobalOrdCardinalityAggregator extends NumericMetricsAggregator.SingleValue {
 
+    // Don't try to dynamically prune fields that have more than 1024 unique terms, there is a chance we never get to 128 unseen terms, and
+    // we'd be paying the overhead of dynamic pruning without getting any benefits.
+    private static final int MAX_FIELD_CARDINALITY_FOR_DYNAMIC_PRUNING = 1024;
+
+    // Only start dynamic pruning when 128 ordinals or less have not been seen yet.
+    private static final int MAX_TERMS_FOR_DYNAMIC_PRUNING = 128;
+
     private final ValuesSource.Bytes.WithOrdinals valuesSource;
+    // The field that this cardinality aggregation runs on, or null if there is no field, or the field doesn't directly map to an index
+    // field.
+    private final String field;
     private final BigArrays bigArrays;
     private final int maxOrd;
     private final int precision;
+    private int dynamicPruningAttempts;
+    private int dynamicPruningSuccess;
+    private int bruteForce;
+    private int noData;
 
     // Build at post-collection phase
     @Nullable
@@ -48,6 +72,7 @@ public class GlobalOrdCardinalityAggregator extends NumericMetricsAggregator.Sin
     public GlobalOrdCardinalityAggregator(
         String name,
         ValuesSource.Bytes.WithOrdinals valuesSource,
+        String field,
         int precision,
         int maxOrd,
         AggregationContext context,
@@ -56,6 +81,7 @@ public GlobalOrdCardinalityAggregator(
     ) throws IOException {
         super(name, context, parent, metadata);
         this.valuesSource = valuesSource;
+        this.field = field;
         this.precision = precision;
         this.maxOrd = maxOrd;
         this.bigArrays = context.bigArrays();
@@ -64,12 +90,182 @@ public GlobalOrdCardinalityAggregator(
 
     @Override
     public ScoreMode scoreMode() {
-        return valuesSource.needsScores() ? ScoreMode.COMPLETE : ScoreMode.COMPLETE_NO_SCORES;
+        if (field != null && valuesSource.needsScores() == false && maxOrd <= MAX_FIELD_CARDINALITY_FOR_DYNAMIC_PRUNING) {
+            return ScoreMode.TOP_DOCS;
+        } else if (valuesSource.needsScores()) {
+            return ScoreMode.COMPLETE;
+        } else {
+            return ScoreMode.COMPLETE_NO_SCORES;
+        }
+    }
+
+    /**
+     * A competitive iterator that helps only collect values that have not been collected so far.
+     */
+    private class CompetitiveIterator extends DocIdSetIterator {
+
+        private final BitArray visitedOrds;
+        private long numNonVisitedOrds;
+        private final TermsEnum indexTerms;
+        private final DocIdSetIterator docsWithField;
+
+        CompetitiveIterator(int numNonVisitedOrds, BitArray visitedOrds, Terms indexTerms, DocIdSetIterator docsWithField)
+            throws IOException {
+            this.visitedOrds = visitedOrds;
+            this.numNonVisitedOrds = numNonVisitedOrds;
+            this.indexTerms = Objects.requireNonNull(indexTerms).iterator();
+            this.docsWithField = docsWithField;
+        }
+
+        private Map<Long, PostingsEnum> nonVisitedOrds;
+        private PriorityQueue<PostingsEnum> nonVisitedPostings;
+
+        private int doc = -1;
+
+        @Override
+        public int docID() {
+            return doc;
+        }
+
+        @Override
+        public int nextDoc() throws IOException {
+            return advance(doc + 1);
+        }
+
+        @Override
+        public int advance(int target) throws IOException {
+            if (nonVisitedPostings == null) {
+                // We haven't started pruning yet, iterate on docs that have a value. This may already help a lot on sparse fields.
+                return doc = docsWithField.advance(target);
+            } else if (nonVisitedPostings.size() == 0) {
+                return doc = DocIdSetIterator.NO_MORE_DOCS;
+            } else {
+                PostingsEnum top = nonVisitedPostings.top();
+                while (top.docID() < target) {
+                    top.advance(target);
+                    top = nonVisitedPostings.updateTop();
+                }
+                return doc = top.docID();
+            }
+        }
+
+        @Override
+        public long cost() {
+            return docsWithField.cost();
+        }
+
+        void startPruning() throws IOException {
+            dynamicPruningSuccess++;
+            nonVisitedOrds = new HashMap<>();
+            // TODO: iterate the bitset using a `nextClearBit` operation?
+            for (long ord = 0; ord < maxOrd; ++ord) {
+                if (visitedOrds.get(ord)) {
+                    continue;
+                }
+                BytesRef term = values.lookupOrd(ord);
+                if (indexTerms.seekExact(term) == false) {
+                    // This global ordinal maps to a value that doesn't exist in this segment
+                    continue;
+                }
+                nonVisitedOrds.put(ord, indexTerms.postings(null, PostingsEnum.NONE));
+            }
+            nonVisitedPostings = new PriorityQueue<>(nonVisitedOrds.size()) {
+                @Override
+                protected boolean lessThan(PostingsEnum a, PostingsEnum b) {
+                    return a.docID() < b.docID();
+                }
+            };
+            for (PostingsEnum pe : nonVisitedOrds.values()) {
+                nonVisitedPostings.add(pe);
+            }
+        }
+
+        void onVisitedOrdinal(long ordinal) throws IOException {
+            numNonVisitedOrds--;
+            if (nonVisitedOrds == null) {
+                if (numNonVisitedOrds <= MAX_TERMS_FOR_DYNAMIC_PRUNING) {
+                    startPruning();
+                }
+            } else {
+                if (nonVisitedOrds.remove(ordinal) != null) {
+                    // Could we make this more efficient?
+                    nonVisitedPostings.clear();
+                    for (PostingsEnum pe : nonVisitedOrds.values()) {
+                        nonVisitedPostings.add(pe);
+                    }
+                }
+            }
+        }
     }
 
     @Override
     public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx, final LeafBucketCollector sub) throws IOException {
         values = valuesSource.globalOrdinalsValues(aggCtx.getLeafReaderContext());
+
+        if (parent == null && field != null) {
+            // This optimization only applies to top-level cardinality aggregations that apply to fields indexed with an inverted index.
+            final Terms indexTerms = aggCtx.getLeafReaderContext().reader().terms(field);
+            if (indexTerms != null) {
+                BitArray bits = visitedOrds.get(0);
+                final int numNonVisitedOrds = maxOrd - (bits == null ? 0 : (int) bits.cardinality());
+                if (maxOrd <= MAX_FIELD_CARDINALITY_FOR_DYNAMIC_PRUNING || numNonVisitedOrds <= MAX_TERMS_FOR_DYNAMIC_PRUNING) {
+                    dynamicPruningAttempts++;
+                    return new LeafBucketCollector() {
+
+                        final BitArray bits;
+                        final CompetitiveIterator competitiveIterator;
+
+                        {
+                            // This optimization only works for top-level cardinality aggregations that collect bucket 0, so we can retrieve
+                            // the appropriate BitArray ahead of time.
+                            visitedOrds = bigArrays.grow(visitedOrds, 1);
+                            BitArray bits = visitedOrds.get(0);
+                            if (bits == null) {
+                                bits = new BitArray(maxOrd, bigArrays);
+                                visitedOrds.set(0, bits);
+                            }
+                            this.bits = bits;
+                            final DocIdSetIterator docsWithField = valuesSource.ordinalsValues(aggCtx.getLeafReaderContext());
+                            competitiveIterator = new CompetitiveIterator(numNonVisitedOrds, bits, indexTerms, docsWithField);
+                            if (numNonVisitedOrds <= MAX_TERMS_FOR_DYNAMIC_PRUNING) {
+                                competitiveIterator.startPruning();
+                            }
+                        }
+
+                        @Override
+                        public void collect(int doc, long bucketOrd) throws IOException {
+                            if (values.advanceExact(doc)) {
+                                for (long ord = values.nextOrd(); ord != SortedSetDocValues.NO_MORE_ORDS; ord = values.nextOrd()) {
+                                    if (bits.getAndSet(ord) == false) {
+                                        competitiveIterator.onVisitedOrdinal(ord);
+                                    }
+                                }
+                            }
+                        }
+
+                        @Override
+                        public CompetitiveIterator competitiveIterator() {
+                            return competitiveIterator;
+                        }
+                    };
+                }
+            } else {
+                final FieldInfo fi = aggCtx.getLeafReaderContext().reader().getFieldInfos().fieldInfo(field);
+                if (fi == null) {
+                    // The field doesn't exist at all, we can skip the segment entirely
+                    noData++;
+                    return LeafBucketCollector.NO_OP_COLLECTOR;
+                } else if (fi.getIndexOptions() != IndexOptions.NONE) {
+                    // The field doesn't have terms while index options are not NONE. This means that this segment doesn't have a single
+                    // value for the field.
+                    noData++;
+                    return LeafBucketCollector.NO_OP_COLLECTOR;
+                }
+                // Otherwise we might be aggregating e.g. an IP field, which indexes data using points rather than an inverted index.
+            }
+        }
+
+        bruteForce++;
         return new LeafBucketCollector() {
             @Override
             public void collect(int doc, long bucketOrd) throws IOException {
@@ -157,4 +353,13 @@ protected void doClose() {
         }
         Releasables.close(visitedOrds, counts);
     }
+
+    @Override
+    public void collectDebugInfo(BiConsumer<String, Object> add) {
+        super.collectDebugInfo(add);
+        add.accept("dynamic_pruning_attempted", dynamicPruningAttempts);
+        add.accept("dynamic_pruning_used", dynamicPruningSuccess);
+        add.accept("brute_force_used", bruteForce);
+        add.accept("skipped_due_to_no_data", noData);
+    }
 }

server/src/test/java/org/elasticsearch/common/util/BitArrayTests.java

@@ -154,4 +154,21 @@ public void testCardinality() {
             }
         }
     }
+
+    public void testGetAndSet() {
+        try (BitArray bitArray = new BitArray(1, BigArrays.NON_RECYCLING_INSTANCE)) {
+            assertFalse(bitArray.getAndSet(100));
+            assertFalse(bitArray.getAndSet(1000));
+            assertTrue(bitArray.getAndSet(100));
+            assertFalse(bitArray.getAndSet(101));
+            assertFalse(bitArray.getAndSet(999));
+            assertTrue(bitArray.getAndSet(1000));
+            assertFalse(bitArray.get(99));
+            assertTrue(bitArray.get(100));
+            assertTrue(bitArray.get(101));
+            assertTrue(bitArray.get(999));
+            assertTrue(bitArray.get(1000));
+            assertFalse(bitArray.get(1001));
+        }
+    }
 }