diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java index 94154247b4b1f..561d61185c9d5 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java @@ -39,6 +39,7 @@ import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.aggregations.AbstractAggregationBuilder; import org.elasticsearch.search.aggregations.AggregationBuilder; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.Aggregations; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.AggregatorBase; @@ -694,7 +695,7 @@ public InternalAggregation buildEmptyAggregation() { public void close() {} @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { return LeafBucketCollector.NO_OP_COLLECTOR; } diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java index 1c78fe67703d1..239f75aae0461 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java @@ -8,7 +8,6 @@ package org.elasticsearch.search.aggregations; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.ScoreMode; import org.elasticsearch.core.CheckedFunction; import org.elasticsearch.search.profile.aggregation.InternalAggregationProfileTree; @@ -77,8 +76,8 @@ public final Aggregator subAggregator(String name) { } @Override - public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { - return delegate.getLeafCollector(ctx); + public final LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + return delegate.getLeafCollector(aggCtx); } @Override diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/AggregationExecutionContext.java b/server/src/main/java/org/elasticsearch/search/aggregations/AggregationExecutionContext.java new file mode 100644 index 0000000000000..eaf82541afbcd --- /dev/null +++ b/server/src/main/java/org/elasticsearch/search/aggregations/AggregationExecutionContext.java @@ -0,0 +1,40 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.search.aggregations; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.CheckedSupplier; + +import java.io.IOException; + +/** + * Used to preserve contextual information during aggregation execution. It can be used by search executors and parent + * aggregations to provide contextual information for the child aggregation during execution such as the currently executed + * time series id or the size of the current date histogram bucket. The information provided by this class is highly contextual and + * only valid during the {@link LeafBucketCollector#collect} call. + */ +public class AggregationExecutionContext { + + private final CheckedSupplier tsidProvider; + private final LeafReaderContext leafReaderContext; + + public AggregationExecutionContext(LeafReaderContext leafReaderContext, CheckedSupplier tsidProvider) { + this.leafReaderContext = leafReaderContext; + this.tsidProvider = tsidProvider; + } + + public LeafReaderContext getLeafReaderContext() { + return leafReaderContext; + } + + public BytesRef getTsid() throws IOException { + return tsidProvider != null ? tsidProvider.get() : null; + } +} diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java b/server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java index 2fcc62d986825..f72ad82d9372f 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java @@ -79,7 +79,7 @@ static void badState() { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext reader) { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) { badState(); assert false; return null; // unreachable but compiler does not agree @@ -201,6 +201,12 @@ public Map metadata() { */ protected abstract LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException; + // TODO: Remove this method in refactoring + protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub, AggregationExecutionContext aggCtx) + throws IOException { + return getLeafCollector(ctx, sub); + } + /** * Collect results for this leaf. *

@@ -210,10 +216,10 @@ public Map metadata() { * for more details on what this does. */ @Override - public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { - preGetSubLeafCollectors(ctx); - final LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(ctx); - return getLeafCollector(ctx, sub); + public final LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + preGetSubLeafCollectors(aggCtx.getLeafReaderContext()); + final LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(aggCtx); + return getLeafCollector(aggCtx.getLeafReaderContext(), sub, aggCtx); } /** diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java b/server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java index 8a5eee01799df..60907169df4ca 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java @@ -22,7 +22,7 @@ public abstract class BucketCollector implements Collector { public static final BucketCollector NO_OP_COLLECTOR = new BucketCollector() { @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext reader) { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) { return LeafBucketCollector.NO_OP_COLLECTOR; } @@ -36,14 +36,18 @@ public void postCollection() throws IOException { // no-op } - @Override public ScoreMode scoreMode() { return ScoreMode.COMPLETE_NO_SCORES; } }; + // TODO: will remove it in a follow up PR @Override - public abstract LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException; + public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + return getLeafCollector(new AggregationExecutionContext(ctx, null)); + } + + public abstract LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException; /** * Pre collection callback. diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java b/server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java index 7fe7b09b8f806..aa811b15d79b2 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java @@ -8,7 +8,6 @@ package org.elasticsearch.search.aggregations; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.CollectionTerminatedException; import org.apache.lucene.search.Collector; import org.apache.lucene.search.LeafCollector; @@ -89,9 +88,9 @@ public void postCollection() throws IOException { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { try { - LeafBucketCollector leafCollector = collector.getLeafCollector(ctx); + LeafBucketCollector leafCollector = collector.getLeafCollector(aggCtx); if (false == leafCollector.isNoop()) { return leafCollector; } @@ -169,11 +168,11 @@ public String toString() { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext context) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { final List leafCollectors = new ArrayList<>(collectors.length); for (BucketCollector collector : collectors) { try { - LeafBucketCollector leafCollector = collector.getLeafCollector(context); + LeafBucketCollector leafCollector = collector.getLeafCollector(aggCtx); if (false == leafCollector.isNoop()) { leafCollectors.add(leafCollector); } diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java index 385c2a5b9cc41..679bdcf5c0f7a 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java @@ -21,6 +21,7 @@ import org.apache.lucene.util.packed.PackedLongValues; import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.LongHash; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.InternalAggregation; @@ -110,7 +111,7 @@ private void clearLeaf() { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { finishLeaf(); return new LeafBucketCollector() { @@ -119,7 +120,7 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOExce @Override public void collect(int doc, long bucket) throws IOException { if (context == null) { - context = ctx; + context = aggCtx.getLeafReaderContext(); docDeltasBuilder = PackedLongValues.packedBuilder(PackedInts.DEFAULT); bucketsBuilder = PackedLongValues.packedBuilder(PackedInts.DEFAULT); } diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java index ded829cb76a44..c13b54b2286d0 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java @@ -8,8 +8,8 @@ package org.elasticsearch.search.aggregations.bucket; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.ScoreMode; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.InternalAggregation; @@ -95,7 +95,7 @@ public void collectDebugInfo(BiConsumer add) { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { throw new IllegalStateException( "Deferred collectors cannot be collected directly. They must be collected through the recording wrapper." ); diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java index abc737afca456..a50db167bf06c 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java @@ -22,6 +22,7 @@ import org.elasticsearch.common.util.ObjectArray; import org.elasticsearch.core.Releasable; import org.elasticsearch.core.Releasables; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; import org.elasticsearch.search.aggregations.MultiBucketCollector; @@ -77,8 +78,8 @@ public void setDeferredCollector(Iterable deferredCollectors) { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { - perSegCollector = new PerSegmentCollects(ctx); + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + perSegCollector = new PerSegmentCollects(aggCtx.getLeafReaderContext()); entries.add(perSegCollector); // Deferring collector diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesAggregator.java index 41b697e025e50..d00f7f07e160f 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesAggregator.java @@ -11,10 +11,8 @@ import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.util.BytesRef; import org.elasticsearch.core.Releasables; -import org.elasticsearch.index.fielddata.IndexFieldData; -import org.elasticsearch.index.fielddata.SortedBinaryDocValues; -import org.elasticsearch.index.fielddata.plain.SortedSetBytesLeafFieldData; import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.AggregatorFactories; import org.elasticsearch.search.aggregations.CardinalityUpperBound; @@ -29,11 +27,9 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; -import java.util.Objects; public class TimeSeriesAggregator extends BucketsAggregator { - private final IndexFieldData tsidFieldData; protected final BytesKeyedBucketOrds bucketOrds; private final boolean keyed; @@ -49,10 +45,6 @@ public TimeSeriesAggregator( ) throws IOException { super(name, factories, context, parent, bucketCardinality, metadata); this.keyed = keyed; - tsidFieldData = (IndexFieldData) Objects.requireNonNull( - context.buildFieldContext("_tsid"), - "Cannot obtain tsid field" - ).indexFieldData(); bucketOrds = BytesKeyedBucketOrds.build(bigArrays(), bucketCardinality); } @@ -98,20 +90,22 @@ protected void doClose() { @Override protected LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector sub) throws IOException { - final SortedBinaryDocValues tsids = tsidFieldData.load(context).getBytesValues(); + // TODO: remove this method in a follow up PR + throw new UnsupportedOperationException("Shouldn't be here"); + } + + protected LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector sub, AggregationExecutionContext aggCtx) + throws IOException { return new LeafBucketCollectorBase(sub, null) { @Override public void collect(int doc, long bucket) throws IOException { - if (tsids.advanceExact(doc)) { - BytesRef newTsid = tsids.nextValue(); - long bucketOrdinal = bucketOrds.add(bucket, newTsid); - if (bucketOrdinal < 0) { // already seen - bucketOrdinal = -1 - bucketOrdinal; - collectExistingBucket(sub, doc, bucketOrdinal); - } else { - collectBucket(sub, doc, bucketOrdinal); - } + long bucketOrdinal = bucketOrds.add(bucket, aggCtx.getTsid()); + if (bucketOrdinal < 0) { // already seen + bucketOrdinal = -1 - bucketOrdinal; + collectExistingBucket(sub, doc, bucketOrdinal); + } else { + collectBucket(sub, doc, bucketOrdinal); } } }; diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcher.java b/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcher.java index 311729cd3d63e..347a33238b751 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcher.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcher.java @@ -14,7 +14,6 @@ import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.Query; import org.apache.lucene.search.Scorer; import org.apache.lucene.search.Weight; @@ -24,6 +23,7 @@ import org.apache.lucene.util.PriorityQueue; import org.elasticsearch.cluster.metadata.DataStream; import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; @@ -60,13 +60,18 @@ public void search(Query query, BucketCollector bucketCollector) throws IOExcept if (++seen % CHECK_CANCELLED_SCORER_INTERVAL == 0) { checkCancelled(); } - LeafBucketCollector leafCollector = bucketCollector.getLeafCollector(leaf); Scorer scorer = weight.scorer(leaf); if (scorer != null) { - LeafWalker leafWalker = new LeafWalker(leaf, scorer, leafCollector); + LeafWalker leafWalker = new LeafWalker(leaf, scorer, bucketCollector, leaf); if (leafWalker.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) { leafWalkers.add(leafWalker); } + } else { + // Even though we will not walk through this aggregation as a part of normal processing + // this is needed to trigger actions in some bucketCollectors that bypass the normal iteration logic + // for example, global aggregator triggers a separate iterator that ignores the query but still needs + // to know all leaves + bucketCollector.getLeafCollector(new AggregationExecutionContext(leaf, null)); } } @@ -148,7 +153,7 @@ private void checkCancelled() { } private static class LeafWalker { - private final LeafCollector collector; + private final LeafBucketCollector collector; private final Bits liveDocs; private final DocIdSetIterator iterator; private final SortedDocValues tsids; @@ -158,8 +163,9 @@ private static class LeafWalker { int tsidOrd; long timestamp; - LeafWalker(LeafReaderContext context, Scorer scorer, LeafCollector collector) throws IOException { - this.collector = collector; + LeafWalker(LeafReaderContext context, Scorer scorer, BucketCollector bucketCollector, LeafReaderContext leaf) throws IOException { + AggregationExecutionContext aggCtx = new AggregationExecutionContext(leaf, scratch::get); + this.collector = bucketCollector.getLeafCollector(aggCtx); liveDocs = context.reader().getLiveDocs(); this.collector.setScorer(scorer); iterator = scorer.iterator(); diff --git a/server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java b/server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java index bc09d1d0577ef..165b4e758150c 100644 --- a/server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java @@ -8,8 +8,8 @@ package org.elasticsearch.search.profile.aggregation; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.ScoreMode; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.InternalAggregation; import org.elasticsearch.search.aggregations.LeafBucketCollector; @@ -87,11 +87,11 @@ public InternalAggregation buildEmptyAggregation() { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { Timer timer = profileBreakdown.getTimer(AggregationTimingType.BUILD_LEAF_COLLECTOR); timer.start(); try { - return new ProfilingLeafBucketCollector(delegate.getLeafCollector(ctx), profileBreakdown); + return new ProfilingLeafBucketCollector(delegate.getLeafCollector(aggCtx), profileBreakdown); } finally { timer.stop(); } diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java index 9f909befba95d..8472e56bc568f 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java @@ -59,11 +59,11 @@ private static class TerminateAfterBucketCollector extends BucketCollector { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext context) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { if (count >= terminateAfter) { return LeafBucketCollector.NO_OP_COLLECTOR; } - final LeafBucketCollector leafCollector = in.getLeafCollector(context); + final LeafBucketCollector leafCollector = in.getLeafCollector(aggCtx); return new LeafBucketCollectorBase(leafCollector, null) { @Override public void collect(int doc, long bucket) throws IOException { @@ -95,7 +95,7 @@ private static class TotalHitCountBucketCollector extends BucketCollector { TotalHitCountBucketCollector() {} @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext context) { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) { return new LeafBucketCollector() { @Override public void collect(int doc, long bucket) throws IOException { @@ -130,8 +130,8 @@ private static class SetScorerBucketCollector extends BucketCollector { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext context) throws IOException { - final LeafBucketCollector leafCollector = in.getLeafCollector(context); + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + final LeafBucketCollector leafCollector = in.getLeafCollector(aggCtx); return new LeafBucketCollectorBase(leafCollector, null) { @Override public void setScorer(Scorable scorer) throws IOException { @@ -252,7 +252,7 @@ public void testSetScorerAfterCollectionTerminated() throws IOException { Collections.shuffle(collectors, random()); BucketCollector collector = MultiBucketCollector.wrap(true, collectors); - LeafBucketCollector leafCollector = collector.getLeafCollector(null); + LeafBucketCollector leafCollector = collector.getLeafCollector((LeafReaderContext) null); leafCollector.setScorer(scorer); assertTrue(setScorerCalled1.get()); assertTrue(setScorerCalled2.get()); diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java index fc683f4b9fd68..f1b6fa6488f5c 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java @@ -14,7 +14,6 @@ import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.Term; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; @@ -26,6 +25,7 @@ import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.RandomIndexWriter; import org.elasticsearch.common.CheckedBiConsumer; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.AggregatorTestCase; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; @@ -101,11 +101,11 @@ public ScoreMode scoreMode() { private BucketCollector bla(Set docIds) { return new BucketCollector() { @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { return new LeafBucketCollector() { @Override public void collect(int doc, long bucket) throws IOException { - docIds.add(ctx.docBase + doc); + docIds.add(aggCtx.getLeafReaderContext().docBase + doc); } }; } @@ -212,8 +212,8 @@ public void preCollection() throws IOException {} public void postCollection() throws IOException {} @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { - LeafBucketCollector delegate = deferringCollector.getLeafCollector(ctx); + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + LeafBucketCollector delegate = deferringCollector.getLeafCollector(aggCtx); return leafCollector.apply(deferringCollector, delegate); } }); @@ -232,7 +232,7 @@ public ScoreMode scoreMode() { } @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { return new LeafBucketCollector() { @Override public void collect(int doc, long owningBucketOrd) throws IOException { diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java index 8dfee6e3ad233..060a3b0befbb0 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java @@ -12,7 +12,6 @@ import org.apache.lucene.document.StringField; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.Term; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.ScoreDoc; @@ -25,6 +24,7 @@ import org.elasticsearch.common.util.MockBigArrays; import org.elasticsearch.common.util.MockPageCacheRecycler; import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.AggregatorTestCase; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; @@ -81,11 +81,11 @@ public void testReplay() throws Exception { private BucketCollector testCollector(Set docIds) { return new BucketCollector() { @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { return new LeafBucketCollector() { @Override public void collect(int doc, long bucket) throws IOException { - docIds.add(ctx.docBase + doc); + docIds.add(aggCtx.getLeafReaderContext().docBase + doc); } }; } diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesCancellationTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesCancellationTests.java index 2cc1e03cd2212..e3c3d8cd30e25 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesCancellationTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesCancellationTests.java @@ -12,7 +12,6 @@ import org.apache.lucene.document.SortedDocValuesField; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.ScoreMode; @@ -24,6 +23,7 @@ import org.elasticsearch.cluster.metadata.DataStream; import org.elasticsearch.core.internal.io.IOUtils; import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; import org.elasticsearch.search.internal.ContextIndexSearcher; @@ -101,7 +101,7 @@ public static class CountingBucketCollector extends BucketCollector { public AtomicInteger count = new AtomicInteger(); @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { return new LeafBucketCollector() { @Override public void collect(int doc, long owningBucketOrd) throws IOException { diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcherTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcherTests.java index 3705ed73a9fe6..0b960c4648962 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcherTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/timeseries/TimeSeriesIndexSearcherTests.java @@ -15,7 +15,6 @@ import org.apache.lucene.index.DocValues; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.NumericDocValues; import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.search.IndexSearcher; @@ -28,6 +27,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.cluster.metadata.DataStream; import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper; +import org.elasticsearch.search.aggregations.AggregationExecutionContext; import org.elasticsearch.search.aggregations.BucketCollector; import org.elasticsearch.search.aggregations.LeafBucketCollector; import org.elasticsearch.test.ESTestCase; @@ -95,9 +95,12 @@ public void testCollectInOrderAcrossSegments() throws IOException, InterruptedEx long total = 0; @Override - public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException { - SortedDocValues tsid = DocValues.getSorted(ctx.reader(), TimeSeriesIdFieldMapper.NAME); - NumericDocValues timestamp = DocValues.getNumeric(ctx.reader(), DataStream.TimestampField.FIXED_TIMESTAMP_FIELD); + public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException { + SortedDocValues tsid = DocValues.getSorted(aggCtx.getLeafReaderContext().reader(), TimeSeriesIdFieldMapper.NAME); + NumericDocValues timestamp = DocValues.getNumeric( + aggCtx.getLeafReaderContext().reader(), + DataStream.TimestampField.FIXED_TIMESTAMP_FIELD + ); return new LeafBucketCollector() { @Override