Star Tree Request/Response structure

server/src/main/java/org/opensearch/index/query/QueryShardContext.java

@@ -56,7 +56,11 @@
 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.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,13 +77,17 @@
 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.SumAggregatorFactory;
 import org.opensearch.search.aggregations.support.AggregationUsageService;
 import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
 import org.opensearch.search.lookup.SearchLookup;
-import org.opensearch.search.query.startree.StarTreeQuery;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+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;
@@ -90,6 +98,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;
@@ -529,6 +538,66 @@ private ParsedQuery toQuery(QueryBuilder queryBuilder, CheckedFunction<QueryBuil
         }
     }
 
+    public ParsedQuery toStarTreeQuery(CompositeIndexFieldInfo starTree, QueryBuilder queryBuilder, Query query) {
+        Map<String, List<Predicate<Long>>> predicateMap = getStarTreePredicates(queryBuilder);
+        StarTreeQuery starTreeQuery = new StarTreeQuery(starTree, predicateMap, null);
+        OriginalOrStarTreeQuery originalOrStarTreeQuery = new OriginalOrStarTreeQuery(starTreeQuery, query);
+        return new ParsedQuery(originalOrStarTreeQuery);
+    }
+
+    /**
+     * Parse query body to star-tree predicates
+     * @param queryBuilder
+     * @return
+     */
+    private Map<String, List<Predicate<Long>>> getStarTreePredicates(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 boolean validateStarTreeMetricSuport(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 false;
+        }
+
+        // TODO: increment supported aggregation type
+        if (aggregatorFactory instanceof SumAggregatorFactory) {
+            field = ((SumAggregatorFactory) aggregatorFactory).getField();
+            if (supportedMetrics.containsKey(field) && supportedMetrics.get(field).contains(MetricStat.SUM)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     public Index index() {
         return indexSettings.getIndex();
     }

server/src/main/java/org/opensearch/search/SearchService.java

@@ -77,12 +77,16 @@
 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;
@@ -97,13 +101,13 @@
 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.startree.StarTreeAggregator;
-import org.opensearch.search.aggregations.startree.StarTreeAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.collapse.CollapseContext;
 import org.opensearch.search.dfs.DfsPhase;
@@ -1349,9 +1353,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);
             }
@@ -1509,35 +1511,40 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
 
         if (canUseStarTree) {
             try {
-                if (setStarTreeQuery(context, queryShardContext, source)) {
-                    logger.info("Star Tree will be used in execution");
-                };
+                setStarTreeQuery(context, queryShardContext, source);
+                logger.info("using star tree");
             } catch (IOException e) {
-                logger.info("Cannot use star-tree");
+                logger.info("not using star tree");
             }
-
         }
     }
 
-    private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source) throws IOException {
-
-        // TODO: (finish)
-        //  1. Check criteria for star-tree query / aggregation formation
-        //  2: Set StarTree Query & Star Tree Aggregator here
-
+    private boolean setStarTreeQuery(SearchContext context, QueryShardContext queryShardContext, SearchSourceBuilder source)
+        throws IOException {
         if (source.aggregations() == null) {
             return false;
         }
 
-        context.parsedQuery(queryShardContext.toStarTreeQuery(null, Set.of("sum_status")));
-
-        StarTreeAggregatorFactory factory = new StarTreeAggregatorFactory("sum_status", queryShardContext, null, null, null, List.of("status"), List.of("sum"));
-        StarTreeAggregatorFactory[] factories = {factory};
-        AggregatorFactories aggregatorFactories = new AggregatorFactories(factories);
+        // TODO: Support for multiple startrees
+        CompositeDataCubeFieldType compositeMappedFieldType = (StarTreeMapper.StarTreeFieldType) context.mapperService()
+            .getCompositeFieldTypes()
+            .iterator()
+            .next();
+        CompositeIndexFieldInfo starTree = new CompositeIndexFieldInfo(
+            compositeMappedFieldType.name(),
+            compositeMappedFieldType.getCompositeIndexType()
+        );
 
-        context.aggregations(new SearchContextAggregations(aggregatorFactories, multiBucketConsumerService.create()));
+        ParsedQuery parsedQuery = queryShardContext.toStarTreeQuery(starTree, source.query(), context.query());
+        AggregatorFactory aggregatorFactory = context.aggregations().factories().getFactories()[0];
+        if (!(aggregatorFactory instanceof ValuesSourceAggregatorFactory
+            && aggregatorFactory.getSubFactories().getFactories().length == 0)) {
+            return false;
+        }
 
-//        StarTreeAggregatorFactory factory = new StarTreeAggregatorFactory()
+        if (queryShardContext.validateStarTreeMetricSuport(compositeMappedFieldType, aggregatorFactory)) {
+            context.parsedQuery(parsedQuery);
+        }
         return true;
     }
 

server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java

@@ -661,4 +661,8 @@ public PipelineTree buildPipelineTree() {
             return new PipelineTree(subTrees, aggregators);
         }
     }
+
+    public AggregatorFactory[] getFactories() {
+        return factories;
+    }
 }

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;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/metrics/NumericMetricsAggregator.java

@@ -31,13 +31,20 @@
 
 package org.opensearch.search.aggregations.metrics;
 
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentReader;
+import org.opensearch.common.lucene.Lucene;
 import org.opensearch.common.util.Comparators;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.CompositeIndexReader;
+import org.opensearch.index.codec.composite.datacube.startree.StarTreeValues;
 import org.opensearch.search.aggregations.Aggregator;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.sort.SortOrder;
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Base class to aggregate all docs into a single numeric metric value.
@@ -107,4 +114,14 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
             return (lhs, rhs) -> Comparators.compareDiscardNaN(metric(key, lhs), metric(key, rhs), order == SortOrder.ASC);
         }
     }
+
+    protected StarTreeValues getStarTreeValues(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+        SegmentReader reader = Lucene.segmentReader(ctx.reader());
+        if (!(reader.getDocValuesReader() instanceof CompositeIndexReader)) return null;
+        CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader();
+        StarTreeValues values = (StarTreeValues) starTreeDocValuesReader.getCompositeIndexValues(starTree);
+        final AtomicReference<StarTreeValues> aggrVal = new AtomicReference<>(null);
+
+        return values;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/metrics/SumAggregator.java

@@ -36,6 +36,8 @@
 import org.opensearch.common.lease.Releasables;
 import org.opensearch.common.util.BigArrays;
 import org.opensearch.common.util.DoubleArray;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.index.codec.composite.datacube.startree.StarTreeValues;
 import org.opensearch.index.fielddata.SortedNumericDoubleValues;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.aggregations.Aggregator;
@@ -45,6 +47,8 @@
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.OriginalOrStarTreeQuery;
+import org.opensearch.search.startree.StarTreeQuery;
 
 import java.io.IOException;
 import java.util.Map;
@@ -56,13 +60,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,
@@ -86,6 +90,14 @@ public ScoreMode scoreMode() {
 
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
+        if (context.query() instanceof OriginalOrStarTreeQuery && ((OriginalOrStarTreeQuery) context.query()).isStarTreeUsed()) {
+            StarTreeQuery starTreeQuery = ((OriginalOrStarTreeQuery) context.query()).getStarTreeQuery();
+            return getStarTreeLeafCollector(ctx, sub, starTreeQuery.getStarTree());
+        }
+        return getDefaultLeafCollector(ctx, sub);
+    }
+
+    private LeafBucketCollector getDefaultLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
         if (valuesSource == null) {
             return LeafBucketCollector.NO_OP_COLLECTOR;
         }
@@ -118,6 +130,28 @@ 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();
+
+        return new LeafBucketCollectorBase(sub, starTreeValues) {
+            @Override
+            public void collect(int doc, long bucket) throws IOException {
+                // TODO: Fix the response for collecting star tree sum
+                sums = bigArrays.grow(sums, bucket + 1);
+                compensations = bigArrays.grow(compensations, bucket + 1);
+                compensations.set(bucket, kahanSummation.delta());
+                sums.set(bucket, kahanSummation.value());
+            }
+        };
+    }
+
     @Override
     public double metric(long owningBucketOrd) {
         if (valuesSource == null || owningBucketOrd >= sums.size()) {

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,

server/src/main/java/org/opensearch/search/aggregations/support/ValuesSource.java

@@ -625,6 +625,10 @@ public SortedNumericDocValues longValues(LeafReaderContext context) {
             public SortedNumericDoubleValues doubleValues(LeafReaderContext context) {
                 return indexFieldData.load(context).getDoubleValues();
             }
+
+            public String getIndexFieldName() {
+                return indexFieldData.getFieldName();
+            }
         }
 
         /**

server/src/main/java/org/opensearch/search/aggregations/support/ValuesSourceAggregatorFactory.java

@@ -172,4 +172,16 @@ protected abstract Aggregator doCreateInternal(
     public String getStatsSubtype() {
         return config.valueSourceType().typeName();
     }
+
+    public String getField() {
+        return config.fieldContext().field();
+    }
+
+    public String getAggregationName() {
+        return name;
+    }
+
+    public ValuesSourceConfig getConfig() {
+        return config;
+    }
 }