Support Nested Aggregations as part of Star-Tree

CHANGELOG.md

@@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added File Cache Pinning ([#17617](https://github.com/opensearch-project/OpenSearch/issues/13648))
 - Support consumer reset in Resume API for pull-based ingestion. This PR includes a breaking change for the experimental pull-based ingestion feature. ([#18332](https://github.com/opensearch-project/OpenSearch/pull/18332))
 - Add FIPS build tooling ([#4254](https://github.com/opensearch-project/security/issues/4254))
+- Support Nested Aggregations as part of Star-Tree ([#18048](https://github.com/opensearch-project/OpenSearch/pull/18048))
 - [Star-Tree] Support for date-range queries with star-tree supported aggregations ([#17855](https://github.com/opensearch-project/OpenSearch/pull/17855)
 
 ### Changed

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java

@@ -990,7 +990,6 @@ private void constructNonStarNodes(InMemoryTreeNode node, int startDocId, int en
             Long dimensionValue = getDimensionValue(i, dimensionId);
             if (Objects.equals(dimensionValue, nodeDimensionValue) == false) {
                 addChildNode(node, i, dimensionId, nodeStartDocId, nodeDimensionValue);
-
                 nodeStartDocId = i;
                 nodeDimensionValue = dimensionValue;
             }

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

@@ -10,8 +10,10 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+import org.opensearch.search.startree.filter.DimensionFilter;
 
 import java.io.IOException;
+import java.util.List;
 
 /**
  * This interface is used to pre-compute the star tree bucket collector for each segment/leaf.
@@ -29,4 +31,18 @@ StarTreeBucketCollector getStarTreeBucketCollector(
         CompositeIndexFieldInfo starTree,
         StarTreeBucketCollector parentCollector
     ) throws IOException;
+
+    /**
+     * Returns the list of dimensions filters involved in this aggregation, which are required for
+     * merging dimension filters during StarTree precomputation. This is specifically needed
+     * for nested bucket aggregations to ensure that the correct dimensions are considered when
+     * constructing or merging filters during StarTree traversal.
+     * For metric aggregations, there is no need to specify dimensions since they operate
+     * purely on values within the buckets formed by parent bucket aggregations.
+     *
+     * @return List of dimension field names involved in the aggregation.
+     */
+    default List<DimensionFilter> getDimensionFilters() {
+        return null;
+    }
 }

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -65,7 +65,7 @@
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
-import org.opensearch.search.startree.StarTreeTraversalUtil;
+import org.opensearch.search.startree.filter.DimensionFilter;
 import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
@@ -283,29 +283,24 @@ private String fetchStarTreeCalendarUnit() {
         return dimensionName;
     }
 
+    @Override
+    public List<DimensionFilter> getDimensionFilters() {
+        return StarTreeQueryHelper.collectDimensionFilters(new MatchAllFilter(fieldName, starTreeDateDimension), subAggregators);
+    }
+
     @Override
     public StarTreeBucketCollector getStarTreeBucketCollector(
         LeafReaderContext ctx,
         CompositeIndexFieldInfo starTree,
         StarTreeBucketCollector parentCollector
     ) throws IOException {
-        assert parentCollector == null;
         StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
         SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
             .getDimensionValuesIterator(starTreeDateDimension);
         SortedNumericStarTreeValuesIterator docCountsIterator = StarTreeQueryHelper.getDocCountsIterator(starTreeValues, starTree);
-
         return new StarTreeBucketCollector(
             starTreeValues,
-            StarTreeTraversalUtil.getStarTreeResult(
-                starTreeValues,
-                StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
-                    context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
-                    fieldName,
-                    List.of(new MatchAllFilter(fieldName, starTreeDateDimension))
-                ),
-                context
-            )
+            parentCollector == null ? StarTreeQueryHelper.getStarTreeResult(starTreeValues, context, getDimensionFilters()) : null
         ) {
             @Override
             public void setSubCollectors() throws IOException {

server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -70,7 +70,7 @@
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
-import org.opensearch.search.startree.StarTreeTraversalUtil;
+import org.opensearch.search.startree.filter.DimensionFilter;
 import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
@@ -380,26 +380,23 @@ private void preComputeWithStarTree(LeafReaderContext ctx, CompositeIndexFieldIn
         }
     }
 
+    @Override
+    public List<DimensionFilter> getDimensionFilters() {
+        return StarTreeQueryHelper.collectDimensionFilters(new MatchAllFilter(fieldName), subAggregators);
+    }
+
     @Override
     public StarTreeBucketCollector getStarTreeBucketCollector(
         LeafReaderContext ctx,
         CompositeIndexFieldInfo starTree,
         StarTreeBucketCollector parentCollector
     ) throws IOException {
-        assert parentCollector == null;
         StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
+
         // TODO: Evaluate optimizing StarTree traversal filter with specific ranges instead of MATCH_ALL_DEFAULT
         return new StarTreeBucketCollector(
             starTreeValues,
-            StarTreeTraversalUtil.getStarTreeResult(
-                starTreeValues,
-                StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
-                    context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
-                    fieldName,
-                    List.of(new MatchAllFilter(fieldName))
-                ),
-                context
-            )
+            parentCollector == null ? StarTreeQueryHelper.getStarTreeResult(starTreeValues, context, getDimensionFilters()) : null
         ) {
             @Override
             public void setSubCollectors() throws IOException {

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -75,7 +75,7 @@
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
-import org.opensearch.search.startree.StarTreeTraversalUtil;
+import org.opensearch.search.startree.filter.DimensionFilter;
 import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
@@ -103,11 +103,11 @@
     private final long valueCount;
     protected final String fieldName;
     private Weight weight;
-    protected final CollectionStrategy collectionStrategy;
+    protected CollectionStrategy collectionStrategy;
     private final SetOnce<SortedSetDocValues> dvs = new SetOnce<>();
     protected int segmentsWithSingleValuedOrds = 0;
     protected int segmentsWithMultiValuedOrds = 0;
-    LongUnaryOperator globalOperator;
+    protected CardinalityUpperBound cardinalityUpperBound;
 
     public GlobalOrdinalsStringTermsAggregator(
         String name,
@@ -127,6 +127,7 @@
         Map<String, Object> metadata
     ) throws IOException {
         super(name, factories, context, parent, order, format, bucketCountThresholds, collectionMode, showTermDocCountError, metadata);
+        this.cardinalityUpperBound = cardinality;
         this.resultStrategy = resultStrategy.apply(this); // ResultStrategy needs a reference to the Aggregator to do its job.
         this.valuesSource = valuesSource;
         final IndexReader reader = context.searcher().getIndexReader();
@@ -240,7 +241,6 @@
     protected boolean tryStarTreePrecompute(LeafReaderContext ctx) throws IOException {
         CompositeIndexFieldInfo supportedStarTree = StarTreeQueryHelper.getSupportedStarTree(this.context.getQueryShardContext());
         if (supportedStarTree != null) {
-            globalOperator = valuesSource.globalOrdinalsMapping(ctx);
             StarTreeBucketCollector starTreeBucketCollector = getStarTreeBucketCollector(ctx, supportedStarTree, null);
             StarTreeQueryHelper.preComputeBucketsWithStarTree(starTreeBucketCollector);
             return true;
@@ -252,7 +252,6 @@
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
         SortedSetDocValues globalOrds = this.getGlobalOrds(ctx);
         collectionStrategy.globalOrdsReady(globalOrds);
-
         SortedDocValues singleValues = DocValues.unwrapSingleton(globalOrds);
         if (singleValues != null) {
             segmentsWithSingleValuedOrds++;
@@ -324,29 +323,36 @@
         });
     }
 
+    @Override
+    public List<DimensionFilter> getDimensionFilters() {
+        return StarTreeQueryHelper.collectDimensionFilters(new MatchAllFilter(fieldName), subAggregators);
+    }
+
     public StarTreeBucketCollector getStarTreeBucketCollector(
         LeafReaderContext ctx,
         CompositeIndexFieldInfo starTree,
         StarTreeBucketCollector parent
     ) throws IOException {
-        assert parent == null;
         StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
         SortedSetStarTreeValuesIterator valuesIterator = (SortedSetStarTreeValuesIterator) starTreeValues.getDimensionValuesIterator(
             fieldName
         );
         SortedNumericStarTreeValuesIterator docCountsIterator = StarTreeQueryHelper.getDocCountsIterator(starTreeValues, starTree);
 
+        /* For nested aggregations, we require the RemapGlobalOrdsStarTree strategy to properly
+        handle global ordinal remapping. This check ensures we don't reinitialize the
+        collectionStrategy again if it's already correctly set. */
+        if (parent != null && !(collectionStrategy instanceof RemapGlobalOrdsStarTree)) {
+            collectionStrategy.close();
+            collectionStrategy = new RemapGlobalOrdsStarTree(this.cardinalityUpperBound);
+            SortedSetDocValues globalOrds = valuesSource.globalOrdinalsValues(ctx);
+            collectionStrategy.globalOrdsReady(globalOrds);
+        }
+
+        LongUnaryOperator globalOperator = valuesSource.globalOrdinalsMapping(ctx);
         return new StarTreeBucketCollector(
             starTreeValues,
-            StarTreeTraversalUtil.getStarTreeResult(
-                starTreeValues,
-                StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
-                    context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
-                    fieldName,
-                    List.of(new MatchAllFilter(fieldName))
-                ),
-                context
-            )
+            parent == null ? StarTreeQueryHelper.getStarTreeResult(starTreeValues, context, getDimensionFilters()) : null
         ) {
             @Override
             public void setSubCollectors() throws IOException {
@@ -363,11 +369,14 @@
                 for (int i = 0, count = valuesIterator.docValueCount(); i < count; i++) {
                     long dimensionValue = valuesIterator.value();
                     long ord = globalOperator.applyAsLong(dimensionValue);
-
                     if (docCountsIterator.advanceExact(starTreeEntry)) {
                         long metricValue = docCountsIterator.nextValue();
-                        long bucketOrd = collectionStrategy.getOrAddBucketOrd(owningBucketOrd, ord);
-                        collectStarTreeBucket(this, metricValue, bucketOrd, starTreeEntry);
+                        if (collectionStrategy instanceof RemapGlobalOrdsStarTree rangeSTGlobalOrds) {
+                            rangeSTGlobalOrds.collectGlobalOrdsForStarTree(owningBucketOrd, starTreeEntry, ord, this, metricValue);
+                        } else {
+                            long bucketOrd = collectionStrategy.getOrAddBucketOrd(owningBucketOrd, ord);
+                            collectStarTreeBucket(this, metricValue, bucketOrd, starTreeEntry);
+                        }
                     }
                 }
             }
@@ -710,7 +719,7 @@
      * less when collecting only a few.
      */
     private class RemapGlobalOrds extends CollectionStrategy {
-        private final LongKeyedBucketOrds bucketOrds;
+        protected final LongKeyedBucketOrds bucketOrds;
 
         private RemapGlobalOrds(CardinalityUpperBound cardinality) {
             bucketOrds = LongKeyedBucketOrds.build(context.bigArrays(), cardinality);
@@ -791,6 +800,28 @@
         }
     }
 
+    private class RemapGlobalOrdsStarTree extends RemapGlobalOrds {
+        private RemapGlobalOrdsStarTree(CardinalityUpperBound cardinality) {
+            super(cardinality);
+        }
+
+        @Override
+        String describe() {
+            return "remapStarTree";
+        }
+
+        void collectGlobalOrdsForStarTree(
+            long owningBucketOrd,
+            int starTreeEntry,
+            long globalOrd,
+            StarTreeBucketCollector collector,
+            long docCount
+        ) throws IOException {
+            long bucketOrd = bucketOrds.add(owningBucketOrd, globalOrd);
+            collectStarTreeBucket(collector, docCount, bucketOrd, starTreeEntry);
+        }
+    }
+
     /**
      * Strategy for building results.
      */

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java

@@ -67,7 +67,7 @@
 import org.opensearch.search.internal.ContextIndexSearcher;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
-import org.opensearch.search.startree.StarTreeTraversalUtil;
+import org.opensearch.search.startree.filter.DimensionFilter;
 import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
@@ -136,7 +136,6 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCol
             public void collect(int doc, long owningBucketOrd) throws IOException {
                 if (values.advanceExact(doc)) {
                     int valuesCount = values.docValueCount();
-
                     long previous = Long.MAX_VALUE;
                     for (int i = 0; i < valuesCount; ++i) {
                         long val = values.nextValue();
@@ -169,28 +168,23 @@ protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws
         return false;
     }
 
+    @Override
+    public List<DimensionFilter> getDimensionFilters() {
+        return StarTreeQueryHelper.collectDimensionFilters(new MatchAllFilter(fieldName), subAggregators);
+    }
+
     public StarTreeBucketCollector getStarTreeBucketCollector(
         LeafReaderContext ctx,
         CompositeIndexFieldInfo starTree,
         StarTreeBucketCollector parent
     ) throws IOException {
-        assert parent == null;
         StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
         SortedNumericStarTreeValuesIterator valuesIterator = (SortedNumericStarTreeValuesIterator) starTreeValues
             .getDimensionValuesIterator(fieldName);
         SortedNumericStarTreeValuesIterator docCountsIterator = StarTreeQueryHelper.getDocCountsIterator(starTreeValues, starTree);
-
         return new StarTreeBucketCollector(
             starTreeValues,
-            StarTreeTraversalUtil.getStarTreeResult(
-                starTreeValues,
-                StarTreeQueryHelper.mergeDimensionFilterIfNotExists(
-                    context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(),
-                    fieldName,
-                    List.of(new MatchAllFilter(fieldName))
-                ),
-                context
-            )
+            parent == null ? StarTreeQueryHelper.getStarTreeResult(starTreeValues, context, getDimensionFilters()) : null
         ) {
             @Override
             public void setSubCollectors() throws IOException {
@@ -215,7 +209,6 @@ public void collectStarTreeEntry(int starTreeEntry, long owningBucketOrd) throws
                 }
 
                 for (int i = 0, count = valuesIterator.entryValueCount(); i < count; i++) {
-
                     if (docCountsIterator.advanceExact(starTreeEntry)) {
                         long metricValue = docCountsIterator.nextValue();
                         long bucketOrd = bucketOrds.add(owningBucketOrd, dimensionValue);
@@ -300,7 +293,6 @@ private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws
             }
 
             buildSubAggs(topBucketsPerOrd);
-
             InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
             for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
                 result[ordIdx] = buildResult(owningBucketOrds[ordIdx], otherDocCounts[ordIdx], topBucketsPerOrd[ordIdx]);