diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bd7f080d9dd7..ddf05b7ca6fa8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 968b73fbb56d6..c10f217cfca0e 100644 --- a/gradle/libs.versions.toml +++ b/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" diff --git a/server/src/main/java/org/apache/lucene/search/DocIdStreamHelper.java b/server/src/main/java/org/apache/lucene/search/DocIdStreamHelper.java new file mode 100644 index 0000000000000..4bcee99e9002c --- /dev/null +++ b/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); + } +} diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index 7fe3cde1e23f2..8f4a0c54e3707 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/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, diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index 14f7b4b321638..45d1fd44b444a 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/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; diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index 9743f5103f4be..f30bf04d0615f 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -281,7 +281,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv public static final Setting 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 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 BUCKET_SELECTION_STRATEGY_FACTOR_SETTING = Setting.intSetting( "search.aggregation.bucket_selection_strategy_factor", diff --git a/server/src/main/java/org/opensearch/search/aggregations/LeafBucketCollector.java b/server/src/main/java/org/opensearch/search/aggregations/LeafBucketCollector.java index 0b34ffc78853a..b89ccd87faceb 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/LeafBucketCollector.java +++ b/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}. + * + *

Note: The provided {@link DocIdStream} may be reused across calls and should be consumed immediately. + * + *

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. + * + *

It is legal for callers to mix calls to {@link #collect(DocIdStream, long)} and {@link #collect(int, long)}. + * + *

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 diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java index dc3411f654022..58ac375c4b938 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java +++ b/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); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java index 0f73d3453b658..603ac101c1015 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FilterAggregator.java +++ b/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); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java index 9ea4c66a9a1c4..f7ea6442b267d 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filter/FiltersAggregator.java +++ b/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); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/rangecollector/SubAggRangeCollector.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/rangecollector/SubAggRangeCollector.java index 5c1f21b22e646..bf3ca0b516b14 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/filterrewrite/rangecollector/SubAggRangeCollector.java +++ b/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); diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java index 8fa9c61821fd8..dfd0b7b3b99ba 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/AutoDateHistogramAggregator.java +++ b/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 diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java index 7ba939f64dbbf..9d691f72d10fe 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java +++ b/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); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/NumericHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/NumericHistogramAggregator.java index 1cd44be431964..8f210c4790591 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/NumericHistogramAggregator.java +++ b/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); + } }; } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/RangeHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/RangeHistogramAggregator.java index fe77033f5a287..1390f51297dbf 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/RangeHistogramAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/RangeHistogramAggregator.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.util.BytesRef; import org.opensearch.index.fielddata.SortedBinaryDocValues; import org.opensearch.index.mapper.RangeFieldMapper; @@ -155,6 +156,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { } } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }; } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java index 087ec7b3d6763..515dac03bfac7 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java @@ -33,6 +33,7 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.FixedBitSet; import org.opensearch.core.ParseField; @@ -353,6 +354,11 @@ public void collect(int doc, long bucket) throws IOException { } } + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } + private int collect(int doc, double value, long owningBucketOrdinal, int lowBound) throws IOException { MatchedRange range = new MatchedRange(ranges, lowBound, value, maxTo); for (int i = range.startLo; i <= range.endHi; ++i) { diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java index 83ca8874d3559..e9bba646a62c5 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java @@ -41,6 +41,7 @@ import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.index.Terms; import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.Weight; import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.BytesRef; @@ -265,6 +266,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol * common and marginally faster. */ return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) { + final int[] buffer = context.cardinalityPrefetchPipeline() ? new int[4096] : null; + @Override public void collect(int doc, long owningBucketOrd) throws IOException { if (false == singleValues.advanceExact(doc)) { @@ -273,6 +276,19 @@ public void collect(int doc, long owningBucketOrd) throws IOException { int globalOrd = singleValues.ordValue(); collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub); } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + if (buffer == null) { + super.collect(stream, owningBucketOrd); + return; + } + int count = stream.intoArray(buffer); + singleValues.prefetchOrdValues(count, buffer); + for (int i = 0; i < count; i++) { + collect(buffer[i], owningBucketOrd); + } + } }); } return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) { @@ -307,6 +323,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub); } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }); } return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, globalOrds) { @@ -324,6 +345,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { collectionStrategy.collectGlobalOrd(owningBucketOrd, doc, globalOrd, sub); } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }); } @@ -542,6 +568,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol if (singleValues != null) { segmentsWithSingleValuedOrds++; return resultStrategy.wrapCollector(new LeafBucketCollectorBase(sub, segmentOrds) { + final int[] buffer = context.cardinalityPrefetchPipeline() ? new int[4096] : null; + @Override public void collect(int doc, long owningBucketOrd) throws IOException { assert owningBucketOrd == 0; @@ -552,6 +580,19 @@ public void collect(int doc, long owningBucketOrd) throws IOException { long docCount = docCountProvider.getDocCount(doc); segmentDocCounts.increment(ord + 1, docCount); } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + if (buffer == null) { + super.collect(stream, owningBucketOrd); + return; + } + int count = stream.intoArray(buffer); + singleValues.prefetchOrdValues(count, buffer); + for (int i = 0; i < count; i++) { + collect(buffer[i], owningBucketOrd); + } + } }); } segmentsWithMultiValuedOrds++; @@ -569,6 +610,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { segmentDocCounts.increment(segmentOrd + 1, docCount); } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }); } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java index 483b678430e41..0feb42b55e733 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/LongRareTermsAggregator.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.opensearch.common.lease.Releasables; import org.opensearch.common.util.LongHash; import org.opensearch.common.util.SetBackedScalingCuckooFilter; @@ -116,6 +117,11 @@ public void collect(int docId, long owningBucketOrd) throws IOException { } } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java index 9b93e6f191055..255cc162b4c0a 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java @@ -35,6 +35,7 @@ import org.apache.lucene.index.IndexReader; 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.NumericUtils; import org.apache.lucene.util.PriorityQueue; @@ -157,6 +158,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { } } } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }); } @@ -715,6 +721,11 @@ public void collect(int doc, long owningBucketOrd) throws IOException { subsetSizes = context.bigArrays().grow(subsetSizes, owningBucketOrd + 1); subsetSizes.increment(owningBucketOrd, 1); } + + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java index 5f99a9cc05558..e91bdbb511b09 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java @@ -31,8 +31,12 @@ package org.opensearch.search.aggregations.metrics; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.FixedBitSet; import org.apache.lucene.util.NumericUtils; @@ -71,6 +75,8 @@ */ class AvgAggregator extends NumericMetricsAggregator.SingleValue implements StarTreePreComputeCollector { + private static final Logger logger = LogManager.getLogger(AvgAggregator.class); + final ValuesSource.Numeric valuesSource; LongArray counts; @@ -128,34 +134,92 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); + final SortedNumericDocValues rawValues = context.cardinalityPrefetchPipeline() ? valuesSource.longValues(ctx) : null; final CompensatedSum kahanSummation = new CompensatedSum(0, 0); + // Prefetch NEXT segment's DV blocks (gives ~1s lead time while current segment collects) + if (rawValues != null) { + rawValues.prefetchRange(0, ctx.reader().maxDoc()); + prefetchNextSegment(ctx, valuesSource); + } return new LeafBucketCollectorBase(sub, values) { + private static final int PREFETCH_WINDOW = 262144; + private static final int BATCH_SIZE = 4096; + private int bufPos = 0; + + @Override public void collect(int doc, long bucket) throws IOException { - counts = bigArrays.grow(counts, bucket + 1); - sums = bigArrays.grow(sums, bucket + 1); - compensations = bigArrays.grow(compensations, bucket + 1); - + if (rawValues != null && ++bufPos >= BATCH_SIZE) { + bufPos = 0; + rawValues.prefetchRange(doc + PREFETCH_WINDOW, PREFETCH_WINDOW); + } if (values.advanceExact(doc)) { - final int valueCount = values.docValueCount(); + int valueCount = values.docValueCount(); + setKahanSummation(bucket); counts.increment(bucket, valueCount); - // Compute the sum of double values with Kahan summation algorithm which is more - // accurate than naive summation. - double sum = sums.get(bucket); - double compensation = compensations.get(bucket); - - kahanSummation.reset(sum, compensation); - for (int i = 0; i < valueCount; i++) { - double value = values.nextValue(); - kahanSummation.add(value); + kahanSummation.add(values.nextValue()); } - sums.set(bucket, kahanSummation.value()); compensations.set(bucket, kahanSummation.delta()); } } + + @Override + public void collect(DocIdStream stream, long bucket) throws IOException { + setKahanSummation(bucket); + final int[] count = { 0 }; + final boolean[] prefetched = { false }; + stream.forEach((doc) -> { + if (!prefetched[0] && rawValues != null) { + prefetched[0] = true; + rawValues.prefetchRange(doc, PREFETCH_WINDOW); + } + if (values.advanceExact(doc)) { + int valueCount = values.docValueCount(); + count[0] += valueCount; + for (int i = 0; i < valueCount; i++) { + kahanSummation.add(values.nextValue()); + } + } + }); + counts.increment(bucket, count[0]); + sums.set(bucket, kahanSummation.value()); + compensations.set(bucket, kahanSummation.delta()); + } + + @Override + public void collectRange(int min, int max) throws IOException { + setKahanSummation(0); + int count = 0; + if (rawValues != null) { + rawValues.prefetchRange(max, max - min); + } + for (int docId = min; docId < max; docId++) { + if (values.advanceExact(docId)) { + int valueCount = values.docValueCount(); + count += valueCount; + for (int i = 0; i < valueCount; i++) { + kahanSummation.add(values.nextValue()); + } + } + } + counts.increment(0, count); + sums.set(0, kahanSummation.value()); + compensations.set(0, kahanSummation.delta()); + } + + private void setKahanSummation(long bucket) { + counts = bigArrays.grow(counts, bucket + 1); + sums = bigArrays.grow(sums, bucket + 1); + compensations = bigArrays.grow(compensations, bucket + 1); + // Compute the sum of double values with Kahan summation algorithm which is more + // accurate than naive summation. + double sum = sums.get(bucket); + double compensation = compensations.get(bucket); + kahanSummation.reset(sum, compensation); + } }; } @@ -275,4 +339,22 @@ public void collectStarTreeEntry(int starTreeEntryBit, long bucket) throws IOExc } }; } + + /** + * Prefetch the NEXT segment's DV blocks so they're loaded while we collect the current segment. + */ + private void prefetchNextSegment(LeafReaderContext ctx, ValuesSource.Numeric valuesSource) { + try { + java.util.List leaves = context.searcher().getIndexReader().leaves(); + int nextOrd = ctx.ord + 1; + if (nextOrd < leaves.size()) { + LeafReaderContext nextCtx = leaves.get(nextOrd); + org.apache.lucene.index.SortedNumericDocValues nextValues = valuesSource.longValues(nextCtx); + nextValues.prefetchRange(0, nextCtx.reader().maxDoc()); + } + } catch (Exception e) { + // Best-effort prefetch + } + } + } 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 index 0cb319b853bce..20a1c5c673960 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/CardinalityAggregator.java @@ -34,7 +34,9 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.DocValues; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.index.SortedSetDocValues; import org.apache.lucene.index.Term; @@ -44,6 +46,7 @@ import org.apache.lucene.search.DisiPriorityQueue; import org.apache.lucene.search.DisiWrapper; import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.Scorer; import org.apache.lucene.search.TermQuery; @@ -142,7 +145,7 @@ private Collector pickCollector(LeafReaderContext ctx) throws IOException { ? MurmurHash3Values.hash(source.doubleValues(ctx)) : MurmurHash3Values.hash(source.longValues(ctx)); numericCollectorsUsed++; - return new DirectCollector(counts, hashValues); + return new DirectCollector(counts, hashValues, context.cardinalityPrefetchPipeline()); } Collector collector = null; @@ -154,7 +157,7 @@ private Collector pickCollector(LeafReaderContext ctx) throws IOException { return new EmptyCollector(); } else if (executionMode == CardinalityAggregatorFactory.ExecutionMode.ORDINALS) { // Force OrdinalsCollector ordinalsCollectorsUsed++; - collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays()); + collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays(), context.cardinalityPrefetchPipeline()); } else if (executionMode == null) { // no hint provided, fall back to heuristics final long ordinalsMemoryUsage = OrdinalsCollector.memoryOverhead(maxOrd); @@ -162,7 +165,7 @@ private Collector pickCollector(LeafReaderContext ctx) throws IOException { // 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()); + collector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays(), context.cardinalityPrefetchPipeline()); } else { ordinalsCollectorsOverheadTooHigh++; } @@ -171,7 +174,7 @@ private Collector pickCollector(LeafReaderContext ctx) throws IOException { if (collector == null) { // not able to build an OrdinalsCollector, or hint is direct stringHashingCollectorsUsed++; - collector = new DirectCollector(counts, MurmurHash3Values.hash(valuesSource.bytesValues(ctx))); + collector = new DirectCollector(counts, MurmurHash3Values.hash(valuesSource.bytesValues(ctx)), context.cardinalityPrefetchPipeline()); } if (canPrune(parent, subAggregators, valuesSourceConfig)) { @@ -513,12 +516,29 @@ public void close() { */ static class DirectCollector extends Collector { + private static final int BATCH_SIZE = 4096; + private static final int PIPELINE_DEPTH = 4; // prefetch this many batches ahead + private final MurmurHash3Values hashes; private final HyperLogLogPlusPlus counts; + private final boolean prefetchPipeline; + private final int[] docBuffer; + // Ring buffer: each slot holds a batch of doc IDs and its count + private final int[][] ring; + private final int[] ringCounts; - DirectCollector(HyperLogLogPlusPlus counts, MurmurHash3Values values) { + DirectCollector(HyperLogLogPlusPlus counts, MurmurHash3Values values, boolean prefetchPipeline) { this.counts = counts; this.hashes = values; + this.prefetchPipeline = prefetchPipeline; + this.docBuffer = new int[BATCH_SIZE]; + if (prefetchPipeline) { + this.ring = new int[PIPELINE_DEPTH][BATCH_SIZE]; + this.ringCounts = new int[PIPELINE_DEPTH]; + } else { + this.ring = null; + this.ringCounts = null; + } } @Override @@ -531,6 +551,74 @@ public void collect(int doc, long bucketOrd) throws IOException { } } + @Override + public void collect(DocIdStream stream, long bucketOrd) throws IOException { + if (prefetchPipeline) { + collectWithPrefetchPipeline(stream, bucketOrd); + } else { + collectSimple(stream, bucketOrd); + } + } + + private void collectSimple(DocIdStream stream, long bucketOrd) throws IOException { + for (int count = stream.intoArray(docBuffer); count > 0; count = stream.intoArray(docBuffer)) { + for (int i = 0; i < count; i++) { + if (hashes.advanceExact(docBuffer[i])) { + final int valueCount = hashes.count(); + for (int v = 0; v < valueCount; ++v) { + counts.collect(bucketOrd, hashes.nextValue()); + } + } + } + } + } + + private void collectWithPrefetchPipeline(DocIdStream stream, long bucketOrd) throws IOException { + // Fill the ring buffer: read up to PIPELINE_DEPTH batches and prefetch all of them + int filled = 0; + for (int slot = 0; slot < PIPELINE_DEPTH; slot++) { + int n = stream.intoArray(ring[slot]); + ringCounts[slot] = n; + if (n > 0) { + hashes.prefetchRange(ring[slot], n); + filled++; + } + if (n == 0 || !stream.mayHaveRemaining()) break; + } + if (filled == 0) return; + + // Process ring: collect slot 0, refill it with new batch, advance + int head = 0; + while (ringCounts[head] > 0) { + int count = ringCounts[head]; + int[] batch = ring[head]; + + // Collect this batch (prefetch was issued PIPELINE_DEPTH batches ago) + for (int i = 0; i < count; i++) { + if (hashes.advanceExact(batch[i])) { + final int valueCount = hashes.count(); + for (int v = 0; v < valueCount; ++v) { + counts.collect(bucketOrd, hashes.nextValue()); + } + } + } + + // Refill this slot with a new batch and prefetch it + if (stream.mayHaveRemaining()) { + int n = stream.intoArray(batch); + ringCounts[head] = n; + if (n > 0) { + hashes.prefetchRange(batch, n); + } + } else { + ringCounts[head] = 0; + } + + // Advance to next slot + head = (head + 1) % PIPELINE_DEPTH; + } + } + @Override public void postCollect() { // no-op @@ -551,6 +639,7 @@ public void close() { static class OrdinalsCollector extends Collector { private static final long SHALLOW_FIXEDBITSET_SIZE = RamUsageEstimator.shallowSizeOfInstance(FixedBitSet.class); + private static final int BATCH_SIZE = 4096; /** * Return an approximate memory overhead per bucket for this collector. @@ -561,11 +650,16 @@ public static long memoryOverhead(long maxOrd) { private final BigArrays bigArrays; private final SortedSetDocValues values; + private final SortedDocValues singleValues; private final int maxOrd; private final HyperLogLogPlusPlus counts; private ObjectArray visitedOrds; + private final int[] docBuffer; + private final int[] prefetchBuffer; + private final int[] ordBuffer; + private final boolean prefetchPipeline; - OrdinalsCollector(HyperLogLogPlusPlus counts, SortedSetDocValues values, BigArrays bigArrays) { + OrdinalsCollector(HyperLogLogPlusPlus counts, SortedSetDocValues values, BigArrays bigArrays, boolean prefetchPipeline) { if (values.getValueCount() > Integer.MAX_VALUE) { throw new IllegalArgumentException(); } @@ -573,6 +667,11 @@ public static long memoryOverhead(long maxOrd) { this.bigArrays = bigArrays; this.counts = counts; this.values = values; + this.singleValues = DocValues.unwrapSingleton(values); + this.docBuffer = new int[BATCH_SIZE]; + this.prefetchBuffer = prefetchPipeline ? new int[BATCH_SIZE] : null; + this.ordBuffer = (singleValues != null) ? new int[BATCH_SIZE] : null; + this.prefetchPipeline = prefetchPipeline; visitedOrds = bigArrays.newObjectArray(1); } @@ -593,6 +692,85 @@ public void collect(int doc, long bucketOrd) throws IOException { } } + @Override + public void collect(DocIdStream stream, long bucketOrd) throws IOException { + visitedOrds = bigArrays.grow(visitedOrds, bucketOrd + 1); + BitArray bits = visitedOrds.get(bucketOrd); + if (bits == null) { + bits = new BitArray(maxOrd, bigArrays); + visitedOrds.set(bucketOrd, bits); + } + if (singleValues != null) { + collectSingleValued(stream, bits); + } else { + collectMultiValued(stream, bits); + } + } + + private void collectSingleValued(DocIdStream stream, BitArray bits) throws IOException { + if (prefetchPipeline == false) { + for (int count = stream.intoArray(docBuffer); count > 0; count = stream.intoArray(docBuffer)) { + singleValues.ordValues(count, docBuffer, ordBuffer, -1); + for (int i = 0; i < count; i++) { + if (ordBuffer[i] != -1) { + bits.set(ordBuffer[i]); + } + } + if (stream.mayHaveRemaining() == false) break; + } + return; + } + + // Pipelined: prefetch batch N+1 while processing batch N + int[] current = docBuffer; + int[] next = prefetchBuffer; + int count = stream.intoArray(current); + if (count <= 0) return; + singleValues.prefetchOrdValues(count, current); + + while (true) { + // Fill next batch and start its prefetch while we process current + int nextCount = stream.mayHaveRemaining() ? stream.intoArray(next) : 0; + if (nextCount > 0) { + singleValues.prefetchOrdValues(nextCount, next); + } + + // Process current batch (IO was prefetched) + singleValues.ordValues(count, current, ordBuffer, -1); + for (int i = 0; i < count; i++) { + if (ordBuffer[i] != -1) { + bits.set(ordBuffer[i]); + } + } + + if (nextCount <= 0) break; + + // Swap buffers + int[] tmp = current; + current = next; + next = tmp; + count = nextCount; + } + } + + private void collectMultiValued(DocIdStream stream, BitArray bits) throws IOException { + for (int count = stream.intoArray(docBuffer); count > 0; count = stream.intoArray(docBuffer)) { + if (prefetchPipeline) { + values.prefetchOrds(docBuffer, count); + } + for (int i = 0; i < count; i++) { + if (values.advanceExact(docBuffer[i])) { + int docValueCount = values.docValueCount(); + long ord; + while ((docValueCount-- > 0) && (ord = values.nextOrd()) != SortedSetDocValues.NO_MORE_DOCS) { + bits.set((int) ord); + } + } + } + if (stream.mayHaveRemaining() == false) break; + } + } + @Override public void postCollect() throws IOException { try (BitArray allVisitedOrds = new BitArray(maxOrd, bigArrays)) { @@ -649,6 +827,9 @@ abstract static class MurmurHash3Values { public abstract long nextValue() throws IOException; + /** Prefetch doc values for the given doc IDs. Default is no-op. */ + public void prefetchRange(int[] docIds, int count) throws IOException {} + /** * Return a {@link MurmurHash3Values} instance that computes hashes on the fly for each double value. */ @@ -697,6 +878,11 @@ public int count() { public long nextValue() throws IOException { return BitMixer.mix64(values.nextValue()); } + + @Override + public void prefetchRange(int[] docIds, int count) throws IOException { + values.prefetchRange(docIds, count); + } } /** diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/CompensatedSum.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/CompensatedSum.java index 4d6d6d880da2e..bae482c9c0726 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/CompensatedSum.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/CompensatedSum.java @@ -112,6 +112,26 @@ public CompensatedSum add(double value, double delta) { return this; } + /** + * Increments the Kahan sum by adding two sums, and updating the correction term for reducing numeric errors. + */ + public void add(double[] values, int count) { + // If the value is Inf or NaN, just add it to the running tally to "convert" to + // Inf/NaN. This keeps the behavior bwc from before kahan summing + double sum = value; + double c = delta; // Compensation for lost low-order bits + + for (int i = 0; i < count; i++) { + double y = values[i] - c; + double t = sum + y; + c = (t - sum) - y; // Calculate the lost part + sum = t; + } + + this.value = sum; + this.delta = c; + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ExtendedStatsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ExtendedStatsAggregator.java index a187e65511864..7df15bf4f2fcf 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ExtendedStatsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ExtendedStatsAggregator.java @@ -32,6 +32,7 @@ package org.opensearch.search.aggregations.metrics; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.opensearch.common.lease.Releasables; import org.opensearch.common.util.BigArrays; @@ -163,6 +164,10 @@ public void collect(int doc, long bucket) throws IOException { } } + @Override + public void collect(DocIdStream stream, long owningBucketOrd) throws IOException { + super.collect(stream, owningBucketOrd); + } }; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java index fbba20d8a6d7d..823cf65d893df 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java @@ -34,7 +34,9 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; +import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.Bits; import org.apache.lucene.util.NumericUtils; @@ -155,15 +157,24 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx); final NumericDoubleValues values = MultiValueMode.MAX.select(allValues); + final SortedNumericDocValues rawValues = context.cardinalityPrefetchPipeline() ? valuesSource.longValues(ctx) : null; + // Prefetch NEXT segment's DV blocks (gives ~1s lead time while current segment collects) + if (rawValues != null) { + rawValues.prefetchRange(0, ctx.reader().maxDoc()); + prefetchNextSegment(ctx, valuesSource); + } return new LeafBucketCollectorBase(sub, allValues) { + private static final int PREFETCH_WINDOW = 262144; + private static final int BATCH_SIZE = 4096; + private int bufPos = 0; @Override public void collect(int doc, long bucket) throws IOException { - if (bucket >= maxes.size()) { - long from = maxes.size(); - maxes = bigArrays.grow(maxes, bucket + 1); - maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY); + if (rawValues != null && ++bufPos >= BATCH_SIZE) { + bufPos = 0; + rawValues.prefetchRange(doc + PREFETCH_WINDOW, PREFETCH_WINDOW); } + growMaxes(bucket); if (values.advanceExact(doc)) { final double value = values.doubleValue(); double max = maxes.get(bucket); @@ -172,6 +183,45 @@ public void collect(int doc, long bucket) throws IOException { } } + @Override + public void collect(DocIdStream stream, long bucket) throws IOException { + growMaxes(bucket); + final double[] maxArr = { maxes.get(bucket) }; + final boolean[] prefetched = { false }; + stream.forEach((doc) -> { + if (!prefetched[0] && rawValues != null) { + prefetched[0] = true; + rawValues.prefetchRange(doc, PREFETCH_WINDOW); + } + if (values.advanceExact(doc)) { + maxArr[0] = Math.max(maxArr[0], values.doubleValue()); + } + }); + maxes.set(bucket, maxArr[0]); + } + + @Override + public void collectRange(int min, int max) throws IOException { + growMaxes(0); + double maximum = maxes.get(0); + if (rawValues != null) { + rawValues.prefetchRange(max, max - min); + } + for (int doc = min; doc < max; doc++) { + if (values.advanceExact(doc)) { + maximum = Math.max(maximum, values.doubleValue()); + } + } + maxes.set(0, maximum); + } + + private void growMaxes(long bucket) { + if (bucket >= maxes.size()) { + long from = maxes.size(); + maxes = bigArrays.grow(maxes, bucket + 1); + maxes.fill(from, maxes.size(), Double.NEGATIVE_INFINITY); + } + } }; } @@ -287,4 +337,22 @@ public void doReset() { public StreamingCostMetrics getStreamingCostMetrics() { return new StreamingCostMetrics(true, 1, 1, 1, 1); } + + /** + * Prefetch the NEXT segment's DV blocks so they're loaded while we collect the current segment. + */ + private void prefetchNextSegment(LeafReaderContext ctx, ValuesSource.Numeric valuesSource) { + try { + java.util.List leaves = context.searcher().getIndexReader().leaves(); + int nextOrd = ctx.ord + 1; + if (nextOrd < leaves.size()) { + LeafReaderContext nextCtx = leaves.get(nextOrd); + org.apache.lucene.index.SortedNumericDocValues nextValues = valuesSource.longValues(nextCtx); + nextValues.prefetchRange(0, nextCtx.reader().maxDoc()); + } + } catch (Exception e) { + // Best-effort prefetch + } + } + } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java index 5c2ed2b240a09..4c0b7b903f48f 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java @@ -34,7 +34,9 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; +import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdStream; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.Bits; import org.apache.lucene.util.NumericUtils; @@ -155,15 +157,24 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx); final NumericDoubleValues values = MultiValueMode.MIN.select(allValues); + final SortedNumericDocValues rawValues = context.cardinalityPrefetchPipeline() ? valuesSource.longValues(ctx) : null; + // Prefetch NEXT segment's DV blocks (gives ~1s lead time while current segment collects) + if (rawValues != null) { + rawValues.prefetchRange(0, ctx.reader().maxDoc()); + prefetchNextSegment(ctx, valuesSource); + } return new LeafBucketCollectorBase(sub, allValues) { + private static final int PREFETCH_WINDOW = 262144; + private static final int BATCH_SIZE = 4096; + private int bufPos = 0; @Override public void collect(int doc, long bucket) throws IOException { - if (bucket >= mins.size()) { - long from = mins.size(); - mins = bigArrays.grow(mins, bucket + 1); - mins.fill(from, mins.size(), Double.POSITIVE_INFINITY); + if (rawValues != null && ++bufPos >= BATCH_SIZE) { + bufPos = 0; + rawValues.prefetchRange(doc + PREFETCH_WINDOW, PREFETCH_WINDOW); } + growMins(bucket); if (values.advanceExact(doc)) { final double value = values.doubleValue(); double min = mins.get(bucket); @@ -171,6 +182,46 @@ public void collect(int doc, long bucket) throws IOException { mins.set(bucket, min); } } + + @Override + public void collect(DocIdStream stream, long bucket) throws IOException { + growMins(bucket); + final double[] minArr = { mins.get(bucket) }; + final boolean[] prefetched = { false }; + stream.forEach((doc) -> { + if (!prefetched[0] && rawValues != null) { + prefetched[0] = true; + rawValues.prefetchRange(doc, PREFETCH_WINDOW); + } + if (values.advanceExact(doc)) { + minArr[0] = Math.min(minArr[0], values.doubleValue()); + } + }); + mins.set(bucket, minArr[0]); + } + + @Override + public void collectRange(int min, int max) throws IOException { + growMins(0); + double minimum = mins.get(0); + if (rawValues != null) { + rawValues.prefetchRange(max, max - min); + } + for (int doc = min; doc < max; doc++) { + if (values.advanceExact(doc)) { + minimum = Math.min(minimum, values.doubleValue()); + } + } + mins.set(0, minimum); + } + + private void growMins(long bucket) { + if (bucket >= mins.size()) { + long from = mins.size(); + mins = bigArrays.grow(mins, bucket + 1); + mins.fill(from, mins.size(), Double.POSITIVE_INFINITY); + } + } }; } @@ -283,4 +334,22 @@ public void doReset() { public StreamingCostMetrics getStreamingCostMetrics() { return new StreamingCostMetrics(true, 1, 1, 1, 1); } + + /** + * Prefetch the NEXT segment's DV blocks so they're loaded while we collect the current segment. + */ + private void prefetchNextSegment(LeafReaderContext ctx, ValuesSource.Numeric valuesSource) { + try { + java.util.List leaves = context.searcher().getIndexReader().leaves(); + int nextOrd = ctx.ord + 1; + if (nextOrd < leaves.size()) { + LeafReaderContext nextCtx = leaves.get(nextOrd); + org.apache.lucene.index.SortedNumericDocValues nextValues = valuesSource.longValues(nextCtx); + nextValues.prefetchRange(0, nextCtx.reader().maxDoc()); + } + } catch (Exception e) { + // Best-effort prefetch + } + } + } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/StatsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/StatsAggregator.java index c52638fe4b9c6..8ddf2f1a4864d 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/StatsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/StatsAggregator.java @@ -32,6 +32,8 @@ package org.opensearch.search.aggregations.metrics; 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.opensearch.common.lease.Releasables; import org.opensearch.common.util.BigArrays; @@ -101,33 +103,32 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); + final SortedNumericDocValues rawValues = context.cardinalityPrefetchPipeline() ? valuesSource.longValues(ctx) : null; + // Prefetch NEXT segment's DV blocks (gives ~1s lead time while current segment collects) + if (rawValues != null) { + rawValues.prefetchRange(0, ctx.reader().maxDoc()); + prefetchNextSegment(ctx, valuesSource); + } final CompensatedSum kahanSummation = new CompensatedSum(0, 0); return new LeafBucketCollectorBase(sub, values) { + private static final int PREFETCH_WINDOW = 262144; + private static final int BATCH_SIZE = 4096; + private int bufPos = 0; + @Override public void collect(int doc, long bucket) throws IOException { - if (bucket >= counts.size()) { - final long from = counts.size(); - final long overSize = BigArrays.overSize(bucket + 1); - counts = bigArrays.resize(counts, overSize); - sums = bigArrays.resize(sums, overSize); - compensations = bigArrays.resize(compensations, overSize); - mins = bigArrays.resize(mins, overSize); - maxes = bigArrays.resize(maxes, overSize); - mins.fill(from, overSize, Double.POSITIVE_INFINITY); - maxes.fill(from, overSize, Double.NEGATIVE_INFINITY); + if (rawValues != null && ++bufPos >= BATCH_SIZE) { + bufPos = 0; + rawValues.prefetchRange(doc + PREFETCH_WINDOW, PREFETCH_WINDOW); } + growStats(bucket); if (values.advanceExact(doc)) { final int valuesCount = values.docValueCount(); counts.increment(bucket, valuesCount); double min = mins.get(bucket); double max = maxes.get(bucket); - // Compute the sum of double values with Kahan summation algorithm which is more - // accurate than naive summation. - double sum = sums.get(bucket); - double compensation = compensations.get(bucket); - kahanSummation.reset(sum, compensation); for (int i = 0; i < valuesCount; i++) { double value = values.nextValue(); @@ -141,6 +142,81 @@ public void collect(int doc, long bucket) throws IOException { maxes.set(bucket, max); } } + + @Override + public void collect(DocIdStream stream, long bucket) throws IOException { + growStats(bucket); + final double[] minArr = { mins.get(bucket) }; + final double[] maxArr = { maxes.get(bucket) }; + final boolean[] prefetched = { false }; + stream.forEach((doc) -> { + if (!prefetched[0] && rawValues != null) { + prefetched[0] = true; + rawValues.prefetchRange(doc, PREFETCH_WINDOW); + } + if (values.advanceExact(doc)) { + final int valuesCount = values.docValueCount(); + counts.increment(bucket, valuesCount); + for (int i = 0; i < valuesCount; i++) { + double value = values.nextValue(); + kahanSummation.add(value); + minArr[0] = Math.min(minArr[0], value); + maxArr[0] = Math.max(maxArr[0], value); + } + } + }); + sums.set(bucket, kahanSummation.value()); + compensations.set(bucket, kahanSummation.delta()); + mins.set(bucket, minArr[0]); + maxes.set(bucket, maxArr[0]); + } + + @Override + public void collectRange(int min, int max) throws IOException { + growStats(0); + + double minimum = mins.get(0); + double maximum = maxes.get(0); + if (rawValues != null) { + rawValues.prefetchRange(max, max - min); + } + for (int doc = min; doc < max; doc++) { + if (values.advanceExact(doc)) { + final int valuesCount = values.docValueCount(); + counts.increment(0, valuesCount); + for (int i = 0; i < valuesCount; i++) { + double value = values.nextValue(); + kahanSummation.add(value); + minimum = Math.min(minimum, value); + maximum = Math.max(maximum, value); + } + } + } + sums.set(0, kahanSummation.value()); + compensations.set(0, kahanSummation.delta()); + mins.set(0, minimum); + maxes.set(0, maximum); + } + + private void growStats(long bucket) { + if (bucket >= counts.size()) { + final long from = counts.size(); + final long overSize = BigArrays.overSize(bucket + 1); + counts = bigArrays.resize(counts, overSize); + sums = bigArrays.resize(sums, overSize); + compensations = bigArrays.resize(compensations, overSize); + mins = bigArrays.resize(mins, overSize); + maxes = bigArrays.resize(maxes, overSize); + mins.fill(from, overSize, Double.POSITIVE_INFINITY); + maxes.fill(from, overSize, Double.NEGATIVE_INFINITY); + } + + // Compute the sum of double values with Kahan summation algorithm which is more + // accurate than naive summation. + double sum = sums.get(bucket); + double compensation = compensations.get(bucket); + kahanSummation.reset(sum, compensation); + } }; } @@ -205,4 +281,22 @@ public InternalAggregation buildEmptyAggregation() { public void doClose() { Releasables.close(counts, maxes, mins, sums, compensations); } + + /** + * Prefetch the NEXT segment's DV blocks so they're loaded while we collect the current segment. + */ + private void prefetchNextSegment(LeafReaderContext ctx, ValuesSource.Numeric valuesSource) { + try { + java.util.List leaves = context.searcher().getIndexReader().leaves(); + int nextOrd = ctx.ord + 1; + if (nextOrd < leaves.size()) { + LeafReaderContext nextCtx = leaves.get(nextOrd); + org.apache.lucene.index.SortedNumericDocValues nextValues = valuesSource.longValues(nextCtx); + nextValues.prefetchRange(0, nextCtx.reader().maxDoc()); + } + } catch (Exception e) { + // Best-effort prefetch + } + } + } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/StreamCardinalityAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/StreamCardinalityAggregator.java index 20b4ac823367c..550e2733a44b2 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/StreamCardinalityAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/StreamCardinalityAggregator.java @@ -76,7 +76,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc streamCollector = new EmptyCollector(); } else { ordinalsCollectorsUsed++; - streamCollector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays()); + streamCollector = new OrdinalsCollector(counts, ordinalValues, context.bigArrays(), context.cardinalityPrefetchPipeline()); } return streamCollector; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java index ba32592f75ea1..1738c961f0acc 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java @@ -32,6 +32,8 @@ package org.opensearch.search.aggregations.metrics; 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.NumericUtils; import org.opensearch.common.lease.Releasables; @@ -118,30 +120,80 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } final BigArrays bigArrays = context.bigArrays(); final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx); + final SortedNumericDocValues rawValues = context.cardinalityPrefetchPipeline() ? valuesSource.longValues(ctx) : null; + // Prefetch NEXT segment's DV blocks (gives ~1s lead time while current segment collects) + if (rawValues != null) { + rawValues.prefetchRange(0, ctx.reader().maxDoc()); + prefetchNextSegment(ctx, valuesSource); + } final CompensatedSum kahanSummation = new CompensatedSum(0, 0); return new LeafBucketCollectorBase(sub, values) { + private static final int PREFETCH_WINDOW = 262144; + private static final int BATCH_SIZE = 4096; + private int bufPos = 0; + @Override public void collect(int doc, long bucket) throws IOException { - sums = bigArrays.grow(sums, bucket + 1); - compensations = bigArrays.grow(compensations, bucket + 1); - + if (rawValues != null && ++bufPos >= BATCH_SIZE) { + bufPos = 0; + rawValues.prefetchRange(doc + PREFETCH_WINDOW, PREFETCH_WINDOW); + } if (values.advanceExact(doc)) { - final int valuesCount = values.docValueCount(); - // Compute the sum of double values with Kahan summation algorithm which is more - // accurate than naive summation. - double sum = sums.get(bucket); - double compensation = compensations.get(bucket); - kahanSummation.reset(sum, compensation); - - for (int i = 0; i < valuesCount; i++) { + setKahanSummation(bucket); + for (int i = 0; i < values.docValueCount(); i++) { double value = values.nextValue(); kahanSummation.add(value); } - compensations.set(bucket, kahanSummation.delta()); sums.set(bucket, kahanSummation.value()); } } + + @Override + public void collect(DocIdStream stream, long bucket) throws IOException { + setKahanSummation(bucket); + final boolean[] prefetched = { false }; + stream.forEach((doc) -> { + if (!prefetched[0] && rawValues != null) { + prefetched[0] = true; + rawValues.prefetchRange(doc, PREFETCH_WINDOW); + } + if (values.advanceExact(doc)) { + for (int i = 0; i < values.docValueCount(); i++) { + kahanSummation.add(values.nextValue()); + } + } + }); + compensations.set(bucket, kahanSummation.delta()); + sums.set(bucket, kahanSummation.value()); + } + + @Override + public void collectRange(int min, int max) throws IOException { + setKahanSummation(0); + if (rawValues != null) { + rawValues.prefetchRange(max, max - min); + } + for (int docId = min; docId < max; docId++) { + if (values.advanceExact(docId)) { + for (int i = 0; i < values.docValueCount(); i++) { + kahanSummation.add(values.nextValue()); + } + } + } + sums.set(0, kahanSummation.value()); + compensations.set(0, kahanSummation.delta()); + } + + private void setKahanSummation(long bucket) { + sums = bigArrays.grow(sums, bucket + 1); + compensations = bigArrays.grow(compensations, bucket + 1); + // Compute the sum of double values with Kahan summation algorithm which is more + // accurate than naive summation. + double sum = sums.get(bucket); + double compensation = compensations.get(bucket); + kahanSummation.reset(sum, compensation); + } }; } @@ -215,4 +267,22 @@ public InternalAggregation buildEmptyAggregation() { public void doClose() { Releasables.close(sums, compensations); } + + /** + * Prefetch the NEXT segment's DV blocks so they're loaded while we collect the current segment. + */ + private void prefetchNextSegment(LeafReaderContext ctx, ValuesSource.Numeric valuesSource) { + try { + java.util.List leaves = context.searcher().getIndexReader().leaves(); + int nextOrd = ctx.ord + 1; + if (nextOrd < leaves.size()) { + LeafReaderContext nextCtx = leaves.get(nextOrd); + org.apache.lucene.index.SortedNumericDocValues nextValues = valuesSource.longValues(nextCtx); + nextValues.prefetchRange(0, nextCtx.reader().maxDoc()); + } + } catch (Exception e) { + // Best-effort prefetch + } + } + } diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java index ac38b364fd36b..4b8b3924d1867 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -563,6 +563,10 @@ public boolean keywordIndexOrDocValuesEnabled() { return false; } + public boolean cardinalityPrefetchPipeline() { + return false; + } + @ExperimentalApi public void setStreamChannelListener(StreamSearchChannelListener listener) { throw new IllegalStateException("Set search channel listener should be implemented for stream search"); diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteSubAggTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteSubAggTests.java index 4221a2837387d..c59fadee03633 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteSubAggTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/filterrewrite/FilterRewriteSubAggTests.java @@ -9,10 +9,14 @@ package org.opensearch.search.aggregations.bucket.filterrewrite; import org.apache.lucene.document.Field; +import org.apache.lucene.document.KeywordField; import org.apache.lucene.document.LongField; +import org.apache.lucene.document.LongPoint; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.BooleanClause; +import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; @@ -22,6 +26,7 @@ import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.indices.breaker.NoneCircuitBreakerService; import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.NumberFieldMapper; import org.opensearch.index.mapper.ParseContext; import org.opensearch.search.aggregations.AggregationBuilder; @@ -38,7 +43,12 @@ import org.opensearch.search.aggregations.bucket.histogram.InternalDateHistogram; import org.opensearch.search.aggregations.bucket.range.InternalRange; import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder; +import org.opensearch.search.aggregations.metrics.InternalAvg; +import org.opensearch.search.aggregations.metrics.InternalCardinality; +import org.opensearch.search.aggregations.metrics.InternalMax; +import org.opensearch.search.aggregations.metrics.InternalMin; import org.opensearch.search.aggregations.metrics.InternalStats; +import org.opensearch.search.aggregations.metrics.InternalSum; import org.opensearch.search.aggregations.pipeline.PipelineAggregator; import org.opensearch.search.internal.SearchContext; @@ -56,26 +66,34 @@ public class FilterRewriteSubAggTests extends AggregatorTestCase { private final String longFieldName = "metric"; private final String dateFieldName = "timestamp"; + private final String nameFieldName = "name"; private final Query matchAllQuery = new MatchAllDocsQuery(); private final NumberFieldMapper.NumberFieldType longFieldType = new NumberFieldMapper.NumberFieldType( longFieldName, NumberFieldMapper.NumberType.LONG ); private final DateFieldMapper.DateFieldType dateFieldType = aggregableDateFieldType(false, true); + private final KeywordFieldMapper.KeywordFieldType nameFieldType = new KeywordFieldMapper.KeywordFieldType(nameFieldName); private final NumberFieldMapper.NumberType numberType = longFieldType.numberType(); private final String rangeAggName = "range"; private final String autoDateAggName = "auto"; private final String dateAggName = "date"; private final String statsAggName = "stats"; + private final String avgAggName = "avg"; + private final String sumAggName = "sum"; + private final String minAggName = "min"; + private final String maxAggName = "max"; + private final String cardinalityAggName = "cardinality"; private final List DEFAULT_DATA = List.of( - new TestDoc(0, Instant.parse("2020-03-01T00:00:00Z")), - new TestDoc(1, Instant.parse("2020-03-01T00:00:00Z")), - new TestDoc(1, Instant.parse("2020-03-01T00:00:01Z")), - new TestDoc(2, Instant.parse("2020-03-01T01:00:00Z")), - new TestDoc(3, Instant.parse("2020-03-01T02:00:00Z")), - new TestDoc(4, Instant.parse("2020-03-01T03:00:00Z")), - new TestDoc(5, Instant.parse("2020-03-01T04:00:00Z")), - new TestDoc(6, Instant.parse("2020-03-01T04:00:00Z")) + new TestDoc(0, Instant.parse("2020-03-01T00:00:00Z"), "abc"), + new TestDoc(1, Instant.parse("2020-03-01T00:00:00Z"), "def"), + new TestDoc(1, Instant.parse("2020-03-01T00:00:01Z"), "ghi"), + new TestDoc(2, Instant.parse("2020-03-01T01:00:00Z"), "jkl"), + new TestDoc(3, Instant.parse("2020-03-01T02:00:00Z"), "jkl"), + new TestDoc(4, Instant.parse("2020-03-01T03:00:00Z"), "mno"), + new TestDoc(4, Instant.parse("2020-03-01T04:00:00Z"), "prq", true), + new TestDoc(5, Instant.parse("2020-03-01T04:00:00Z"), "stu"), + new TestDoc(6, Instant.parse("2020-03-01T04:00:00Z"), "stu") ); public void testRange() throws IOException { @@ -107,12 +125,157 @@ public void testRange() throws IOException { assertEquals(3, thirdAuto.getBuckets().size()); } + public void testRangeWithAvgAndSum() throws IOException { + // Test for sum metric aggregation + RangeAggregationBuilder rangeAggregationBuilder = new RangeAggregationBuilder(rangeAggName).field(longFieldName) + .addRange(1, 2) + .addRange(2, 4) + .addRange(4, 6) + .subAggregation(AggregationBuilders.sum(sumAggName).field(longFieldName)); + + InternalRange result = executeAggregation(DEFAULT_DATA, rangeAggregationBuilder, true); + + // Verify results + List buckets = result.getBuckets(); + assertEquals(3, buckets.size()); + + InternalRange.Bucket firstBucket = buckets.get(0); + assertEquals(2, firstBucket.getDocCount()); + InternalSum firstSum = firstBucket.getAggregations().get(sumAggName); + assertEquals(2, firstSum.getValue(), 0); + + InternalRange.Bucket secondBucket = buckets.get(1); + assertEquals(2, secondBucket.getDocCount()); + InternalSum secondSum = secondBucket.getAggregations().get(sumAggName); + assertEquals(5, secondSum.getValue(), 0); + + InternalRange.Bucket thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.getDocCount()); + InternalSum thirdSum = thirdBucket.getAggregations().get(sumAggName); + assertEquals(9, thirdSum.getValue(), 0); + + // Test for average metric aggregation now + rangeAggregationBuilder = new RangeAggregationBuilder(rangeAggName).field(longFieldName) + .addRange(1, 2) + .addRange(2, 4) + .addRange(4, 6) + .subAggregation(AggregationBuilders.avg(avgAggName).field(longFieldName)); + + result = executeAggregation(DEFAULT_DATA, rangeAggregationBuilder, true); + + // Verify results + buckets = result.getBuckets(); + assertEquals(3, buckets.size()); + + firstBucket = buckets.get(0); + assertEquals(2, firstBucket.getDocCount()); + InternalAvg firstAvg = firstBucket.getAggregations().get(avgAggName); + assertEquals(1, firstAvg.getValue(), 0); + + secondBucket = buckets.get(1); + assertEquals(2, secondBucket.getDocCount()); + InternalAvg secondAvg = secondBucket.getAggregations().get(avgAggName); + assertEquals(2.5, secondAvg.getValue(), 0); + + thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.getDocCount()); + InternalAvg thirdAvg = thirdBucket.getAggregations().get(avgAggName); + assertEquals(4.5, thirdAvg.getValue(), 0); + } + + public void testRangeWithMinAndMax() throws IOException { + // Test for min metric aggregation + RangeAggregationBuilder rangeAggregationBuilder = new RangeAggregationBuilder(rangeAggName).field(longFieldName) + .addRange(1, 2) + .addRange(2, 4) + .addRange(4, 6) + .subAggregation(AggregationBuilders.min(minAggName).field(longFieldName)); + + InternalRange result = executeAggregation(DEFAULT_DATA, rangeAggregationBuilder, true); + + // Verify results + List buckets = result.getBuckets(); + assertEquals(3, buckets.size()); + + InternalRange.Bucket firstBucket = buckets.get(0); + assertEquals(2, firstBucket.getDocCount()); + InternalMin firstMin = firstBucket.getAggregations().get(minAggName); + assertEquals(1, firstMin.getValue(), 0); + + InternalRange.Bucket secondBucket = buckets.get(1); + assertEquals(2, secondBucket.getDocCount()); + InternalMin secondMin = secondBucket.getAggregations().get(minAggName); + assertEquals(2, secondMin.getValue(), 0); + + InternalRange.Bucket thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.getDocCount()); + InternalMin thirdMin = thirdBucket.getAggregations().get(minAggName); + assertEquals(4, thirdMin.getValue(), 0); + + // Test for max metric aggregation now + rangeAggregationBuilder = new RangeAggregationBuilder(rangeAggName).field(longFieldName) + .addRange(1, 2) + .addRange(2, 4) + .addRange(4, 6) + .subAggregation(AggregationBuilders.max(maxAggName).field(longFieldName)); + + result = executeAggregation(DEFAULT_DATA, rangeAggregationBuilder, true); + + // Verify results + buckets = result.getBuckets(); + assertEquals(3, buckets.size()); + + firstBucket = buckets.get(0); + assertEquals(2, firstBucket.getDocCount()); + InternalMax firstMax = firstBucket.getAggregations().get(maxAggName); + assertEquals(1, firstMax.getValue(), 0); + + secondBucket = buckets.get(1); + assertEquals(2, secondBucket.getDocCount()); + InternalMax secondMax = secondBucket.getAggregations().get(maxAggName); + assertEquals(3, secondMax.getValue(), 0); + + thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.getDocCount()); + InternalMax thirdMax = thirdBucket.getAggregations().get(maxAggName); + assertEquals(5, thirdMax.getValue(), 0); + } + + public void testRangeWithCard() throws IOException { + RangeAggregationBuilder rangeAggregationBuilder = new RangeAggregationBuilder(rangeAggName).field(longFieldName) + .addRange(1, 2) + .addRange(2, 4) + .addRange(4, 6) + .subAggregation(AggregationBuilders.cardinality(cardinalityAggName).field(nameFieldName).executionHint("ordinals")); + + InternalRange result = executeAggregation(DEFAULT_DATA, rangeAggregationBuilder, true); + + // Verify results + List buckets = result.getBuckets(); + assertEquals(3, buckets.size()); + + InternalRange.Bucket firstBucket = buckets.get(0); + assertEquals(2, firstBucket.getDocCount()); + InternalCardinality firstCardinality = firstBucket.getAggregations().get(cardinalityAggName); + assertEquals(2, firstCardinality.getValue(), 0); + + InternalRange.Bucket secondBucket = buckets.get(1); + assertEquals(2, secondBucket.getDocCount()); + InternalCardinality secondCardinality = secondBucket.getAggregations().get(cardinalityAggName); + assertEquals(1, secondCardinality.getValue(), 0); + + InternalRange.Bucket thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.getDocCount()); + InternalCardinality thirdCardinality = thirdBucket.getAggregations().get(cardinalityAggName); + assertEquals(2, thirdCardinality.getValue(), 0); + } + public void testDateHisto() throws IOException { DateHistogramAggregationBuilder dateHistogramAggregationBuilder = new DateHistogramAggregationBuilder(dateAggName).field( dateFieldName ).calendarInterval(DateHistogramInterval.HOUR).subAggregation(AggregationBuilders.stats(statsAggName).field(longFieldName)); - InternalDateHistogram result = executeAggregation(DEFAULT_DATA, dateHistogramAggregationBuilder, true); + InternalDateHistogram result = executeAggregation(DEFAULT_DATA, dateHistogramAggregationBuilder, false); // Verify results List buckets = result.getBuckets(); @@ -125,6 +288,7 @@ public void testDateHisto() throws IOException { assertEquals(3, firstStats.getCount()); assertEquals(1, firstStats.getMax(), 0); assertEquals(0, firstStats.getMin(), 0); + assertEquals(2, firstStats.getSum(), 0); InternalDateHistogram.Bucket secondBucket = buckets.get(1); assertEquals("2020-03-01T01:00:00.000Z", secondBucket.getKeyAsString()); @@ -133,6 +297,7 @@ public void testDateHisto() throws IOException { assertEquals(1, secondStats.getCount()); assertEquals(2, secondStats.getMax(), 0); assertEquals(2, secondStats.getMin(), 0); + assertEquals(2, secondStats.getSum(), 0); InternalDateHistogram.Bucket thirdBucket = buckets.get(2); assertEquals("2020-03-01T02:00:00.000Z", thirdBucket.getKeyAsString()); @@ -141,6 +306,7 @@ public void testDateHisto() throws IOException { assertEquals(1, thirdStats.getCount()); assertEquals(3, thirdStats.getMax(), 0); assertEquals(3, thirdStats.getMin(), 0); + assertEquals(3, thirdStats.getSum(), 0); InternalDateHistogram.Bucket fourthBucket = buckets.get(3); assertEquals("2020-03-01T03:00:00.000Z", fourthBucket.getKeyAsString()); @@ -149,6 +315,7 @@ public void testDateHisto() throws IOException { assertEquals(1, fourthStats.getCount()); assertEquals(4, fourthStats.getMax(), 0); assertEquals(4, fourthStats.getMin(), 0); + assertEquals(4, fourthStats.getSum(), 0); InternalDateHistogram.Bucket fifthBucket = buckets.get(4); assertEquals("2020-03-01T04:00:00.000Z", fifthBucket.getKeyAsString()); @@ -157,6 +324,7 @@ public void testDateHisto() throws IOException { assertEquals(2, fifthStats.getCount()); assertEquals(6, fifthStats.getMax(), 0); assertEquals(5, fifthStats.getMin(), 0); + assertEquals(11, fifthStats.getSum(), 0); } public void testAutoDateHisto() throws IOException { @@ -337,11 +505,35 @@ private Directory setupIndex(List docs, boolean random) throws IOExcept for (TestDoc doc : docs) { indexWriter.addDocument(doc.toDocument()); } + + indexWriter.commit(); + } + + try (IndexWriter indexWriter = new IndexWriter(directory, new IndexWriterConfig().setCodec(TestUtil.getDefaultCodec()))) { + for (TestDoc doc : docs) { + if (doc.deleted) { + BooleanQuery.Builder booleanQueryBuilder = new BooleanQuery.Builder(); + booleanQueryBuilder.add(LongPoint.newRangeQuery(longFieldName, doc.metric, doc.metric), BooleanClause.Occur.MUST); + booleanQueryBuilder.add( + LongField.newRangeQuery( + dateFieldName, + dateFieldType.parse(doc.timestamp.toString()), + dateFieldType.parse(doc.timestamp.toString()) + ), + BooleanClause.Occur.MUST + ); + indexWriter.deleteDocuments(booleanQueryBuilder.build()); + } + } + + indexWriter.commit(); } } else { try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) { for (TestDoc doc : docs) { - indexWriter.addDocument(doc.toDocument()); + if (!doc.deleted) { + indexWriter.addDocument(doc.toDocument()); + } } } } @@ -361,7 +553,8 @@ private IA executeAggregationOnReader( matchAllQuery, bucketConsumer, longFieldType, - dateFieldType + dateFieldType, + nameFieldType ); Aggregator aggregator = createAggregator(aggregationBuilder, searchContext); CountingAggregator countingAggregator = new CountingAggregator(new AtomicInteger(), aggregator); @@ -413,10 +606,22 @@ private InternalAggregation.ReduceContext createReduceContext( private class TestDoc { private final long metric; private final Instant timestamp; + private final String name; + private final boolean deleted; public TestDoc(long metric, Instant timestamp) { + this(metric, timestamp, "abc", false); + } + + public TestDoc(long metric, Instant timestamp, String name) { + this(metric, timestamp, name, false); + } + + public TestDoc(long metric, Instant timestamp, String name, boolean deleted) { this.metric = metric; this.timestamp = timestamp; + this.name = name; + this.deleted = deleted; } public ParseContext.Document toDocument() { @@ -426,6 +631,7 @@ public ParseContext.Document toDocument() { for (Field fld : fieldList) doc.add(fld); doc.add(new LongField(dateFieldName, dateFieldType.parse(timestamp.toString()), Field.Store.NO)); + doc.add(new KeywordField(nameFieldName, name, Field.Store.NO)); return doc; }