Remove intermediate SearchLookup classes (#68052)

SearchLookup has two intermediate classes, DocMap and StoredFieldsLookup, that
are simple factories for their Leaf implementations. They are never accessed
outside SearchLookup, with the exception of two calls on DocMap that can be
easily refactored. This commit removes them, making SearchLookup.getLeafSearchLookup
directly responsible for creating the leaf lookups.

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java

@@ -453,13 +453,13 @@ public class ExpressionScriptEngine implements ScriptEngine {
         }
 
         String fieldname = parts[1].text;
-        MappedFieldType fieldType = lookup.doc().fieldType(fieldname);
+        MappedFieldType fieldType = lookup.fieldType(fieldname);
 
         if (fieldType == null) {
             throw new ParseException("Field [" + fieldname + "] does not exist in mappings", 5);
         }
 
-        IndexFieldData<?> fieldData = lookup.doc().getForField(fieldType);
+        IndexFieldData<?> fieldData = lookup.getForField(fieldType);
         final DoubleValuesSource valueSource;
         if (fieldType instanceof GeoPointFieldType) {
             // geo

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchDocValuesPhase.java

@@ -58,7 +58,7 @@ public final class FetchDocValuesPhase implements FetchSubPhase {
             }
             ValueFetcher fetcher = new DocValueFetcher(
                 ft.docValueFormat(fieldAndFormat.format, null),
-                context.searchLookup().doc().getForField(ft)
+                context.searchLookup().getForField(ft)
             );
             fields.add(new DocValueField(fieldAndFormat.field, fetcher));
         }

server/src/main/java/org/elasticsearch/search/lookup/DocLookup.java

@@ -1,48 +0,0 @@
-/*
- * Licensed to Elasticsearch under one or more contributor
- * license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright
- * ownership. Elasticsearch licenses this file to you under
- * the Apache License, Version 2.0 (the "License"); you may
- * not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-package org.elasticsearch.search.lookup;
-
-import org.apache.lucene.index.LeafReaderContext;
-import org.elasticsearch.index.fielddata.IndexFieldData;
-import org.elasticsearch.index.mapper.MappedFieldType;
-
-import java.util.function.Function;
-
-public final class DocLookup {
-
-    private final Function<String, MappedFieldType> fieldTypeLookup;
-    private final Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup;
-
-    DocLookup(Function<String, MappedFieldType> fieldTypeLookup, Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup) {
-        this.fieldTypeLookup = fieldTypeLookup;
-        this.fieldDataLookup = fieldDataLookup;
-    }
-
-    public MappedFieldType fieldType(String fieldName) {
-        return fieldTypeLookup.apply(fieldName);
-    }
-
-    public IndexFieldData<?> getForField(MappedFieldType fieldType) {
-        return fieldDataLookup.apply(fieldType);
-    }
-
-    public LeafDocLookup getLeafDocLookup(LeafReaderContext context) {
-        return new LeafDocLookup(fieldTypeLookup, fieldDataLookup, context);
-    }
-}

server/src/main/java/org/elasticsearch/search/lookup/LeafStoredFieldsLookup.java

@@ -18,8 +18,9 @@
  */
 package org.elasticsearch.search.lookup;
 
-import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.StoredFieldVisitor;
 import org.elasticsearch.ElasticsearchParseException;
+import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.index.fieldvisitor.SingleFieldsVisitor;
 import org.elasticsearch.index.mapper.MappedFieldType;
 
@@ -38,13 +39,14 @@ import static java.util.Collections.singletonMap;
 public class LeafStoredFieldsLookup implements Map<Object, Object> {
 
     private final Function<String, MappedFieldType> fieldTypeLookup;
-    private final LeafReader reader;
+    private final CheckedBiConsumer<Integer, StoredFieldVisitor, IOException> reader;
 
     private int docId = -1;
 
     private final Map<String, FieldLookup> cachedFieldData = new HashMap<>();
 
-    LeafStoredFieldsLookup(Function<String, MappedFieldType> fieldTypeLookup, LeafReader reader) {
+    LeafStoredFieldsLookup(Function<String, MappedFieldType> fieldTypeLookup,
+                           CheckedBiConsumer<Integer, StoredFieldVisitor, IOException> reader) {
         this.fieldTypeLookup = fieldTypeLookup;
         this.reader = reader;
     }
@@ -136,7 +138,7 @@ public class LeafStoredFieldsLookup implements Map<Object, Object> {
             List<Object> values = new ArrayList<>(2);
             SingleFieldsVisitor visitor = new SingleFieldsVisitor(data.fieldType(), values);
             try {
-                reader.document(docId, visitor);
+                reader.accept(docId, visitor);
             } catch (IOException e) {
                 throw new ElasticsearchParseException("failed to load field [{}]", e, name);
             }

server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java

@@ -50,9 +50,7 @@ public class SearchLookup {
      * {@code b}'s chain will contain {@code ["a", "b"]}.
      */
     private final Set<String> fieldChain;
-    private final DocLookup docMap;
     private final SourceLookup sourceLookup;
-    private final StoredFieldsLookup storedFieldsLookup;
     private final Function<String, MappedFieldType> fieldTypeLookup;
     private final BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup;
 
@@ -64,10 +62,7 @@ public class SearchLookup {
                         BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup) {
         this.fieldTypeLookup = fieldTypeLookup;
         this.fieldChain = Collections.emptySet();
-        docMap = new DocLookup(fieldTypeLookup,
+        this.sourceLookup = new SourceLookup();
-            fieldType -> fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name())));
-        sourceLookup = new SourceLookup();
-        storedFieldsLookup = new StoredFieldsLookup(fieldTypeLookup);
         this.fieldDataLookup = fieldDataLookup;
     }
 
@@ -80,10 +75,7 @@ public class SearchLookup {
      */
     private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
         this.fieldChain = Collections.unmodifiableSet(fieldChain);
-        this.docMap = new DocLookup(searchLookup.fieldTypeLookup,
-            fieldType -> searchLookup.fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name())));
         this.sourceLookup = searchLookup.sourceLookup;
-        this.storedFieldsLookup = searchLookup.storedFieldsLookup;
         this.fieldTypeLookup = searchLookup.fieldTypeLookup;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
     }
@@ -111,13 +103,17 @@ public class SearchLookup {
 
     public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) {
         return new LeafSearchLookup(context,
-                docMap.getLeafDocLookup(context),
+                new LeafDocLookup(fieldTypeLookup, this::getForField, context),
                 sourceLookup,
-                storedFieldsLookup.getLeafFieldsLookup(context));
+                new LeafStoredFieldsLookup(fieldTypeLookup, (doc, visitor) -> context.reader().document(doc, visitor)));
     }
 
-    public DocLookup doc() {
+    public MappedFieldType fieldType(String fieldName) {
-        return docMap;
+        return fieldTypeLookup.apply(fieldName);
+    }
+
+    public IndexFieldData<?> getForField(MappedFieldType fieldType) {
+        return fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()));
     }
 
     public SourceLookup source() {

server/src/main/java/org/elasticsearch/search/lookup/StoredFieldsLookup.java

@@ -1,37 +0,0 @@
-/*
- * Licensed to Elasticsearch under one or more contributor
- * license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright
- * ownership. Elasticsearch licenses this file to you under
- * the Apache License, Version 2.0 (the "License"); you may
- * not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-package org.elasticsearch.search.lookup;
-
-import org.apache.lucene.index.LeafReaderContext;
-import org.elasticsearch.index.mapper.MappedFieldType;
-
-import java.util.function.Function;
-
-public class StoredFieldsLookup {
-
-    private final Function<String, MappedFieldType> fieldTypeLookup;
-
-    StoredFieldsLookup(Function<String, MappedFieldType> fieldTypeLookup) {
-        this.fieldTypeLookup = fieldTypeLookup;
-    }
-
-    LeafStoredFieldsLookup getLeafFieldsLookup(LeafReaderContext context) {
-        return new LeafStoredFieldsLookup(fieldTypeLookup, context.reader());
-    }
-}

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

@@ -560,7 +560,7 @@ public class SearchExecutionContextTests extends ESTestCase {
                 if (randomBoolean()) {
                     indexFieldData = searchExecutionContext.getForField(fieldType);
                 } else {
-                    indexFieldData = searchExecutionContext.lookup().doc().getForField(fieldType);
+                    indexFieldData = searchExecutionContext.lookup().getForField(fieldType);
                 }
                 searcher.search(query, new Collector() {
                     @Override
@@ -579,7 +579,7 @@ public class SearchExecutionContextTests extends ESTestCase {
                             public void collect(int doc) throws IOException {
                                 ScriptDocValues<?> scriptDocValues;
                                 if(randomBoolean()) {
-                                    LeafDocLookup leafDocLookup = searchExecutionContext.lookup().doc().getLeafDocLookup(context);
+                                    LeafDocLookup leafDocLookup = searchExecutionContext.lookup().getLeafSearchLookup(context).doc();
                                     leafDocLookup.setDocument(doc);
                                     scriptDocValues = leafDocLookup.get(field);
                                 } else {

server/src/test/java/org/elasticsearch/search/lookup/LeafStoredFieldsLookupTests.java

@@ -21,8 +21,6 @@ package org.elasticsearch.search.lookup;
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.IndexOptions;
-import org.apache.lucene.index.LeafReader;
-import org.apache.lucene.index.StoredFieldVisitor;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.test.ESTestCase;
 import org.junit.Before;
@@ -30,10 +28,7 @@ import org.junit.Before;
 import java.util.Collections;
 import java.util.List;
 
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyObject;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -53,15 +48,8 @@ public class LeafStoredFieldsLookupTests extends ESTestCase {
         FieldInfo mockFieldInfo = new FieldInfo("field", 1, false, false, true,
             IndexOptions.NONE, DocValuesType.NONE, -1, Collections.emptyMap(), 0, 0, 0, false);
 
-        LeafReader leafReader = mock(LeafReader.class);
+        fieldsLookup = new LeafStoredFieldsLookup(field -> field.equals("field") || field.equals("alias") ? fieldType : null,
-        doAnswer(invocation -> {
+            (doc, visitor) -> visitor.doubleField(mockFieldInfo, 2.718));
-            Object[] args = invocation.getArguments();
-            StoredFieldVisitor visitor = (StoredFieldVisitor) args[1];
-            visitor.doubleField(mockFieldInfo, 2.718);
-            return null;
-        }).when(leafReader).document(anyInt(), any(StoredFieldVisitor.class));
-
-        fieldsLookup = new LeafStoredFieldsLookup(field -> field.equals("field") || field.equals("alias") ? fieldType : null, leafReader);
     }
 
     public void testBasicLookup() {

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java

@@ -315,7 +315,7 @@ public abstract class MapperTestCase extends MapperServiceTestCase {
             iw.addDocument(mapperService.documentMapper().parse(source(b -> b.field(ft.name(), sourceValue))).rootDoc());
         }, iw -> {
             SearchLookup lookup = new SearchLookup(mapperService::fieldType, fieldDataLookup);
-            ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.doc().getForField(ft));
+            ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.getForField(ft));
             IndexSearcher searcher = newSearcher(iw);
             LeafReaderContext context = searcher.getIndexReader().leaves().get(0);
             lookup.source().setSegmentAndDocument(context, 0);

test/framework/src/main/java/org/elasticsearch/script/ScoreAccessor.java

@@ -20,14 +20,13 @@
 package org.elasticsearch.script;
 
 import org.apache.lucene.search.Scorable;
-import org.elasticsearch.search.lookup.DocLookup;
 
 import java.io.IOException;
 
 /**
  * A float encapsulation that dynamically accesses the score of a document.
  *
- * The provided {@link DocLookup} is used to retrieve the score
+ * The provided {@link Scorable} is used to retrieve the score
  * for the current document.
  */
 public final class ScoreAccessor extends Number {

x-pack/plugin/mapper-flattened/src/test/java/org/elasticsearch/index/mapper/FlattenedFieldLookupTests.java

@@ -163,7 +163,7 @@ public class FlattenedFieldLookupTests extends ESTestCase {
             }
             return null;
         }, fieldDataSupplier);
-        LeafDocLookup docLookup = searchLookup.doc().getLeafDocLookup(null);
+        LeafDocLookup docLookup = searchLookup.getLeafSearchLookup(null).doc();
 
         assertEquals(docValues1, docLookup.get("json.key1"));
         assertEquals(docValues2, docLookup.get("json.key2"));