Skip to content

Commit

Permalink
temp
Browse files Browse the repository at this point in the history
Signed-off-by: Sandesh Kumar <[email protected]>
  • Loading branch information
sandeshkr419 committed Sep 19, 2024
1 parent 7af6aa6 commit 07abd51
Show file tree
Hide file tree
Showing 16 changed files with 309 additions and 416 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, true, Property.NodeScope);

/**
* Gates the functionality of application based configuration templates.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,29 +10,36 @@

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;
import org.opensearch.index.compositeindex.datacube.Dimension;
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;

/**
Expand All @@ -43,6 +50,8 @@
*/
public class StarTreeQueryHelper {

private static Map<LeafReaderContext, StarTreeValues> starTreeValuesMap = new HashMap<>();

/**
* Checks if the search context can be supported by star-tree
*/
Expand All @@ -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()
Expand All @@ -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;
}

Expand All @@ -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
Expand All @@ -110,7 +128,7 @@ private static StarTreeQuery toStarTreeQuery(
return null;
}

return new StarTreeQuery(starTree, queryMap);
return new StarTreeQueryContext(starTree, queryMap);
}

/**
Expand Down Expand Up @@ -151,18 +169,56 @@ 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;
}
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<Long> 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();
}
};
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -270,6 +276,7 @@ final class DefaultSearchContext extends SearchContext {
this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
this.starTreeValuesMap = new HashMap<>();
}

@Override
Expand Down Expand Up @@ -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;
}
}
9 changes: 4 additions & 5 deletions server/src/main/java/org/opensearch/search/SearchService.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down
Loading

0 comments on commit 07abd51

Please sign in to comment.