From 4f6cc675920cc8c2b4844a382840c8612ee15070 Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Wed, 18 Sep 2024 12:54:27 +0530 Subject: [PATCH] star tree file formats refactoring and fixing offset bug Signed-off-by: Bharathwaj G --- .../opensearch/common/util/FeatureFlags.java | 2 +- .../startree/utils/StarTreeQueryHelper.java | 114 ++++++++-------- .../SortedNumericStarTreeValuesIterator.java | 4 + .../iterator/StarTreeValuesIterator.java | 2 +- .../search/DefaultSearchContext.java | 22 ++++ .../org/opensearch/search/SearchService.java | 5 +- .../aggregations/metrics/AvgAggregator.java | 2 +- .../aggregations/metrics/MaxAggregator.java | 82 +++++++----- .../aggregations/metrics/MinAggregator.java | 2 +- .../metrics/NumericMetricsAggregator.java | 3 + .../aggregations/metrics/SumAggregator.java | 105 ++------------- .../metrics/ValueCountAggregator.java | 36 ++---- .../search/internal/SearchContext.java | 9 +- .../startree/OriginalOrStarTreeQuery.java | 63 --------- .../search/startree/StarTreeFilter.java | 73 ++--------- .../search/startree/StarTreeQuery.java | 122 ------------------ 16 files changed, 187 insertions(+), 459 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java delete mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index a5acea004c3b2..bbee8329a4821 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -107,7 +107,7 @@ public class FeatureFlags { * aggregations. */ public static final String STAR_TREE_INDEX = "opensearch.experimental.feature.composite_index.star_tree.enabled"; - public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope); + public static final Setting STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope); /** * Gates the functionality of application based configuration templates. diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java index 80bca749ace9c..02a618b3379b2 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java @@ -10,6 +10,9 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.util.NumericUtils; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.codec.composite.CompositeIndexReader; @@ -17,23 +20,30 @@ import org.opensearch.index.compositeindex.datacube.Metric; import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; +import org.opensearch.index.fielddata.SortedNumericDoubleValues; import org.opensearch.index.mapper.CompositeDataCubeFieldType; import org.opensearch.index.mapper.StarTreeMapper; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.search.aggregations.AggregatorFactory; +import org.opensearch.search.aggregations.LeafBucketCollector; +import org.opensearch.search.aggregations.LeafBucketCollectorBase; import org.opensearch.search.aggregations.metrics.MetricAggregatorFactory; +import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.internal.SearchContext; -import org.opensearch.search.startree.OriginalOrStarTreeQuery; -import org.opensearch.search.startree.StarTreeQuery; +import org.opensearch.search.startree.StarTreeFilter; import org.opensearch.search.startree.StarTreeQueryContext; import java.io.IOException; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import java.util.stream.Collectors; /** @@ -44,6 +54,9 @@ */ public class StarTreeQueryHelper { + + private static Map starTreeValuesMap = new HashMap<>(); + /** * Checks if the search context can be supported by star-tree */ @@ -65,30 +78,7 @@ public static boolean isStarTreeSupported(SearchContext context, boolean trackTo * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder. * Returns null if the query cannot be supported. */ - public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext context, SearchSourceBuilder source) throws IOException { - // Current implementation assumes only single star-tree is supported - CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService() - .getCompositeFieldTypes() - .iterator() - .next(); - CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo( - compositeMappedFieldType.name(), - compositeMappedFieldType.getCompositeIndexType() - ); - - StarTreeQuery starTreeQuery = StarTreeQueryHelper.toStarTreeQuery(starTree, compositeMappedFieldType, source.query()); - if (starTreeQuery == null) { - return null; - } - for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) { - if (validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) { - return null; - } - } - - return new OriginalOrStarTreeQuery(starTreeQuery, context.query()); - } /** * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder. @@ -123,30 +113,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context return starTreeQueryContext; } - private static StarTreeQuery toStarTreeQuery( - CompositeIndexFieldInfo starTree, - CompositeDataCubeFieldType compositeIndexFieldInfo, - QueryBuilder queryBuilder - ) { - Map queryMap; - if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) { - queryMap = null; - } else if (queryBuilder instanceof TermQueryBuilder) { - List supportedDimensions = compositeIndexFieldInfo.getDimensions() - .stream() - .map(Dimension::getField) - .collect(Collectors.toList()); - queryMap = getStarTreePredicates(queryBuilder, supportedDimensions); - if (queryMap == null) { - return null; - } - } else { - return null; - } - - return new StarTreeQuery(starTree, queryMap); - } - private static StarTreeQueryContext toStarTreeQueryContext( CompositeIndexFieldInfo starTree, CompositeDataCubeFieldType compositeIndexFieldInfo, @@ -209,13 +175,11 @@ private static boolean validateStarTreeMetricSuport( } public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context) { - if (context.query() instanceof StarTreeQuery) { - return ((StarTreeQuery) context.query()).getStarTree(); - } - return null; + StarTreeQueryContext starTreeQueryContext = context.getStarTreeQueryContext(); + return (starTreeQueryContext != null) ? starTreeQueryContext.getStarTree() : null; } - public static StarTreeValues getStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException { + public static StarTreeValues computeStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException { SegmentReader reader = Lucene.segmentReader(context.reader()); if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) { return null; @@ -223,4 +187,46 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader(); return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree); } + + public static LeafBucketCollector getStarTreeLeafCollector( + SearchContext context, + ValuesSource.Numeric valuesSource, + LeafReaderContext ctx, + LeafBucketCollector sub, + CompositeIndexFieldInfo starTree, + String metric, + Consumer valueConsumer, + Runnable finalConsumer + ) throws IOException { + StarTreeValues starTreeValues = context.getStarTreeValues(ctx, starTree); + String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); + String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + starTree.getField(), + fieldName, + metric + ); + + assert starTreeValues != null; + SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(metricName); + StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); + StarTreeValuesIterator result = filter.getStarTreeResult(); + + int entryId; + while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) { + if (valuesIterator.advance(entryId) > 0) { + int count = valuesIterator.docValueCount(); + for (int i = 0; i < count; i++) { + long value = valuesIterator.nextValue(); + valueConsumer.accept(value); // Apply the operation (max, sum, etc.) + } + } + } + finalConsumer.run(); + return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) { + @Override + public void collect(int doc, long bucket) { + throw new CollectionTerminatedException(); + } + }; + } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java index 27afdf1479b4e..ea92e0541ba3f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java @@ -29,4 +29,8 @@ public SortedNumericStarTreeValuesIterator(DocIdSetIterator docIdSetIterator) { public long nextValue() throws IOException { return ((SortedNumericDocValues) docIdSetIterator).nextValue(); } + + public int docValueCount() throws IOException { + return ((SortedNumericDocValues) docIdSetIterator).docValueCount(); + } } diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java index 32866f3e50092..454e5b393973f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java +++ b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java @@ -21,7 +21,7 @@ * @opensearch.experimental */ @ExperimentalApi -public abstract class StarTreeValuesIterator { +public class StarTreeValuesIterator { public static final int NO_MORE_ENTRIES = Integer.MAX_VALUE; protected final DocIdSetIterator docIdSetIterator; diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index ab6fcbfa6a58f..6facbfeaa94f4 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -34,6 +34,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.BooleanClause.Occur; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.BoostQuery; @@ -56,6 +57,8 @@ import org.opensearch.index.IndexService; import org.opensearch.index.IndexSettings; import org.opensearch.index.cache.bitset.BitsetFilterCache; +import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; +import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.engine.Engine; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; @@ -116,6 +119,7 @@ import java.util.function.Function; import java.util.function.LongSupplier; +import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.computeStarTreeValues; import static org.opensearch.search.SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD; import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_MODE; import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING; @@ -272,6 +276,7 @@ final class DefaultSearchContext extends SearchContext { this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold(); this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories; this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled(); + this.starTreeValuesMap = new HashMap<>(); } @Override @@ -1151,8 +1156,25 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() { return false; } + @Override public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { this.starTreeQueryContext = starTreeQueryContext; return this; } + + @Override + public StarTreeQueryContext getStarTreeQueryContext() { + return this.starTreeQueryContext; + } + + + @Override + public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException { + if (this.starTreeValuesMap.containsKey(ctx)) { + return starTreeValuesMap.get(ctx); + } + StarTreeValues starTreeValues = computeStarTreeValues(ctx, starTree); + starTreeValuesMap.put(ctx, starTreeValues); + return starTreeValues; + } } diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index 17235f94ed483..0eea3da46885f 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -139,7 +139,6 @@ import org.opensearch.search.sort.SortAndFormats; import org.opensearch.search.sort.SortBuilder; import org.opensearch.search.sort.SortOrder; -import org.opensearch.search.startree.OriginalOrStarTreeQuery; import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.Suggest; import org.opensearch.search.suggest.completion.CompletionSuggestion; @@ -1548,9 +1547,9 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled() && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) { try { - OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source); +// OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source); StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source); - if (parsedQuery != null) { + if (starTreeQueryContext != null) { // context.parsedQuery(new ParsedQuery(parsedQuery)); context.starTreeQueryContext(starTreeQueryContext); logger.debug("can use star tree"); 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 93c72b88f221d..e452d7ba92e95 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 @@ -56,7 +56,7 @@ import java.io.IOException; import java.util.Map; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; +//import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** 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 3152eb49f315b..719ea99e40bcd 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 @@ -31,11 +31,14 @@ package org.opensearch.search.aggregations.metrics; +import org.apache.lucene.index.DocValues; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.PointValues; +import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.CollectionTerminatedException; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.Bits; import org.apache.lucene.util.NumericUtils; @@ -46,6 +49,8 @@ import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.fielddata.NumericDoubleValues; import org.opensearch.index.fielddata.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; @@ -57,13 +62,14 @@ import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.aggregations.support.ValuesSourceConfig; import org.opensearch.search.internal.SearchContext; +import org.opensearch.search.startree.StarTreeFilter; import java.io.IOException; import java.util.Arrays; import java.util.Map; import java.util.function.Function; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; +//import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** @@ -162,38 +168,48 @@ public void collect(int doc, long bucket) throws IOException { }; } -// private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) -// throws IOException { -// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); -// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); -// String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( -// starTree.getField(), -// fieldName, -// MetricStat.MAX.getTypeName() -// ); -// assert starTreeValues != null; -// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName); -// -// final BigArrays bigArrays = context.bigArrays(); -// final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx); -// return new LeafBucketCollectorBase(sub, allValues) { -// -// @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 (values.advanceExact(doc)) { -// final double value = NumericUtils.sortableLongToDouble(values.nextValue()); -// double max = maxes.get(bucket); -// max = Math.max(max, value); -// maxes.set(bucket, max); -// } -// } -// }; -// } + private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + throws IOException { + StarTreeValues starTreeValues = context.getStarTreeValues(ctx, starTree); + String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); + String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + starTree.getField(), + fieldName, + MetricStat.MAX.getTypeName() + ); + assert starTreeValues != null; + SortedNumericStarTreeValuesIterator values = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(metricName); + + final SortedNumericDoubleValues allValues = valuesSource.doubleValues(ctx); + + StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); + StarTreeValuesIterator result = filter.getStarTreeResult(); + + double max = Double.NEGATIVE_INFINITY; + int docID; + while ((docID = result.nextEntry()) != DocIdSetIterator.NO_MORE_DOCS) { + // Move to the document in the SortedNumericDocValues + if (values.advance(docID) > 0) { + int count = values.docValueCount(); // Get the number of values for the document + + for (int i = 0; i < count; i++) { + long rawValue = values.nextValue(); + double value = NumericUtils.sortableLongToDouble(rawValue); + + // Update the max value if the current value is greater + max = Math.max(max, value); + } + } + } + maxes.set(0, max); + + return new LeafBucketCollectorBase(sub, allValues) { + @Override + public void collect(int doc, long bucket) throws IOException { + throw new CollectionTerminatedException(); + } + }; + } @Override public double metric(long owningBucketOrd) { 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 fd5a1d8bdd650..918cdea6d6cb0 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 @@ -62,7 +62,7 @@ import java.util.Map; import java.util.function.Function; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; +//import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java index f90e5a092385f..3f08635afc427 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java @@ -31,7 +31,9 @@ package org.opensearch.search.aggregations.metrics; +import org.apache.lucene.index.LeafReaderContext; import org.opensearch.common.util.Comparators; +import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.search.aggregations.Aggregator; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.sort.SortOrder; @@ -64,6 +66,7 @@ protected SingleValue(String name, SearchContext context, Aggregator parent, Map public abstract double metric(long owningBucketOrd); + @Override public BucketComparator bucketComparator(String key, SortOrder order) { if (key != null && false == "value".equals(key)) { 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 22f3ad6a7e65a..eca282ae255f4 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 @@ -31,8 +31,9 @@ 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.CollectionTerminatedException; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.ScoreMode; @@ -43,6 +44,7 @@ import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.MetricStat; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; +import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper; import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; @@ -56,12 +58,12 @@ import org.opensearch.search.aggregations.support.ValuesSourceConfig; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.startree.StarTreeFilter; -import org.opensearch.search.startree.StarTreeQueryContext; import java.io.IOException; import java.util.Map; +import java.util.function.BiConsumer; +import java.util.function.Consumer; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** @@ -106,13 +108,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); -// if (supportedStarTree != null) { -// return getStarTreeLeafCollector(ctx, sub, supportedStarTree); -// } - StarTreeQueryContext starTreeQueryContext = this.context().getStarTreeQueryContext(); - if (starTreeQueryContext != null) { - return getAltStarTreeCollector(ctx, sub, supportedStarTree); - + if (supportedStarTree != null) { + return getStarTreeCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); } @@ -147,90 +144,14 @@ public void collect(int doc, long bucket) throws IOException { }; } -// private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) -// throws IOException { -// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); -// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); -// String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( -// starTree.getField(), -// fieldName, -// MetricStat.SUM.getTypeName() -// ); -// assert starTreeValues != null; -// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName); -// -// final BigArrays bigArrays = context.bigArrays(); -// final CompensatedSum kahanSummation = new CompensatedSum(0, 0); -// -// return new LeafBucketCollectorBase(sub, values) { -// @Override -// public void collect(int doc, long bucket) throws IOException { -// sums = bigArrays.grow(sums, bucket + 1); -// compensations = bigArrays.grow(compensations, bucket + 1); -// -// if (values.advanceExact(doc)) { -// final int valuesCount = values.docValueCount(); -// double sum = sums.get(bucket); -// double compensation = compensations.get(bucket); -// kahanSummation.reset(sum, compensation); -// -// for (int i = 0; i < valuesCount; i++) { -// double value = NumericUtils.sortableLongToDouble(values.nextValue()); -// kahanSummation.add(value); -// } -// -// compensations.set(bucket, kahanSummation.delta()); -// sums.set(bucket, kahanSummation.value()); -// } -// } -// }; -// } - - - public LeafBucketCollector getAltStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) - throws IOException { - StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); - String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); - String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( - starTree.getField(), - fieldName, - MetricStat.SUM.getTypeName() - ); - assert starTreeValues != null; - SortedNumericStarTreeValuesIterator values = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(metricName); - - final BigArrays bigArrays = context.bigArrays(); + public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException { final CompensatedSum kahanSummation = new CompensatedSum(0, 0); - - StarTreeFilter filter = new StarTreeFilter(starTreeValues, context().getStarTreeQueryContext().getQueryMap()); - StarTreeValuesIterator result = filter.getStarTreeResult(); - - int docID; - while ((docID = result.nextEntry()) != DocIdSetIterator.NO_MORE_DOCS) { - // Move to the document in the SortedNumericDocValues - if (values.advance(docID) > 0) { - // Retrieve the number of values for this document - int count = values.docValueCount(); - - // Traverse the values for the current document - for (int i = 0; i < count; i++) { - double value = NumericUtils.sortableLongToDouble(values.nextValue()); - // Process the value, e.g., adding it to the sum - kahanSummation.add(value); - } - } - } - - // Now you can use kahanSummation to get the final sum or perform further processing -// double finalSum = kahanSummation.value(); - - return new LeafBucketCollectorBase(sub, values) { - @Override - public void collect(int doc, long bucket) { - throw new CollectionTerminatedException(); - } - }; + return StarTreeQueryHelper.getStarTreeLeafCollector(context, valuesSource, ctx, sub, starTree, MetricStat.SUM.getTypeName(), + value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), + () -> sums.set(0, kahanSummation.value()) + ); } + @Override public double metric(long owningBucketOrd) { if (valuesSource == null || owningBucketOrd >= sums.size()) { diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java index 217085c3e9f57..e6a1200b57d8c 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/metrics/ValueCountAggregator.java @@ -39,8 +39,7 @@ import org.opensearch.common.util.LongArray; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.MetricStat; -import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; -import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeUtils; +import org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper; import org.opensearch.index.fielddata.MultiGeoPointValues; import org.opensearch.index.fielddata.SortedBinaryDocValues; import org.opensearch.search.aggregations.Aggregator; @@ -54,7 +53,7 @@ import java.io.IOException; import java.util.Map; -import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; +//import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; /** @@ -98,7 +97,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { -// return getStarTreeLeafCollector(ctx, sub, supportedStarTree); + return getStarTreeCollector(ctx, sub, supportedStarTree); } final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx); @@ -140,29 +139,12 @@ public void collect(int doc, long bucket) throws IOException { }; } -// private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) -// throws IOException { -// StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); -// String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); -// String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( -// starTree.getField(), -// fieldName, -// MetricStat.VALUE_COUNT.getTypeName() -// ); -// assert starTreeValues != null; -// SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName); -// final BigArrays bigArrays = context.bigArrays(); -// -// return new LeafBucketCollectorBase(sub, values) { -// @Override -// public void collect(int doc, long bucket) throws IOException { -// counts = bigArrays.grow(counts, bucket + 1); -// if (values.advanceExact(doc)) { -// counts.increment(bucket, values.nextValue()); -// } -// } -// }; -// } + public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) throws IOException { + return StarTreeQueryHelper.getStarTreeLeafCollector(context, (ValuesSource.Numeric) valuesSource, ctx, sub, starTree, MetricStat.VALUE_COUNT.getTypeName(), + value -> counts.increment(0, value), + () -> {} + ); + } @Override public double metric(long owningBucketOrd) { 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 6cb5382ba6cee..ed3cbde61ce74 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -31,6 +31,7 @@ package org.opensearch.search.internal; +import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.Collector; import org.apache.lucene.search.CollectorManager; import org.apache.lucene.search.FieldDoc; @@ -44,6 +45,8 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.BigArrays; import org.opensearch.index.cache.bitset.BitsetFilterCache; +import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; +import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.ObjectMapper; @@ -79,6 +82,7 @@ import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestionSearchContext; +import java.io.IOException; import java.util.Collection; import java.util.HashMap; import java.util.List; @@ -125,7 +129,7 @@ public List toInternalAggregations(Collection co private final List releasables = new CopyOnWriteArrayList<>(); private final AtomicBoolean closed = new AtomicBoolean(false); private InnerHitsContext innerHitsContext; - + protected volatile Map starTreeValuesMap; private volatile boolean searchTimedOut; protected SearchContext() {} @@ -540,4 +544,7 @@ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryCont return this; } + public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException { + return null; + } } diff --git a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java deleted file mode 100644 index 4806265888a2c..0000000000000 --- a/server/src/main/java/org/opensearch/search/startree/OriginalOrStarTreeQuery.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.opensearch.search.startree; - -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.QueryVisitor; - -import java.io.IOException; -import java.util.Objects; - -/** - * Preserves star-tree queries which can be used along with original query - * Decides which star-tree query to use (or not) based on cost factors - * - * @opensearch.experimental - */ -public class OriginalOrStarTreeQuery extends Query { - - private final StarTreeQuery starTreeQuery; - private final Query originalQuery; - - public OriginalOrStarTreeQuery(StarTreeQuery starTreeQuery, Query originalQuery) { - this.starTreeQuery = starTreeQuery; - this.originalQuery = originalQuery; - } - - @Override - public String toString(String s) { - return originalQuery.toString(s); - } - - @Override - public void visit(QueryVisitor queryVisitor) {} - - @Override - public boolean equals(Object o) { - return sameClassAs(o) && equalsTo(getClass().cast(o)); - } - - private boolean equalsTo(OriginalOrStarTreeQuery other) { - return starTreeQuery.equals(other.starTreeQuery) && originalQuery.equals(other.originalQuery); - } - - @Override - public int hashCode() { - return Objects.hash(classHash(), starTreeQuery, originalQuery, starTreeQuery); - } - - @Override - public Query rewrite(IndexSearcher indexSearcher) throws IOException { - if (indexSearcher.getIndexReader().hasDeletions()) { - return originalQuery; - } - return starTreeQuery; - } -} diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java index ff3cd4a345c5b..6c0f5d3278076 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java @@ -55,18 +55,18 @@ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map pr /** *
    *
  • First go over the star tree and try to match as many dimensions as possible - *
  • For the remaining columns, use doc values indexes to match them + *
  • For the remaining columns, use star-tree doc values to match them *
*/ public StarTreeValuesIterator getStarTreeResult() throws IOException { StarTreeResult starTreeResult = traverseStarTree(); List andIterators = new ArrayList<>(); andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator())); - StarTreeValuesIterator docIdSetIterator = andIterators.get(0); + StarTreeValuesIterator starTreeValuesIterator = andIterators.get(0); // No matches, return if (starTreeResult.maxMatchedDoc == -1) { - return docIdSetIterator; + return starTreeValuesIterator; } for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) { logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc); @@ -74,77 +74,30 @@ public StarTreeValuesIterator getStarTreeResult() throws IOException { SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) this.starTreeValues.getDimensionValuesIterator( remainingPredicateColumn ); - List docIds = new ArrayList<>(); + List entryIds = new ArrayList<>(); long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly - while (docIdSetIterator.nextEntry() != NO_MORE_DOCS) { - int docID = docIdSetIterator.entryId(); - if (ndv.advance(docID) > 0) { + while (starTreeValuesIterator.nextEntry() != NO_MORE_DOCS) { + int entryId = starTreeValuesIterator.entryId(); + if (ndv.advance(entryId) > 0) { final int valuesCount = ndv.docValueCount(); for (int i = 0; i < valuesCount; i++) { long value = ndv.nextValue(); // Directly compare value with queryValue if (value == queryValue) { - docIds.add(docID); + entryIds.add(entryId); break; } } } } - DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size()); - for (int docID : docIds) { - adder.add(docID); + DocIdSetBuilder.BulkAdder adder = builder.grow(entryIds.size()); + for (int entryId : entryIds) { + adder.add(entryId); } - docIdSetIterator = new StarTreeValuesIterator(builder.build().iterator()); - } - return docIdSetIterator; - } - - /** - *
    - *
  • First go over the star tree and try to match as many dimensions as possible - *
  • For the remaining columns, use doc values indexes to match them - *
- */ - public void getStarTreeResultNew() throws IOException { - StarTreeResult starTreeResult = traverseStarTree(); - List andIterators = new ArrayList<>(); - andIterators.add(starTreeResult._matchedDocIds.build().iterator()); - DocIdSetIterator docIdSetIterator = andIterators.get(0); - - // No matches, return - if (starTreeResult.maxMatchedDoc == -1) { - return; - } - for (String remainingPredicateColumn : starTreeResult._remainingPredicateColumns) { - logger.debug("remainingPredicateColumn : {}, maxMatchedDoc : {} ", remainingPredicateColumn, starTreeResult.maxMatchedDoc); - DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1); - SortedNumericStarTreeValuesIterator ndv = (SortedNumericStarTreeValuesIterator) this.starTreeValues.getDimensionValuesIterator( - remainingPredicateColumn - ); - List docIds = new ArrayList<>(); - long queryValue = queryMap.get(remainingPredicateColumn); // Get the query value directly - - while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) { - int docID = docIdSetIterator.docID(); - if (ndv.advance(docID) > 0) { - final int valuesCount = ndv.docValueCount(); - for (int i = 0; i < valuesCount; i++) { - long value = ndv.nextValue(); - // Directly compare value with queryValue - if (value == queryValue) { - docIds.add(docID); - break; - } - } - } - } - DocIdSetBuilder.BulkAdder adder = builder.grow(docIds.size()); - for (int docID : docIds) { - adder.add(docID); - } - docIdSetIterator = builder.build().iterator(); + starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator()); } + return starTreeValuesIterator; } /** diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java deleted file mode 100644 index a51a12fd99a63..0000000000000 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java +++ /dev/null @@ -1,122 +0,0 @@ -/* - * 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.opensearch.search.startree; - -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.SegmentReader; -import org.apache.lucene.search.ConstantScoreScorer; -import org.apache.lucene.search.ConstantScoreWeight; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.QueryVisitor; -import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.search.Scorer; -import org.apache.lucene.search.Weight; -import org.opensearch.common.lucene.Lucene; -import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; -import org.opensearch.index.codec.composite.CompositeIndexReader; -import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; - -import java.io.IOException; -import java.util.List; -import java.util.Map; -import java.util.Objects; - -/** - * Query class for querying star tree data structure. - * - * @opensearch.experimental - */ -public class StarTreeQuery extends Query { - - /** - * Star tree field info - * This is used to get the star tree data structure - */ - private final CompositeIndexFieldInfo starTree; - - /** - * Map of field name to a value to be queried for that field - * This is used to filter the data based on the query - */ - private final Map queryMap; - - public StarTreeQuery(CompositeIndexFieldInfo starTree, Map queryMap) { - this.starTree = starTree; - this.queryMap = queryMap; - } - - @Override - public void visit(QueryVisitor visitor) {} - - @Override - public boolean equals(Object obj) { - return sameClassAs(obj) && equalsTo(getClass().cast(obj)); - } - - private boolean equalsTo(StarTreeQuery other) { - return starTree.equals(other.starTree) && queryMap != null && queryMap.equals(other.queryMap); - } - - @Override - public int hashCode() { - return Objects.hash(classHash(), starTree, queryMap); - } - - @Override - public String toString(String field) { - StringBuilder sb = new StringBuilder(); - sb.append(getClass().getSimpleName()); - sb.append("("); - sb.append(this.starTree); - if (queryMap != null) { - sb.append(", "); - sb.append(queryMap); - sb.append(")"); - } - return sb.toString(); - } - - @Override - public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public Scorer scorer(LeafReaderContext context) throws IOException { - StarTreeValues starTreeValues = getStarTreeValues(context); - if (starTreeValues == null) { - return null; - } - StarTreeFilter filter = new StarTreeFilter(starTreeValues, queryMap); - DocIdSetIterator result = null; - return new ConstantScoreScorer(this, score(), scoreMode, result); - } - - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; - } - - private StarTreeValues getStarTreeValues(LeafReaderContext ctx) throws IOException { - SegmentReader reader = Lucene.segmentReader(ctx.reader()); - CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader(); - List compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields(); - if (compositeIndexFields != null && !compositeIndexFields.isEmpty()) { - return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree); - } else { - return null; - } - } - }; - } - - public CompositeIndexFieldInfo getStarTree() { - return starTree; - } -}