Add cardinality prefetch pipeline with depth-4 ring buffer

CHANGELOG.md

@@ -5,6 +5,34 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 3.3.x]
 ### Added
+- Allow setting index.creation_date on index creation and restore for plugin compatibility and migrations ([#19931](https://github.com/opensearch-project/OpenSearch/pull/19931))
+- Add support for a ForkJoinPool type ([#19008](https://github.com/opensearch-project/OpenSearch/pull/19008))
+- Add seperate shard limit validation for local and remote indices ([#19532](https://github.com/opensearch-project/OpenSearch/pull/19532))
+- Use Lucene `pack` method for `half_float` and `usigned_long` when using `ApproximatePointRangeQuery`.
+- Add a mapper for context aware segments grouping criteria ([#19233](https://github.com/opensearch-project/OpenSearch/pull/19233))
+- Return full error for GRPC error response ([#19568](https://github.com/opensearch-project/OpenSearch/pull/19568))
+- Add support for repository with Server side encryption enabled and client side encryption as well based on a flag. ([#19630)](https://github.com/opensearch-project/OpenSearch/pull/19630))
+- Add pluggable gRPC interceptors with explicit ordering([#19005](https://github.com/opensearch-project/OpenSearch/pull/19005))
+- Add BindableServices extension point to transport-grpc-spi ([#19304](https://github.com/opensearch-project/OpenSearch/pull/19304))
+- Add metrics for the merged segment warmer feature ([#18929](https://github.com/opensearch-project/OpenSearch/pull/18929))
+- Handle deleted documents for filter rewrite sub-aggregation optimization ([#19643](https://github.com/opensearch-project/OpenSearch/pull/19643))
+- Add bulk collect API for filter rewrite sub-aggregation optimization ([#19933](https://github.com/opensearch-project/OpenSearch/pull/19933))
+- Allow collectors take advantage of preaggregated data using collectRange API ([#20009](https://github.com/opensearch-project/OpenSearch/pull/20009))
+- Bulk collection logic for metrics and cardinality aggregations ([#20067](https://github.com/opensearch-project/OpenSearch/pull/20067))
+- Add pointer based lag metric in pull-based ingestion ([#19635](https://github.com/opensearch-project/OpenSearch/pull/19635))
+- Introduced internal API for retrieving metadata about requested indices from transport actions  ([#18523](https://github.com/opensearch-project/OpenSearch/pull/18523))
+- Add cluster defaults for merge autoThrottle, maxMergeThreads, and maxMergeCount; Add segment size filter to the merged segment warmer ([#19629](https://github.com/opensearch-project/OpenSearch/pull/19629))
+- Add build-tooling to run in FIPS environment ([#18921](https://github.com/opensearch-project/OpenSearch/pull/18921))
+- Add SMILE/CBOR/YAML document format support to Bulk GRPC endpoint ([#19744](https://github.com/opensearch-project/OpenSearch/pull/19744))
+- Implement GRPC Search params `Highlight`and `Sort` ([#19868](https://github.com/opensearch-project/OpenSearch/pull/19868))
+- Implement GRPC ConstantScoreQuery, FuzzyQuery, MatchBoolPrefixQuery, MatchPhrasePrefix, PrefixQuery, MatchQuery ([#19854](https://github.com/opensearch-project/OpenSearch/pull/19854))
+- Add async periodic flush task support for pull-based ingestion ([#19878](https://github.com/opensearch-project/OpenSearch/pull/19878))
+- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098))
+- Implement GRPC FunctionScoreQuery ([#19888](https://github.com/opensearch-project/OpenSearch/pull/19888))
+- Implement error_trace parameter for bulk requests ([#19985](https://github.com/opensearch-project/OpenSearch/pull/19985))
+- Allow the truncate filter in normalizers ([#19778](https://github.com/opensearch-project/OpenSearch/issues/19778))
+- Support pull-based ingestion message mappers and raw payload support ([#19765](https://github.com/opensearch-project/OpenSearch/pull/19765))
+- Support dynamic consumer configuration update in pull-based ingestion ([#19963](https://github.com/opensearch-project/OpenSearch/pull/19963))
 
 ### Changed
 

gradle/libs.versions.toml

@@ -1,6 +1,6 @@
 [versions]
 opensearch        = "3.3.2"
-lucene            = "10.3.1"
+lucene            = "10.3.1-prefetch-SNAPSHOT"
 
 bundled_jdk_vendor = "adoptium"
 bundled_jdk = "24.0.2+12"

server/src/main/java/org/apache/lucene/search/DocIdStreamHelper.java

@@ -0,0 +1,21 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.apache.lucene.search;
+
+import org.apache.lucene.util.FixedBitSet;
+
+/*
+Need this helper class for initializing BitSetDocIdStream as it is
+package-private class in Lucene
+ */
+public class DocIdStreamHelper {
+    public static DocIdStream getDocIdStream(FixedBitSet fixedBitSet) {
+        return new BitSetDocIdStream(fixedBitSet, 0);
+    }
+}

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -575,6 +575,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.SEARCH_MAX_QUERY_STRING_LENGTH,
                 SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD,
                 SearchService.KEYWORD_INDEX_OR_DOC_VALUES_ENABLED,
+                SearchService.CARDINALITY_PREFETCH_PIPELINE,
                 CreatePitController.PIT_INIT_KEEP_ALIVE,
                 Node.WRITE_PORTS_FILE_SETTING,
                 Node.NODE_NAME_SETTING,

server/src/main/java/org/opensearch/search/DefaultSearchContext.java

@@ -222,6 +222,7 @@ final class DefaultSearchContext extends SearchContext {
     private final int cardinalityAggregationPruningThreshold;
     private final int bucketSelectionStrategyFactor;
     private final boolean keywordIndexOrDocValuesEnabled;
+    private final boolean cardinalityPrefetchPipeline;
 
     private boolean isStreamSearch;
     private StreamSearchChannelListener listener;
@@ -290,6 +291,7 @@ final class DefaultSearchContext extends SearchContext {
         this.bucketSelectionStrategyFactor = evaluateBucketSelectionStrategyFactor();
         this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
         this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
+        this.cardinalityPrefetchPipeline = evaluateCardinalityPrefetchPipeline();
         this.isStreamSearch = isStreamSearch;
     }
 
@@ -1269,6 +1271,18 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() {
         return false;
     }
 
+    @Override
+    public boolean cardinalityPrefetchPipeline() {
+        return cardinalityPrefetchPipeline;
+    }
+
+    private boolean evaluateCardinalityPrefetchPipeline() {
+        if (clusterService != null) {
+            return clusterService.getClusterSettings().get(SearchService.CARDINALITY_PREFETCH_PIPELINE);
+        }
+        return false;
+    }
+
     public void setStreamChannelListener(StreamSearchChannelListener listener) {
         assert isStreamSearch() : "Stream search not enabled";
         this.listener = listener;

server/src/main/java/org/opensearch/search/SearchService.java

@@ -281,7 +281,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
 
     public static final Setting<Boolean> QUERY_REWRITING_ENABLED_SETTING = Setting.boolSetting(
         "search.query_rewriting.enabled",
-        false,
+        true,
         Property.Dynamic,
         Property.NodeScope
     );
@@ -402,6 +402,13 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> CARDINALITY_PREFETCH_PIPELINE = Setting.boolSetting(
+        "search.cardinality_prefetch_pipeline",
+        false,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final int DEFAULT_BUCKET_SELECTION_STRATEGY_FACTOR = 5;
     public static final Setting<Integer> BUCKET_SELECTION_STRATEGY_FACTOR_SETTING = Setting.intSetting(
         "search.aggregation.bucket_selection_strategy_factor",

server/src/main/java/org/opensearch/search/aggregations/LeafBucketCollector.java

@@ -32,8 +32,10 @@
 
 package org.opensearch.search.aggregations;
 
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Scorable;
+import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.search.aggregations.bucket.terms.LongKeyedBucketOrds;
 
 import java.io.IOException;
@@ -123,6 +125,31 @@ public void collect(int doc) throws IOException {
         collect(doc, 0);
     }
 
+    @Override
+    public void collect(DocIdStream stream) throws IOException {
+        collect(stream, 0);
+    }
+
+    /**
+     * Bulk-collect doc IDs within {@code owningBucketOrd}.
+     *
+     * <p>Note: The provided {@link DocIdStream} may be reused across calls and should be consumed immediately.
+     *
+     * <p>Note: The provided DocIdStream typically only holds a small subset of query matches. This method may be called multiple times per segment.
+     * Like collect(int), it is guaranteed that doc IDs get collected in order, ie. doc IDs are collected in order within a DocIdStream, and if
+     * called twice, all doc IDs from the second DocIdStream will be greater than all doc IDs from the first DocIdStream.
+     *
+     * <p>It is legal for callers to mix calls to {@link #collect(DocIdStream, long)} and {@link #collect(int, long)}.
+     *
+     * <p>The default implementation calls {@code stream.forEach(doc -> collect(doc, owningBucketOrd))}.
+     */
+    @ExperimentalApi
+    public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+        // Different aggregator implementations should override this method even if to just delegate to super for
+        // helping the performance: when the super call inlines, calls to #collect(int, long) become monomorphic.
+        stream.forEach((doc) -> collect(doc, owningBucketOrd));
+    }
+
     @Override
     public void setScorer(Scorable scorer) throws IOException {
         // no-op by default

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -42,6 +42,7 @@
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.DocIdSet;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.LeafFieldComparator;
@@ -636,6 +637,11 @@ public void collect(int doc, long zeroBucket) throws IOException {
                         assert zeroBucket == 0L;
                         inner.collect(doc);
                     }
+
+                    @Override
+                    public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                        super.collect(stream, owningBucketOrd);
+                    }
                 };
             }
         }
@@ -665,6 +671,11 @@ public void collect(int doc, long bucket) throws IOException {
                     throw exc;
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
@@ -724,6 +735,11 @@ public void collect(int doc, long zeroBucket) throws IOException {
                     subCollector.collect(doc, slot);
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java

@@ -32,6 +32,7 @@
 package org.opensearch.search.aggregations.bucket.filter;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.lucene.Lucene;
@@ -82,6 +83,11 @@ public void collect(int doc, long bucket) throws IOException {
                     collectBucket(sub, doc, bucket);
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java

@@ -33,6 +33,7 @@
 package org.opensearch.search.aggregations.bucket.filter;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.lucene.Lucene;
@@ -191,6 +192,11 @@ public void collect(int doc, long bucket) throws IOException {
                     collectBucket(sub, doc, bucketOrd(bucket, bits.length));
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/rangecollector/SubAggRangeCollector.java

@@ -12,7 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.util.BitDocIdSet;
+import org.apache.lucene.search.DocIdStreamHelper;
 import org.apache.lucene.util.FixedBitSet;
 import org.opensearch.search.aggregations.BucketCollector;
 import org.opensearch.search.aggregations.LeafBucketCollector;
@@ -23,8 +23,6 @@
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 
-import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
-
 /**
  * Range collector implementation that supports sub-aggregations by collecting doc IDs.
  */
@@ -38,7 +36,6 @@ public class SubAggRangeCollector extends SimpleRangeCollector {
     private final LeafReaderContext leafCtx;
 
     private final FixedBitSet bitSet;
-    private final BitDocIdSet bitDocIdSet;
 
     public SubAggRangeCollector(
         Ranges ranges,
@@ -53,9 +50,7 @@ public SubAggRangeCollector(
         this.getBucketOrd = getBucketOrd;
         this.collectableSubAggregators = subAggCollectorParam.collectableSubAggregators();
         this.leafCtx = subAggCollectorParam.leafCtx();
-        int numDocs = leafCtx.reader().maxDoc();
-        bitSet = new FixedBitSet(numDocs);
-        bitDocIdSet = new BitDocIdSet(bitSet);
+        bitSet = new FixedBitSet(leafCtx.reader().maxDoc());
     }
 
     @Override
@@ -82,13 +77,9 @@ public void finalizePreviousRange() {
 
         // trigger the sub agg collection for this range
         try {
-            DocIdSetIterator iterator = bitDocIdSet.iterator();
             // build a new leaf collector for each bucket
             LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(leafCtx);
-            while (iterator.nextDoc() != NO_MORE_DOCS) {
-                int currentDoc = iterator.docID();
-                sub.collect(currentDoc, bucketOrd);
-            }
+            sub.collect(DocIdStreamHelper.getDocIdStream(bitSet), bucketOrd);
             logger.trace("collected sub aggregation for bucket {}", bucketOrd);
         } catch (IOException e) {
             throw new RuntimeException(e);

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java

@@ -33,6 +33,7 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.CollectionUtil;
 import org.opensearch.common.Rounding;
@@ -267,6 +268,11 @@ public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBuc
             public void collect(int doc, long owningBucketOrd) throws IOException {
                 iteratingCollector.collect(doc, owningBucketOrd);
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 
@@ -414,6 +420,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     }
                 }
 
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
+
                 private void collectValue(int doc, long rounded) throws IOException {
                     long bucketOrd = bucketOrds.add(0, rounded);
                     if (bucketOrd < 0) { // already seen
@@ -663,6 +674,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     }
                 }
 
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
+
                 private int collectValue(long owningBucketOrd, int roundingIdx, int doc, long rounded) throws IOException {
                     long bucketOrd = bucketOrds.add(owningBucketOrd, rounded);
                     if (bucketOrd < 0) { // already seen

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -251,6 +251,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                         collectValue(sub, doc, owningBucketOrd, preparedRounding.round(value));
                     }
                 }
+
+                @Override
+                public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                    super.collect(stream, owningBucketOrd);
+                }
             };
         }
 
@@ -274,6 +279,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     }
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/NumericHistogramAggregator.java

@@ -33,6 +33,7 @@
 package org.opensearch.search.aggregations.bucket.histogram;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.ScoreMode;
 import org.opensearch.index.fielddata.SortedNumericDoubleValues;
 import org.opensearch.search.aggregations.Aggregator;
@@ -137,6 +138,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
                     }
                 }
             }
+
+            @Override
+            public void collect(DocIdStream stream, long owningBucketOrd) throws IOException {
+                super.collect(stream, owningBucketOrd);
+            }
         };
     }
 }