From 1087f86c027f4b95678861c6954807ba57dfdb73 Mon Sep 17 00:00:00 2001 From: kkewwei Date: Mon, 22 Apr 2024 16:24:32 +0800 Subject: [PATCH] fix: The aggs result of NestedAggregator with sub NestedAggregator may be not accurately --- .../nested/NestedAggregationBuilder.java | 5 +- .../bucket/nested/NestedAggregator.java | 59 ++-- .../nested/NestedAggregatorFactory.java | 5 +- .../bucket/nested/NestedAggregatorTests.java | 305 ++++++++++++++++-- .../aggregations/AggregatorTestCase.java | 16 + 5 files changed, 344 insertions(+), 46 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregationBuilder.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregationBuilder.java index 0d98fdbf15f78..cb0e96e2a8101 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregationBuilder.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregationBuilder.java @@ -115,17 +115,16 @@ protected AggregatorFactory doBuild(QueryShardContext queryShardContext, Aggrega ObjectMapper childObjectMapper = queryShardContext.getObjectMapper(path); if (childObjectMapper == null) { // in case the path has been unmapped: - return new NestedAggregatorFactory(name, null, null, queryShardContext, parent, subFactoriesBuilder, metadata); + return new NestedAggregatorFactory(name, null, queryShardContext, parent, subFactoriesBuilder, metadata); } if (childObjectMapper.nested().isNested() == false) { throw new AggregationExecutionException("[nested] nested path [" + path + "] is not nested"); } try { - ObjectMapper parentObjectMapper = queryShardContext.nestedScope().nextLevel(childObjectMapper); + queryShardContext.nestedScope().nextLevel(childObjectMapper); return new NestedAggregatorFactory( name, - parentObjectMapper, childObjectMapper, queryShardContext, parent, diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregator.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregator.java index cfa1d32a52501..6737a21feacbe 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregator.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregator.java @@ -43,6 +43,7 @@ import org.apache.lucene.search.Weight; import org.apache.lucene.search.join.BitSetProducer; import org.apache.lucene.util.BitSet; +import org.opensearch.common.collect.Tuple; import org.opensearch.common.lucene.search.Queries; import org.opensearch.core.ParseField; import org.opensearch.index.mapper.ObjectMapper; @@ -79,7 +80,6 @@ public class NestedAggregator extends BucketsAggregator implements SingleBucketA NestedAggregator( String name, AggregatorFactories factories, - ObjectMapper parentObjectMapper, ObjectMapper childObjectMapper, SearchContext context, Aggregator parent, @@ -88,8 +88,7 @@ public class NestedAggregator extends BucketsAggregator implements SingleBucketA ) throws IOException { super(name, factories, context, parent, cardinality, metadata); - Query parentFilter = parentObjectMapper != null ? parentObjectMapper.nestedTypeFilter() : Queries.newNonNestedFilter(); - this.parentFilter = context.bitsetFilterCache().getBitSetProducer(parentFilter); + this.parentFilter = context.bitsetFilterCache().getBitSetProducer(Queries.newNonNestedFilter()); this.childFilter = childObjectMapper.nestedTypeFilter(); this.collectsFromSingleBucket = cardinality.map(estimate -> estimate < 2); } @@ -108,19 +107,16 @@ public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final L return new LeafBucketCollectorBase(sub, null) { @Override public void collect(int parentDoc, long bucket) throws IOException { - // if parentDoc is 0 then this means that this parent doesn't have child docs (b/c these appear always before the parent - // doc), so we can skip: - if (parentDoc == 0 || parentDocs == null || childDocs == null) { + // parentDoc can be 0 when searching: + if (parentDocs == null || childDocs == null) { return; } - final int prevParentDoc = parentDocs.prevSetBit(parentDoc - 1); - int childDocId = childDocs.docID(); - if (childDocId <= prevParentDoc) { - childDocId = childDocs.advance(prevParentDoc + 1); - } + Tuple res = getChildAndRootParent(parentDocs, childDocs, parentDoc); + int currentRootDoc = res.v1(); + int childDocId = res.v2(); - for (; childDocId < parentDoc; childDocId = childDocs.nextDoc()) { + for (; childDocId < currentRootDoc; childDocId = childDocs.nextDoc()) { collectBucket(sub, childDocId, bucket); } } @@ -130,6 +126,30 @@ public void collect(int parentDoc, long bucket) throws IOException { } } + static Tuple getChildAndRootParent(BitSet parentDocs, DocIdSetIterator childDocs, int parentDoc + ) throws IOException { + int currentRootDoc; + int prevParentDoc = parentDocs.prevSetBit(parentDoc); + if (prevParentDoc == -1) { + currentRootDoc = parentDocs.nextSetBit(0); + } else if (prevParentDoc == parentDoc) { + currentRootDoc = parentDoc; + if (currentRootDoc == 0) { + prevParentDoc = -1; + } else { + prevParentDoc = parentDocs.prevSetBit(currentRootDoc - 1); + } + } else { + currentRootDoc = parentDocs.nextSetBit(prevParentDoc + 1); + } + + int childDocId = childDocs.docID(); + if (childDocId <= prevParentDoc) { + childDocId = childDocs.advance(prevParentDoc + 1); + } + return Tuple.tuple(currentRootDoc, childDocId); + } + @Override protected void preGetSubLeafCollectors(LeafReaderContext ctx) throws IOException { super.preGetSubLeafCollectors(ctx); @@ -191,9 +211,8 @@ public void setScorer(Scorable scorer) throws IOException { @Override public void collect(int parentDoc, long bucket) throws IOException { - // if parentDoc is 0 then this means that this parent doesn't have child docs (b/c these appear always before the parent - // doc), so we can skip: - if (parentDoc == 0 || parentDocs == null || childDocs == null) { + // parentDoc can be 0 when searching: + if (parentDocs == null || childDocs == null) { return; } @@ -214,13 +233,11 @@ void processBufferedChildBuckets() throws IOException { return; } - final int prevParentDoc = parentDocs.prevSetBit(currentParentDoc - 1); - int childDocId = childDocs.docID(); - if (childDocId <= prevParentDoc) { - childDocId = childDocs.advance(prevParentDoc + 1); - } + Tuple res = getChildAndRootParent(parentDocs, childDocs, currentParentDoc); + int currentRootDoc = res.v1(); + int childDocId = res.v2(); - for (; childDocId < currentParentDoc; childDocId = childDocs.nextDoc()) { + for (; childDocId < currentRootDoc; childDocId = childDocs.nextDoc()) { cachedScorer.doc = childDocId; for (var bucket : bucketBuffer) { collectBucket(sub, childDocId, bucket); diff --git a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java index a43d41882e475..372ba638031eb 100644 --- a/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java +++ b/server/src/main/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorFactory.java @@ -52,12 +52,10 @@ */ public class NestedAggregatorFactory extends AggregatorFactory { - private final ObjectMapper parentObjectMapper; private final ObjectMapper childObjectMapper; NestedAggregatorFactory( String name, - ObjectMapper parentObjectMapper, ObjectMapper childObjectMapper, QueryShardContext queryShardContext, AggregatorFactory parent, @@ -65,7 +63,6 @@ public class NestedAggregatorFactory extends AggregatorFactory { Map metadata ) throws IOException { super(name, queryShardContext, parent, subFactories, metadata); - this.parentObjectMapper = parentObjectMapper; this.childObjectMapper = childObjectMapper; } @@ -79,7 +76,7 @@ public Aggregator createInternal( if (childObjectMapper == null) { return new Unmapped(name, searchContext, parent, factories, metadata); } - return new NestedAggregator(name, factories, parentObjectMapper, childObjectMapper, searchContext, parent, cardinality, metadata); + return new NestedAggregator(name, factories, childObjectMapper, searchContext, parent, cardinality, metadata); } /** diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorTests.java index 406c411494d60..5ea6c4f2d5fe8 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/bucket/nested/NestedAggregatorTests.java @@ -34,6 +34,7 @@ import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; +import org.apache.lucene.document.LongPoint; import org.apache.lucene.document.SortedDocValuesField; import org.apache.lucene.document.SortedNumericDocValuesField; import org.apache.lucene.document.SortedSetDocValuesField; @@ -45,23 +46,37 @@ import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.ConstantScoreQuery; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.TermQuery; +import org.apache.lucene.search.join.ScoreMode; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.util.BitSet; import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.FixedBitSet; +import org.mockito.Mockito; import org.opensearch.common.CheckedConsumer; import org.opensearch.common.collect.Tuple; import org.opensearch.common.lucene.search.Queries; import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.cache.bitset.BitsetFilterCache; +import org.opensearch.index.mapper.ContentPath; import org.opensearch.index.mapper.IdFieldMapper; import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.index.mapper.Mapper; import org.opensearch.index.mapper.NestedPathFieldMapper; import org.opensearch.index.mapper.NumberFieldMapper; +import org.opensearch.index.mapper.ObjectMapper; import org.opensearch.index.mapper.SeqNoFieldMapper; import org.opensearch.index.mapper.Uid; import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.NestedQueryBuilder; +import org.opensearch.index.query.QueryShardContext; +import org.opensearch.index.query.TermsQueryBuilder; +import org.opensearch.index.query.support.NestedScope; import org.opensearch.script.MockScriptEngine; import org.opensearch.script.Script; import org.opensearch.script.ScriptEngine; @@ -105,20 +120,32 @@ import java.util.stream.LongStream; import static java.util.stream.Collectors.toList; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.common.lucene.search.Queries.newNonNestedFilter; import static org.opensearch.search.aggregations.AggregationBuilders.max; import static org.opensearch.search.aggregations.AggregationBuilders.nested; import static org.hamcrest.Matchers.equalTo; +import static org.opensearch.search.aggregations.bucket.nested.NestedAggregator.getChildAndRootParent; +import static org.opensearch.test.InternalAggregationTestCase.DEFAULT_MAX_BUCKETS; public class NestedAggregatorTests extends AggregatorTestCase { private static final String VALUE_FIELD_NAME = "number"; + private static final String VALUE_FIELD_NAME2 = "number2"; + private static final String NESTED_OBJECT = "nested_object"; private static final String NESTED_OBJECT2 = "nested_object2"; private static final String NESTED_AGG = "nestedAgg"; private static final String MAX_AGG_NAME = "maxAgg"; private static final String SUM_AGG_NAME = "sumAgg"; private static final String INVERSE_SCRIPT = "inverse"; - + private static final String OUT_NESTED = "outNested"; + private static final String OUT_TERMS = "outTerms"; + private static final String INNER_NESTED = "innerNested"; + private static final String INNER_TERMS = "innerTerms"; private static final SeqNoFieldMapper.SequenceIDFields sequenceIDFields = SeqNoFieldMapper.SequenceIDFields.emptySeqID(); /** @@ -201,17 +228,20 @@ public void testSingleNestingMax() throws IOException { } try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) { NestedAggregationBuilder nestedBuilder = new NestedAggregationBuilder(NESTED_AGG, NESTED_OBJECT); - MaxAggregationBuilder maxAgg = new MaxAggregationBuilder(MAX_AGG_NAME).field(VALUE_FIELD_NAME); + MaxAggregationBuilder maxAgg = new MaxAggregationBuilder(MAX_AGG_NAME).field(NESTED_OBJECT + "." + VALUE_FIELD_NAME); nestedBuilder.subAggregation(maxAgg); - MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); + MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NESTED_OBJECT + "." + VALUE_FIELD_NAME, + NumberFieldMapper.NumberType.LONG); InternalNested nested = searchAndReduce( + createIndexSettings(), newSearcher(indexReader, false, true), new MatchAllDocsQuery(), nestedBuilder, + DEFAULT_MAX_BUCKETS, + true, fieldType ); - assertEquals(expectedNestedDocs, nested.getDocCount()); assertEquals(NESTED_AGG, nested.getName()); assertEquals(expectedNestedDocs, nested.getDocCount()); @@ -240,7 +270,7 @@ public void testDoubleNestingMax() throws IOException { int numNestedDocs = randomIntBetween(0, 20); expectedMaxValue = Math.max( expectedMaxValue, - generateMaxDocs(documents, numNestedDocs, i, NESTED_OBJECT + "." + NESTED_OBJECT2, VALUE_FIELD_NAME) + generateMaxDocs(documents, numNestedDocs, i, NESTED_OBJECT, VALUE_FIELD_NAME) ); expectedNestedDocs += numNestedDocs; @@ -253,19 +283,22 @@ public void testDoubleNestingMax() throws IOException { iw.commit(); } try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) { - NestedAggregationBuilder nestedBuilder = new NestedAggregationBuilder(NESTED_AGG, NESTED_OBJECT + "." + NESTED_OBJECT2); - MaxAggregationBuilder maxAgg = new MaxAggregationBuilder(MAX_AGG_NAME).field(VALUE_FIELD_NAME); + NestedAggregationBuilder nestedBuilder = new NestedAggregationBuilder(NESTED_AGG, NESTED_OBJECT); + MaxAggregationBuilder maxAgg = new MaxAggregationBuilder(MAX_AGG_NAME).field(NESTED_OBJECT+ "." + VALUE_FIELD_NAME); nestedBuilder.subAggregation(maxAgg); - MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); + MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NESTED_OBJECT+ "." + VALUE_FIELD_NAME, + NumberFieldMapper.NumberType.LONG); InternalNested nested = searchAndReduce( + createIndexSettings(), newSearcher(indexReader, false, true), new MatchAllDocsQuery(), nestedBuilder, + DEFAULT_MAX_BUCKETS, + true, fieldType ); - assertEquals(expectedNestedDocs, nested.getDocCount()); assertEquals(NESTED_AGG, nested.getName()); assertEquals(expectedNestedDocs, nested.getDocCount()); @@ -310,17 +343,19 @@ public void testOrphanedDocs() throws IOException { } try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) { NestedAggregationBuilder nestedBuilder = new NestedAggregationBuilder(NESTED_AGG, NESTED_OBJECT); - SumAggregationBuilder sumAgg = new SumAggregationBuilder(SUM_AGG_NAME).field(VALUE_FIELD_NAME); + SumAggregationBuilder sumAgg = new SumAggregationBuilder(SUM_AGG_NAME).field(NESTED_OBJECT + "." + VALUE_FIELD_NAME); nestedBuilder.subAggregation(sumAgg); - MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); + MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NESTED_OBJECT + "." + VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); InternalNested nested = searchAndReduce( + createIndexSettings(), newSearcher(indexReader, false, true), new MatchAllDocsQuery(), nestedBuilder, + DEFAULT_MAX_BUCKETS, + true, fieldType ); - assertEquals(expectedNestedDocs, nested.getDocCount()); assertEquals(NESTED_AGG, nested.getName()); assertEquals(expectedNestedDocs, nested.getDocCount()); @@ -393,7 +428,7 @@ public void testResetRootDocId() throws Exception { MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); BooleanQuery.Builder bq = new BooleanQuery.Builder(); - bq.add(Queries.newNonNestedFilter(), BooleanClause.Occur.MUST); + bq.add(newNonNestedFilter(), BooleanClause.Occur.MUST); bq.add(new TermQuery(new Term(IdFieldMapper.NAME, Uid.encodeId("2"))), BooleanClause.Occur.MUST_NOT); InternalNested nested = searchAndReduce( @@ -685,7 +720,7 @@ public void testPreGetChildLeafCollectors() throws IOException { Filter filter = searchAndReduce( newSearcher(indexReader, false, true), - Queries.newNonNestedFilter(), + newNonNestedFilter(), filterAggregationBuilder, fieldType1, fieldType2 @@ -747,8 +782,17 @@ public void testFieldAlias() throws IOException { max(MAX_AGG_NAME).field(VALUE_FIELD_NAME + "-alias") ); - InternalNested nested = searchAndReduce(newSearcher(indexReader, false, true), new MatchAllDocsQuery(), agg, fieldType); - Nested aliasNested = searchAndReduce(newSearcher(indexReader, false, true), new MatchAllDocsQuery(), aliasAgg, fieldType); + InternalNested nested = searchAndReduce(createIndexSettings(), + newSearcher(indexReader, false, true), new MatchAllDocsQuery(), agg, + DEFAULT_MAX_BUCKETS,true, fieldType); + Nested aliasNested = searchAndReduce( + createIndexSettings(), + newSearcher(indexReader, false, true), + new MatchAllDocsQuery(), + aliasAgg, + DEFAULT_MAX_BUCKETS, + true, + fieldType); assertEquals(nested, aliasNested); assertEquals(expectedNestedDocs, nested.getDocCount()); @@ -796,13 +840,15 @@ public void testNestedWithPipeline() throws IOException { MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); InternalNested nested = searchAndReduce( + createIndexSettings(), newSearcher(indexReader, false, true), new MatchAllDocsQuery(), nestedBuilder, + DEFAULT_MAX_BUCKETS, + true, fieldType ); - assertEquals(expectedNestedDocs, nested.getDocCount()); assertEquals(NESTED_AGG, nested.getName()); assertEquals(expectedNestedDocs, nested.getDocCount()); @@ -853,6 +899,218 @@ public void testNestedUnderTerms() throws IOException { }, resellersMappedFields()); } + public void testBufferingNestedLeafBucketCollector() throws IOException { + int numRootDocs = scaledRandomIntBetween(2, 200); + int expectedNestedDocs; + String[] bucketKeys; + try (Directory directory = newDirectory()) { + try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) { + for (int i = 0; i < numRootDocs; i++) { + + List documents = new ArrayList<>(); + if (randomBoolean()) { + generateDocument(documents, i, NESTED_OBJECT, VALUE_FIELD_NAME, 1); + generateDocument(documents, i, NESTED_OBJECT2, VALUE_FIELD_NAME2, i); + } else { + generateDocument(documents, i, NESTED_OBJECT2, VALUE_FIELD_NAME2, i); + generateDocument(documents, i, NESTED_OBJECT, VALUE_FIELD_NAME, 1); + } + Document document = new Document(); + document.add(new Field(IdFieldMapper.NAME, Uid.encodeId(Integer.toString(i)), IdFieldMapper.Defaults.FIELD_TYPE)); + document.add(sequenceIDFields.primaryTerm); + documents.add(document); + iw.addDocuments(documents); + } + iw.commit(); + } + try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) { + IndexSettings indexSettings = createIndexSettings(); + MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NESTED_OBJECT + "." +VALUE_FIELD_NAME, NumberFieldMapper.NumberType.LONG); + MappedFieldType fieldType1 = new NumberFieldMapper.NumberFieldType(NESTED_OBJECT2 + "." + VALUE_FIELD_NAME2, NumberFieldMapper.NumberType.LONG); + QueryShardContext queryShardContext = createQueryShardContext(NESTED_OBJECT2, indexSettings, fieldType1); + // query + expectedNestedDocs = numRootDocs/2; + bucketKeys = new String[expectedNestedDocs]; + BytesRef [] values = new BytesRef[numRootDocs/2]; + for (int i = 0; i < numRootDocs/2; i++) { + bucketKeys[i] = "" + (i * 2); + values[i] = new BytesRef(bucketKeys[i]); + } + TermsQueryBuilder termsQueryBuilder = new TermsQueryBuilder(NESTED_OBJECT2 + "." + VALUE_FIELD_NAME2, values); + NestedQueryBuilder nestedQueryBuilder = new NestedQueryBuilder(NESTED_OBJECT2, termsQueryBuilder, ScoreMode.None); + + //out nested aggs + NestedAggregationBuilder outNestedBuilder = new NestedAggregationBuilder(OUT_NESTED, NESTED_OBJECT); + TermsAggregationBuilder outTermsAggregator = new TermsAggregationBuilder(OUT_TERMS) + .field(NESTED_OBJECT + "." + VALUE_FIELD_NAME) + .size(100); + outNestedBuilder.subAggregation(outTermsAggregator); + + //inner nested aggs + NestedAggregationBuilder innerNestedBuilder = new NestedAggregationBuilder(INNER_NESTED, NESTED_OBJECT2); + TermsAggregationBuilder innerTermsAggregator = new TermsAggregationBuilder(INNER_TERMS) + .field(NESTED_OBJECT2 + "." + VALUE_FIELD_NAME2) + .size(100); + innerNestedBuilder.subAggregation(innerTermsAggregator); + outTermsAggregator.subAggregation(innerNestedBuilder); + System.out.println(outNestedBuilder); + + InternalNested nested = searchAndReduce( + indexSettings, + newSearcher(indexReader, false, true), + nestedQueryBuilder.toQuery(queryShardContext), + outNestedBuilder, + DEFAULT_MAX_BUCKETS, + true, + fieldType, + fieldType1); + + assertEquals(OUT_NESTED, nested.getName()); + assertEquals(expectedNestedDocs, nested.getDocCount()); + + LongTerms outTerms = (LongTerms) nested.getProperty(OUT_TERMS); + assertEquals(1, outTerms.getBuckets().size()); + + InternalNested internalNested = (InternalNested) (((Object[])outTerms.getProperty(INNER_NESTED))[0]); + assertEquals(expectedNestedDocs, internalNested.getDocCount()); + + LongTerms innerTerms = (LongTerms) internalNested.getProperty(INNER_TERMS); + assertEquals(bucketKeys.length, innerTerms.getBuckets().size()); + for (int i = 0; i bits.length || bits[index] != 1) { + return -1; + } + return index; + } + + @Override + public int nextDoc() { + for (int i = index; i < bits.length; i++) { + if (bits[i] == 1) { + index = i; + return index; + } + } + index = bits.length; + return -1; + } + + @Override + public int advance(int target) { + for (int i = target; i < bits.length; i++) { + if (bits[i] == 1) { + index = i; + return index; + } + } + index = bits.length; + return -1; + } + + @Override + public long cost() { + return bits.length; + } + }; + } + + public void testGetChildAndRootParent() throws IOException { + { + // p: parent c: child + // [p1], [c2,p3], [c4,p6] [p8] + BitSet parentDocs = new FixedBitSet(20); + parentDocs.set(0); + parentDocs.set(3); + parentDocs.set(6); + parentDocs.set(8); + DocIdSetIterator childDocs = getDocIdSetIterator(new int[]{2,4}); + Tuple res = getChildAndRootParent(parentDocs, childDocs, 0); + assertEquals(0, res.v1().intValue()); + assertEquals(2, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 3); + assertEquals(3, res.v1().intValue()); + assertEquals(2, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 4); + assertEquals(6, res.v1().intValue()); + assertEquals(4, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 8); + assertEquals(8, res.v1().intValue()); + assertEquals(-1, res.v2().intValue()); + } + + { + // p: parent c: child1 d: child2 + // [c1,d2,p3], [d4,c5,p6], [c7,d8,p9] + BitSet parentDocs = new FixedBitSet(20); + parentDocs.set(3); + parentDocs.set(6); + parentDocs.set(9); + { + DocIdSetIterator childDocs = getDocIdSetIterator(new int[]{1, 5, 7}); + Tuple res = getChildAndRootParent(parentDocs, childDocs, 2); + assertEquals(3, res.v1().intValue()); + assertEquals(1, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 4); + assertEquals(6, res.v1().intValue()); + assertEquals(5, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 8); + assertEquals(9, res.v1().intValue()); + assertEquals(7, res.v2().intValue()); + } + + { + DocIdSetIterator childDocs = getDocIdSetIterator(new int[]{1, 7}); + Tuple res = getChildAndRootParent(parentDocs, childDocs, 2); + assertEquals(3, res.v1().intValue()); + assertEquals(1, res.v2().intValue()); + + res = getChildAndRootParent(parentDocs, childDocs, 8); + assertEquals(9, res.v1().intValue()); + assertEquals(7, res.v2().intValue()); + } + } + } + + protected QueryShardContext createQueryShardContext(String fieldName, IndexSettings indexSettings, MappedFieldType fieldType) { + QueryShardContext queryShardContext = mock(QueryShardContext.class); + when(queryShardContext.nestedScope()).thenReturn(new NestedScope(indexSettings)); + + BitsetFilterCache bitsetFilterCache = new BitsetFilterCache(indexSettings, Mockito.mock(BitsetFilterCache.Listener.class)); + when(queryShardContext.bitsetFilter(any())).thenReturn(bitsetFilterCache.getBitSetProducer(Queries.newNonNestedFilter())); + when(queryShardContext.fieldMapper(anyString())).thenReturn(fieldType); + when(queryShardContext.getSearchQuoteAnalyzer(any())).thenCallRealMethod(); + when(queryShardContext.getSearchAnalyzer(any())).thenCallRealMethod(); + when(queryShardContext.getIndexSettings()).thenReturn(indexSettings); + when(queryShardContext.getObjectMapper(anyString())).thenAnswer(invocation -> { + Mapper.BuilderContext context = new Mapper.BuilderContext(indexSettings.getSettings(), new ContentPath()); + return new ObjectMapper.Builder<>(fieldName).nested(ObjectMapper.Nested.newNested()).build(context); + }); + when(queryShardContext.allowExpensiveQueries()).thenReturn(true); + return queryShardContext; + } + public static CheckedConsumer buildResellerData(int numProducts, int numResellers) { return iw -> { for (int p = 0; p < numProducts; p++) { @@ -893,13 +1151,24 @@ private static double[] generateDocuments(List documents, int numNeste document.add(new Field(IdFieldMapper.NAME, Uid.encodeId(Integer.toString(id)), IdFieldMapper.Defaults.NESTED_FIELD_TYPE)); document.add(new Field(NestedPathFieldMapper.NAME, path, NestedPathFieldMapper.Defaults.FIELD_TYPE)); long value = randomNonNegativeLong() % 10000; - document.add(new SortedNumericDocValuesField(fieldName, value)); + document.add(new SortedNumericDocValuesField(path+ "." + fieldName, value)); documents.add(document); values[nested] = value; } return values; } + private static void generateDocument(List documents, int id, + String path, String fieldName, long vales ) { + Document document = new Document(); + document.add(new Field(IdFieldMapper.NAME, Uid.encodeId(Integer.toString(id)), IdFieldMapper.Defaults.NESTED_FIELD_TYPE)); + document.add(new Field(NestedPathFieldMapper.NAME, path, NestedPathFieldMapper.Defaults.FIELD_TYPE)); + document.add(new SortedNumericDocValuesField(path + "." + fieldName, vales)); + document.add(new LongPoint(path + "." + fieldName, vales)); + documents.add(document); + + } + private List generateBook(String id, String[] authors, int[] numPages) { List documents = new ArrayList<>(); 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 e538dede07fc8..e2b7e3a5779d9 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 @@ -69,6 +69,7 @@ import org.opensearch.common.lease.Releasable; import org.opensearch.common.lease.Releasables; import org.opensearch.common.lucene.index.OpenSearchDirectoryReader; +import org.opensearch.common.lucene.search.Queries; import org.opensearch.common.network.NetworkAddress; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.BigArrays; @@ -519,6 +520,17 @@ protected A searchAndReduc return searchAndReduce(createIndexSettings(), searcher, query, builder, maxBucket, fieldTypes); } + protected A searchAndReduce( + IndexSettings indexSettings, + IndexSearcher searcher, + Query query, + AggregationBuilder builder, + int maxBucket, + MappedFieldType... fieldTypes + ) throws IOException { + return searchAndReduce(indexSettings, searcher, query, builder, maxBucket, false, fieldTypes); + } + /** * Collects all documents that match the provided query {@link Query} and * returns the reduced {@link InternalAggregation}. @@ -533,11 +545,15 @@ protected A searchAndReduc Query query, AggregationBuilder builder, int maxBucket, + boolean hasNested, MappedFieldType... fieldTypes ) throws IOException { final IndexReaderContext ctx = searcher.getTopReaderContext(); final PipelineTree pipelines = builder.buildPipelineTree(); List aggs = new ArrayList<>(); + if (hasNested) { + query = Queries.filtered(query, Queries.newNonNestedFilter()); + } Query rewritten = searcher.rewrite(query); MultiBucketConsumer bucketConsumer = new MultiBucketConsumer( maxBucket,