Skip to content

Commit

Permalink
star tree file formats refactoring and fixing offset bug
Browse files Browse the repository at this point in the history 
Signed-off-by: Bharathwaj G <[email protected]>
  • Loading branch information
@bharath-techie @sandeshkr419
bharath-techie authored and sandeshkr419 committed Sep 19, 2024
1 parent 4962594 commit 4f6cc67
Show file tree
Hide file tree
Showing 16 changed files with 187 additions and 459 deletions.
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/common/util/FeatureFlags.java
View file
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
114 changes: 60 additions & 54 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 @@ -10,30 +10,40 @@

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.CollectionTerminatedException;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.NumericUtils;
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.fielddata.SortedNumericDoubleValues;
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.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
Expand All @@ -44,6 +54,9 @@
*/
public class StarTreeQueryHelper {


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

/**
* Checks if the search context can be supported by star-tree
*/
Expand All @@ -65,30 +78,7 @@ 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 {
// 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()
);

StarTreeQuery starTreeQuery = StarTreeQueryHelper.toStarTreeQuery(starTree, compositeMappedFieldType, source.query());
if (starTreeQuery == null) {
return null;
}

for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
if (validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory) == false) {
return null;
}
}

return new OriginalOrStarTreeQuery(starTreeQuery, context.query());
}

/**
* Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
Expand Down Expand Up @@ -123,30 +113,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
return starTreeQueryContext;
}

private static StarTreeQuery toStarTreeQuery(
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 StarTreeQuery(starTree, queryMap);
}

private static StarTreeQueryContext toStarTreeQueryContext(
CompositeIndexFieldInfo starTree,
CompositeDataCubeFieldType compositeIndexFieldInfo,
Expand Down Expand Up @@ -209,18 +175,58 @@ 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) > 0) {
int count = valuesIterator.docValueCount();
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();
}
};
}
}
4 changes: 4 additions & 0 deletions .../compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
View file
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 docValueCount() throws IOException {
return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
}
}
2 changes: 1 addition & 1 deletion ...nsearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
View file
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
22 changes: 22 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 @@ -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 @@ -116,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 @@ -272,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 @@ -1151,8 +1156,25 @@ 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)) {
return starTreeValuesMap.get(ctx);
}
StarTreeValues starTreeValues = computeStarTreeValues(ctx, starTree);
starTreeValuesMap.put(ctx, starTreeValues);
return starTreeValues;
}
}
5 changes: 2 additions & 3 deletions server/src/main/java/org/opensearch/search/SearchService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,6 @@
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;
Expand Down Expand Up @@ -1548,9 +1547,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);
// OriginalOrStarTreeQuery parsedQuery = StarTreeQueryHelper.getOriginalOrStarTreeQuery(context, source);
StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source);
if (parsedQuery != null) {
if (starTreeQueryContext != null) {
// context.parsedQuery(new ParsedQuery(parsedQuery));
context.starTreeQueryContext(starTreeQueryContext);
logger.debug("can use star tree");
Expand Down
2 changes: 1 addition & 1 deletion 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 @@ -56,7 +56,7 @@
import java.io.IOException;
import java.util.Map;

import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
//import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getStarTreeValues;
import static org.opensearch.index.compositeindex.datacube.startree.utils.StarTreeQueryHelper.getSupportedStarTree;

/**
Expand Down
Loading

0 comments on commit 4f6cc67

Please sign in to comment.