Refactor DocLookup to use field def lookup (#656)

* Refactor DocLookup to use field def lookup

src/main/java/com/yelp/nrtsearch/server/luceneserver/IndexState.java

@@ -95,7 +95,7 @@ public abstract class IndexState implements Closeable {
    * Index level doc values lookup. Generates {@link
    * com.yelp.nrtsearch.server.luceneserver.doc.SegmentDocLookup} for a given lucene segment.
    */
-  public final DocLookup docLookup = new DocLookup(this);
+  public final DocLookup docLookup = new DocLookup(this, this::getField);
 
   /** Search-time analyzer. */
   public final Analyzer searchAnalyzer =

src/main/java/com/yelp/nrtsearch/server/luceneserver/QueryNodeMapper.java

@@ -32,6 +32,7 @@
 import com.yelp.nrtsearch.server.grpc.RangeQuery;
 import com.yelp.nrtsearch.server.grpc.RewriteMethod;
 import com.yelp.nrtsearch.server.luceneserver.analysis.AnalyzerCreator;
+import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.TextBaseFieldDef;
@@ -97,7 +98,12 @@ public static QueryNodeMapper getInstance() {
               MatchOperator.MUST, BooleanClause.Occur.MUST));
 
   public Query getQuery(com.yelp.nrtsearch.server.grpc.Query query, IndexState state) {
-    Query queryNode = getQueryNode(query, state);
+    return getQuery(query, state, state.docLookup);
+  }
+
+  public Query getQuery(
+      com.yelp.nrtsearch.server.grpc.Query query, IndexState state, DocLookup docLookup) {
+    Query queryNode = getQueryNode(query, state, docLookup);
 
     if (query.getBoost() < 0) {
       throw new IllegalArgumentException("Boost must be a positive number");
@@ -126,20 +132,21 @@ public Query getNestedPathQuery(IndexState indexState, String path) {
     return new TermQuery(new Term(IndexState.NESTED_PATH, indexState.resolveQueryNestedPath(path)));
   }
 
-  private Query getQueryNode(com.yelp.nrtsearch.server.grpc.Query query, IndexState state) {
+  private Query getQueryNode(
+      com.yelp.nrtsearch.server.grpc.Query query, IndexState state, DocLookup docLookup) {
     switch (query.getQueryNodeCase()) {
       case BOOLEANQUERY:
-        return getBooleanQuery(query.getBooleanQuery(), state);
+        return getBooleanQuery(query.getBooleanQuery(), state, docLookup);
       case PHRASEQUERY:
         return getPhraseQuery(query.getPhraseQuery());
       case FUNCTIONSCOREQUERY:
-        return getFunctionScoreQuery(query.getFunctionScoreQuery(), state);
+        return getFunctionScoreQuery(query.getFunctionScoreQuery(), state, docLookup);
       case TERMQUERY:
         return getTermQuery(query.getTermQuery(), state);
       case TERMINSETQUERY:
         return getTermInSetQuery(query.getTermInSetQuery(), state);
       case DISJUNCTIONMAXQUERY:
-        return getDisjunctionMaxQuery(query.getDisjunctionMaxQuery(), state);
+        return getDisjunctionMaxQuery(query.getDisjunctionMaxQuery(), state, docLookup);
       case MATCHQUERY:
         return getMatchQuery(query.getMatchQuery(), state);
       case MATCHPHRASEQUERY:
@@ -153,7 +160,7 @@ private Query getQueryNode(com.yelp.nrtsearch.server.grpc.Query query, IndexStat
       case GEOPOINTQUERY:
         return getGeoPointQuery(query.getGeoPointQuery(), state);
       case NESTEDQUERY:
-        return getNestedQuery(query.getNestedQuery(), state);
+        return getNestedQuery(query.getNestedQuery(), state, docLookup);
       case EXISTSQUERY:
         return getExistsQuery(query.getExistsQuery(), state);
       case GEORADIUSQUERY:
@@ -169,7 +176,7 @@ private Query getQueryNode(com.yelp.nrtsearch.server.grpc.Query query, IndexStat
       case PREFIXQUERY:
         return getPrefixQuery(query.getPrefixQuery(), state);
       case CONSTANTSCOREQUERY:
-        return getConstantScoreQuery(query.getConstantScoreQuery(), state);
+        return getConstantScoreQuery(query.getConstantScoreQuery(), state, docLookup);
       case GEOPOLYGONQUERY:
         return getGeoPolygonQuery(query.getGeoPolygonQuery(), state);
       case QUERYNODE_NOT_SET:
@@ -206,8 +213,10 @@ private Query getCompletionQuery(
   }
 
   private Query getNestedQuery(
-      com.yelp.nrtsearch.server.grpc.NestedQuery nestedQuery, IndexState state) {
-    Query childRawQuery = getQuery(nestedQuery.getQuery(), state);
+      com.yelp.nrtsearch.server.grpc.NestedQuery nestedQuery,
+      IndexState state,
+      DocLookup docLookup) {
+    Query childRawQuery = getQuery(nestedQuery.getQuery(), state, docLookup);
     Query childQuery =
         new BooleanQuery.Builder()
             .add(getNestedPathQuery(state, nestedQuery.getPath()), BooleanClause.Occur.FILTER)
@@ -237,7 +246,9 @@ private ScoreMode getScoreMode(com.yelp.nrtsearch.server.grpc.NestedQuery nested
   }
 
   private BooleanQuery getBooleanQuery(
-      com.yelp.nrtsearch.server.grpc.BooleanQuery booleanQuery, IndexState state) {
+      com.yelp.nrtsearch.server.grpc.BooleanQuery booleanQuery,
+      IndexState state,
+      DocLookup docLookup) {
     BooleanQuery.Builder builder =
         new BooleanQuery.Builder()
             .setMinimumNumberShouldMatch(booleanQuery.getMinimumNumberShouldMatch());
@@ -252,7 +263,7 @@ private BooleanQuery getBooleanQuery(
         .forEach(
             clause -> {
               com.yelp.nrtsearch.server.grpc.BooleanClause.Occur occur = clause.getOccur();
-              builder.add(getQuery(clause.getQuery(), state), occurMapping.get(occur));
+              builder.add(getQuery(clause.getQuery(), state, docLookup), occurMapping.get(occur));
               if (occur != com.yelp.nrtsearch.server.grpc.BooleanClause.Occur.MUST_NOT) {
                 allMustNot.set(false);
               }
@@ -273,15 +284,17 @@ private PhraseQuery getPhraseQuery(com.yelp.nrtsearch.server.grpc.PhraseQuery ph
   }
 
   private FunctionScoreQuery getFunctionScoreQuery(
-      com.yelp.nrtsearch.server.grpc.FunctionScoreQuery functionScoreQuery, IndexState state) {
+      com.yelp.nrtsearch.server.grpc.FunctionScoreQuery functionScoreQuery,
+      IndexState state,
+      DocLookup docLookup) {
     ScoreScript.Factory scriptFactory =
         ScriptService.getInstance().compile(functionScoreQuery.getScript(), ScoreScript.CONTEXT);
 
     Map<String, Object> params =
         ScriptParamsUtils.decodeParams(functionScoreQuery.getScript().getParamsMap());
     return new FunctionScoreQuery(
-        getQuery(functionScoreQuery.getQuery(), state),
-        scriptFactory.newFactory(params, state.docLookup));
+        getQuery(functionScoreQuery.getQuery(), state, docLookup),
+        scriptFactory.newFactory(params, docLookup));
   }
 
   private FunctionMatchQuery getFunctionFilterQuery(
@@ -333,10 +346,12 @@ private Query getTermInSetQuery(
   }
 
   private DisjunctionMaxQuery getDisjunctionMaxQuery(
-      com.yelp.nrtsearch.server.grpc.DisjunctionMaxQuery disjunctionMaxQuery, IndexState state) {
+      com.yelp.nrtsearch.server.grpc.DisjunctionMaxQuery disjunctionMaxQuery,
+      IndexState state,
+      DocLookup docLookup) {
     List<Query> disjuncts =
         disjunctionMaxQuery.getDisjunctsList().stream()
-            .map(query -> getQuery(query, state))
+            .map(query -> getQuery(query, state, docLookup))
             .collect(Collectors.toList());
     return new DisjunctionMaxQuery(disjuncts, disjunctionMaxQuery.getTieBreakerMultiplier());
   }
@@ -623,8 +638,9 @@ private static MultiTermQuery.RewriteMethod getRewriteMethod(
 
   private Query getConstantScoreQuery(
       com.yelp.nrtsearch.server.grpc.ConstantScoreQuery constantScoreQueryGrpc,
-      IndexState indexState) {
-    Query filterQuery = getQuery(constantScoreQueryGrpc.getFilter(), indexState);
+      IndexState indexState,
+      DocLookup docLookup) {
+    Query filterQuery = getQuery(constantScoreQueryGrpc.getFilter(), indexState, docLookup);
     return new ConstantScoreQuery(filterQuery);
   }
 }

src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/DocLookup.java

@@ -16,6 +16,8 @@
 package com.yelp.nrtsearch.server.luceneserver.doc;
 
 import com.yelp.nrtsearch.server.luceneserver.IndexState;
+import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
+import java.util.function.Function;
 import org.apache.lucene.index.LeafReaderContext;
 
 /**
@@ -24,9 +26,11 @@
  */
 public class DocLookup {
   private final IndexState indexState;
+  private final Function<String, FieldDef> fieldDefLookup;
 
-  public DocLookup(IndexState indexState) {
+  public DocLookup(IndexState indexState, Function<String, FieldDef> fieldDefLookup) {
     this.indexState = indexState;
+    this.fieldDefLookup = fieldDefLookup;
   }
 
   /**
@@ -36,7 +40,7 @@ public DocLookup(IndexState indexState) {
    * @return lookup accessor for given segment context
    */
   public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
-    return new SegmentDocLookup(indexState, context);
+    return new SegmentDocLookup(fieldDefLookup, context);
   }
 
   /**
@@ -47,4 +51,14 @@ public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
   public IndexState getIndexState() {
     return indexState;
   }
+
+  /**
+   * Get the field definition for the given field name.
+   *
+   * @param fieldName field name
+   * @return field definition
+   */
+  public FieldDef getFieldDef(String fieldName) {
+    return fieldDefLookup.apply(fieldName);
+  }
 }

src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/SegmentDocLookup.java

@@ -15,7 +15,6 @@
  */
 package com.yelp.nrtsearch.server.luceneserver.doc;
 
-import com.yelp.nrtsearch.server.luceneserver.IndexState;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
 import java.io.IOException;
@@ -24,6 +23,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
 import org.apache.lucene.index.LeafReaderContext;
 
 /**
@@ -35,14 +35,14 @@
  */
 public class SegmentDocLookup implements Map<String, LoadedDocValues<?>> {
 
-  private final IndexState indexState;
+  private final Function<String, FieldDef> fieldDefLookup;
   private final LeafReaderContext context;
   private final Map<String, LoadedDocValues<?>> loaderCache = new HashMap<>();
 
   private int docId = -1;
 
-  public SegmentDocLookup(IndexState indexState, LeafReaderContext context) {
-    this.indexState = indexState;
+  public SegmentDocLookup(Function<String, FieldDef> fieldDefLookup, LeafReaderContext context) {
+    this.fieldDefLookup = fieldDefLookup;
     this.context = context;
   }
 
@@ -80,7 +80,7 @@ public boolean containsKey(Object key) {
     }
     String fieldName = key.toString();
     try {
-      FieldDef field = indexState.getField(fieldName);
+      FieldDef field = fieldDefLookup.apply(fieldName);
       return field instanceof IndexableFieldDef && ((IndexableFieldDef) field).hasDocValues();
     } catch (Exception ignored) {
       return false;
@@ -108,7 +108,7 @@ public LoadedDocValues<?> get(Object key) {
     String fieldName = key.toString();
     LoadedDocValues<?> docValues = loaderCache.get(fieldName);
     if (docValues == null) {
-      FieldDef fieldDef = indexState.getField(fieldName);
+      FieldDef fieldDef = fieldDefLookup.apply(fieldName);
       if (fieldDef == null) {
         throw new IllegalArgumentException("Field does not exist: " + fieldName);
       }

src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchContext.java

@@ -18,6 +18,7 @@
 import com.yelp.nrtsearch.server.grpc.SearchResponse;
 import com.yelp.nrtsearch.server.luceneserver.IndexState;
 import com.yelp.nrtsearch.server.luceneserver.ShardState;
+import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup;
 import com.yelp.nrtsearch.server.luceneserver.doc.SharedDocContext;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreTask;
@@ -32,6 +33,7 @@
 /** Search context class to provide all the information to perform a search. */
 public class SearchContext implements FieldFetchContext {
   private final IndexState indexState;
+  private final DocLookup docLookup;
   private final ShardState shardState;
   private final SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy;
   private final SearchResponse.Builder responseBuilder;
@@ -53,6 +55,7 @@ public class SearchContext implements FieldFetchContext {
 
   private SearchContext(Builder builder, boolean validate) {
     this.indexState = builder.indexState;
+    this.docLookup = builder.docLookup;
     this.shardState = builder.shardState;
     this.searcherAndTaxonomy = builder.searcherAndTaxonomy;
     this.responseBuilder = builder.responseBuilder;
@@ -80,6 +83,11 @@ public IndexState getIndexState() {
     return indexState;
   }
 
+  /** Get DocLookup. */
+  public DocLookup getDocLookup() {
+    return docLookup;
+  }
+
   /** Get query shard state. */
   public ShardState getShardState() {
     return shardState;
@@ -186,6 +194,7 @@ private void validate() {
     Objects.requireNonNull(fetchTasks);
     Objects.requireNonNull(rescorers);
     Objects.requireNonNull(sharedDocContext);
+    Objects.requireNonNull(docLookup);
 
     if (timestampSec < 0) {
       throw new IllegalStateException("Invalid timestamp value: " + timestampSec);
@@ -208,6 +217,7 @@ public SearchContext getSearchContext() {
   public static class Builder {
 
     private IndexState indexState;
+    private DocLookup docLookup;
     private ShardState shardState;
     private SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy;
     private SearchResponse.Builder responseBuilder;
@@ -234,6 +244,12 @@ public Builder setIndexState(IndexState indexState) {
       return this;
     }
 
+    /** Set doclookup. */
+    public Builder setDocLookup(DocLookup docLookup) {
+      this.docLookup = docLookup;
+      return this;
+    }
+
     /** Set query shard state. */
     public Builder setShardState(ShardState shardState) {
       this.shardState = shardState;

src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java

@@ -29,6 +29,7 @@
 import com.yelp.nrtsearch.server.luceneserver.QueryNodeMapper;
 import com.yelp.nrtsearch.server.luceneserver.ShardState;
 import com.yelp.nrtsearch.server.luceneserver.doc.DefaultSharedDocContext;
+import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.RuntimeFieldDef;
@@ -63,6 +64,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.apache.lucene.facet.DrillDownQuery;
 import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager;
@@ -138,6 +140,10 @@ public static SearchContext buildContextForRequest(
         getRetrieveFields(searchRequest.getRetrieveFieldsList(), queryFields);
     contextBuilder.setRetrieveFields(Collections.unmodifiableMap(retrieveFields));
 
+    Function<String, FieldDef> fieldDefLookup = (String s) -> queryFields.get(s);
+    DocLookup docLookup = new DocLookup(indexState, fieldDefLookup);
+    contextBuilder.setDocLookup(docLookup);
+
     String rootQueryNestedPath =
         indexState.resolveQueryNestedPath(searchRequest.getQueryNestedPath());
     contextBuilder.setQueryNestedPath(rootQueryNestedPath);
@@ -146,7 +152,8 @@ public static SearchContext buildContextForRequest(
             indexState,
             searchRequest.getQueryText(),
             searchRequest.getQuery(),
-            rootQueryNestedPath);
+            rootQueryNestedPath,
+            docLookup);
     if (profileResult != null) {
       profileResult.setParsedQuery(query.toString());
     }
@@ -184,7 +191,8 @@ public static SearchContext buildContextForRequest(
                     rootQueryNestedPath,
                     entry.getKey(),
                     entry.getValue(),
-                    searchRequest.getExplain())));
+                    searchRequest.getExplain(),
+                    docLookup)));
       }
     }
 
@@ -361,7 +369,8 @@ private static Query extractQuery(
       IndexState state,
       String queryText,
       com.yelp.nrtsearch.server.grpc.Query query,
-      String queryNestedPath) {
+      String queryNestedPath,
+      DocLookup docLookup) {
     Query q;
     if (!queryText.isEmpty()) {
       QueryBuilder queryParser = createQueryParser(state, null);
@@ -373,7 +382,7 @@ private static Query extractQuery(
             String.format("could not parse queryText: %s", queryText));
       }
     } else {
-      q = QUERY_NODE_MAPPER.getQuery(query, state);
+      q = QUERY_NODE_MAPPER.getQuery(query, state, docLookup);
     }
 
     if (state.hasNestedChildFields()) {
@@ -509,10 +518,11 @@ private static InnerHitContext buildInnerHitContext(
       String parentQueryNestedPath,
       String innerHitName,
       InnerHit innerHit,
-      boolean explain) {
+      boolean explain,
+      DocLookup docLookup) {
     // Do not apply nestedPath here. This is query is used to create a shared
     // weight.
-    Query childQuery = extractQuery(indexState, "", innerHit.getInnerQuery(), null);
+    Query childQuery = extractQuery(indexState, "", innerHit.getInnerQuery(), null, docLookup);
     return InnerHitContextBuilder.Builder()
         .withInnerHitName(innerHitName)
         .withQuery(childQuery)