[8.16] Allow for queries on _tier to skip shards during coordinator rewrite (elastic#114990) (elastic#115513)

* Allow for queries on _tier to skip shards during coordinator rewrite (elastic#114990)

The `_tier` metadata field was not used on the  coordinator when
rewriting queries in order to exclude shards that don't match. This lead
to queries in the following form to continue to report failures even
though the only unavailable shards were in the  tier that was excluded
from search (frozen tier in this example):

```
POST testing/_search
{
  "query": {
    "bool": {
      "must_not": [
        {
          "term": {
            "_tier": "data_frozen"
          }
        }
      ]
    }
  }
}
```

This PR addresses this by having the queries that can execute on `_tier`
(term, match, query string, simple query string, prefix, wildcard)
execute a coordinator rewrite to  exclude the indices that don't match
the `_tier` query  **before** attempting to reach to the shards (shards,
that might not be available and raise errors). 

Fixes elastic#114910

* Don't use getFirst

* Test compile

Author: andreidan

docs/changelog/114990.yaml

@@ -0,0 +1,6 @@
+pr: 114990
+summary: Allow for querries on `_tier` to skip shards in the `can_match` phase
+area: Search
+type: bug
+issues:
+ - 114910

server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java

@@ -9,17 +9,23 @@
 
 package org.elasticsearch.index.query;
 
+import org.apache.lucene.search.Query;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.index.mapper.ConstantFieldType;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MappingLookup;
+import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.shard.IndexLongFieldRange;
 import org.elasticsearch.indices.DateFieldRangeInfo;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 
 import java.util.Collections;
+import java.util.Map;
 import java.util.function.LongSupplier;
 
 /**
@@ -30,20 +36,57 @@
  * and skip the shards that don't hold queried data. See IndexMetadata for more details.
  */
 public class CoordinatorRewriteContext extends QueryRewriteContext {
+
+    public static final String TIER_FIELD_NAME = "_tier";
+
+    private static final ConstantFieldType TIER_FIELD_TYPE = new ConstantFieldType(TIER_FIELD_NAME, Map.of()) {
+        @Override
+        public ValueFetcher valueFetcher(SearchExecutionContext context, String format) {
+            throw new UnsupportedOperationException("fetching field values is not supported on the coordinator node");
+        }
+
+        @Override
+        public String typeName() {
+            return TIER_FIELD_NAME;
+        }
+
+        @Override
+        protected boolean matches(String pattern, boolean caseInsensitive, QueryRewriteContext context) {
+            if (caseInsensitive) {
+                pattern = Strings.toLowercaseAscii(pattern);
+            }
+
+            String tierPreference = context.getTierPreference();
+            if (tierPreference == null) {
+                return false;
+            }
+            return Regex.simpleMatch(pattern, tierPreference);
+        }
+
+        @Override
+        public Query existsQuery(SearchExecutionContext context) {
+            throw new UnsupportedOperationException("field exists query is not supported on the coordinator node");
+        }
+    };
+
     private final DateFieldRangeInfo dateFieldRangeInfo;
+    private final String tier;
 
     /**
      * Context for coordinator search rewrites based on time ranges for the @timestamp field and/or 'event.ingested' field
+     *
      * @param parserConfig
      * @param client
      * @param nowInMillis
      * @param dateFieldRangeInfo range and field type info for @timestamp and 'event.ingested'
+     * @param tier               the configured data tier (via the _tier_preference setting) for the index
      */
     public CoordinatorRewriteContext(
         XContentParserConfiguration parserConfig,
         Client client,
         LongSupplier nowInMillis,
-        DateFieldRangeInfo dateFieldRangeInfo
+        DateFieldRangeInfo dateFieldRangeInfo,
+        String tier
     ) {
         super(
             parserConfig,
@@ -63,10 +106,12 @@ public CoordinatorRewriteContext(
             null
         );
         this.dateFieldRangeInfo = dateFieldRangeInfo;
+        this.tier = tier;
     }
 
     /**
-     * @param fieldName Must be one of DataStream.TIMESTAMP_FIELD_FIELD or IndexMetadata.EVENT_INGESTED_FIELD_NAME
+     * @param fieldName Must be one of DataStream.TIMESTAMP_FIELD_FIELD, IndexMetadata.EVENT_INGESTED_FIELD_NAME, or
+     *                  DataTierFiledMapper.NAME
      * @return MappedField with type for the field. Returns null if fieldName is not one of the allowed field names.
      */
     @Nullable
@@ -75,6 +120,8 @@ public MappedFieldType getFieldType(String fieldName) {
             return dateFieldRangeInfo.timestampFieldType();
         } else if (IndexMetadata.EVENT_INGESTED_FIELD_NAME.equals(fieldName)) {
             return dateFieldRangeInfo.eventIngestedFieldType();
+        } else if (TIER_FIELD_NAME.equals(fieldName)) {
+            return TIER_FIELD_TYPE;
         } else {
             return null;
         }
@@ -99,4 +146,18 @@ public IndexLongFieldRange getFieldRange(String fieldName) {
     public CoordinatorRewriteContext convertToCoordinatorRewriteContext() {
         return this;
     }
+
+    @Override
+    public String getTierPreference() {
+        // dominant branch first (tier preference is configured)
+        return tier.isEmpty() == false ? tier : null;
+    }
+
+    /**
+     * We're holding on to the index tier in the context as otherwise we'd need
+     * to re-parse it from the index settings when evaluating the _tier field.
+     */
+    public String tier() {
+        return tier;
+    }
 }

server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContextProvider.java

@@ -52,6 +52,12 @@ public CoordinatorRewriteContext getCoordinatorRewriteContext(Index index) {
             return null;
         }
         DateFieldRangeInfo dateFieldRangeInfo = mappingSupplier.apply(index);
+        // we've now added a coordinator rewrite based on the _tier field so the requirement
+        // for the timestamps fields to be present is artificial (we could do a coordinator
+        // rewrite only based on the _tier field) and we might decide to remove this artificial
+        // limitation to enable coordinator rewrites based on _tier for hot and warm indices
+        // (currently the _tier coordinator rewrite is only available for mounted and partially mounted
+        // indices)
         if (dateFieldRangeInfo == null) {
             return null;
         }
@@ -74,7 +80,8 @@ public CoordinatorRewriteContext getCoordinatorRewriteContext(Index index) {
             parserConfig,
             client,
             nowInMillis,
-            new DateFieldRangeInfo(timestampFieldType, timestampRange, dateFieldRangeInfo.eventIngestedFieldType(), eventIngestedRange)
+            new DateFieldRangeInfo(timestampFieldType, timestampRange, dateFieldRangeInfo.eventIngestedFieldType(), eventIngestedRange),
+            indexMetadata.getTierPreference().isEmpty() == false ? indexMetadata.getTierPreference().get(0) : ""
         );
     }
 }

server/src/main/java/org/elasticsearch/index/query/PrefixQueryBuilder.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.mapper.ConstantFieldType;
 import org.elasticsearch.index.mapper.MappedFieldType;
@@ -189,11 +190,24 @@ public String getWriteableName() {
     }
 
     @Override
-    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throws IOException {
+    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) {
         MappedFieldType fieldType = context.getFieldType(this.fieldName);
         if (fieldType == null) {
             return new MatchNoneQueryBuilder("The \"" + getName() + "\" query is against a field that does not exist");
-        } else if (fieldType instanceof ConstantFieldType constantFieldType) {
+        }
+        return maybeRewriteBasedOnConstantFields(fieldType, context);
+    }
+
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(this.fieldName);
+        // we don't rewrite a null field type to `match_none` on the coordinator because the coordinator has access
+        // to only a subset of fields see {@link CoordinatorRewriteContext#getFieldType}
+        return maybeRewriteBasedOnConstantFields(fieldType, coordinatorRewriteContext);
+    }
+
+    private QueryBuilder maybeRewriteBasedOnConstantFields(@Nullable MappedFieldType fieldType, QueryRewriteContext context) {
+        if (fieldType instanceof ConstantFieldType constantFieldType) {
             // This logic is correct for all field types, but by only applying it to constant
             // fields we also have the guarantee that it doesn't perform I/O, which is important
             // since rewrites might happen on a network thread.

server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -11,9 +11,12 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ResolvedIndices;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.cluster.routing.allocation.DataTier;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.regex.Regex;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.CountDown;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
@@ -406,4 +409,22 @@ public ResolvedIndices getResolvedIndices() {
     public PointInTimeBuilder getPointInTimeBuilder() {
         return pit;
     }
+
+    /**
+     * Retrieve the first tier preference from the index setting. If the setting is not
+     * present, then return null.
+     */
+    @Nullable
+    public String getTierPreference() {
+        Settings settings = getIndexSettings().getSettings();
+        String value = DataTier.TIER_PREFERENCE_SETTING.get(settings);
+
+        if (Strings.hasText(value) == false) {
+            return null;
+        }
+
+        // Tier preference can be a comma-delimited list of tiers, ordered by preference
+        // It was decided we should only test the first of these potentially multiple preferences.
+        return value.split(",")[0].trim();
+    }
 }

server/src/main/java/org/elasticsearch/index/query/TermQueryBuilder.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.mapper.ConstantFieldType;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.xcontent.ParseField;
@@ -170,11 +171,24 @@ protected void addExtraXContent(XContentBuilder builder, Params params) throws I
     }
 
     @Override
-    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throws IOException {
+    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) {
         MappedFieldType fieldType = context.getFieldType(this.fieldName);
         if (fieldType == null) {
             return new MatchNoneQueryBuilder("The \"" + getName() + "\" query is against a field that does not exist");
-        } else if (fieldType instanceof ConstantFieldType constantFieldType) {
+        }
+        return maybeRewriteBasedOnConstantFields(fieldType, context);
+    }
+
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(this.fieldName);
+        // we don't rewrite a null field type to `match_none` on the coordinator because the coordinator has access
+        // to only a subset of fields see {@link CoordinatorRewriteContext#getFieldType}
+        return maybeRewriteBasedOnConstantFields(fieldType, coordinatorRewriteContext);
+    }
+
+    private QueryBuilder maybeRewriteBasedOnConstantFields(@Nullable MappedFieldType fieldType, QueryRewriteContext context) {
+        if (fieldType instanceof ConstantFieldType constantFieldType) {
             // This logic is correct for all field types, but by only applying it to constant
             // fields we also have the guarantee that it doesn't perform I/O, which is important
             // since rewrites might happen on a network thread.

server/src/main/java/org/elasticsearch/index/query/TermsQueryBuilder.java

@@ -393,11 +393,24 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
     }
 
     @Override
-    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throws IOException {
+    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) {
         MappedFieldType fieldType = context.getFieldType(this.fieldName);
         if (fieldType == null) {
             return new MatchNoneQueryBuilder("The \"" + getName() + "\" query is against a field that does not exist");
-        } else if (fieldType instanceof ConstantFieldType constantFieldType) {
+        }
+        return maybeRewriteBasedOnConstantFields(fieldType, context);
+    }
+
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(this.fieldName);
+        // we don't rewrite a null field type to `match_none` on the coordinator because the coordinator has access
+        // to only a subset of fields see {@link CoordinatorRewriteContext#getFieldType}
+        return maybeRewriteBasedOnConstantFields(fieldType, coordinatorRewriteContext);
+    }
+
+    private QueryBuilder maybeRewriteBasedOnConstantFields(@Nullable MappedFieldType fieldType, QueryRewriteContext context) {
+        if (fieldType instanceof ConstantFieldType constantFieldType) {
             // This logic is correct for all field types, but by only applying it to constant
             // fields we also have the guarantee that it doesn't perform I/O, which is important
             // since rewrites might happen on a network thread.

server/src/main/java/org/elasticsearch/index/query/WildcardQueryBuilder.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.mapper.ConstantFieldType;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.support.QueryParsers;
@@ -200,11 +201,24 @@ public static WildcardQueryBuilder fromXContent(XContentParser parser) throws IO
     }
 
     @Override
-    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throws IOException {
+    protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) {
         MappedFieldType fieldType = context.getFieldType(this.fieldName);
         if (fieldType == null) {
-            return new MatchNoneQueryBuilder("The \"" + getName() + "\" query  is against a field that does not exist");
-        } else if (fieldType instanceof ConstantFieldType constantFieldType) {
+            return new MatchNoneQueryBuilder("The \"" + getName() + "\" query is against a field that does not exist");
+        }
+        return maybeRewriteBasedOnConstantFields(fieldType, context);
+    }
+
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(this.fieldName);
+        // we don't rewrite a null field type to `match_none` on the coordinator because the coordinator has access
+        // to only a subset of fields see {@link CoordinatorRewriteContext#getFieldType}
+        return maybeRewriteBasedOnConstantFields(fieldType, coordinatorRewriteContext);
+    }
+
+    private QueryBuilder maybeRewriteBasedOnConstantFields(@Nullable MappedFieldType fieldType, QueryRewriteContext context) {
+        if (fieldType instanceof ConstantFieldType constantFieldType) {
             // This logic is correct for all field types, but by only applying it to constant
             // fields we also have the guarantee that it doesn't perform I/O, which is important
             // since rewrites might happen on a network thread.

server/src/test/java/org/elasticsearch/index/query/PrefixQueryBuilderTests.java

@@ -17,7 +17,9 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.core.Strings;
+import org.elasticsearch.index.mapper.DateFieldMapper;
 import org.elasticsearch.test.AbstractQueryTestCase;
+import org.hamcrest.CoreMatchers;
 import org.hamcrest.Matchers;
 
 import java.io.IOException;
@@ -175,4 +177,37 @@ public void testMustRewrite() throws IOException {
         IllegalStateException e = expectThrows(IllegalStateException.class, () -> queryBuilder.toQuery(context));
         assertEquals("Rewrite first", e.getMessage());
     }
+
+    public void testCoordinatorTierRewriteToMatchAll() throws IOException {
+        QueryBuilder query = new PrefixQueryBuilder("_tier", "data_fro");
+        final String timestampFieldName = "@timestamp";
+        long minTimestamp = 1685714000000L;
+        long maxTimestamp = 1685715000000L;
+        final CoordinatorRewriteContext coordinatorRewriteContext = createCoordinatorRewriteContext(
+            new DateFieldMapper.DateFieldType(timestampFieldName),
+            minTimestamp,
+            maxTimestamp,
+            "data_frozen"
+        );
+
+        QueryBuilder rewritten = query.rewrite(coordinatorRewriteContext);
+        assertThat(rewritten, CoreMatchers.instanceOf(MatchAllQueryBuilder.class));
+    }
+
+    public void testCoordinatorTierRewriteToMatchNone() throws IOException {
+        QueryBuilder query = QueryBuilders.boolQuery().mustNot(new PrefixQueryBuilder("_tier", "data_fro"));
+        final String timestampFieldName = "@timestamp";
+        long minTimestamp = 1685714000000L;
+        long maxTimestamp = 1685715000000L;
+        final CoordinatorRewriteContext coordinatorRewriteContext = createCoordinatorRewriteContext(
+            new DateFieldMapper.DateFieldType(timestampFieldName),
+            minTimestamp,
+            maxTimestamp,
+            "data_frozen"
+        );
+
+        QueryBuilder rewritten = query.rewrite(coordinatorRewriteContext);
+        assertThat(rewritten, CoreMatchers.instanceOf(MatchNoneQueryBuilder.class));
+    }
+
 }