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 a1ec6cf7712c0..04a98e82a97fd 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 @@ -173,7 +173,7 @@ public static CompositeIndexFieldInfo getSupportedStarTree(SearchContext context return (starTreeQueryContext != null) ? starTreeQueryContext.getStarTree() : null; } - public static StarTreeValues computeStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException { + public static StarTreeValues getStarTreeValues(LeafReaderContext context, CompositeIndexFieldInfo starTree) throws IOException { SegmentReader reader = Lucene.segmentReader(context.reader()); if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) { return null; @@ -192,7 +192,7 @@ public static LeafBucketCollector getStarTreeLeafCollector( Consumer valueConsumer, Runnable finalConsumer ) throws IOException { - StarTreeValues starTreeValues = context.getStarTreeValues(ctx, starTree); + StarTreeValues starTreeValues = getStarTreeValues(ctx, starTree); String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); String metricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues(starTree.getField(), fieldName, metric); @@ -200,8 +200,7 @@ public static LeafBucketCollector getStarTreeLeafCollector( SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( metricName ); - StarTreeFilter filter = new StarTreeFilter(starTreeValues, context.getStarTreeQueryContext().getQueryMap()); - StarTreeValuesIterator result = filter.getStarTreeResult(); + StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues); int entryId; while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) { diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java index 204e52c32342c..fd53fa4153b1e 100644 --- a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java +++ b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java @@ -43,6 +43,7 @@ import org.apache.lucene.search.FieldDoc; import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; +import org.apache.lucene.util.BitSet; import org.opensearch.Version; import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchType; @@ -59,6 +60,8 @@ 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.compositeindex.datacube.startree.utils.StarTreeQueryHelper; +import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.engine.Engine; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; @@ -101,6 +104,7 @@ import org.opensearch.search.rescore.RescoreContext; import org.opensearch.search.slice.SliceBuilder; import org.opensearch.search.sort.SortAndFormats; +import org.opensearch.search.startree.StarTreeFilter; import org.opensearch.search.startree.StarTreeQueryContext; import org.opensearch.search.suggest.SuggestionSearchContext; @@ -119,7 +123,6 @@ 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; @@ -1167,17 +1170,15 @@ public StarTreeQueryContext getStarTreeQueryContext() { } @Override - public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException { + public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { if (this.starTreeValuesMap.containsKey(ctx)) { - logger.info("Used cached values"); - return starTreeValuesMap.get(ctx); - } else { - logger.info("not using cache"); + return starTreeValuesMap.get(ctx); } + StarTreeFilter filter = new StarTreeFilter(starTreeValues, this.getStarTreeQueryContext().getQueryMap()); + StarTreeValuesIterator result = filter.getStarTreeResult(); - StarTreeValues starTreeValues = computeStarTreeValues(ctx, starTree); - starTreeValuesMap.put(ctx, starTreeValues); - return starTreeValues; + starTreeValuesMap.put(ctx, result); + return result; } } 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 017bb91871b0e..f1b32fe8c9602 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,12 +32,20 @@ package org.opensearch.search.aggregations.metrics; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.CollectionTerminatedException; 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.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; import org.opensearch.index.fielddata.SortedNumericDoubleValues; import org.opensearch.search.DocValueFormat; import org.opensearch.search.aggregations.Aggregator; @@ -47,9 +55,11 @@ 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.Map; +import java.util.function.Consumer; import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree; @@ -138,57 +148,146 @@ 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()); - // } - // } - // }; - // } + public LeafBucketCollector getStarTreeCollector(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + throws IOException { + final CompensatedSum kahanSummation = new CompensatedSum(0, 0); + return StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.SUM.getTypeName(), + value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), + () -> sums.set(0, kahanSummation.value()) + ); + } + +// private LeafBucketCollector getStarTreeLeafCollector1(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()); +// } +// } +// }; +// } + + + public LeafBucketCollector getStarTreeCollector2(LeafReaderContext ctx, LeafBucketCollector sub, CompositeIndexFieldInfo starTree) + throws IOException { + final CompensatedSum kahanSummation = new CompensatedSum(0, 0); + StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.SUM.getTypeName(), + value -> kahanSummation.add(NumericUtils.sortableLongToDouble(value)), + () -> sums.set(0, kahanSummation.value()) + ); + StarTreeQueryHelper.getStarTreeLeafCollector( + context, + valuesSource, + ctx, + sub, + starTree, + MetricStat.VALUE_COUNT.getTypeName(), + value -> counts.increment(0, value), + () -> {} + ); + return LeafBucketCollector.NO_OP_COLLECTOR; + } + + public LeafBucketCollector getStarTreeLeafCollector( + LeafReaderContext ctx, + LeafBucketCollector sub, + CompositeIndexFieldInfo starTree + ) throws IOException { + StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree); + String fieldName = ((ValuesSource.Numeric.FieldData) valuesSource).getIndexFieldName(); + String sumMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + starTree.getField(), + fieldName, + MetricStat.SUM.getTypeName() + ); + String countMetricName = StarTreeUtils.fullyQualifiedFieldNameForStarTreeMetricsDocValues( + starTree.getField(), + fieldName, + MetricStat.VALUE_COUNT.getTypeName() + ); + + assert starTreeValues != null; + + final CompensatedSum kahanSummation = new CompensatedSum(0, 0); + SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( + sumMetricName + ); + SortedNumericStarTreeValuesIterator countValueIterator = (SortedNumericStarTreeValuesIterator) starTreeValues.getMetricValuesIterator( + countMetricName + ); + StarTreeValuesIterator result = context.getStarTreeFilteredValues(ctx, starTreeValues); + + int entryId; + while ((entryId = result.nextEntry()) != StarTreeValuesIterator.NO_MORE_ENTRIES) { + if (sumValuesIterator.advance(entryId) != StarTreeValuesIterator.NO_MORE_ENTRIES) { + int count = sumValuesIterator.valuesCount(); + for (int i = 0; i < count; i++) { + kahanSummation.add(NumericUtils.sortableLongToDouble(sumValuesIterator.nextValue())); + counts.increment(0, countValueIterator.nextValue()); + } + } + } + sums.set(0, kahanSummation.value()); + return new LeafBucketCollectorBase(sub, valuesSource.doubleValues(ctx)) { + @Override + public void collect(int doc, long bucket) { + throw new CollectionTerminatedException(); + } + }; + } @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 c7e697ad45d8f..e68dcc4acb614 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 @@ -130,7 +130,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - System.out.println("max star tree"); return getStarTreeCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); 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 188c615a26ff2..11f38020aa8a1 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 @@ -129,7 +129,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBuc CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context); if (supportedStarTree != null) { - System.out.println("min star tree"); return getStarTreeCollector(ctx, sub, supportedStarTree); } return getDefaultLeafCollector(ctx, sub); 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 ed3cbde61ce74..66b2ba9c98431 100644 --- a/server/src/main/java/org/opensearch/search/internal/SearchContext.java +++ b/server/src/main/java/org/opensearch/search/internal/SearchContext.java @@ -47,6 +47,7 @@ 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.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.mapper.ObjectMapper; @@ -129,7 +130,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; + protected volatile Map starTreeValuesMap; private volatile boolean searchTimedOut; protected SearchContext() {} @@ -544,7 +545,7 @@ public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryCont return this; } - public StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException { + public StarTreeValuesIterator getStarTreeFilteredValues(LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException { return null; } } 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 493bdf3ca143e..fee32ba48e1f5 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java @@ -11,6 +11,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.util.DocIdSetBuilder; +import org.apache.lucene.util.FixedBitSet; 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; @@ -62,6 +63,7 @@ public StarTreeValuesIterator getStarTreeResult() throws IOException { andIterators.add(new StarTreeValuesIterator(starTreeResult._matchedDocIds.build().iterator())); StarTreeValuesIterator starTreeValuesIterator = andIterators.get(0); + int length = 0; // No matches, return if (starTreeResult.maxMatchedDoc == -1) { return starTreeValuesIterator; @@ -93,6 +95,7 @@ public StarTreeValuesIterator getStarTreeResult() throws IOException { for (int entryId : entryIds) { adder.add(entryId); } + length = entryIds.size(); starTreeValuesIterator = new StarTreeValuesIterator(builder.build().iterator()); } return starTreeValuesIterator;