Replace the SearchContext with QueryShardContext when building collapsing context (#46543)

This commit replaces the `SearchContext` with the `QueryShardContext` when building collapsing conteext
Collapse context is part of the `SearchContext` so it shouldn't require a `SearchContext` to create one.

Relates #46523

Author: jimczi

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -901,7 +901,16 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         }
 
         if (source.collapse() != null) {
-            final CollapseContext collapseContext = source.collapse().build(context);
+            if (context.scrollContext() != null) {
+                throw new SearchContextException(context, "cannot use `collapse` in a scroll context");
+            }
+            if (context.searchAfter() != null) {
+                throw new SearchContextException(context, "cannot use `collapse` in conjunction with `search_after`");
+            }
+            if (context.rescore() != null && context.rescore().isEmpty() == false) {
+                throw new SearchContextException(context, "cannot use `collapse` in conjunction with `rescore`");
+            }
+            final CollapseContext collapseContext = source.collapse().build(queryShardContext);
             context.collapse(collapseContext);
         }
     }

server/src/main/java/org/elasticsearch/search/collapse/CollapseBuilder.java

@@ -34,8 +34,7 @@
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.index.query.InnerHitBuilder;
-import org.elasticsearch.search.SearchContextException;
-import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.index.query.QueryShardContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -200,32 +199,22 @@ public int hashCode() {
         return result;
     }
 
-    public CollapseContext build(SearchContext context) {
-        if (context.scrollContext() != null) {
-            throw new SearchContextException(context, "cannot use `collapse` in a scroll context");
-        }
-        if (context.searchAfter() != null) {
-            throw new SearchContextException(context, "cannot use `collapse` in conjunction with `search_after`");
-        }
-        if (context.rescore() != null && context.rescore().isEmpty() == false) {
-            throw new SearchContextException(context, "cannot use `collapse` in conjunction with `rescore`");
-        }
-
-        MappedFieldType fieldType = context.getQueryShardContext().fieldMapper(field);
+    public CollapseContext build(QueryShardContext queryShardContext) {
+        MappedFieldType fieldType = queryShardContext.fieldMapper(field);
         if (fieldType == null) {
-            throw new SearchContextException(context, "no mapping found for `" + field + "` in order to collapse on");
+            throw new IllegalArgumentException("no mapping found for `" + field + "` in order to collapse on");
         }
         if (fieldType instanceof KeywordFieldMapper.KeywordFieldType == false &&
             fieldType instanceof NumberFieldMapper.NumberFieldType == false) {
-            throw new SearchContextException(context, "unknown type for collapse field `" + field +
+            throw new IllegalArgumentException("unknown type for collapse field `" + field +
                 "`, only keywords and numbers are accepted");
         }
 
         if (fieldType.hasDocValues() == false) {
-            throw new SearchContextException(context, "cannot collapse on field `" + field + "` without `doc_values`");
+            throw new IllegalArgumentException("cannot collapse on field `" + field + "` without `doc_values`");
         }
         if (fieldType.indexOptions() == IndexOptions.NONE && (innerHits != null && !innerHits.isEmpty())) {
-            throw new SearchContextException(context, "cannot expand `inner_hits` for collapse field `"
+            throw new IllegalArgumentException("cannot expand `inner_hits` for collapse field `"
                 + field + "`, " + "only indexed field can retrieve `inner_hits`");
         }
 

server/src/test/java/org/elasticsearch/search/collapse/CollapseBuilderTests.java

@@ -37,9 +37,7 @@
 import org.elasticsearch.index.query.InnerHitBuilder;
 import org.elasticsearch.index.query.InnerHitBuilderTests;
 import org.elasticsearch.index.query.QueryShardContext;
-import org.elasticsearch.search.SearchContextException;
 import org.elasticsearch.search.SearchModule;
-import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.test.AbstractSerializingTestCase;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -138,59 +136,49 @@ protected NamedXContentRegistry xContentRegistry() {
         return xContentRegistry;
     }
 
-    private SearchContext mockSearchContext() {
-        SearchContext context = mock(SearchContext.class);
-        QueryShardContext shardContext = mock(QueryShardContext.class);
-        when(context.getQueryShardContext()).thenReturn(shardContext);
-        when(context.scrollContext()).thenReturn(null);
-        when(context.rescore()).thenReturn(null);
-        when(context.searchAfter()).thenReturn(null);
-        return context;
-    }
-
     public void testBuild() throws IOException {
         Directory dir = new RAMDirectory();
         try (IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig(new MockAnalyzer(random())))) {
             writer.commit();
         }
-        SearchContext searchContext = mockSearchContext();
+        QueryShardContext shardContext = mock(QueryShardContext.class);
         try (IndexReader reader = DirectoryReader.open(dir)) {
-            when(searchContext.getQueryShardContext().getIndexReader()).thenReturn(reader);
+            when(shardContext.getIndexReader()).thenReturn(reader);
             MappedFieldType numberFieldType =
                 new NumberFieldMapper.NumberFieldType(NumberFieldMapper.NumberType.LONG);
             MappedFieldType keywordFieldType =
                 new KeywordFieldMapper.KeywordFieldType();
             for (MappedFieldType fieldType : new MappedFieldType[] {numberFieldType, keywordFieldType}) {
                 fieldType.setName("field");
                 fieldType.setHasDocValues(true);
-                when(searchContext.getQueryShardContext().fieldMapper("field")).thenReturn(fieldType);
+                when(shardContext.fieldMapper("field")).thenReturn(fieldType);
                 CollapseBuilder builder = new CollapseBuilder("field");
-                CollapseContext collapseContext = builder.build(searchContext);
+                CollapseContext collapseContext = builder.build(shardContext);
                 assertEquals(collapseContext.getFieldType(), fieldType);
 
                 fieldType.setIndexOptions(IndexOptions.NONE);
-                collapseContext = builder.build(searchContext);
+                collapseContext = builder.build(shardContext);
                 assertEquals(collapseContext.getFieldType(), fieldType);
 
                 fieldType.setHasDocValues(false);
-                SearchContextException exc = expectThrows(SearchContextException.class, () -> builder.build(searchContext));
+                IllegalArgumentException exc = expectThrows(IllegalArgumentException.class, () -> builder.build(shardContext));
                 assertEquals(exc.getMessage(), "cannot collapse on field `field` without `doc_values`");
 
                 fieldType.setHasDocValues(true);
                 builder.setInnerHits(new InnerHitBuilder());
-                exc = expectThrows(SearchContextException.class, () -> builder.build(searchContext));
+                exc = expectThrows(IllegalArgumentException.class, () -> builder.build(shardContext));
                 assertEquals(exc.getMessage(),
                     "cannot expand `inner_hits` for collapse field `field`, " +
                         "only indexed field can retrieve `inner_hits`");
             }
         }
     }
 
-    public void testBuildWithSearchContextExceptions() {
-        SearchContext context = mockSearchContext();
+    public void testBuildWithExceptions() {
+        QueryShardContext shardContext = mock(QueryShardContext.class);
         {
             CollapseBuilder builder = new CollapseBuilder("unknown_field");
-            SearchContextException exc = expectThrows(SearchContextException.class, () -> builder.build(context));
+            IllegalArgumentException exc = expectThrows(IllegalArgumentException.class, () -> builder.build(shardContext));
             assertEquals(exc.getMessage(), "no mapping found for `unknown_field` in order to collapse on");
         }
 
@@ -217,9 +205,9 @@ public Query existsQuery(QueryShardContext context) {
             };
             fieldType.setName("field");
             fieldType.setHasDocValues(true);
-            when(context.getQueryShardContext().fieldMapper("field")).thenReturn(fieldType);
+            when(shardContext.fieldMapper("field")).thenReturn(fieldType);
             CollapseBuilder builder = new CollapseBuilder("field");
-            SearchContextException exc = expectThrows(SearchContextException.class, () -> builder.build(context));
+            IllegalArgumentException exc = expectThrows(IllegalArgumentException.class, () -> builder.build(shardContext));
             assertEquals(exc.getMessage(), "unknown type for collapse field `field`, only keywords and numbers are accepted");
         }
     }