From 07abd51dc64ac8e813b207a8d400b11855164b90 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Wed, 18 Sep 2024 17:32:55 -0700 Subject: [PATCH] temp Signed-off-by: Sandesh Kumar --- .../opensearch/common/util/FeatureFlags.java | 2 +- .../startree/utils/StarTreeQueryHelper.java | 82 ++++++++++-- .../SortedNumericStarTreeValuesIterator.java | 4 + .../iterator/StarTreeValuesIterator.java | 2 +- .../search/DefaultSearchContext.java | 33 +++++ .../org/opensearch/search/SearchService.java | 9 +- .../aggregations/metrics/AvgAggregator.java | 109 ++++++++-------- .../aggregations/metrics/MaxAggregator.java | 51 +++----- .../aggregations/metrics/MinAggregator.java | 51 +++----- .../aggregations/metrics/SumAggregator.java | 53 ++------ .../metrics/ValueCountAggregator.java | 36 ++---- .../search/internal/SearchContext.java | 18 ++- .../startree/OriginalOrStarTreeQuery.java | 63 --------- .../search/startree/StarTreeFilter.java | 42 +++--- .../search/startree/StarTreeQuery.java | 122 ------------------ .../search/startree/StarTreeQueryContext.java | 48 +++++++ 16 files changed, 309 insertions(+), 416 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 create mode 100644 server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.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 351f38afff7fd..a1ec6cf7712c0 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,7 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.search.CollectionTerminatedException; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.codec.composite.CompositeIndexReader; @@ -17,22 +18,28 @@ 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.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.Consumer; import java.util.stream.Collectors; /** @@ -43,6 +50,8 @@ */ public class StarTreeQueryHelper { + private static Map starTreeValuesMap = new HashMap<>(); + /** * Checks if the search context can be supported by star-tree */ @@ -64,7 +73,12 @@ 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 { + + /** + * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder. + * Returns null if the query cannot be supported. + */ + public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException { // Current implementation assumes only single star-tree is supported CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService() .getCompositeFieldTypes() @@ -75,8 +89,12 @@ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext c compositeMappedFieldType.getCompositeIndexType() ); - StarTreeQuery starTreeQuery = StarTreeQueryHelper.toStarTreeQuery(starTree, compositeMappedFieldType, source.query()); - if (starTreeQuery == null) { + StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.toStarTreeQueryContext( + starTree, + compositeMappedFieldType, + source.query() + ); + if (starTreeQueryContext == null) { return null; } @@ -86,10 +104,10 @@ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext c } } - return new OriginalOrStarTreeQuery(starTreeQuery, context.query()); + return starTreeQueryContext; } - private static StarTreeQuery toStarTreeQuery( + private static StarTreeQueryContext toStarTreeQueryContext( CompositeIndexFieldInfo starTree, CompositeDataCubeFieldType compositeIndexFieldInfo, QueryBuilder queryBuilder @@ -110,7 +128,7 @@ private static StarTreeQuery toStarTreeQuery( return null; } - return new StarTreeQuery(starTree, queryMap); + return new StarTreeQueryContext(starTree, queryMap); } /** @@ -151,13 +169,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; @@ -165,4 +181,44 @@ 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) != StarTreeValuesIterator.NO_MORE_ENTRIES) { + int count = valuesIterator.valuesCount(); + 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..44f9545ce4f7f 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 valuesCount() 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 74a7482d975df..204e52c32342c 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; @@ -98,6 +101,7 @@ import org.opensearch.search.rescore.RescoreContext; import org.opensearch.search.slice.SliceBuilder; import org.opensearch.search.sort.SortAndFormats; +import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestionSearchContext; import java.io.IOException; @@ -115,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; @@ -176,6 +181,7 @@ final class DefaultSearchContext extends SearchContext { private SliceBuilder sliceBuilder; private SearchShardTask task; private final Version minNodeVersion; + private StarTreeQueryContext starTreeQueryContext; /** * The original query as sent by the user without the types and aliases @@ -270,6 +276,7 @@ final class DefaultSearchContext extends SearchContext { this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold(); this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories; this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled(); + this.starTreeValuesMap = new HashMap<>(); } @Override @@ -1147,4 +1154,30 @@ 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)) { + logger.info("Used cached values"); + return starTreeValuesMap.get(ctx); + + } else { + logger.info("not using cache"); + } + + 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 ac08ca0cc072c..d38d1280c592e 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -84,7 +84,6 @@ import org.opensearch.index.query.InnerHitContextBuilder; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.MatchNoneQueryBuilder; -import org.opensearch.index.query.ParsedQuery; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryRewriteContext; import org.opensearch.index.query.QueryShardContext; @@ -139,7 +138,7 @@ 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; import org.opensearch.tasks.TaskResourceTrackingService; @@ -1547,9 +1546,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); - if (parsedQuery != null) { - context.parsedQuery(new ParsedQuery(parsedQuery)); + StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source); + if (starTreeQueryContext != null) { + context.starTreeQueryContext(starTreeQueryContext); logger.debug("can use star tree"); } else { logger.debug("cannot 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 d028f38be1b61..017bb91871b0e 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 @@ -32,17 +32,12 @@ package org.opensearch.search.aggregations.metrics; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.util.NumericUtils; import org.opensearch.common.lease.Releasables; import org.opensearch.common.util.BigArrays; import org.opensearch.common.util.DoubleArray; 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.fielddata.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; @@ -56,7 +51,6 @@ 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.getSupportedStarTree; /** @@ -104,7 +98,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - return getStarTreeLeafCollector(ctx, sub, supportedStarTree); + // return getStarTreeLeafCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); } @@ -144,56 +138,57 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) - throws IOException { - final BigArrays bigArrays = context.bigArrays(); - final CompensatedSum kahanSummation = new CompensatedSum(0, 0); - - StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); - String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); - String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( - starTree.getField(), - fieldName, - MetricStat.SUM.getTypeName() - ); - assert starTreeValues != null; - SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName); - - String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( - starTree.getField(), - fieldName, - MetricStat.VALUE_COUNT.getTypeName() - ); - SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName); - - return new LeafBucketCollectorBase(sub, values) { - @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 (values.advanceExact(doc) && countValues.advanceExact(doc)) { - final long valueCount = values.docValueCount(); - counts.increment(bucket, countValues.nextValue()); - // 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 = NumericUtils.sortableLongToDouble(values.nextValue()); - kahanSummation.add(value); - } - - sums.set(bucket, kahanSummation.value()); - compensations.set(bucket, kahanSummation.delta()); - } - } - }; - } + // private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo + // starTree) + // throws IOException { + // final BigArrays bigArrays = context.bigArrays(); + // final CompensatedSum kahanSummation = new CompensatedSum(0, 0); + // + // StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); + // String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); + // String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + // starTree.getField(), + // fieldName, + // MetricStat.SUM.getTypeName() + // ); + // assert starTreeValues != null; + // SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(sumMetricName); + // + // String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + // starTree.getField(), + // fieldName, + // MetricStat.VALUE_COUNT.getTypeName() + // ); + // SortedNumericDocValues countValues = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(countMetricName); + // + // return new LeafBucketCollectorBase(sub, values) { + // @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 (values.advanceExact(doc) && countValues.advanceExact(doc)) { + // final long valueCount = values.docValueCount(); + // counts.increment(bucket, countValues.nextValue()); + // // 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 = NumericUtils.sortableLongToDouble(values.nextValue()); + // kahanSummation.add(value); + // } + // + // sums.set(bucket, kahanSummation.value()); + // compensations.set(bucket, kahanSummation.delta()); + // } + // } + // }; + // } @Override public double metric(long 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 5f0c5e242a453..c7e697ad45d8f 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,6 @@ 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.ScoreMode; import org.apache.lucene.util.Bits; @@ -44,8 +43,7 @@ import org.opensearch.common.util.DoubleArray; 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.NumericDoubleValues; import org.opensearch.index.fielddata.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; @@ -61,9 +59,9 @@ import java.io.IOException; import java.util.Arrays; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; 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.getSupportedStarTree; /** @@ -132,7 +130,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - return getStarTreeLeafCollector(ctx, sub, supportedStarTree); + System.out.println("max star tree"); + return getStarTreeCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); } @@ -162,37 +161,21 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + public LeafBucketCollector getStarTreeCollector(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() + AtomicReference max = new AtomicReference<>(Double.NEGATIVE_INFINITY); + return StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.SUM.getTypeName(), + value -> { + max.set(Math.max(max.get(), (NumericUtils.sortableLongToDouble(value)))); + }, + () -> maxes.set(0, max.get()) ); - 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); - } - } - }; } @Override 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 81246a192555c..188c615a26ff2 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,6 @@ 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.ScoreMode; import org.apache.lucene.util.Bits; @@ -44,8 +43,7 @@ import org.opensearch.common.util.DoubleArray; 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.NumericDoubleValues; import org.opensearch.index.fielddata.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; @@ -60,9 +58,9 @@ import java.io.IOException; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; 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.getSupportedStarTree; /** @@ -131,7 +129,8 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - return getStarTreeLeafCollector(ctx, sub, supportedStarTree); + System.out.println("min star tree"); + return getStarTreeCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); } @@ -159,37 +158,21 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + public LeafBucketCollector getStarTreeCollector(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.MIN.getTypeName() + AtomicReference min = new AtomicReference<>(Double.POSITIVE_INFINITY); + return StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.SUM.getTypeName(), + value -> { + min.set(Math.min(min.get(), (NumericUtils.sortableLongToDouble(value)))); + }, + () -> mins.set(0, min.get()) ); - 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 >= mins.size()) { - long from = mins.size(); - mins = bigArrays.grow(mins, bucket + 1); - mins.fill(from, mins.size(), Double.POSITIVE_INFINITY); - } - if (values.advanceExact(doc)) { - final double value = NumericUtils.sortableLongToDouble(values.nextValue()); - double min = mins.get(bucket); - min = Math.min(min, value); - mins.set(bucket, min); - } - } - }; } @Override 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 c205c9abb7bee..617ee6a939e6c 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,7 +32,6 @@ package org.opensearch.search.aggregations.metrics; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.util.NumericUtils; import org.opensearch.common.lease.Releasables; @@ -40,8 +39,7 @@ import org.opensearch.common.util.DoubleArray; 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.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; @@ -55,7 +53,6 @@ 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.getSupportedStarTree; /** @@ -101,7 +98,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); } return getDefaultLeafCollector(ctx, sub); } @@ -136,43 +133,19 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + public LeafBucketCollector getStarTreeCollector(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()); - } - } - }; + return StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.SUM.getTypeName(), + value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), + () -> sums.set(0, kahanSummation.value()) + ); } @Override 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 0d9a145b2b8ad..a156ec49983fa 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,6 @@ 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.getSupportedStarTree; /** @@ -98,7 +96,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,28 +138,18 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + public LeafBucketCollector getStarTreeCollector(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() + return StarTreeQueryHelper.getStarTreeLeafCollector( + context, + (ValuesSource.Numeric) valuesSource, + ctx, + sub, + starTree, + MetricStat.VALUE_COUNT.getTypeName(), + value -> counts.increment(0, value), + () -> {} ); - 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()); - } - } - }; } @Override 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 5357206e8c117..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; @@ -76,8 +79,10 @@ import org.opensearch.search.query.ReduceableSearchResult; import org.opensearch.search.rescore.RescoreContext; import org.opensearch.search.sort.SortAndFormats; +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; @@ -124,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() {} @@ -531,4 +536,15 @@ public boolean keywordIndexOrDocValuesEnabled() { return false; } + public StarTreeQueryContext getStarTreeQueryContext() { + return null; + } + + public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { + 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 7b48ecd4561c7..493bdf3ca143e 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java @@ -10,13 +10,13 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.apache.lucene.index.SortedNumericDocValues; -import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.util.DocIdSetBuilder; import org.opensearch.index.compositeindex.datacube.Dimension; import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues; import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode; 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 java.io.IOException; import java.util.ArrayDeque; @@ -38,7 +38,7 @@ * @opensearch.experimental * @opensearch.internal */ -class StarTreeFilter { +public class StarTreeFilter { private static final Logger logger = LogManager.getLogger(StarTreeFilter.class); private final Map queryMap; @@ -53,49 +53,49 @@ 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 DocIdSetIterator getStarTreeResult() throws IOException { + public StarTreeValuesIterator getStarTreeResult() throws IOException { StarTreeResult starTreeResult = traverseStarTree(); - List andIterators = new ArrayList<>(); - andIterators.add(starTreeResult._matchedDocIds.build().iterator()); - DocIdSetIterator docIdSetIterator = andIterators.get(0); + List andIterators = new ArrayList<>(); + andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator())); + 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); DocIdSetBuilder builder = new DocIdSetBuilder(starTreeResult.maxMatchedDoc + 1); - SortedNumericDocValues ndv = (SortedNumericDocValues) this.starTreeValues.getDimensionDocIdSetIterator( + 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.nextDoc() != NO_MORE_DOCS) { - int docID = docIdSetIterator.docID(); - if (ndv.advanceExact(docID)) { - final int valuesCount = ndv.docValueCount(); + while (starTreeValuesIterator.nextEntry() != NO_MORE_DOCS) { + int entryId = starTreeValuesIterator.entryId(); + if (ndv.advance(entryId) > 0) { + final int valuesCount = ndv.valuesCount(); 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 = builder.build().iterator(); + starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator()); } - return docIdSetIterator; + 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 e8354b0415ac1..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 = filter.getStarTreeResult(); - 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; - } -} diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java new file mode 100644 index 0000000000000..2b12c9bba604f --- /dev/null +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -0,0 +1,48 @@ +/* + * 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.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; + +import java.util.Map; + +/** + * Query class for querying star tree data structure. + * + * @opensearch.experimental + */ +@ExperimentalApi +public class StarTreeQueryContext { + + /** + * 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 StarTreeQueryContext(CompositeIndexFieldInfo starTree, Map queryMap) { + this.starTree = starTree; + this.queryMap = queryMap; + } + + public CompositeIndexFieldInfo getStarTree() { + return starTree; + } + + public Map getQueryMap() { + return queryMap; + } +}