Skip to content

Commit

Permalink
temp temp
Browse files Browse the repository at this point in the history 
Signed-off-by: Sandesh Kumar <[email protected]>
  • Loading branch information
@sandeshkr419
sandeshkr419 committed Sep 19, 2024
1 parent 07abd51 commit afbc128
Show file tree
Hide file tree
Showing 7 changed files with 169 additions and 68 deletions.
7 changes: 3 additions & 4 deletions ...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -192,16 +192,15 @@ public static LeafBucketCollector getStarTreeLeafCollector(
Consumer<Long> 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);

assert starTreeValues != null;
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) {
Expand Down
19 changes: 10 additions & 9 deletions server/src/main/java/org/opensearch/search/DefaultSearchContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
}
201 changes: 150 additions & 51 deletions server/src/main/java/org/opensearch/search/aggregations/metrics/AvgAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand Down Expand Up @@ -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) {
Expand Down
1 change: 0 additions & 1 deletion server/src/main/java/org/opensearch/search/aggregations/metrics/MaxAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
1 change: 0 additions & 1 deletion server/src/main/java/org/opensearch/search/aggregations/metrics/MinAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
5 changes: 3 additions & 2 deletions server/src/main/java/org/opensearch/search/internal/SearchContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -129,7 +130,7 @@ public List<InternalAggregation> toInternalAggregations(Collection<Collector> co
private final List<Releasable> releasables = new CopyOnWriteArrayList<>();
private final AtomicBoolean closed = new AtomicBoolean(false);
private InnerHitsContext innerHitsContext;
protected volatile Map<LeafReaderContext, StarTreeValues> starTreeValuesMap;
protected volatile Map<LeafReaderContext, StarTreeValuesIterator> starTreeValuesMap;
private volatile boolean searchTimedOut;

protected SearchContext() {}
Expand Down Expand Up @@ -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;
}
}
3 changes: 3 additions & 0 deletions server/src/main/java/org/opensearch/search/startree/StarTreeFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit afbc128

Please sign in to comment.