diff --git a/CHANGELOG.md b/CHANGELOG.md index f9c2f920267f1..86950c46fd39f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java index 935c490b5a4dc..e6e8332b0723f 100644 --- a/server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java +++ b/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; } diff --git a/server/src/main/java/org/opensearch/search/aggregations/StarTreePreComputeCollector.java b/server/src/main/java/org/opensearch/search/aggregations/StarTreePreComputeCollector.java index c2f2017997c4d..42f5253d74c22 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/StarTreePreComputeCollector.java +++ b/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 getDimensionFilters() { + return null; + } } diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java index 383b5e1ecf0b8..f734646a7b8dd 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java +++ b/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 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 { diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java index b3d5254d7a186..18ca3bb047b2f 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java +++ b/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 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 { diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java index b5972b2145c54..95f8a249a7e55 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java +++ b/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 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr private final long valueCount; protected final String fieldName; private Weight weight; - protected final CollectionStrategy collectionStrategy; + protected CollectionStrategy collectionStrategy; private final SetOnce dvs = new SetOnce<>(); protected int segmentsWithSingleValuedOrds = 0; protected int segmentsWithMultiValuedOrds = 0; - LongUnaryOperator globalOperator; + protected CardinalityUpperBound cardinalityUpperBound; public GlobalOrdinalsStringTermsAggregator( String name, @@ -127,6 +127,7 @@ public GlobalOrdinalsStringTermsAggregator( Map 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 tryPrecomputeAggregationForLeaf(LeafReaderContext ctx) throws 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 @@ protected boolean tryStarTreePrecompute(LeafReaderContext ctx) throws IOExceptio 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 @@ public void collect(int doc, long owningBucketOrd) throws IOException { }); } + @Override + public List 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 @@ public void collectStarTreeEntry(int starTreeEntry, long owningBucketOrd) throws 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 @@ public void close() {} * 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 @@ public void close() { } } + 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. */ diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java index bf89a5c4ccb8e..7ca69553158a8 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java +++ b/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 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]); diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java index 53a5a7e007417..9128087a6fb4a 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryContext.java @@ -10,6 +10,7 @@ import org.apache.lucene.util.FixedBitSet; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; import org.opensearch.index.compositeindex.datacube.DateDimension; import org.opensearch.index.compositeindex.datacube.Dimension; @@ -107,23 +108,7 @@ public void maybeSetCachedNodeIdsForSegment(int key, FixedBitSet values) { public boolean consolidateAllFilters(SearchContext context) { // Validate the fields and metrics required by aggregations are supported in star tree for (AggregatorFactory aggregatorFactory : context.aggregations().factories().getFactories()) { - // first check for aggregation is a metric aggregation - if (validateStarTreeMetricSupport(compositeMappedFieldType, aggregatorFactory)) { - continue; - } - - // if not a metric aggregation, check for applicable date histogram shape - if (validateDateHistogramSupport(compositeMappedFieldType, aggregatorFactory)) { - continue; - } - - // validation for terms aggregation - if (validateKeywordTermsAggregationSupport(compositeMappedFieldType, aggregatorFactory)) { - continue; - } - - // validation for range aggregation - if (validateRangeAggregationSupport(compositeMappedFieldType, aggregatorFactory)) { + if (validateNestedAggregationStructure(compositeMappedFieldType, aggregatorFactory)) { continue; } // invalid query shape @@ -181,12 +166,6 @@ private static boolean validateKeywordTermsAggregationSupport( return false; } - // Validate all sub-factories - for (AggregatorFactory subFactory : aggregatorFactory.getSubFactories().getFactories()) { - if (!validateStarTreeMetricSupport(compositeIndexFieldInfo, subFactory)) { - return false; - } - } return true; } @@ -208,12 +187,6 @@ private static boolean validateRangeAggregationSupport( return false; } - // Validate all sub-factories - for (AggregatorFactory subFactory : aggregatorFactory.getSubFactories().getFactories()) { - if (!validateStarTreeMetricSupport(compositeIndexFieldInfo, subFactory)) { - return false; - } - } return true; } @@ -270,12 +243,46 @@ private static boolean validateDateHistogramSupport( return false; } - // Validate all sub-factories + return true; + } + + private static boolean validateNestedAggregationStructure( + CompositeDataCubeFieldType compositeIndexFieldInfo, + AggregatorFactory aggregatorFactory + ) { + boolean isValid; + boolean isFeatureFlagEnabled = FeatureFlags.isEnabled(FeatureFlags.STAR_TREE_INDEX_SETTING); + + switch (aggregatorFactory) { + case TermsAggregatorFactory termsAggregatorFactory -> isValid = validateKeywordTermsAggregationSupport( + compositeIndexFieldInfo, + termsAggregatorFactory + ) && isFeatureFlagEnabled; + case DateHistogramAggregatorFactory dateHistogramAggregatorFactory -> isValid = validateDateHistogramSupport( + compositeIndexFieldInfo, + dateHistogramAggregatorFactory + ); + case RangeAggregatorFactory rangeAggregatorFactory -> isValid = validateRangeAggregationSupport( + compositeIndexFieldInfo, + rangeAggregatorFactory + ) && isFeatureFlagEnabled; + case MetricAggregatorFactory metricAggregatorFactory -> { + isValid = validateStarTreeMetricSupport(compositeIndexFieldInfo, metricAggregatorFactory); + return isValid && metricAggregatorFactory.getSubFactories().getFactories().length == 0; + } + case null, default -> { + return false; + } + } + + if (isValid == false) return false; + for (AggregatorFactory subFactory : aggregatorFactory.getSubFactories().getFactories()) { - if (!validateStarTreeMetricSupport(compositeIndexFieldInfo, subFactory)) { + if (!validateNestedAggregationStructure(compositeIndexFieldInfo, subFactory)) { return false; } } + return true; } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java index 68a613a373edf..59ebd9afe1dc7 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeQueryHelper.java @@ -22,13 +22,17 @@ import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator; import org.opensearch.index.mapper.DocCountFieldMapper; import org.opensearch.index.query.QueryShardContext; +import org.opensearch.search.aggregations.Aggregator; import org.opensearch.search.aggregations.StarTreeBucketCollector; +import org.opensearch.search.aggregations.StarTreePreComputeCollector; import org.opensearch.search.aggregations.support.ValuesSource; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.startree.filter.DimensionFilter; import org.opensearch.search.startree.filter.StarTreeFilter; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -226,4 +230,35 @@ public static StarTreeFilter mergeDimensionFilterIfNotExists( return new StarTreeFilter(dimensionFilterMap); } + public static FixedBitSet getStarTreeResult( + StarTreeValues starTreeValues, + SearchContext context, + List dimensionFiltersToMerge + ) throws IOException { + StarTreeFilter starTreeFilter = context.getQueryShardContext().getStarTreeQueryContext().getBaseQueryStarTreeFilter(); + for (DimensionFilter dimensionFilter : dimensionFiltersToMerge) { + starTreeFilter = StarTreeQueryHelper.mergeDimensionFilterIfNotExists( + starTreeFilter, + dimensionFilter.getMatchingDimension(), + List.of(dimensionFilter) + ); + } + + return StarTreeTraversalUtil.getStarTreeResult(starTreeValues, starTreeFilter, context); + } + + public static List collectDimensionFilters(DimensionFilter initialDimensionFilter, Aggregator[] subAggregators) { + List dimensionFiltersToMerge = new ArrayList<>(); + dimensionFiltersToMerge.add(initialDimensionFilter); + + for (Aggregator subAgg : subAggregators) { + if (subAgg instanceof StarTreePreComputeCollector collector) { + List childFilters = collector.getDimensionFilters(); + dimensionFiltersToMerge.addAll(childFilters != null ? childFilters : Collections.emptyList()); + } + } + + return dimensionFiltersToMerge; + } + } diff --git a/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java b/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java index 92935c5243ad1..661cc6fedbf2d 100644 --- a/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java +++ b/server/src/main/java/org/opensearch/search/startree/StarTreeTraversalUtil.java @@ -21,6 +21,7 @@ import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.startree.filter.DimensionFilter; +import org.opensearch.search.startree.filter.MatchAllFilter; import org.opensearch.search.startree.filter.StarTreeFilter; import java.io.IOException; @@ -93,6 +94,17 @@ public static FixedBitSet getStarTreeResult(StarTreeValues starTreeValues, StarT // Clear the temporary bit set before reuse tempBitSet.clear(0, starTreeResult.maxMatchedDoc + 1); + // Skip filtering if a MatchAllFilter is present for this dimension, since it implies all values match and no further filtering + // is needed + boolean isMatchAllFilterPresent = false; + for (DimensionFilter dimensionFilter : dimensionFilters) { + if (dimensionFilter instanceof MatchAllFilter) { + isMatchAllFilterPresent = true; + break; + } + } + if (isMatchAllFilterPresent) continue; + if (bitSet.length() > 0) { // Iterate over the current set of matched document IDs for (int entryId = bitSet.nextSetBit(0); entryId != DocIdSetIterator.NO_MORE_DOCS; entryId = (entryId + 1 < bitSet.length()) diff --git a/server/src/main/java/org/opensearch/search/startree/filter/MatchAllFilter.java b/server/src/main/java/org/opensearch/search/startree/filter/MatchAllFilter.java index 5fde506993116..9ff2baed4dc2d 100644 --- a/server/src/main/java/org/opensearch/search/startree/filter/MatchAllFilter.java +++ b/server/src/main/java/org/opensearch/search/startree/filter/MatchAllFilter.java @@ -46,7 +46,8 @@ public void matchStarTreeNodes(StarTreeNode parentNode, StarTreeValues starTreeV if (parentNode != null) { for (Iterator it = parentNode.getChildrenIterator(); it.hasNext();) { StarTreeNode starTreeNode = it.next(); - if (starTreeNode.getStarTreeNodeType() == StarTreeNodeType.DEFAULT.getValue()) { + if (starTreeNode.getStarTreeNodeType() == StarTreeNodeType.DEFAULT.getValue() + || starTreeNode.getStarTreeNodeType() == StarTreeNodeType.NULL.getValue()) { collector.collectStarTreeNode(starTreeNode); } } @@ -66,5 +67,4 @@ public String getDimensionName() { public String getSubDimensionName() { return subDimensionName; } - } diff --git a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java index 2e1a98767922a..d65edf52928ec 100644 --- a/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java +++ b/server/src/test/java/org/opensearch/search/SearchServiceStarTreeTests.java @@ -17,6 +17,7 @@ import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.Rounding; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.common.Strings; import org.opensearch.index.IndexService; @@ -44,6 +45,7 @@ import org.opensearch.index.shard.IndexShard; import org.opensearch.indices.IndicesService; import org.opensearch.search.aggregations.AggregationBuilders; +import org.opensearch.search.aggregations.Aggregator; import org.opensearch.search.aggregations.AggregatorFactories; import org.opensearch.search.aggregations.AggregatorFactory; import org.opensearch.search.aggregations.SearchContextAggregations; @@ -70,15 +72,19 @@ import java.time.Instant; import java.time.format.DateTimeFormatter; import java.time.temporal.ChronoUnit; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Set; import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX; +import static org.opensearch.search.aggregations.AggregationBuilders.count; import static org.opensearch.search.aggregations.AggregationBuilders.dateHistogram; import static org.opensearch.search.aggregations.AggregationBuilders.max; import static org.opensearch.search.aggregations.AggregationBuilders.medianAbsoluteDeviation; +import static org.opensearch.search.aggregations.AggregationBuilders.min; import static org.opensearch.search.aggregations.AggregationBuilders.range; import static org.opensearch.search.aggregations.AggregationBuilders.sum; import static org.opensearch.search.aggregations.AggregationBuilders.terms; @@ -94,6 +100,7 @@ */ public class SearchServiceStarTreeTests extends OpenSearchSingleNodeTestCase { + private static final String FIELD_NAME = "status"; private static final String TIMESTAMP_FIELD = "@timestamp"; private static final String STATUS = "status"; private static final String SIZE = "size"; @@ -249,6 +256,135 @@ public void testQueryParsingForMetricAggregations() throws IOException { searchContext.close(); } + public void testStarTreeNestedAggregations() throws IOException { + FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build()); + setStarTreeIndexSetting("true"); + + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1) + .put(StarTreeIndexSettings.IS_COMPOSITE_INDEX_SETTING.getKey(), true) + .put(IndexMetadata.INDEX_APPEND_ONLY_ENABLED_SETTING.getKey(), true) + .build(); + CreateIndexRequestBuilder builder = client().admin() + .indices() + .prepareCreate("test") + .setSettings(settings) + .setMapping(NumericTermsAggregatorTests.getExpandedMapping(1, false)); + + createIndex("test", builder); + + IndicesService indicesService = getInstanceFromNode(IndicesService.class); + IndexService indexService = indicesService.indexServiceSafe(resolveIndex("test")); + IndexShard indexShard = indexService.getShard(0); + ShardSearchRequest request = new ShardSearchRequest( + OriginalIndices.NONE, + new SearchRequest().allowPartialSearchResults(true), + indexShard.shardId(), + 1, + new AliasFilter(null, Strings.EMPTY_ARRAY), + 1.0f, + -1, + null, + null + ); + + QueryBuilder baseQuery; + SearchContext searchContext = createSearchContext(indexService); + StarTreeFieldConfiguration starTreeFieldConfiguration = new StarTreeFieldConfiguration( + 1, + Collections.emptySet(), + StarTreeFieldConfiguration.StarTreeBuildMode.ON_HEAP + ); + + List>> aggregationSuppliers = getAggregationSuppliers(); + + ValuesSourceAggregationBuilder[] aggBuilders = { + sum("_sum").field(FIELD_NAME), + max("_max").field(FIELD_NAME), + min("_min").field(FIELD_NAME), + count("_count").field(FIELD_NAME), }; + + // 3-LEVELS [BUCKET -> BUCKET -> METRIC] + for (Supplier> firstSupplier : aggregationSuppliers) { + for (Supplier> secondSupplier : aggregationSuppliers) { + for (ValuesSourceAggregationBuilder metricAgg : aggBuilders) { + + ValuesSourceAggregationBuilder secondBucket = secondSupplier.get().subAggregation(metricAgg); + ValuesSourceAggregationBuilder firstBucket = firstSupplier.get().subAggregation(secondBucket); + + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().size(0) + .query(new MatchAllQueryBuilder()) + .aggregation(firstBucket); + + MetricStat stat = getMetricStatFromAgg(metricAgg); + List metrics = List.of(new Metric(FIELD_NAME, List.of(stat))); + + assertStarTreeContext( + request, + sourceBuilder, + getStarTreeQueryContext( + searchContext, + starTreeFieldConfiguration, + "startree1", + -1, + getDimensions(aggregationSuppliers.indexOf(firstSupplier), aggregationSuppliers.indexOf(secondSupplier)), + metrics, + new MatchAllQueryBuilder(), + sourceBuilder, + true + ), + -1 + ); + } + } + } + + // 4-LEVELS [BUCKET -> BUCKET -> BUCKET -> METRIC] + for (Supplier> firstSupplier : aggregationSuppliers) { + for (Supplier> secondSupplier : aggregationSuppliers) { + for (Supplier> thirdSupplier : aggregationSuppliers) { + for (ValuesSourceAggregationBuilder metricAgg : aggBuilders) { + + ValuesSourceAggregationBuilder thirdBucket = thirdSupplier.get().subAggregation(metricAgg); + ValuesSourceAggregationBuilder secondBucket = secondSupplier.get().subAggregation(thirdBucket); + ValuesSourceAggregationBuilder firstBucket = firstSupplier.get().subAggregation(secondBucket); + + SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().size(0) + .query(new MatchAllQueryBuilder()) + .aggregation(firstBucket); + + MetricStat stat = getMetricStatFromAgg(metricAgg); + List metrics = List.of(new Metric(FIELD_NAME, List.of(stat))); + + assertStarTreeContext( + request, + sourceBuilder, + getStarTreeQueryContext( + searchContext, + starTreeFieldConfiguration, + "startree1", + -1, + getDimensions( + aggregationSuppliers.indexOf(firstSupplier), + aggregationSuppliers.indexOf(secondSupplier), + aggregationSuppliers.indexOf(thirdSupplier) + ), + metrics, + new MatchAllQueryBuilder(), + sourceBuilder, + true + ), + -1 + ); + } + } + } + } + + setStarTreeIndexSetting(null); + } + /** * Test query parsing for date histogram aggregations, with/without numeric term query */ @@ -1021,6 +1157,52 @@ private StarTreeQueryContext getStarTreeQueryContext( return starTreeQueryContext; } + private static List getDimensions(int... indices) { + return Arrays.stream(indices).mapToObj(SearchServiceStarTreeTests::getDimensionByIndex).toList(); + } + + private static List>> getAggregationSuppliers() { + String TIMESTAMP_FIELD = "timestamp"; + String KEYWORD_FIELD = "clientip"; + String SIZE = "size"; + String STATUS = "status"; + + return List.of( + () -> terms("term_size").field(SIZE), + () -> terms("term_status").field(STATUS), + () -> dateHistogram("by_day").field(TIMESTAMP_FIELD).calendarInterval(DateHistogramInterval.DAY), + () -> range("range").field(STATUS).addRange(0, 10), + () -> terms("term_keyword").field(KEYWORD_FIELD).collectMode(Aggregator.SubAggCollectionMode.BREADTH_FIRST) + ); + } + + private static Dimension getDimensionByIndex(int index) { + String TIMESTAMP_FIELD = "timestamp"; + String KEYWORD_FIELD = "clientip"; + String SIZE = "size"; + String STATUS = "status"; + + return switch (index) { + case 0 -> new NumericDimension(SIZE); + case 2 -> new DateDimension( + TIMESTAMP_FIELD, + List.of(new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH)), + DateFieldMapper.Resolution.MILLISECONDS + ); + case 4 -> new OrdinalDimension(KEYWORD_FIELD); + default -> new NumericDimension(STATUS); + }; + } + + private MetricStat getMetricStatFromAgg(ValuesSourceAggregationBuilder agg) { + String name = agg.getName(); + if (name.contains("sum")) return MetricStat.SUM; + else if (name.contains("max")) return MetricStat.MAX; + else if (name.contains("min")) return MetricStat.MIN; + else if (name.contains("count")) return MetricStat.VALUE_COUNT; + throw new IllegalArgumentException("Unknown metric aggregation: " + name); + } + public void indexRandomDocuments(int totalDocs) throws IOException { long nowMillis = Instant.now().toEpochMilli(); diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/KeywordTermsAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/KeywordTermsAggregatorTests.java index d0d27d51535ce..011c6f274aeb3 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/startree/KeywordTermsAggregatorTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/KeywordTermsAggregatorTests.java @@ -170,7 +170,6 @@ public void testStarTreeKeywordTerms() throws IOException { for (ValuesSourceAggregationBuilder aggregationBuilder : aggBuilders) { query = new MatchAllDocsQuery(); queryBuilder = null; - termsAggregationBuilder = terms("terms_agg").field(CLIENTIP).subAggregation(aggregationBuilder); testCase(indexSearcher, query, queryBuilder, termsAggregationBuilder, starTree, supportedDimensions); diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeNestedAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeNestedAggregatorTests.java new file mode 100644 index 0000000000000..0d738b4b2119e --- /dev/null +++ b/server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeNestedAggregatorTests.java @@ -0,0 +1,439 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.search.aggregations.startree; + +import com.carrotsearch.randomizedtesting.RandomizedTest; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.lucene101.Lucene101Codec; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.LongPoint; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.document.SortedSetDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SegmentReader; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.util.BytesRef; +import org.opensearch.common.Rounding; +import org.opensearch.common.lucene.Lucene; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.index.codec.composite.CompositeIndexFieldInfo; +import org.opensearch.index.codec.composite.CompositeIndexReader; +import org.opensearch.index.codec.composite.composite101.Composite101Codec; +import org.opensearch.index.codec.composite912.datacube.startree.StarTreeDocValuesFormatTests; +import org.opensearch.index.compositeindex.datacube.DateDimension; +import org.opensearch.index.compositeindex.datacube.Dimension; +import org.opensearch.index.compositeindex.datacube.NumericDimension; +import org.opensearch.index.compositeindex.datacube.OrdinalDimension; +import org.opensearch.index.compositeindex.datacube.startree.utils.date.DateTimeUnitAdapter; +import org.opensearch.index.mapper.DateFieldMapper; +import org.opensearch.index.mapper.KeywordFieldMapper; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.MapperService; +import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.search.aggregations.Aggregator; +import org.opensearch.search.aggregations.InternalAggregation; +import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder; +import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregatorTestCase; +import org.opensearch.search.aggregations.bucket.histogram.DateHistogramInterval; +import org.opensearch.search.aggregations.bucket.histogram.InternalDateHistogram; +import org.opensearch.search.aggregations.bucket.range.InternalRange; +import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder; +import org.opensearch.search.aggregations.bucket.terms.InternalTerms; +import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; +import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Random; +import java.util.function.Function; +import java.util.function.Supplier; + +import static org.opensearch.common.util.FeatureFlags.STAR_TREE_INDEX; +import static org.opensearch.search.aggregations.AggregationBuilders.count; +import static org.opensearch.search.aggregations.AggregationBuilders.dateHistogram; +import static org.opensearch.search.aggregations.AggregationBuilders.max; +import static org.opensearch.search.aggregations.AggregationBuilders.min; +import static org.opensearch.search.aggregations.AggregationBuilders.range; +import static org.opensearch.search.aggregations.AggregationBuilders.sum; +import static org.opensearch.search.aggregations.AggregationBuilders.terms; +import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS; + +public class StarTreeNestedAggregatorTests extends DateHistogramAggregatorTestCase { + private static final String TIMESTAMP_FIELD = "@timestamp"; + private static final MappedFieldType TIMESTAMP_FIELD_TYPE = new DateFieldMapper.DateFieldType(TIMESTAMP_FIELD); + + private static final String KEYWORD_FIELD = "clientip"; + MappedFieldType KEYWORD_FIELD_TYPE = new KeywordFieldMapper.KeywordFieldType(KEYWORD_FIELD); + + final static String STATUS = "status"; + final static String SIZE = "size"; + private static final MappedFieldType STATUS_FIELD_TYPE = new NumberFieldMapper.NumberFieldType( + STATUS, + NumberFieldMapper.NumberType.LONG + ); + private static final MappedFieldType SIZE_FIELD_TYPE = new NumberFieldMapper.NumberFieldType(SIZE, NumberFieldMapper.NumberType.LONG); + private static FeatureFlags.TestUtils.FlagWriteLock fflock = null; + + @Before + public void setup() { + fflock = new FeatureFlags.TestUtils.FlagWriteLock(STAR_TREE_INDEX); + // FeatureFlags.initializeFeatureFlags(Settings.builder().put(FeatureFlags.STAR_TREE_INDEX, true).build()); + } + + @After + public void teardown() throws IOException { + fflock.close(); + // FeatureFlags.initializeFeatureFlags(Settings.EMPTY); + } + + protected Codec getCodec() { + final Logger testLogger = LogManager.getLogger(MetricAggregatorTests.class); + MapperService mapperService; + try { + mapperService = StarTreeDocValuesFormatTests.createMapperService(NumericTermsAggregatorTests.getExpandedMapping(1, false)); + } catch (IOException e) { + throw new RuntimeException(e); + } + return new Composite101Codec(Lucene101Codec.Mode.BEST_SPEED, mapperService, testLogger); + } + + public void testStarTreeNestedAggregations() throws IOException { + Directory directory = newDirectory(); + IndexWriterConfig conf = newIndexWriterConfig(null); + conf.setCodec(getCodec()); + conf.setMergePolicy(newLogMergePolicy()); + RandomIndexWriter iw = new RandomIndexWriter(random(), directory, conf); + + Random random = RandomizedTest.getRandom(); + int totalDocs = 100; + + long val; + long date; + List docs = new ArrayList<>(); + // Index 100 random documents + for (int i = 0; i < totalDocs; i++) { + Document doc = new Document(); + if (randomBoolean()) { + val = random.nextInt(100); // Random int between 0 and 99 for status + doc.add(new SortedNumericDocValuesField(STATUS, val)); + } + if (randomBoolean()) { + val = random.nextInt(100); + doc.add(new SortedNumericDocValuesField(SIZE, val)); + } + if (randomBoolean()) { + val = random.nextInt(10); // Random strings for int between 0 and 9 for keyword terms + doc.add(new SortedSetDocValuesField(KEYWORD_FIELD, new BytesRef(String.valueOf(val)))); + } + if (randomBoolean()) { + date = random.nextInt(180) * 24 * 60 * 60 * 1000L; // Random date within 180 days + doc.add(new SortedNumericDocValuesField(TIMESTAMP_FIELD, date)); + doc.add(new LongPoint(TIMESTAMP_FIELD, date)); + } + + iw.addDocument(doc); + docs.add(doc); + } + + if (randomBoolean()) { + iw.forceMerge(1); + } + iw.close(); + DirectoryReader ir = DirectoryReader.open(directory); + LeafReaderContext context = ir.leaves().get(0); + + SegmentReader reader = Lucene.segmentReader(context.reader()); + IndexSearcher indexSearcher = newSearcher(wrapInMockESDirectoryReader(ir), false, false); + CompositeIndexReader starTreeDocValuesReader = (CompositeIndexReader) reader.getDocValuesReader(); + + List compositeIndexFields = starTreeDocValuesReader.getCompositeIndexFields(); + CompositeIndexFieldInfo starTree = compositeIndexFields.get(0); + + LinkedHashMap supportedDimensions = new LinkedHashMap<>(); + supportedDimensions.put(new NumericDimension(STATUS), STATUS_FIELD_TYPE); + supportedDimensions.put(new NumericDimension(SIZE), SIZE_FIELD_TYPE); + supportedDimensions.put( + new DateDimension( + TIMESTAMP_FIELD, + List.of( + new DateTimeUnitAdapter(Rounding.DateTimeUnit.MONTH_OF_YEAR), + new DateTimeUnitAdapter(Rounding.DateTimeUnit.DAY_OF_MONTH) + ), + DateFieldMapper.Resolution.MILLISECONDS + ), + new DateFieldMapper.DateFieldType(TIMESTAMP_FIELD) + ); + supportedDimensions.put(new OrdinalDimension(KEYWORD_FIELD), KEYWORD_FIELD_TYPE); + + Query query = new MatchAllDocsQuery(); + QueryBuilder queryBuilder = null; + + ValuesSourceAggregationBuilder[] aggBuilders = { + sum("_sum").field(STATUS), + max("_max").field(STATUS), + min("_min").field(STATUS), + count("_count").field(STATUS) }; + + List>> aggregationSuppliers = List.of( + () -> terms("term_size").field(SIZE), + () -> terms("term_status").field(STATUS), + () -> range("range_agg").field(STATUS).addRange(10, 30).addRange(30, 50), + () -> terms("term_keyword").field(KEYWORD_FIELD).collectMode(Aggregator.SubAggCollectionMode.DEPTH_FIRST), + () -> terms("term_keyword").field(KEYWORD_FIELD).collectMode(Aggregator.SubAggCollectionMode.BREADTH_FIRST), + () -> dateHistogram("by_day").field(TIMESTAMP_FIELD).calendarInterval(DateHistogramInterval.DAY) + ); + // 3-LEVELS [BUCKET -> BUCKET -> METRIC] + for (ValuesSourceAggregationBuilder aggregationBuilder : aggBuilders) { + query = new MatchAllDocsQuery(); + queryBuilder = null; + for (Supplier> outerSupplier : aggregationSuppliers) { + for (Supplier> innerSupplier : aggregationSuppliers) { + + ValuesSourceAggregationBuilder inner = innerSupplier.get().subAggregation(aggregationBuilder); + ValuesSourceAggregationBuilder outer = outerSupplier.get().subAggregation(inner); + + // Skipping [DateHistogramAggregationBuilder + RangeAggregationBuilder] combinations for a ReducedMultiBucketConsumer + // assertion in + // searchAndReduceStarTree + boolean skipReducedMultiBucketConsumerAssertion = (inner instanceof RangeAggregationBuilder + && outer instanceof DateHistogramAggregationBuilder); + testCase( + indexSearcher, + query, + queryBuilder, + outer, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion + ); + + // Numeric-terms query with numeric terms aggregation + for (int cases = 0; cases < 5; cases++) { + String queryField; + long queryValue; + if (randomBoolean()) { + queryField = STATUS; + } else { + queryField = SIZE; + } + queryValue = random.nextInt(10); + query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue); + queryBuilder = new TermQueryBuilder(queryField, queryValue); + testCase( + indexSearcher, + query, + queryBuilder, + outer, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion + ); + } + } + } + } + + // 4-LEVELS [BUCKET -> BUCKET -> BUCKET -> METRIC] + for (ValuesSourceAggregationBuilder aggregationBuilder : aggBuilders) { + query = new MatchAllDocsQuery(); + queryBuilder = null; + for (Supplier> outermostSupplier : aggregationSuppliers) { + for (Supplier> middleSupplier : aggregationSuppliers) { + for (Supplier> innerSupplier : aggregationSuppliers) { + + ValuesSourceAggregationBuilder innermost = innerSupplier.get().subAggregation(aggregationBuilder); + ValuesSourceAggregationBuilder middle = middleSupplier.get().subAggregation(innermost); + ValuesSourceAggregationBuilder outermost = outermostSupplier.get().subAggregation(middle); + + // Skipping [DateHistogramAggregationBuilder + RangeAggregationBuilder] combinations for a + // ReducedMultiBucketConsumer assertion in + // searchAndReduceStarTree + boolean skipReducedMultiBucketConsumerAssertion = (middle instanceof RangeAggregationBuilder + && outermost instanceof DateHistogramAggregationBuilder) + || (middle instanceof DateHistogramAggregationBuilder && innermost instanceof RangeAggregationBuilder); + testCase( + indexSearcher, + query, + queryBuilder, + outermost, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion + ); + + // Numeric-terms query with numeric terms aggregation + for (int cases = 0; cases < 5; cases++) { + String queryField; + long queryValue; + if (randomBoolean()) { + queryField = STATUS; + } else { + queryField = SIZE; + } + queryValue = random.nextInt(10); + query = SortedNumericDocValuesField.newSlowExactQuery(queryField, queryValue); + queryBuilder = new TermQueryBuilder(queryField, queryValue); + testCase( + indexSearcher, + query, + queryBuilder, + outermost, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion + ); + } + + } + } + } + } + + ir.close(); + directory.close(); + + } + + private void testCase( + IndexSearcher indexSearcher, + Query query, + QueryBuilder queryBuilder, + ValuesSourceAggregationBuilder aggregationBuilder, + CompositeIndexFieldInfo starTree, + LinkedHashMap supportedDimensions, + boolean skipReducedMultiBucketConsumerAssertion + ) throws IOException { + + if (aggregationBuilder instanceof TermsAggregationBuilder) { + assertEqualStarTreeAggregation( + InternalTerms.class, + InternalTerms::getBuckets, + aggregationBuilder, + indexSearcher, + query, + queryBuilder, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion, + STATUS_FIELD_TYPE, + SIZE_FIELD_TYPE, + TIMESTAMP_FIELD_TYPE, + KEYWORD_FIELD_TYPE + ); + + } else if (aggregationBuilder instanceof DateHistogramAggregationBuilder) { + assertEqualStarTreeAggregation( + InternalDateHistogram.class, + InternalDateHistogram::getBuckets, + aggregationBuilder, + indexSearcher, + query, + queryBuilder, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion, + STATUS_FIELD_TYPE, + TIMESTAMP_FIELD_TYPE, + SIZE_FIELD_TYPE, + KEYWORD_FIELD_TYPE + ); + + } else if (aggregationBuilder instanceof RangeAggregationBuilder) { + assertEqualStarTreeAggregation( + InternalRange.class, + InternalRange::getBuckets, + aggregationBuilder, + indexSearcher, + query, + queryBuilder, + starTree, + supportedDimensions, + skipReducedMultiBucketConsumerAssertion, + STATUS_FIELD_TYPE, + SIZE_FIELD_TYPE, + TIMESTAMP_FIELD_TYPE, + KEYWORD_FIELD_TYPE + ); + } + } + + private void assertEqualStarTreeAggregation( + Class clazz, + Function> getBuckets, + ValuesSourceAggregationBuilder aggregationBuilder, + IndexSearcher indexSearcher, + Query query, + QueryBuilder queryBuilder, + CompositeIndexFieldInfo starTree, + LinkedHashMap supportedDimensions, + boolean skipReducedMultiBucketConsumerAssertion, + MappedFieldType... fieldTypes + ) throws IOException { + + T defaultAgg = clazz.cast( + searchAndReduceStarTree( + createIndexSettings(), + indexSearcher, + query, + queryBuilder, + aggregationBuilder, + null, + null, + null, + DEFAULT_MAX_BUCKETS, + false, + null, + false, + skipReducedMultiBucketConsumerAssertion, + fieldTypes + ) + ); + + T starTreeAgg = clazz.cast( + searchAndReduceStarTree( + createIndexSettings(), + indexSearcher, + query, + queryBuilder, + aggregationBuilder, + starTree, + supportedDimensions, + null, + DEFAULT_MAX_BUCKETS, + false, + null, + true, + skipReducedMultiBucketConsumerAssertion, + fieldTypes + ) + ); + + List defaultBuckets = getBuckets.apply(defaultAgg); + List starTreeBuckets = getBuckets.apply(starTreeAgg); + + assertEquals(defaultBuckets.size(), starTreeBuckets.size()); + assertEquals(defaultBuckets, starTreeBuckets); + } + +} diff --git a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java index df982d4f0c7f3..d9348aacd7a11 100644 --- a/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java +++ b/test/framework/src/main/java/org/opensearch/search/aggregations/AggregatorTestCase.java @@ -775,6 +775,40 @@ protected A searchAndReduc AggregatorFactory aggregatorFactory, boolean assertCollectorEarlyTermination, MappedFieldType... fieldTypes + ) throws IOException { + return searchAndReduceStarTree( + indexSettings, + searcher, + query, + queryBuilder, + builder, + compositeIndexFieldInfo, + supportedDimensions, + supportedMetrics, + maxBucket, + hasNested, + aggregatorFactory, + assertCollectorEarlyTermination, + false, + fieldTypes + ); + } + + protected A searchAndReduceStarTree( + IndexSettings indexSettings, + IndexSearcher searcher, + Query query, + QueryBuilder queryBuilder, + AggregationBuilder builder, + CompositeIndexFieldInfo compositeIndexFieldInfo, + LinkedHashMap supportedDimensions, + List supportedMetrics, + int maxBucket, + boolean hasNested, + AggregatorFactory aggregatorFactory, + boolean assertCollectorEarlyTermination, + boolean skipReducedMultiBucketConsumerAssertion, + MappedFieldType... fieldTypes ) throws IOException { query = query.rewrite(searcher); final IndexReaderContext ctx = searcher.getTopReaderContext(); @@ -823,7 +857,7 @@ protected A searchAndReduc @SuppressWarnings("unchecked") A internalAgg = (A) aggs.get(0).reduce(aggs, context); - doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer); + if (skipReducedMultiBucketConsumerAssertion == false) doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer); return internalAgg; }