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..80bca749ace9c 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 @@ -28,6 +28,7 @@ import org.opensearch.search.internal.SearchContext; import org.opensearch.search.startree.OriginalOrStarTreeQuery; import org.opensearch.search.startree.StarTreeQuery; +import org.opensearch.search.startree.StarTreeQueryContext; import java.io.IOException; import java.util.HashMap; @@ -89,6 +90,39 @@ public static OriginalOrStarTreeQuery getOriginalOrStarTreeQuery(SearchContext c return new OriginalOrStarTreeQuery(starTreeQuery, context.query()); } + /** + * 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() + .iterator() + .next(); + CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo( + compositeMappedFieldType.name(), + compositeMappedFieldType.getCompositeIndexType() + ); + + StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.toStarTreeQueryContext( + starTree, + compositeMappedFieldType, + source.query() + ); + if (starTreeQueryContext == null) { + return null; + } + + for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) { + if (validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) { + return null; + } + } + + return starTreeQueryContext; + } + private static StarTreeQuery toStarTreeQuery( CompositeIndexFieldInfo starTree, CompositeDataCubeFieldType compositeIndexFieldInfo, @@ -113,6 +147,30 @@ private static StarTreeQuery toStarTreeQuery( return new StarTreeQuery(starTree, queryMap); } + private static StarTreeQueryContext toStarTreeQueryContext( + 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 StarTreeQueryContext(starTree, queryMap); + } + /** * Parse query body to star-tree predicates * @param queryBuilder diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index 74a7482d975df..ab6fcbfa6a58f 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -98,6 +98,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; @@ -176,6 +177,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 @@ -744,6 +746,7 @@ public SearchContext parsedQuery(ParsedQuery query) { return this; } + @Override public ParsedQuery parsedQuery() { return this.originalQuery; @@ -1147,4 +1150,9 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() { } return false; } + + public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { + this.starTreeQueryContext = starTreeQueryContext; + return this; + } } diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index ac08ca0cc072c..17235f94ed483 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -140,6 +140,7 @@ 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; @@ -1548,8 +1549,10 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) { try { OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source); + StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source); if (parsedQuery != null) { - context.parsedQuery(new ParsedQuery(parsedQuery)); +// context.parsedQuery(new ParsedQuery(parsedQuery)); + 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..93c72b88f221d 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 @@ -104,7 +104,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 +144,56 @@ 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..3152eb49f315b 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 @@ -132,7 +132,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); } @@ -162,38 +162,38 @@ 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 = 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); +// } +// } +// }; +// } @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 81246a192555c..fd5a1d8bdd650 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 @@ -131,7 +131,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); } @@ -159,38 +159,38 @@ 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.MIN.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 >= 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); - } - } - }; - } +// 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.MIN.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 >= 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 public double metric(long owningBucketOrd) { 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..22f3ad6a7e65a 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 @@ -33,6 +33,8 @@ 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; import org.apache.lucene.util.NumericUtils; import org.opensearch.common.lease.Releasables; @@ -42,6 +44,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.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; @@ -51,6 +55,8 @@ 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 org.opensearch.search.startree.StarTreeQueryContext; import java.io.IOException; import java.util.Map; @@ -100,8 +106,13 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc } CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); - if (supportedStarTree != null) { - return getStarTreeLeafCollector(ctx, sub, supportedStarTree); +// if (supportedStarTree != null) { +// return getStarTreeLeafCollector(ctx, sub, supportedStarTree); +// } + StarTreeQueryContext starTreeQueryContext = this.context().getStarTreeQueryContext(); + if (starTreeQueryContext != null) { + return getAltStarTreeCollector(ctx, sub, supportedStarTree); + } return getDefaultLeafCollector(ctx, sub); } @@ -136,7 +147,47 @@ public void collect(int doc, long bucket) throws IOException { }; } - private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) +// 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(); @@ -146,35 +197,40 @@ private LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, Leaf MetricStat.SUM.getTypeName() ); assert starTreeValues != null; - SortedNumericDocValues values = (SortedNumericDocValues) starTreeValues.getMetricDocIdSetIterator(metricName); + SortedNumericStarTreeValuesIterator values = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator(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); + StarTreeFilter filter = new StarTreeFilter(starTreeValues, context().getStarTreeQueryContext().getQueryMap()); + StarTreeValuesIterator result = filter.getStarTreeResult(); - for (int i = 0; i < valuesCount; i++) { - double value = NumericUtils.sortableLongToDouble(values.nextValue()); - kahanSummation.add(value); - } + 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(); - compensations.set(bucket, kahanSummation.delta()); - sums.set(bucket, kahanSummation.value()); + // 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(); + } }; } - @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 0d9a145b2b8ad..217085c3e9f57 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 @@ -98,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); } final SortedNumericDocValues values = ((ValuesSource.Numeric) valuesSource).longValues(ctx); @@ -140,29 +140,29 @@ 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()); - } - } - }; - } +// 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()); +// } +// } +// }; +// } @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 5357206e8c117..6cb5382ba6cee 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -76,6 +76,7 @@ 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.util.Collection; @@ -531,4 +532,12 @@ public boolean keywordIndexOrDocValuesEnabled() { return false; } + public StarTreeQueryContext getStarTreeQueryContext() { + return null; + } + + public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) { + return this; + } + } 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..ff3cd4a345c5b 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java @@ -17,6 +17,8 @@ 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 +40,7 @@ * @opensearch.experimental * @opensearch.internal */ -class StarTreeFilter { +public class StarTreeFilter { private static final Logger logger = LogManager.getLogger(StarTreeFilter.class); private final Map queryMap; @@ -56,7 +58,55 @@ public StarTreeFilter(StarTreeValues starTreeAggrStructure, Map pr *
  • For the remaining columns, use doc values indexes to match them * */ - public DocIdSetIterator getStarTreeResult() throws IOException { + 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); + + // No matches, return + if (starTreeResult.maxMatchedDoc == -1) { + return docIdSetIterator; + } + 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.nextEntry() != NO_MORE_DOCS) { + int docID = docIdSetIterator.entryId(); + 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 = 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()); @@ -64,12 +114,12 @@ public DocIdSetIterator getStarTreeResult() throws IOException { // No matches, return if (starTreeResult.maxMatchedDoc == -1) { - return docIdSetIterator; + return; } 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<>(); @@ -77,7 +127,7 @@ public DocIdSetIterator getStarTreeResult() throws IOException { while (docIdSetIterator.nextDoc() != NO_MORE_DOCS) { int docID = docIdSetIterator.docID(); - if (ndv.advanceExact(docID)) { + if (ndv.advance(docID) > 0) { final int valuesCount = ndv.docValueCount(); for (int i = 0; i < valuesCount; i++) { long value = ndv.nextValue(); @@ -95,7 +145,6 @@ public DocIdSetIterator getStarTreeResult() throws IOException { } docIdSetIterator = builder.build().iterator(); } - return docIdSetIterator; } /** diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java index e8354b0415ac1..a51a12fd99a63 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQuery.java @@ -94,7 +94,7 @@ public Scorer scorer(LeafReaderContext context) throws IOException { return null; } StarTreeFilter filter = new StarTreeFilter(starTreeValues, queryMap); - DocIdSetIterator result = filter.getStarTreeResult(); + DocIdSetIterator result = null; return new ConstantScoreScorer(this, score(), scoreMode, result); } 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; + } +}