[RollupV2] Implement search resolution

server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java

@@ -72,7 +72,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
     private final BiFunction<String, String, Transport.Connection> nodeIdToConnection;
     private final SearchTask task;
     protected final SearchPhaseResults<Result> results;
-    private final ClusterState clusterState;
+    protected final ClusterState clusterState;
     private final Map<String, AliasFilter> aliasFilter;
     private final Map<String, Float> concreteIndexBoosts;
     private final SetOnce<AtomicArray<ShardSearchFailure>> shardFailures = new SetOnce<>();

server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java

@@ -11,10 +11,13 @@
 import org.apache.lucene.util.FixedBitSet;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexAbstraction;
 import org.elasticsearch.cluster.routing.GroupShardsIterator;
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.index.query.CoordinatorRewriteContext;
 import org.elasticsearch.index.query.CoordinatorRewriteContextProvider;
+import org.elasticsearch.rollup.RollupV2;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.SearchService.CanMatchResponse;
 import org.elasticsearch.search.SearchShardTarget;
@@ -27,6 +30,7 @@
 import org.elasticsearch.transport.Transport;
 
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -107,6 +111,11 @@ private GroupShardsIterator<SearchShardIterator> getIterator(CanMatchSearchPhase
             }
             possibleMatches.set(shardIndexToQuery);
         }
+
+        if (RollupV2.isEnabled()) {
+            possibleMatches = chooseOptimalRollupShards(results, shardsIts, possibleMatches);
+        }
+
         SearchSourceBuilder source = getRequest().source();
         int i = 0;
         for (SearchShardIterator iter : shardsIts) {
@@ -123,6 +132,69 @@ private GroupShardsIterator<SearchShardIterator> getIterator(CanMatchSearchPhase
         return new GroupShardsIterator<>(sortShards(shardsIts, results.minAndMaxes, fieldSort.order()));
     }
 
+    /**
+     * Iterate over the shard iterator and search for overlapping rollup indices. If rollup indices exist,
+     * the method includes the shards from the optimal rollup index that can match the query and ignores
+     * the shards from other overlapping indices.
+     *
+     * @return A {@link FixedBitSet} with the shards of the optimal indices.
+     */
+    private FixedBitSet chooseOptimalRollupShards(CanMatchSearchPhaseResults results,
+                                                  GroupShardsIterator<SearchShardIterator> shardsIts,
+                                                  FixedBitSet possibleMatches) {
+        // Map with key the index UUID that will be replaced by an optimal rollup index. The optimal index
+        // will be its substitute. The value is a tuple containing the name of the substitute index
+        // and its priority. In the end all indices contained in the keys will be ignored.
+        Map<String, Tuple<String, Long>> indexSubstitute = new HashMap<>(shardsIts.size());
+
+        int i = 0;
+        for (SearchShardIterator iter : shardsIts) {
+            CanMatchResponse result = results.getResult(i++);
+            if (result.canMatch()) {
+                String indexName = iter.shardId().getIndexName();
+                String indexUuid = iter.shardId().getIndex().getUUID();
+                Long priority = result.priority();
+                String sourceIndexUuid = result.sourceIndexUuid();
+                IndexAbstraction originalIndex = clusterState.getMetadata().getIndicesLookup().get(indexName);
+
+                // If index is not a member of a data stream or is not a rollup index, it will not be replaced
+                // Also, if an index has already been marked to be replaced, it should be skipped.
+                if (originalIndex.getParentDataStream() == null || sourceIndexUuid == null || priority == null
+                    || indexSubstitute.containsKey(indexUuid)) {
+                    continue;
+                }
+
+                if (indexSubstitute.containsKey(sourceIndexUuid)) {
+                    // Retrieve the previously optimal index and compare it with the current index
+                    // Find a new optimal index and replace source index and suboptimal rollup index
+                    // with the new optimal index.
+                    Tuple<String, Long> previousOptimalIndex = indexSubstitute.get(sourceIndexUuid);
+                    String newOptimalIndex = previousOptimalIndex.v2() >= priority ? previousOptimalIndex.v1() : indexUuid;
+                    Tuple<String, Long> optimalTuple = Tuple.tuple(newOptimalIndex, Math.max(priority, previousOptimalIndex.v2()));
+
+                    // Replace original index with the new optimal index
+                    indexSubstitute.put(sourceIndexUuid, optimalTuple);
+
+                    // Replace suboptimal index with the new optimal index
+                    String suboptimalIndex = indexUuid.equals(newOptimalIndex) == false ? indexUuid : previousOptimalIndex.v1();
+                    indexSubstitute.put(suboptimalIndex, optimalTuple);
+                } else {
+                    indexSubstitute.put(sourceIndexUuid, Tuple.tuple(indexUuid, priority));
+                }
+            }
+        }
+
+        FixedBitSet newMatches = possibleMatches.clone();
+        i = 0;
+        for (SearchShardIterator iter : shardsIts) {
+            if (newMatches.get(i) && indexSubstitute.containsKey(iter.shardId().getIndex().getUUID())) {
+                newMatches.clear(i);
+            }
+            i++;
+        }
+        return newMatches;
+    }
+
     @Override
     protected void performPhaseOnShard(int shardIndex, SearchShardIterator shardIt, SearchShardTarget shard) {
         CoordinatorRewriteContext coordinatorRewriteContext =
@@ -193,17 +265,19 @@ private static final class CanMatchSearchPhaseResults extends SearchPhaseResults
         private final FixedBitSet possibleMatches;
         private final MinAndMax<?>[] minAndMaxes;
         private int numPossibleMatches;
+        private final CanMatchResponse[] results;
 
         CanMatchSearchPhaseResults(int size) {
             super(size);
             possibleMatches = new FixedBitSet(size);
             minAndMaxes = new MinAndMax[size];
+            results = new CanMatchResponse[size];
         }
 
         @Override
         void consumeResult(CanMatchResponse result, Runnable next) {
             try {
-                consumeResult(result.getShardIndex(), result.canMatch(), result.estimatedMinAndMax());
+                consumeResult(result.getShardIndex(), result);
             } finally {
                 next.run();
             }
@@ -217,15 +291,16 @@ boolean hasResult(int shardIndex) {
         @Override
         void consumeShardFailure(int shardIndex) {
             // we have to carry over shard failures in order to account for them in the response.
-            consumeResult(shardIndex, true, null);
+            consumeResult(shardIndex, new CanMatchResponse(true, null, null, null));
         }
 
-        synchronized void consumeResult(int shardIndex, boolean canMatch, MinAndMax<?> minAndMax) {
-            if (canMatch) {
+        synchronized void consumeResult(int shardIndex, CanMatchResponse result) {
+            if (result.canMatch()) {
                 possibleMatches.set(shardIndex);
                 numPossibleMatches++;
             }
-            minAndMaxes[shardIndex] = minAndMax;
+            minAndMaxes[shardIndex] = result.estimatedMinAndMax();
+            results[shardIndex] = result;
         }
 
         synchronized int getNumPossibleMatches() {
@@ -240,5 +315,9 @@ synchronized FixedBitSet getPossibleMatches() {
         Stream<CanMatchResponse> getSuccessfulResults() {
             return Stream.empty();
         }
+
+        synchronized CanMatchResponse getResult(int shardIndex) {
+            return results[shardIndex];
+        }
     }
 }

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -20,6 +20,8 @@
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
+import org.elasticsearch.cluster.metadata.DataStream;
+import org.elasticsearch.cluster.metadata.IndexAbstraction;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -44,6 +46,7 @@
 import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.rollup.RollupV2;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.SearchShardTarget;
@@ -710,9 +713,26 @@ static boolean shouldPreFilterSearchShards(ClusterState clusterState,
         } else if (preFilterShardSize == null) {
             preFilterShardSize = SearchRequest.DEFAULT_PRE_FILTER_SHARD_SIZE;
         }
+        if (RollupV2.isEnabled() && hasRollupDatastream(indices, searchRequest.indices(), clusterState)) {
+            return true;
+        }
         return searchRequest.searchType() == QUERY_THEN_FETCH // we can't do this for DFS it needs to fan out to all shards all the time
-                    && (SearchService.canRewriteToMatchNone(source) || hasPrimaryFieldSort(source))
-                    && preFilterShardSize < numShards;
+            && (SearchService.canRewriteToMatchNone(source) || hasPrimaryFieldSort(source))
+            && preFilterShardSize < numShards;
+    }
+
+    private static boolean hasRollupDatastream(String[] indices, String[] requestIndices, ClusterState clusterState) {
+        Set<String> requestIndicesSet = Set.of(requestIndices);
+        for (String index : indices) {
+            IndexAbstraction originalIndex = clusterState.getMetadata().getIndicesLookup().get(index);
+            DataStream datastream = originalIndex.getParentDataStream() != null
+                ? originalIndex.getParentDataStream().getDataStream() : null;
+            IndexMetadata indexMetadata = clusterState.getMetadata().index(index);
+            if (datastream != null && indexMetadata.isRollupIndex() && requestIndicesSet.contains(index) == false) {
+                return true;
+            }
+        }
+        return false;
     }
 
     private static boolean hasReadOnlyIndices(String[] indices, ClusterState clusterState) {

server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java

@@ -577,6 +577,10 @@ public Index getResizeSourceIndex() {
     public static final Setting<String> INDEX_ROLLUP_SOURCE_NAME = Setting.simpleString(INDEX_ROLLUP_SOURCE_NAME_KEY,
         Property.IndexScope, Property.PrivateIndex);
 
+    public boolean isRollupIndex() {
+        return INDEX_ROLLUP_SOURCE_UUID.exists(settings);
+    }
+
     ImmutableOpenMap<String, DiffableStringMap> getCustomData() {
         return this.customData;
     }