refactoring

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · Oct 1, 2024 · 4c35a89 · 4c35a89
1 parent a18249d
commit 4c35a89
Show file tree

Hide file tree

Showing 8 changed files with 162 additions and 268 deletions.
diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -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, true, Property.NodeScope);
+    public static final Setting<Boolean> STAR_TREE_INDEX_SETTING = Setting.boolSetting(STAR_TREE_INDEX, false, Property.NodeScope);
 
     /**
      * Gates the functionality of application based configuration templates.

diff --git a/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java b/...java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeQueryHelper.java
@@ -54,24 +54,19 @@ public class StarTreeQueryHelper {
     /**
      * Checks if the search context can be supported by star-tree
      */
-    public static boolean isStarTreeSupported(SearchContext context, boolean trackTotalHits) {
-        boolean canUseStarTree = context.aggregations() != null
-            && context.size() == 0
+    public static boolean isStarTreeSupported(SearchContext context) {
+        return context.aggregations() != null
             && context.mapperService().isCompositeIndexPresent()
             && context.parsedPostFilter() == null
             && context.innerHits().getInnerHits().isEmpty()
             && context.sort() == null
-            && (!trackTotalHits || context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_DISABLED)
             && context.trackScores() == false
-            && context.minimumScore() == null
-            && context.terminateAfter() == 0;
-        return canUseStarTree;
+            && context.minimumScore() == null;
     }
 
-
     /**
-     * Gets a parsed OriginalOrStarTreeQuery from the search context and source builder.
-     * Returns null if the query cannot be supported.
+     * Gets StarTreeQueryContext from the search context and source builder.
+     * Returns null if the query & aggregation cannot be supported.
      */
     public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context, SearchSourceBuilder source) throws IOException {
         // Current implementation assumes only single star-tree is supported
@@ -93,7 +88,6 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
             return null;
         }
 
-
         for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) {
             MetricStat metricStat = validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory);
             if (metricStat == null) {
@@ -104,20 +98,22 @@ public static StarTreeQueryContext getStarTreeQueryContext(SearchContext context
         if (context.aggregations().factories().getFactories().length > 1) {
             context.initializeStarTreeValuesMap();
         }
-
         return starTreeQueryContext;
     }
 
+    /**
+     * Uses query builder & composite index info to form star-tree query context
+     */
     private static StarTreeQueryContext toStarTreeQueryContext(
-        CompositeIndexFieldInfo starTree,
-        CompositeDataCubeFieldType compositeIndexFieldInfo,
+        CompositeIndexFieldInfo compositeIndexFieldInfo,
+        CompositeDataCubeFieldType compositeFieldType,
         QueryBuilder queryBuilder
     ) {
         Map<String, Long> queryMap;
         if (queryBuilder == null || queryBuilder instanceof MatchAllQueryBuilder) {
             queryMap = null;
         } else if (queryBuilder instanceof TermQueryBuilder) {
-            List<String> supportedDimensions = compositeIndexFieldInfo.getDimensions()
+            List<String> supportedDimensions = compositeFieldType.getDimensions()
                 .stream()
                 .map(Dimension::getField)
                 .collect(Collectors.toList());
@@ -128,13 +124,12 @@ private static StarTreeQueryContext toStarTreeQueryContext(
         } else {
             return null;
         }
-
-        return new StarTreeQueryContext(starTree, queryMap);
+        return new StarTreeQueryContext(compositeIndexFieldInfo, queryMap);
     }
 
     /**
      * Parse query body to star-tree predicates
-     * @param queryBuilder to match supported query shape
+     * @param queryBuilder to match star-tree supported query shape
      * @return predicates to match
      */
     private static Map<String, Long> getStarTreePredicates(QueryBuilder queryBuilder, List<String> supportedDimensions) {
@@ -185,6 +180,10 @@ public static StarTreeValues getStarTreeValues(LeafReaderContext context, Compos
         return (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
     }
 
+    /**
+     * Get the star-tree leaf collector
+     * This collector computes the aggregation prematurely and invokes an early termination collector
+     */
     public static LeafBucketCollector getStarTreeLeafCollector(
         SearchContext context,
         ValuesSource.Numeric valuesSource,
@@ -225,7 +224,6 @@ public static LeafBucketCollector getStarTreeLeafCollector(
             }
         }
 
-
         // Call the final consumer after processing all entries
         finalConsumer.run();
 
@@ -238,7 +236,11 @@ public void collect(int doc, long bucket) {
         };
     }
 
-    public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues) throws IOException {
+    /**
+     * Get the filtered values for the star-tree query
+     */
+    public static FixedBitSet getStarTreeFilteredValues(SearchContext context, LeafReaderContext ctx, StarTreeValues starTreeValues)
+        throws IOException {
         if (context.getStarTreeValuesMap() != null && context.getStarTreeValuesMap().containsKey(ctx)) {
             return context.getStarTreeValuesMap().get(ctx);
         }

diff --git a/server/src/main/java/org/opensearch/search/DefaultSearchContext.java b/server/src/main/java/org/opensearch/search/DefaultSearchContext.java
@@ -34,7 +34,6 @@
 
 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;
@@ -43,7 +42,6 @@
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.util.FixedBitSet;
 import org.opensearch.Version;
 import org.opensearch.action.search.SearchShardTask;
 import org.opensearch.action.search.SearchType;
@@ -58,7 +56,6 @@
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.cache.bitset.BitsetFilterCache;
-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;
@@ -101,8 +98,6 @@
 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;
 
 import java.io.IOException;
@@ -275,7 +270,6 @@ final class DefaultSearchContext extends SearchContext {
         this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
         this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
         this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
-        this.starTreeValuesMap = new HashMap<>();
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -1544,7 +1544,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
 
         if (this.indicesService.getCompositeIndexSettings() != null
             && this.indicesService.getCompositeIndexSettings().isStarTreeIndexCreationEnabled()
-            && StarTreeQueryHelper.isStarTreeSupported(context, source.trackTotalHitsUpTo() != null)) {
+            && StarTreeQueryHelper.isStarTreeSupported(context)) {
             try {
                 StarTreeQueryContext starTreeQueryContext = StarTreeQueryHelper.getStarTreeQueryContext(context, source);
                 if (starTreeQueryContext != null) {

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
@@ -108,7 +108,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);
     }
@@ -164,7 +164,6 @@ public LeafBucketCollector getStarTreeLeafCollector(LeafReaderContext ctx, LeafB
             MetricStat.VALUE_COUNT.getTypeName()
         );
 
-
         final CompensatedSum kahanSummation = new CompensatedSum(sums.get(0), 0);
         SortedNumericStarTreeValuesIterator sumValuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
             .getMetricValuesIterator(sumMetricName);

diff --git a/server/src/main/java/org/opensearch/search/internal/SearchContext.java b/server/src/main/java/org/opensearch/search/internal/SearchContext.java
@@ -46,7 +46,6 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.BigArrays;
 import org.opensearch.index.cache.bitset.BitsetFilterCache;
-import org.opensearch.index.compositeindex.datacube.startree.index.StarTreeValues;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.ObjectMapper;
@@ -79,11 +78,9 @@
 import org.opensearch.search.query.ReduceableSearchResult;
 import org.opensearch.search.rescore.RescoreContext;
 import org.opensearch.search.sort.SortAndFormats;
-import org.opensearch.search.startree.StarTreeFilter;
 import org.opensearch.search.startree.StarTreeQueryContext;
 import org.opensearch.search.suggest.SuggestionSearchContext;
 
-import java.io.IOException;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -538,7 +535,6 @@ public boolean keywordIndexOrDocValuesEnabled() {
         return false;
     }
 
-
     public SearchContext starTreeQueryContext(StarTreeQueryContext starTreeQueryContext) {
         this.starTreeQueryContext = starTreeQueryContext;
         return this;