Star tree P0 changes

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · Aug 10, 2024 · bbcb8bf · bbcb8bf
1 parent 781a2d4
commit bbcb8bf
Show file tree

Hide file tree

Showing 22 changed files with 1,790 additions and 11 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
@@ -105,7 +105,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);
 
     private static final List<Setting<Boolean>> ALL_FEATURE_FLAG_SETTINGS = List.of(
         REMOTE_STORE_MIGRATION_EXPERIMENTAL_SETTING,

diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -56,7 +56,12 @@
 import org.opensearch.index.IndexSortConfig;
 import org.opensearch.index.analysis.IndexAnalyzers;
 import org.opensearch.index.cache.bitset.BitsetFilterCache;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.compositeindex.datacube.Metric;
+import org.opensearch.index.compositeindex.datacube.MetricStat;
+import org.opensearch.index.compositeindex.datacube.startree.aggregators.MetricAggregatorInfo;
 import org.opensearch.index.fielddata.IndexFieldData;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
 import org.opensearch.index.mapper.ContentPath;
 import org.opensearch.index.mapper.DerivedFieldResolver;
 import org.opensearch.index.mapper.DerivedFieldResolverFactory;
@@ -73,12 +78,18 @@
 import org.opensearch.script.ScriptContext;
 import org.opensearch.script.ScriptFactory;
 import org.opensearch.script.ScriptService;
+import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.metrics.MetricsAggregator;
+import org.opensearch.search.aggregations.metrics.SumAggregatorFactory;
 import org.opensearch.search.aggregations.support.AggregationUsageService;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
 import org.opensearch.search.lookup.SearchLookup;
+import org.opensearch.search.startree.StarTreeQuery;
 import org.opensearch.transport.RemoteClusterAware;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -89,6 +100,7 @@
 import java.util.function.LongSupplier;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptyMap;
@@ -522,6 +534,63 @@ private ParsedQuery toQuery(QueryBuilder queryBuilder, CheckedFunction<QueryBuil
         }
     }
 
+    public ParsedQuery toStarTreeQuery(CompositeIndexFieldInfo starTree, QueryBuilder queryBuilder, Query query) {
+        Map<String, List<Predicate<Long>>> predicateMap = getFilterMap(queryBuilder);
+        StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, predicateMap, null);
+        return new ParsedQuery(starTreeQuery);
+    }
+
+    private Map<String, List<Predicate<Long>>> getFilterMap(QueryBuilder queryBuilder) {
+        // Assuming the following variables have been initialized:
+        Map<String, List<Predicate<Long>>> predicateMap = new HashMap<>();
+
+
+        // Check if the query builder is an instance of TermQueryBuilder
+        if (queryBuilder instanceof TermQueryBuilder) {
+            TermQueryBuilder tq = (TermQueryBuilder) queryBuilder;
+            String field = tq.fieldName();
+            long inputQueryVal = Long.parseLong(tq.value().toString());
+
+            // Get or create the list of predicates for the given field
+            List<Predicate<Long>> predicates = predicateMap.getOrDefault(field, new ArrayList<>());
+
+            // Create a predicate to match the input query value
+            Predicate<Long> predicate = dimVal -> dimVal == inputQueryVal;
+            predicates.add(predicate);
+
+            // Put the predicates list back into the map
+            predicateMap.put(field, predicates);
+        } else {
+            throw new IllegalArgumentException("The query is not a term query");
+        }
+        return predicateMap;
+
+    }
+
+    public String getMetricKey(CompositeDataCubeFieldType compositeIndexFieldInfo, AggregatorFactory aggregatorFactory) {
+        String field = null;
+        Map<String, List<MetricStat>> supportedMetrics = compositeIndexFieldInfo.getMetrics().stream()
+            .collect(Collectors.toMap(Metric::getField, Metric::getMetrics));
+
+        // Existing support only for MetricAggregators without sub-aggregations
+        if (aggregatorFactory.getSubFactories().getFactories().length != 0) {
+            return null;
+        }
+
+        if (aggregatorFactory instanceof SumAggregatorFactory) {
+            field = ((SumAggregatorFactory)aggregatorFactory).getField();
+            if (!(supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.SUM))) {
+                return null;
+            }
+        }
+
+        return field != null
+            ? compositeIndexFieldInfo.name() + "_" + field + "_" + "sum_metric"
+            : null;
+    }
+
+
+
     public Index index() {
         return indexSettings.getIndex();
     }

diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java
@@ -35,6 +35,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.TopDocs;
 import org.opensearch.OpenSearchException;
@@ -77,16 +78,21 @@
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
 import org.opensearch.index.engine.Engine;
+import org.opensearch.index.mapper.CompositeDataCubeFieldType;
 import org.opensearch.index.mapper.DerivedFieldResolver;
 import org.opensearch.index.mapper.DerivedFieldResolverFactory;
+import org.opensearch.index.mapper.StarTreeMapper;
 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;
 import org.opensearch.index.query.Rewriteable;
+import org.opensearch.index.query.TermsQueryBuilder;
 import org.opensearch.index.shard.IndexEventListener;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.shard.SearchOperationListener;
@@ -97,11 +103,14 @@
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.aggregations.AggregationInitializationException;
 import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.AggregatorFactory;
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.InternalAggregation.ReduceContext;
 import org.opensearch.search.aggregations.MultiBucketConsumerService;
 import org.opensearch.search.aggregations.SearchContextAggregations;
 import org.opensearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.collapse.CollapseContext;
 import org.opensearch.search.dfs.DfsPhase;
@@ -136,6 +145,7 @@
 import org.opensearch.search.sort.SortAndFormats;
 import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.search.sort.SortOrder;
+import org.opensearch.search.startree.StarTreeAggregatorFactory;
 import org.opensearch.search.suggest.Suggest;
 import org.opensearch.search.suggest.completion.CompletionSuggestion;
 import org.opensearch.tasks.TaskResourceTrackingService;
@@ -1314,6 +1324,11 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
             context.evaluateRequestShouldUseConcurrentSearch();
             return;
         }
+        // Can be marked false for majority cases for which star-tree cannot be used
+        // Will save checking the criteria later and we can have a limit on what search requests are supported
+        // As we increment the cases where star-tree can be used, this can be set back to true
+        boolean canUseStarTree = context.mapperService().isCompositeIndexPresent();
+
         SearchShardTarget shardTarget = context.shardTarget();
         QueryShardContext queryShardContext = context.getQueryShardContext();
         context.from(source.from());
@@ -1339,9 +1354,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         if (source.sorts() != null) {
             try {
                 Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(source.sorts(), context.getQueryShardContext());
-                if (optionalSort.isPresent()) {
-                    context.sort(optionalSort.get());
-                }
+                optionalSort.ifPresent(context::sort);
             } catch (IOException e) {
                 throw new SearchException(shardTarget, "failed to create sort elements", e);
             }
@@ -1496,6 +1509,44 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         if (source.profile()) {
             context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch()));
         }
+
+        if (canUseStarTree) {
+            try {
+                setStarTreeQuery(context, queryShardContext, source);
+                logger.info("using star tree");
+            }
+            catch (IOException e) {
+               logger.info("not using star tree");
+            }
+        }
+    }
+
+    private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source) throws IOException {
+
+        if (source.aggregations() == null) {
+            return false;
+        }
+
+        // TODO: Support for multiple startrees
+        CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService().getCompositeFieldTypes().iterator().next();
+        CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo(compositeMappedFieldType.name(), compositeMappedFieldType.getCompositeIndexType());
+
+        ParsedQuery parsedQuery = queryShardContext.toStarTreeQuery(starTree, source.query(), context.query());
+        AggregatorFactory aggregatorFactory = context.aggregations().factories().getFactories()[0];
+        if (!(aggregatorFactory instanceof ValuesSourceAggregatorFactory)) {
+            return false;
+        }
+        ValuesSourceConfig config = ((ValuesSourceAggregatorFactory) aggregatorFactory).getConfig();
+        String metricKey = queryShardContext.getMetricKey(compositeMappedFieldType, aggregatorFactory);
+        StarTreeAggregatorFactory factory = new StarTreeAggregatorFactory(aggregatorFactory.name(), queryShardContext, config, List.of(metricKey));
+
+
+        AggregatorFactories aggregatorFactories = new AggregatorFactories(new StarTreeAggregatorFactory[]{factory});
+
+        context.parsedQuery(parsedQuery)
+            .aggregations(new SearchContextAggregations(aggregatorFactories, multiBucketConsumerService.create()));
+
+        return false;
     }
 
     /**

diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
@@ -255,7 +255,7 @@ public static Builder builder() {
         return new Builder();
     }
 
-    private AggregatorFactories(AggregatorFactory[] factories) {
+    public AggregatorFactories(AggregatorFactory[] factories) {
         this.factories = factories;
     }
 
@@ -661,4 +661,8 @@ public PipelineTree buildPipelineTree() {
             return new PipelineTree(subTrees, aggregators);
         }
     }
+
+    public AggregatorFactory[] getFactories() {
+        return factories;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactory.java
@@ -127,4 +127,8 @@ protected boolean supportsConcurrentSegmentSearch() {
     public boolean evaluateChildFactories() {
         return factories.allFactoriesSupportConcurrentSearch();
     }
+
+    public AggregatorFactories getSubFactories() {
+        return factories;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java
@@ -56,13 +56,13 @@
  */
 public class SumAggregator extends NumericMetricsAggregator.SingleValue {
 
-    private final ValuesSource.Numeric valuesSource;
-    private final DocValueFormat format;
+    protected final ValuesSource.Numeric valuesSource;
+    protected final DocValueFormat format;
 
-    private DoubleArray sums;
-    private DoubleArray compensations;
+    protected DoubleArray sums;
+    protected DoubleArray compensations;
 
-    SumAggregator(
+    public SumAggregator(
         String name,
         ValuesSourceConfig valuesSourceConfig,
         SearchContext context,

diff --git a/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregatorFactory.java
@@ -52,7 +52,7 @@
  *
  * @opensearch.internal
  */
-class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
+public class SumAggregatorFactory extends ValuesSourceAggregatorFactory {
 
     SumAggregatorFactory(
         String name,