More queries and tests

andreidan · Oct 18, 2024 · a1d3bfd · a1d3bfd
1 parent 8510222
commit a1d3bfd
Show file tree

Hide file tree

Showing 9 changed files with 164 additions and 15 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java
@@ -10,6 +10,7 @@
 package org.elasticsearch.index.query;
 
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.NamedMatches;
 import org.apache.lucene.search.Query;
@@ -21,6 +22,8 @@
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.SuggestingErrorOnUnknown;
+import org.elasticsearch.index.mapper.DataTierFieldMapper;
+import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.xcontent.AbstractObjectParser;
 import org.elasticsearch.xcontent.FilterXContentParser;
 import org.elasticsearch.xcontent.FilterXContentParserWrapper;
@@ -311,6 +314,9 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
         }
         final QueryRewriteContext context = queryRewriteContext.convertToIndexMetadataContext();
         if (context != null) {
+            // Even though we're still on the coordinator, this executes at the shard level.
+            // Any optimisations around skipping unavailable shards should be performed in {@link AbstractQueryBuilder#doCoordinatorRewrite}
+            // as that will execute for unavailable shards as well (which can be skipped if needed)
             return doIndexMetadataRewrite(context);
         }
         return this;
@@ -326,6 +332,47 @@ protected QueryBuilder doCoordinatorRewrite(final CoordinatorRewriteContext coor
         return this;
     }
 
+    /**
+     * Enables a coordinator rewrite based on the _tier field for queries that can be expressed
+     * as a term query (i.e. match, term, simple query etc).
+     * @param query the source query to rewrite
+     * @param caseInsensitive some queries are optionally case sensitive, this controls if the underlying
+     *                        term query we execute on _tier will be case sensitive or insensitive
+     * @param queryFieldName the name of the field the source query is meant to execute against
+     * @param queryFieldValue the value of the source query field
+     * @param coordinatorRewriteContext the coordinator rewrite context
+     * @return the rewritten query if applicable
+     */
+    protected static QueryBuilder tierFieldTermQueryCoordinatorRewriteIfPresent(
+        QueryBuilder query,
+        boolean caseInsensitive,
+        String queryFieldName,
+        Object queryFieldValue,
+        CoordinatorRewriteContext coordinatorRewriteContext
+    ) {
+        if (queryFieldName.equals(DataTierFieldMapper.NAME) == false) {
+            return query;
+        }
+        final MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(DataTierFieldMapper.NAME);
+        if (fieldType instanceof final DataTierFieldMapper.DataTierFieldType tierFieldType) {
+            Query tierFieldQuery;
+            if (caseInsensitive) {
+                tierFieldQuery = tierFieldType.internalTermQueryCaseInsensitive(queryFieldValue, coordinatorRewriteContext);
+            } else {
+                tierFieldQuery = tierFieldType.internalTermQuery(queryFieldValue, coordinatorRewriteContext);
+            }
+
+            if (tierFieldQuery instanceof MatchNoDocsQuery) {
+                return new MatchNoneQueryBuilder("The \"" + query.getName() + "\" query was rewritten to a \"match_none\" query.");
+            } else if (tierFieldQuery instanceof MatchAllDocsQuery) {
+                return new MatchAllQueryBuilder();
+            } else {
+                assert false : "Constant fields must produce match-all or match-none queries, got " + query;
+            }
+        }
+        return query;
+    }
+
     /**
      * @param searchExecutionContext A {@link QueryRewriteContext} that enables full rewrite capabilities
      *                               happening on the data node with all information available for rewriting.

diff --git a/server/src/main/java/org/elasticsearch/index/query/MatchQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/MatchQueryBuilder.java
@@ -385,6 +385,11 @@ protected QueryBuilder doIndexMetadataRewrite(QueryRewriteContext context) throw
         return this;
     }
 
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        return tierFieldTermQueryCoordinatorRewriteIfPresent(this, true, fieldName, value, coordinatorRewriteContext);
+    }
+
     private NamedAnalyzer configuredAnalyzer(QueryRewriteContext context) {
         if (analyzer != null) {
             return context.getIndexAnalyzers().get(analyzer);

diff --git a/server/src/main/java/org/elasticsearch/index/query/PrefixQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/PrefixQueryBuilder.java
@@ -22,6 +22,7 @@
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.mapper.ConstantFieldType;
+import org.elasticsearch.index.mapper.DataTierFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.support.QueryParsers;
 import org.elasticsearch.xcontent.ParseField;
@@ -234,6 +235,25 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return fieldType.prefixQuery(value, method, caseInsensitive, context);
     }
 
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        if (fieldName.equals(DataTierFieldMapper.NAME) == false) {
+            return this;
+        }
+        final MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(DataTierFieldMapper.NAME);
+        if (fieldType instanceof final DataTierFieldMapper.DataTierFieldType tierFieldType) {
+            Query tierFieldQuery = tierFieldType.prefixQuery(value, true, coordinatorRewriteContext);
+            if (tierFieldQuery instanceof MatchNoDocsQuery) {
+                return new MatchNoneQueryBuilder("The \"" + getName() + "\" query was rewritten to a \"match_none\" query.");
+            } else if (tierFieldQuery instanceof MatchAllDocsQuery) {
+                return new MatchAllQueryBuilder();
+            } else {
+                assert false : "Constant fields must produce match-all or match-none queries, got " + tierFieldQuery;
+            }
+        }
+        return this;
+    }
+
     @Override
     protected final int doHashCode() {
         return Objects.hash(fieldName, value, rewrite, caseInsensitive);

diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryStringQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/QueryStringQueryBuilder.java
@@ -23,6 +23,7 @@
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.index.analysis.NamedAnalyzer;
+import org.elasticsearch.index.mapper.DataTierFieldMapper;
 import org.elasticsearch.index.query.support.QueryParsers;
 import org.elasticsearch.index.search.QueryParserHelper;
 import org.elasticsearch.index.search.QueryStringQueryParser;
@@ -968,6 +969,22 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return query;
     }
 
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        if ((fieldsAndWeights.isEmpty() == false && fieldsAndWeights.containsKey(DataTierFieldMapper.NAME))
+            || (defaultField != null && (defaultField.equals(DataTierFieldMapper.NAME) || Regex.isMatchAllPattern(defaultField)))) {
+            return tierFieldTermQueryCoordinatorRewriteIfPresent(
+                this,
+                true,
+                DataTierFieldMapper.NAME,
+                queryString,
+                coordinatorRewriteContext
+            );
+        } else {
+            return this;
+        }
+    }
+
     @Override
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.ZERO;

diff --git a/server/src/main/java/org/elasticsearch/index/query/SimpleQueryStringBuilder.java b/server/src/main/java/org/elasticsearch/index/query/SimpleQueryStringBuilder.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.index.mapper.DataTierFieldMapper;
 import org.elasticsearch.index.search.QueryParserHelper;
 import org.elasticsearch.index.search.SimpleQueryStringQueryParser;
 import org.elasticsearch.index.search.SimpleQueryStringQueryParser.Settings;
@@ -609,6 +610,21 @@ public static SimpleQueryStringBuilder fromXContent(XContentParser parser) throw
         return qb;
     }
 
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        if (fieldsAndWeights.containsKey(DataTierFieldMapper.NAME)) {
+            return tierFieldTermQueryCoordinatorRewriteIfPresent(
+                this,
+                true,
+                DataTierFieldMapper.NAME,
+                queryText,
+                coordinatorRewriteContext
+            );
+        } else {
+            return this;
+        }
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;

diff --git a/server/src/main/java/org/elasticsearch/index/query/TermQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/TermQueryBuilder.java
@@ -18,7 +18,6 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.index.mapper.ConstantFieldType;
-import org.elasticsearch.index.mapper.DataTierFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -221,19 +220,7 @@ protected final int doHashCode() {
 
     @Override
     protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
-        if (fieldName.equals(DataTierFieldMapper.NAME) == false) {
-            return this;
-        }
-        final MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(DataTierFieldMapper.NAME);
-        if (fieldType instanceof final DataTierFieldMapper.DataTierFieldType tierFieldType) {
-            Query tierFieldQuery = tierFieldType.internalTermQueryCaseInsensitive(value, coordinatorRewriteContext);
-            if (tierFieldQuery instanceof MatchNoDocsQuery) {
-                return new MatchNoneQueryBuilder("The \"" + getName() + "\" query was rewritten to a \"match_none\" query.");
-            } else if (tierFieldQuery instanceof MatchAllDocsQuery) {
-                return new MatchAllQueryBuilder();
-            }
-        }
-        return this;
+        return tierFieldTermQueryCoordinatorRewriteIfPresent(this, caseInsensitive, fieldName, value, coordinatorRewriteContext);
     }
 
     @Override

diff --git a/server/src/main/java/org/elasticsearch/index/query/TermsQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/TermsQueryBuilder.java
@@ -361,6 +361,8 @@ protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinato
                     return new MatchNoneQueryBuilder("The \"" + getName() + "\" query was rewritten to a \"match_none\" query.");
                 } else if (tierFieldQuery instanceof MatchAllDocsQuery) {
                     return new MatchAllQueryBuilder();
+                } else {
+                    assert false : "Constant fields must produce match-all or match-none queries, got " + tierFieldQuery;
                 }
             }
         }

diff --git a/server/src/main/java/org/elasticsearch/index/query/WildcardQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/WildcardQueryBuilder.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.index.mapper.ConstantFieldType;
+import org.elasticsearch.index.mapper.DataTierFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.support.QueryParsers;
 import org.elasticsearch.xcontent.ParseField;
@@ -232,6 +233,25 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
         return fieldType.wildcardQuery(value, method, caseInsensitive, context);
     }
 
+    @Override
+    protected QueryBuilder doCoordinatorRewrite(CoordinatorRewriteContext coordinatorRewriteContext) {
+        if (fieldName.equals(DataTierFieldMapper.NAME) == false) {
+            return this;
+        }
+        final MappedFieldType fieldType = coordinatorRewriteContext.getFieldType(DataTierFieldMapper.NAME);
+        if (fieldType instanceof final DataTierFieldMapper.DataTierFieldType tierFieldType) {
+            Query tierFieldQuery = tierFieldType.wildcardQuery(value, true, coordinatorRewriteContext);
+            if (tierFieldQuery instanceof MatchNoDocsQuery) {
+                return new MatchNoneQueryBuilder("The \"" + getName() + "\" query was rewritten to a \"match_none\" query.");
+            } else if (tierFieldQuery instanceof MatchAllDocsQuery) {
+                return new MatchAllQueryBuilder();
+            } else {
+                assert false : "Constant fields must produce match-all or match-none queries, got " + tierFieldQuery;
+            }
+        }
+        return this;
+    }
+
     @Override
     protected int doHashCode() {
         return Objects.hash(fieldName, value, rewrite, caseInsensitive);

diff --git a/...csearch/xpack/searchablesnapshots/SearchableSnapshotsCanMatchOnCoordinatorIntegTests.java b/...csearch/xpack/searchablesnapshots/SearchableSnapshotsCanMatchOnCoordinatorIntegTests.java
@@ -1059,7 +1059,7 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
         }
 
         {
-            // bool query
+            // bool term query
             BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().mustNot(QueryBuilders.termQuery("_tier", "data_frozen"));
             List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
             SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
@@ -1072,6 +1072,41 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
                 assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) numDocsRegularIndex));
             });
         }
+
+        {
+            // bool match query
+            BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().mustNot(QueryBuilders.matchQuery("_tier", "data_frozen"));
+            List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
+            SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                .source(new SearchSourceBuilder().query(boolQueryBuilder));
+
+            assertResponse(client().search(request), searchResponse -> {
+                // as we excluded the frozen tier we shouldn't get any failures
+                assertThat(searchResponse.getFailedShards(), equalTo(0));
+                // we should be receiving all the hits from the index that's in the data_content tier
+                assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) numDocsRegularIndex));
+            });
+        }
+
+        {
+            // bool query string, prefix, wildcard, or simple query string
+            BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().mustNot(randomFrom(
+                QueryBuilders.queryStringQuery("data_frozen").field("_tier"),
+                QueryBuilders.simpleQueryStringQuery("data_frozen").field("_tier"),
+                QueryBuilders.wildcardQuery("_tier", "dat*ozen"),
+                QueryBuilders.prefixQuery("_tier", "data_fro")
+            ));
+            List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
+            SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                .source(new SearchSourceBuilder().query(boolQueryBuilder));
+
+            assertResponse(client().search(request), searchResponse -> {
+                // as we excluded the frozen tier we shouldn't get any failures
+                assertThat(searchResponse.getFailedShards(), equalTo(0));
+                // we should be receiving all the hits from the index that's in the data_content tier
+                assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) numDocsRegularIndex));
+            });
+        }
     }
 
     private void createIndexWithTimestampAndEventIngested(String indexName, int numShards, Settings extraSettings) throws IOException {