Skip to content

Commit

Permalink
temp
Browse files Browse the repository at this point in the history 
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
  • Loading branch information
@sandeshkr419
sandeshkr419 committed Sep 19, 2024
1 parent 7af6aa6 commit 4962594
Show file tree
Hide file tree
Showing 12 changed files with 402 additions and 171 deletions.
58 changes: 58 additions & 0 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 @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -113,6 +147,30 @@ private static StarTreeQuery toStarTreeQuery(
return new StarTreeQuery(starTree, queryMap);
}

private static StarTreeQueryContext toStarTreeQueryContext(
CompositeIndexFieldInfo starTree,
CompositeDataCubeFieldType compositeIndexFieldInfo,
QueryBuilder queryBuilder
) {
Map<String, Long> queryMap;
if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
queryMap = null;
} else if (queryBuilder instanceof TermQueryBuilder) {
List<String> 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
Expand Down
8 changes: 8 additions & 0 deletions server/src/main/java/org/opensearch/search/DefaultSearchContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -744,6 +746,7 @@ public SearchContext parsedQuery(ParsedQuery query) {
return this;
}


@Override
public ParsedQuery parsedQuery() {
return this.originalQuery;
Expand Down Expand Up @@ -1147,4 +1150,9 @@ public boolean evaluateKeywordIndexOrDocValuesEnabled() {
}
return false;
}

public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
this.starTreeQueryContext = starTreeQueryContext;
return this;
}
}
5 changes: 4 additions & 1 deletion server/src/main/java/org/opensearch/search/SearchService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down
102 changes: 51 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 @@ -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);
}
Expand Down Expand Up @@ -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) {
Expand Down
66 changes: 33 additions & 33 deletions 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 @@ -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);
}
Expand Down Expand Up @@ -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) {
Expand Down
Loading

0 comments on commit 4962594

Please sign in to comment.