Modernize cardinality agg tests (elastic#90114)

This provides support for routing more cardinality tests through the "standard" aggregations testing infrastructure.  Ultimately, this will help with improving our memory access testing around aggregation context, although this PR doesn't implement that.

Author: not-napoleon

Diff for: server/src/test/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorTests.java

@@ -14,14 +14,9 @@
 import org.apache.lucene.document.SortedDocValuesField;
 import org.apache.lucene.document.SortedNumericDocValuesField;
 import org.apache.lucene.document.SortedSetDocValuesField;
-import org.apache.lucene.index.DirectoryReader;
-import org.apache.lucene.index.IndexReader;
-import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.search.FieldExistsQuery;
-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.elasticsearch.common.geo.GeoPoint;
@@ -47,7 +42,6 @@
 import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
 import org.elasticsearch.search.aggregations.bucket.terms.Terms;
 import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
-import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.AggregationInspectionHelper;
 import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.aggregations.support.ValuesSourceType;
@@ -410,40 +404,20 @@ public void testUnmappedMissingNumber() throws IOException {
     }
 
     public void testSingleValuedFieldPartiallyUnmapped() throws IOException {
-        final Directory directory = newDirectory();
-        final RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory);
-        final int numDocs = 10;
-        for (int i = 0; i < numDocs; i++) {
-            indexWriter.addDocument(singleton(new NumericDocValuesField("number", i + 1)));
-        }
-        indexWriter.close();
-
-        final Directory unmappedDirectory = newDirectory();
-        final RandomIndexWriter unmappedIndexWriter = new RandomIndexWriter(random(), unmappedDirectory);
-        unmappedIndexWriter.close();
-
-        final IndexReader indexReader = DirectoryReader.open(directory);
-        final IndexReader unamappedIndexReader = DirectoryReader.open(unmappedDirectory);
-        final MultiReader multiReader = new MultiReader(indexReader, unamappedIndexReader);
-        final IndexSearcher indexSearcher = newSearcher(multiReader, true, true);
-
         final MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.INTEGER);
         final AggregationBuilder aggregationBuilder = new CardinalityAggregationBuilder("cardinality").field("number");
 
-        final CardinalityAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
-        aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
-        aggregator.postCollection();
-
-        final InternalCardinality cardinality = (InternalCardinality) aggregator.buildAggregation(0L);
-
-        assertEquals(10.0, cardinality.getValue(), 0);
-        assertEquals("cardinality", cardinality.getName());
-        assertTrue(AggregationInspectionHelper.hasValue(cardinality));
-
-        multiReader.close();
-        directory.close();
-        unmappedDirectory.close();
+        multiIndexTestCase(aggregationBuilder, new MatchAllDocsQuery(), List.of((unmappedIndexWriter) -> {}, (indexWriter) -> {
+            final int numDocs = 10;
+            for (int i = 0; i < numDocs; i++) {
+                indexWriter.addDocument(singleton(new NumericDocValuesField("number", i + 1)));
+            }
+        }), (internalAggregation) -> {
+            InternalCardinality cardinality = (InternalCardinality) internalAggregation;
+            assertEquals(10.0, cardinality.getValue(), 0);
+            assertEquals("cardinality", cardinality.getName());
+            assertTrue(AggregationInspectionHelper.hasValue(cardinality));
+        }, fieldType);
     }
 
     public void testSingleValuedNumericValueScript() throws IOException {
@@ -618,47 +592,6 @@ public void testAsSubAggregation() throws IOException {
         }, mappedFieldTypes);
     }
 
-    public void testCacheAggregation() throws IOException {
-        final Directory directory = newDirectory();
-        final RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory);
-        final int numDocs = 10;
-        for (int i = 0; i < numDocs; i++) {
-            indexWriter.addDocument(singleton(new NumericDocValuesField("number", i + 1)));
-        }
-        indexWriter.close();
-
-        final Directory unmappedDirectory = newDirectory();
-        final RandomIndexWriter unmappedIndexWriter = new RandomIndexWriter(random(), unmappedDirectory);
-        unmappedIndexWriter.close();
-
-        final IndexReader indexReader = DirectoryReader.open(directory);
-        final IndexReader unamappedIndexReader = DirectoryReader.open(unmappedDirectory);
-        final MultiReader multiReader = new MultiReader(indexReader, unamappedIndexReader);
-        final IndexSearcher indexSearcher = newSearcher(multiReader, true, true);
-
-        final MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.INTEGER);
-        final CardinalityAggregationBuilder aggregationBuilder = new CardinalityAggregationBuilder("cardinality").field("number");
-
-        final AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
-        final CardinalityAggregator aggregator = createAggregator(aggregationBuilder, context);
-        aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
-        aggregator.postCollection();
-
-        final InternalCardinality cardinality = (InternalCardinality) aggregator.buildAggregation(0L);
-
-        assertEquals(10.0, cardinality.getValue(), 0);
-        assertEquals("cardinality", cardinality.getName());
-        assertTrue(AggregationInspectionHelper.hasValue(cardinality));
-
-        // Test that an aggregation not using a script does get cached
-        assertTrue(context.isCacheable());
-
-        multiReader.close();
-        directory.close();
-        unmappedDirectory.close();
-    }
-
     private void testAggregation(
         Query query,
         CheckedConsumer<RandomIndexWriter, IOException> buildIndex,

Diff for: test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java

@@ -21,6 +21,7 @@
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.NoMergePolicy;
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.search.Collector;
@@ -56,6 +57,7 @@
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.MockPageCacheRecycler;
 import org.elasticsearch.core.CheckedConsumer;
+import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.Index;
@@ -451,6 +453,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
                     aggTestConfig.builder(),
                     aggTestConfig.maxBuckets(),
                     aggTestConfig.splitLeavesIntoSeparateAggregators(),
+                    aggTestConfig.shouldBeCached(),
                     crankyService,
                     aggTestConfig.fieldTypes()
                 );
@@ -469,6 +472,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
             aggTestConfig.builder(),
             aggTestConfig.maxBuckets(),
             aggTestConfig.splitLeavesIntoSeparateAggregators(),
+            aggTestConfig.shouldBeCached(),
             breakerService,
             aggTestConfig.fieldTypes()
         );
@@ -482,6 +486,7 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
         AggregationBuilder builder,
         int maxBucket,
         boolean splitLeavesIntoSeparateAggregators,
+        boolean shouldBeCached,
         CircuitBreakerService breakerService,
         MappedFieldType... fieldTypes
     ) throws IOException {
@@ -535,6 +540,7 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
             root.postCollection();
             aggs.add(root.buildTopLevel());
         }
+        assertEquals(shouldBeCached, context.isCacheable());
         assertRoundTrip(aggs);
         if (randomBoolean() && aggs.size() > 1) {
             // sometimes do an incremental reduce
@@ -624,6 +630,55 @@ protected <T extends AggregationBuilder, V extends InternalAggregation> void tes
         }
     }
 
+    protected <T extends AggregationBuilder, V extends InternalAggregation> void multiIndexTestCase(
+        T aggregationBuilder,
+        Query query,
+        List<CheckedConsumer<RandomIndexWriter, IOException>> indexBuilders,
+        Consumer<V> verify,
+        MappedFieldType... fieldTypes
+    ) throws IOException {
+        Directory[] directories = new Directory[indexBuilders.size()];
+        boolean timeSeries = aggregationBuilder.isInSortOrderExecutionRequired();
+        try {
+            for (int i = 0; i < indexBuilders.size(); i++) {
+                directories[i] = newDirectory();
+                IndexWriterConfig config = LuceneTestCase.newIndexWriterConfig(random(), new MockAnalyzer(random()));
+                if (timeSeries) {
+                    Sort sort = new Sort(
+                        new SortField(TimeSeriesIdFieldMapper.NAME, SortField.Type.STRING, false),
+                        new SortedNumericSortField(DataStreamTimestampFieldMapper.DEFAULT_PATH, SortField.Type.LONG, true)
+                    );
+                    config.setIndexSort(sort);
+                }
+                RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directories[i], config);
+                indexBuilders.get(i).accept(indexWriter);
+                indexWriter.close();
+            }
+            // construct the multi-reader
+            List<DirectoryReader> directoryReaders = new ArrayList<>();
+            try {
+                for (Directory directory : directories) {
+                    DirectoryReader open = DirectoryReader.open(directory);
+                    directoryReaders.add(open);
+                }
+                DirectoryReader[] readers = directoryReaders.toArray(new DirectoryReader[0]);
+                try (MultiReader multiReader = new MultiReader(readers)) {
+                    IndexSearcher indexSearcher = newIndexSearcher(multiReader);
+
+                    V agg = searchAndReduce(new AggTestConfig(indexSearcher, query, aggregationBuilder, fieldTypes));
+                    verify.accept(agg);
+
+                    verifyOutputFieldNames(aggregationBuilder, agg);
+                }
+            } finally {
+                IOUtils.close(directoryReaders);
+            }
+        } finally {
+            IOUtils.close(directories);
+        }
+
+    }
+
     protected void withIndex(
         CheckedConsumer<RandomIndexWriter, IOException> buildIndex,
         CheckedConsumer<IndexSearcher, IOException> consume
@@ -1421,26 +1476,31 @@ public record AggTestConfig(
         AggregationBuilder builder,
         int maxBuckets,
         boolean splitLeavesIntoSeparateAggregators,
+        boolean shouldBeCached,
         MappedFieldType... fieldTypes
     ) {
         public AggTestConfig(IndexSearcher searcher, AggregationBuilder builder, MappedFieldType... fieldTypes) {
-            this(searcher, new MatchAllDocsQuery(), builder, DEFAULT_MAX_BUCKETS, randomBoolean(), fieldTypes);
+            this(searcher, new MatchAllDocsQuery(), builder, DEFAULT_MAX_BUCKETS, randomBoolean(), true, fieldTypes);
         }
 
         public AggTestConfig(IndexSearcher searcher, Query query, AggregationBuilder builder, MappedFieldType... fieldTypes) {
-            this(searcher, query, builder, DEFAULT_MAX_BUCKETS, randomBoolean(), fieldTypes);
+            this(searcher, query, builder, DEFAULT_MAX_BUCKETS, randomBoolean(), true, fieldTypes);
         }
 
         public AggTestConfig withQuery(Query query) {
-            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, fieldTypes);
+            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, shouldBeCached, fieldTypes);
         }
 
         public AggTestConfig withSplitLeavesIntoSeperateAggregators(boolean splitLeavesIntoSeparateAggregators) {
-            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, fieldTypes);
+            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, shouldBeCached, fieldTypes);
+        }
+
+        public AggTestConfig withShouldBeCached(boolean shouldBeCached) {
+            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, shouldBeCached, fieldTypes);
         }
 
         public AggTestConfig withMaxBuckets(int maxBuckets) {
-            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, fieldTypes);
+            return new AggTestConfig(searcher, query, builder, maxBuckets, splitLeavesIntoSeparateAggregators, shouldBeCached, fieldTypes);
         }
     }
 }