multi terms aggregation changes (opensearch-project#19284)

Signed-off-by: Sandesh Kumar <[email protected]>

Author: sandeshkr419

CHANGELOG.md

@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Implement GRPC Boolean query and inject registry for all internal query converters ([#19391](https://github.com/opensearch-project/OpenSearch/pull/19391))
 - Implement GRPC Script query ([#19455](https://github.com/opensearch-project/OpenSearch/pull/19455))
 - [Search Stats] Add search & star-tree search query failure count metrics ([#19210](https://github.com/opensearch-project/OpenSearch/issues/19210))
+- [Star-tree] Support for multi-terms aggregation ([#18398](https://github.com/opensearch-project/OpenSearch/issues/18398))
 
 ### Changed
 - Refactor `if-else` chains to use `Java 17 pattern matching switch expressions`(([#18965](https://github.com/opensearch-project/OpenSearch/pull/18965))

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java

@@ -35,7 +35,7 @@ public long nextValue() throws IOException {
         return ((SortedNumericDocValues) docIdSetIterator).nextValue();
     }
 
-    public int entryValueCount() throws IOException {
+    public int entryValueCount() {
         return ((SortedNumericDocValues) docIdSetIterator).docValueCount();
     }
 

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/SortedSetStarTreeValuesIterator.java

@@ -44,7 +44,7 @@ public long nextOrd() throws IOException {
         return ((SortedSetDocValues) docIdSetIterator).nextOrd();
     }
 
-    public int docValueCount() {
+    public int entryValueCount() {
         return ((SortedSetDocValues) docIdSetIterator).docValueCount();
     }
 

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java

@@ -42,6 +42,8 @@ public int advance(int target) throws IOException {
         return docIdSetIterator.advance(target);
     }
 
+    public abstract int entryValueCount();
+
     public long cost() {
         return docIdSetIterator.cost();
     }

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

@@ -366,7 +366,7 @@ public void collectStarTreeEntry(int starTreeEntry, long owningBucketOrd) throws
                 if (valuesIterator.advanceExact(starTreeEntry) == false) {
                     return;
                 }
-                for (int i = 0, count = valuesIterator.docValueCount(); i < count; i++) {
+                for (int i = 0, count = valuesIterator.entryValueCount(); i < count; i++) {
                     long dimensionValue = valuesIterator.value();
                     long ord = globalOperator.applyAsLong(dimensionValue);
                     if (docCountsIterator.advanceExact(starTreeEntry)) {

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

@@ -48,6 +48,7 @@ public class MultiTermsAggregationFactory extends AggregatorFactory {
     private final Aggregator.SubAggCollectionMode collectMode;
     private final TermsAggregator.BucketCountThresholds bucketCountThresholds;
     private final boolean showTermDocCountError;
+    private final List<String> requestFields;
 
     public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
         builder.register(REGISTRY_KEY, List.of(CoreValuesSourceType.BYTES, CoreValuesSourceType.IP), config -> {
@@ -116,6 +117,7 @@ public MultiTermsAggregationFactory(
                 )
             )
             .collect(Collectors.toList());
+        this.requestFields = multiTermConfigs.stream().map(MultiTermsValuesSourceConfig::getFieldName).toList();
         this.formats = this.configs.stream().map(c -> c.v1().format()).collect(Collectors.toList());
         this.order = order;
         this.collectMode = collectMode;
@@ -138,14 +140,17 @@ protected Aggregator createInternal(
             // counting
             bucketCountThresholds.setShardSize(BucketUtils.suggestShardSideQueueSize(bucketCountThresholds.getRequiredSize()));
         }
+        // TODO: Optimize passing too many value source config derived objects to aggregator
         bucketCountThresholds.ensureValidity();
         return new MultiTermsAggregator(
             name,
             factories,
             showTermDocCountError,
+            configs.stream().map(config -> config.v1().getValuesSource()).toList(),
             configs.stream()
                 .map(config -> queryShardContext.getValuesSourceRegistry().getAggregator(REGISTRY_KEY, config.v1()).build(config))
                 .collect(Collectors.toList()),
+            this.getRequestFields(),
             configs.stream().map(c -> c.v1().format()).collect(Collectors.toList()),
             order,
             collectMode,
@@ -157,6 +162,10 @@ protected Aggregator createInternal(
         );
     }
 
+    public List<String> getRequestFields() {
+        return requestFields;
+    }
+
     @Override
     protected boolean supportsConcurrentSegmentSearch() {
         return true;

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

@@ -23,8 +23,13 @@
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.index.codec.composite.CompositeIndexFieldInfo;
+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.SortedBinaryDocValues;
 import org.opensearch.index.fielddata.SortedNumericDoubleValues;
+import org.opensearch.index.mapper.NumberFieldMapper;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.aggregations.Aggregator;
 import org.opensearch.search.aggregations.AggregatorFactories;
@@ -33,12 +38,17 @@
 import org.opensearch.search.aggregations.InternalAggregation;
 import org.opensearch.search.aggregations.InternalOrder;
 import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.aggregations.StarTreeBucketCollector;
+import org.opensearch.search.aggregations.StarTreePreComputeCollector;
 import org.opensearch.search.aggregations.bucket.BucketsAggregator;
 import org.opensearch.search.aggregations.bucket.DeferableBucketAggregator;
 import org.opensearch.search.aggregations.bucket.LocalBucketCountThresholds;
 import org.opensearch.search.aggregations.support.AggregationPath;
 import org.opensearch.search.aggregations.support.ValuesSource;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.startree.StarTreeQueryHelper;
+import org.opensearch.search.startree.filter.DimensionFilter;
+import org.opensearch.search.startree.filter.MatchAllFilter;
 
 import java.io.IOException;
 import java.math.BigInteger;
@@ -50,32 +60,38 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import static org.opensearch.search.aggregations.InternalOrder.isKeyOrder;
 import static org.opensearch.search.aggregations.bucket.terms.TermsAggregator.descendsFromNestedAggregator;
+import static org.opensearch.search.startree.StarTreeQueryHelper.getSupportedStarTree;
 
 /**
  * An aggregator that aggregate with multi_terms.
  *
  * @opensearch.internal
  */
-public class MultiTermsAggregator extends DeferableBucketAggregator {
+public class MultiTermsAggregator extends DeferableBucketAggregator implements StarTreePreComputeCollector {
 
     private final BytesKeyedBucketOrds bucketOrds;
     private final MultiTermsValuesSource multiTermsValue;
     private final boolean showTermDocCountError;
     private final List<DocValueFormat> formats;
+    private final List<String> fields;
     private final TermsAggregator.BucketCountThresholds bucketCountThresholds;
     private final BucketOrder order;
     private final Comparator<InternalMultiTerms.Bucket> partiallyBuiltBucketComparator;
     private final SubAggCollectionMode collectMode;
     private final Set<Aggregator> aggsUsedForSorting = new HashSet<>();
+    private final BytesStreamOutput starTreeScratch = new BytesStreamOutput();
 
     public MultiTermsAggregator(
         String name,
         AggregatorFactories factories,
         boolean showTermDocCountError,
+        List<ValuesSource> rawValuesSources,
         List<InternalValuesSource> internalValuesSources,
+        List<String> fields,
         List<DocValueFormat> formats,
         BucketOrder order,
         SubAggCollectionMode collectMode,
@@ -87,7 +103,7 @@ public MultiTermsAggregator(
     ) throws IOException {
         super(name, factories, context, parent, metadata);
         this.bucketOrds = BytesKeyedBucketOrds.build(context.bigArrays(), cardinality);
-        this.multiTermsValue = new MultiTermsValuesSource(internalValuesSources);
+        this.multiTermsValue = new MultiTermsValuesSource(rawValuesSources, internalValuesSources);
         this.showTermDocCountError = showTermDocCountError;
         this.formats = formats;
         this.bucketCountThresholds = bucketCountThresholds;
@@ -104,12 +120,12 @@ public MultiTermsAggregator(
         } else {
             this.collectMode = collectMode;
         }
+        this.fields = fields;
         // Don't defer any child agg if we are dependent on it for pruning results
         if (order instanceof InternalOrder.Aggregation) {
             AggregationPath path = ((InternalOrder.Aggregation) order).path();
             aggsUsedForSorting.add(path.resolveTopmostAggregator(this));
-        } else if (order instanceof InternalOrder.CompoundOrder) {
-            InternalOrder.CompoundOrder compoundOrder = (InternalOrder.CompoundOrder) order;
+        } else if (order instanceof InternalOrder.CompoundOrder compoundOrder) {
             for (BucketOrder orderElement : compoundOrder.orderElements()) {
                 if (orderElement instanceof InternalOrder.Aggregation) {
                     AggregationPath path = ((InternalOrder.Aggregation) orderElement).path();
@@ -226,6 +242,142 @@ public void collect(int doc, long owningBucketOrd) throws IOException {
         };
     }
 
+    @Override
+    protected boolean tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws IOException {
+        CompositeIndexFieldInfo supportedStarTree = getSupportedStarTree(this.context.getQueryShardContext());
+        if (supportedStarTree != null) {
+            preComputeWithStarTree(ctx, supportedStarTree);
+            return true;
+        }
+        return false;
+    }
+
+    private void preComputeWithStarTree(LeafReaderContext ctx, CompositeIndexFieldInfo starTree) throws IOException {
+        StarTreeBucketCollector starTreeBucketCollector = getStarTreeBucketCollector(ctx, starTree, null);
+        StarTreeQueryHelper.preComputeBucketsWithStarTree(starTreeBucketCollector);
+    }
+
+    /**
+     * Creates a {@link StarTreeBucketCollector} for pre-aggregating with a star-tree index.
+     * This collector generates the cartesian product of dimension values within a single star-tree entry
+     * to form the composite keys for the multi-terms aggregation.
+     */
+    public StarTreeBucketCollector getStarTreeBucketCollector(
+        LeafReaderContext ctx,
+        CompositeIndexFieldInfo starTree,
+        StarTreeBucketCollector parent
+    ) throws IOException {
+        StarTreeValues starTreeValues = StarTreeQueryHelper.getStarTreeValues(ctx, starTree);
+        assert starTreeValues != null;
+        SortedNumericStarTreeValuesIterator docCountsIterator = StarTreeQueryHelper.getDocCountsIterator(starTreeValues, starTree);
+
+        // Get an iterator for each field (dimension) in the multi-terms aggregation.
+        final List<StarTreeValuesIterator> dimensionIterators = new ArrayList<>();
+        // We also need a way to convert the raw long values from the iterators into the correct TermValue type.
+        final List<Function<Long, TermValue<?>>> termValueBuilders = new ArrayList<>();
+
+        for (int i = 0; i < fields.size(); i++) {
+            String fieldName = fields.get(i);
+            dimensionIterators.add(starTreeValues.getDimensionValuesIterator(fieldName));
+            ValuesSource vs = multiTermsValue.rawValueSources.get(i);
+
+            if (vs instanceof ValuesSource.Bytes.WithOrdinals vsBytes) {
+                termValueBuilders.add(ord -> {
+                    try {
+                        return TermValue.of(vsBytes.globalOrdinalsValues(ctx).lookupOrd(ord));
+                    } catch (IOException e) {
+                        throw new RuntimeException(e);
+                    }
+
+                });
+            } else if (vs instanceof ValuesSource.Numeric numericSource) {
+                if (numericSource.isFloatingPoint()) {
+                    NumberFieldMapper.NumberFieldType numberFieldType = ((NumberFieldMapper.NumberFieldType) context.mapperService()
+                        .fieldType(fieldName));
+                    termValueBuilders.add(val -> TermValue.of(numberFieldType.toDoubleValue(val)));
+                } else {
+                    termValueBuilders.add(TermValue::of);
+                }
+            } else {
+                throw new IllegalStateException("Unsupported ValuesSource type for star-tree: " + vs.getClass().getName());
+            }
+
+        }
+
+        return new StarTreeBucketCollector(
+            starTreeValues,
+            parent == null ? StarTreeQueryHelper.getStarTreeResult(starTreeValues, context, getDimensionFilters()) : null
+        ) {
+            @Override
+            public void setSubCollectors() throws IOException {
+                for (Aggregator aggregator : subAggregators) {
+                    if (aggregator instanceof StarTreePreComputeCollector collector) {
+                        this.subCollectors.add(collector.getStarTreeBucketCollector(ctx, starTree, this));
+                    }
+                }
+            }
+
+            @Override
+            public void collectStarTreeEntry(int starTreeEntry, long owningBucketOrd) throws IOException {
+                if (docCountsIterator.advanceExact(starTreeEntry) == false) {
+                    return; // No documents in this star-tree entry.
+                }
+                long docCountMetric = docCountsIterator.nextValue();
+
+                List<List<TermValue<?>>> collectedValues = new ArrayList<>();
+                for (int i = 0; i < dimensionIterators.size(); i++) {
+                    StarTreeValuesIterator dimIterator = dimensionIterators.get(i);
+                    if (!dimIterator.advanceExact(starTreeEntry)) {
+                        // If any dimension is missing for this entry, the cartesian product is empty.
+                        return;
+                    }
+
+                    List<TermValue<?>> valuesForDim = new ArrayList<>();
+                    Function<Long, TermValue<?>> builder = termValueBuilders.get(i);
+                    for (int j = 0; j < dimIterator.entryValueCount(); j++) {
+                        valuesForDim.add(builder.apply(dimIterator.value()));
+                    }
+                    collectedValues.add(valuesForDim);
+                }
+
+                starTreeScratch.seek(0);
+                starTreeScratch.writeVInt(dimensionIterators.size());
+                generateAndCollectFromStarTree(collectedValues, 0, owningBucketOrd, starTreeEntry, docCountMetric);
+            }
+
+            private void generateAndCollectFromStarTree(
+                List<List<TermValue<?>>> collectedValues,
+                int index,
+                long owningBucketOrd,
+                int starTreeEntry,
+                long docCountMetric
+            ) throws IOException {
+                if (index == collectedValues.size()) {
+                    // A full composite key is in the buffer, add it to bucketOrds.
+                    long bucketOrd = bucketOrds.add(owningBucketOrd, starTreeScratch.bytes().toBytesRef());
+                    collectStarTreeBucket(this, docCountMetric, bucketOrd, starTreeEntry);
+                    return;
+                }
+
+                long position = starTreeScratch.position();
+                List<TermValue<?>> values = collectedValues.get(index);
+                for (TermValue<?> value : values) {
+                    value.writeTo(starTreeScratch);
+                    generateAndCollectFromStarTree(collectedValues, index + 1, owningBucketOrd, starTreeEntry, docCountMetric);
+                    starTreeScratch.seek(position);
+                }
+            }
+        };
+    }
+
+    @Override
+    public List<DimensionFilter> getDimensionFilters() {
+        return StarTreeQueryHelper.collectDimensionFilters(
+            fields.stream().map(a -> (DimensionFilter) new MatchAllFilter(a)).toList(),
+            subAggregators
+        );
+    }
+
     @Override
     protected void doClose() {
         Releasables.close(bucketOrds, multiTermsValue);
@@ -347,10 +499,12 @@ public static TermValue<Double> of(Double value) {
      * @opensearch.internal
      */
     static class MultiTermsValuesSource implements Releasable {
+        private final List<ValuesSource> rawValueSources;
         private final List<InternalValuesSource> valuesSources;
         private final BytesStreamOutput scratch = new BytesStreamOutput();
 
-        public MultiTermsValuesSource(List<InternalValuesSource> valuesSources) {
+        public MultiTermsValuesSource(List<ValuesSource> rawValueSources, List<InternalValuesSource> valuesSources) {
+            this.rawValueSources = rawValueSources;
             this.valuesSources = valuesSources;
         }
 
@@ -416,8 +570,7 @@ private void generateAndCollectCompositeKeys(
                     int numIterations = values.size();
                     // For each loop is not done to reduce the allocations done for Iterator objects
                     // once for every field in every doc.
-                    for (int i = 0; i < numIterations; i++) {
-                        TermValue<?> value = values.get(i);
+                    for (TermValue<?> value : values) {
                         value.writeTo(scratch); // encode the value
                         generateAndCollectCompositeKeys(collectedValues, index + 1, owningBucketOrd, doc); // dfs
                         scratch.seek(position); // backtrack