Remove IndexState from DocLookup (#742)

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

@@ -44,7 +44,6 @@
 import java.util.regex.Pattern;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.AnalyzerWrapper;
-import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.facet.FacetsConfig;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.search.similarities.PerFieldSimilarityWrapper;
@@ -95,7 +94,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, this::getField);
+  public final DocLookup docLookup = new DocLookup(this::getField);
 
   /** Search-time analyzer. */
   public final Analyzer searchAnalyzer =
@@ -362,9 +361,6 @@ public abstract void start(
   /** Get fields with doc values that do eager global ordinal building. */
   public abstract Map<String, GlobalOrdinalable> getEagerFieldGlobalOrdinalFields();
 
-  /** Get field bindings to use for javascript expressions. */
-  public abstract Bindings getExpressionBindings();
-
   /** Verifies if it has nested child object fields. */
   public abstract boolean hasNestedChildFields();
 

src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/DocLookup.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 java.util.function.Function;
 import org.apache.lucene.index.LeafReaderContext;
@@ -25,11 +24,9 @@
  * SegmentDocLookup} bound to single lucene segment.
  */
 public class DocLookup {
-  private final IndexState indexState;
   private final Function<String, FieldDef> fieldDefLookup;
 
-  public DocLookup(IndexState indexState, Function<String, FieldDef> fieldDefLookup) {
-    this.indexState = indexState;
+  public DocLookup(Function<String, FieldDef> fieldDefLookup) {
     this.fieldDefLookup = fieldDefLookup;
   }
 
@@ -43,15 +40,6 @@ public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
     return new SegmentDocLookup(fieldDefLookup, context);
   }
 
-  /**
-   * Get the state information associated with this index.
-   *
-   * @return index state
-   */
-  public IndexState getIndexState() {
-    return indexState;
-  }
-
   /**
    * Get the field definition for the given field name.
    *

src/main/java/com/yelp/nrtsearch/server/luceneserver/field/FieldDefBindings.java

@@ -16,7 +16,7 @@
 package com.yelp.nrtsearch.server.luceneserver.field;
 
 import com.yelp.nrtsearch.server.luceneserver.field.properties.Bindable;
-import java.util.Map;
+import java.util.function.Function;
 import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.expressions.js.VariableContext;
 import org.apache.lucene.expressions.js.VariableContext.Type;
@@ -25,11 +25,11 @@
 /** Implements {@link Bindings} on top of the registered fields. */
 public final class FieldDefBindings extends Bindings {
 
-  private final Map<String, FieldDef> fields;
+  private final Function<String, FieldDef> fieldDefLookup;
 
   /** Sole constructor. */
-  public FieldDefBindings(Map<String, FieldDef> fields) {
-    this.fields = fields;
+  public FieldDefBindings(Function<String, FieldDef> fieldDefLookup) {
+    this.fieldDefLookup = fieldDefLookup;
   }
 
   /**
@@ -55,7 +55,7 @@ public DoubleValuesSource getDoubleValuesSource(String name) {
     if (name.equals("_score")) {
       return DoubleValuesSource.SCORES;
     }
-    FieldDef fd = fields.get(name);
+    FieldDef fd = fieldDefLookup.apply(name);
     String property = Bindable.VALUE_PROPERTY;
     String fieldName = name;
     if (fd == null) {
@@ -72,7 +72,7 @@ public DoubleValuesSource getDoubleValuesSource(String name) {
             "Invalid field binding format: " + name + ", expected: doc['field_name'].property");
       }
       fieldName = parsed[1].text;
-      fd = fields.get(fieldName);
+      fd = fieldDefLookup.apply(fieldName);
       if (fd == null) {
         throw new IllegalArgumentException("Unknown field to bind: " + fieldName);
       }

src/main/java/com/yelp/nrtsearch/server/luceneserver/index/FieldAndFacetState.java

@@ -17,7 +17,6 @@
 
 import com.yelp.nrtsearch.server.grpc.Field;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
-import com.yelp.nrtsearch.server.luceneserver.field.FieldDefBindings;
 import com.yelp.nrtsearch.server.luceneserver.field.IdFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.ObjectFieldDef;
@@ -31,7 +30,6 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.facet.FacetsConfig;
 import org.apache.lucene.facet.FacetsConfig.DimConfig;
 
@@ -48,7 +46,6 @@ public class FieldAndFacetState {
   private final List<String> indexedAnalyzedFields;
   private final Map<String, FieldDef> eagerGlobalOrdinalFields;
   private final Map<String, GlobalOrdinalable> eagerFieldGlobalOrdinalFields;
-  private final Bindings exprBindings;
 
   // facet
   private final FacetsConfig facetsConfig;
@@ -62,7 +59,6 @@ public FieldAndFacetState() {
     indexedAnalyzedFields = Collections.emptyList();
     eagerGlobalOrdinalFields = Collections.emptyMap();
     eagerFieldGlobalOrdinalFields = Collections.emptyMap();
-    exprBindings = new FieldDefBindings(fields);
 
     facetsConfig = new FacetsConfig();
     internalFacetFieldNames = Collections.emptySet();
@@ -81,7 +77,6 @@ public FieldAndFacetState() {
     eagerGlobalOrdinalFields = Collections.unmodifiableMap(builder.eagerGlobalOrdinalFields);
     eagerFieldGlobalOrdinalFields =
         Collections.unmodifiableMap(builder.eagerFieldGlobalOrdinalFields);
-    exprBindings = builder.exprBindings;
 
     facetsConfig = builder.facetsConfig;
     internalFacetFieldNames = Collections.unmodifiableSet(builder.internalFacetFieldNames);
@@ -117,11 +112,6 @@ public Map<String, GlobalOrdinalable> getFieldEagerGlobalOrdinalFields() {
     return eagerFieldGlobalOrdinalFields;
   }
 
-  /** Get field expression {@link Bindings} used for js scripting language. */
-  public Bindings getExprBindings() {
-    return exprBindings;
-  }
-
   /** Get facet config. */
   public FacetsConfig getFacetsConfig() {
     return facetsConfig;
@@ -145,7 +135,6 @@ public static class Builder {
     private final List<String> indexedAnalyzedFields;
     private final Map<String, FieldDef> eagerGlobalOrdinalFields;
     private final Map<String, GlobalOrdinalable> eagerFieldGlobalOrdinalFields;
-    private final Bindings exprBindings;
 
     // facet
     private final FacetsConfig facetsConfig;
@@ -158,7 +147,6 @@ private Builder(FieldAndFacetState initial) {
       this.indexedAnalyzedFields = new ArrayList<>(initial.indexedAnalyzedFields);
       this.eagerGlobalOrdinalFields = new HashMap<>(initial.eagerGlobalOrdinalFields);
       this.eagerFieldGlobalOrdinalFields = new HashMap<>(initial.eagerFieldGlobalOrdinalFields);
-      this.exprBindings = new FieldDefBindings(this.fields);
 
       this.facetsConfig = new FacetsConfig();
       this.internalFacetFieldNames = new HashSet<>();
@@ -174,8 +162,13 @@ private Builder(FieldAndFacetState initial) {
       }
     }
 
-    public Bindings getBindings() {
-      return exprBindings;
+    /**
+     * Get fields currently registered with the builder.
+     *
+     * @return map of registered fields
+     */
+    public Map<String, FieldDef> getFields() {
+      return fields;
     }
 
     /**

src/main/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexState.java

@@ -54,7 +54,6 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.facet.FacetsConfig;
 import org.apache.lucene.index.ConcurrentMergeScheduler;
 import org.apache.lucene.index.IndexWriterConfig;
@@ -469,11 +468,6 @@ public Map<String, GlobalOrdinalable> getEagerFieldGlobalOrdinalFields() {
     return fieldAndFacetState.getFieldEagerGlobalOrdinalFields();
   }
 
-  @Override
-  public Bindings getExpressionBindings() {
-    return fieldAndFacetState.getExprBindings();
-  }
-
   @Override
   public boolean hasNestedChildFields() {
     return fieldAndFacetState.getHasNestedChildFields();

src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java

@@ -20,6 +20,7 @@
 import com.yelp.nrtsearch.server.grpc.FieldDefResponse;
 import com.yelp.nrtsearch.server.grpc.FieldType;
 import com.yelp.nrtsearch.server.luceneserver.IndexState;
+import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDefCreator;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
@@ -28,7 +29,6 @@
 import com.yelp.nrtsearch.server.luceneserver.index.IndexStateManager;
 import com.yelp.nrtsearch.server.luceneserver.script.ScoreScript;
 import com.yelp.nrtsearch.server.luceneserver.script.ScriptService;
-import com.yelp.nrtsearch.server.luceneserver.script.js.JsScriptEngine;
 import com.yelp.nrtsearch.server.utils.ScriptParamsUtils;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -179,21 +179,8 @@ public static void parseVirtualField(Field field, FieldAndFacetState.Builder fie
     ScoreScript.Factory factory =
         ScriptService.getInstance().compile(field.getScript(), ScoreScript.CONTEXT);
     Map<String, Object> params = ScriptParamsUtils.decodeParams(field.getScript().getParamsMap());
-    // Workaround for the fact that the javascript expression may need bindings to other fields in
-    // this request.
-    // Build the complete bindings and pass it as a script parameter. We might want to think about a
-    // better way of
-    // doing this (or maybe updating index state in general).
-    if (field.getScript().getLang().equals(JsScriptEngine.LANG)) {
-      params = new HashMap<>(params);
-      params.put("bindings", fieldStateBuilder.getBindings());
-    } else {
-      // TODO fix this, by removing DocLookup dependency on IndexState. Should be possible to just
-      // use the fields from the field state builder
-      throw new IllegalArgumentException("Only js lang supported for index virtual fields");
-    }
-    // js scripts use Bindings instead of DocLookup
-    DoubleValuesSource values = factory.newFactory(params, null);
+    DoubleValuesSource values =
+        factory.newFactory(params, new DocLookup(fieldStateBuilder.getFields()::get));
 
     FieldDef virtualFieldDef = new VirtualFieldDef(field.getName(), values);
     fieldStateBuilder.addField(virtualFieldDef, field);

src/main/java/com/yelp/nrtsearch/server/luceneserver/script/js/JsScriptEngine.java

@@ -15,12 +15,11 @@
  */
 package com.yelp.nrtsearch.server.luceneserver.script.js;
 
+import com.yelp.nrtsearch.server.luceneserver.field.FieldDefBindings;
 import com.yelp.nrtsearch.server.luceneserver.script.ScoreScript;
 import com.yelp.nrtsearch.server.luceneserver.script.ScriptContext;
 import com.yelp.nrtsearch.server.luceneserver.script.ScriptEngine;
 import java.text.ParseException;
-import java.util.HashMap;
-import java.util.Map;
 import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.expressions.Expression;
 import org.apache.lucene.expressions.js.JavascriptCompiler;
@@ -71,22 +70,8 @@ public <T> T compile(String source, ScriptContext<T> context) {
     }
     ScoreScript.Factory factory =
         ((params, docLookup) -> {
-          Map<String, Object> scriptParams;
-          Bindings fieldBindings;
-          Object bindingsParam = params.get("bindings");
-          if (bindingsParam instanceof Bindings) {
-            fieldBindings = (Bindings) bindingsParam;
-
-            // we do not want the bindings to be used as an expression parameter, so remove it.
-            // the extra copy may not be absolutely needed, but this only happens when a new
-            // virtual field is added to the index, and this keeps the code thread safe.
-            scriptParams = new HashMap<>(params);
-            scriptParams.remove("bindings");
-          } else {
-            fieldBindings = docLookup.getIndexState().getExpressionBindings();
-            scriptParams = params;
-          }
-          return expr.getDoubleValuesSource(new JsScriptBindings(fieldBindings, scriptParams));
+          Bindings fieldBindings = new FieldDefBindings(docLookup::getFieldDef);
+          return expr.getDoubleValuesSource(new JsScriptBindings(fieldBindings, params));
         });
     return context.factoryClazz.cast(factory);
   }

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

@@ -68,7 +68,6 @@
 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;
@@ -143,8 +142,7 @@ 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);
+    DocLookup docLookup = new DocLookup(queryFields::get);
     contextBuilder.setDocLookup(docLookup);
 
     String rootQueryNestedPath =
@@ -418,8 +416,8 @@ private static Map<String, FieldDef> getRetrieveFields(
   /**
    * Add index fields to given query fields map.
    *
-   * @param indexState state for query index
    * @param queryFields mutable current map of query fields
+   * @param otherFields fields to add to query fields
    * @throws IllegalArgumentException if any index field already exists
    */
   private static void addToQueryFields(

src/test/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexStateTest.java

@@ -72,7 +72,6 @@
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import org.apache.lucene.expressions.Bindings;
 import org.apache.lucene.facet.FacetsConfig;
 import org.apache.lucene.index.ConcurrentMergeScheduler;
 import org.apache.lucene.index.IndexWriterConfig;
@@ -1030,17 +1029,6 @@ public void testGetEagerGlobalOrdinalFields() throws IOException {
     verifyNoMoreInteractions(mockFieldState);
   }
 
-  @Test
-  public void testGetExpressionBindings() throws IOException {
-    FieldAndFacetState mockFieldState = mock(FieldAndFacetState.class);
-    Bindings mockBindings = mock(Bindings.class);
-    when(mockFieldState.getExprBindings()).thenReturn(mockBindings);
-    ImmutableIndexState indexState = getIndexState(getEmptyState(), mockFieldState);
-    assertSame(mockBindings, indexState.getExpressionBindings());
-    verify(mockFieldState, times(1)).getExprBindings();
-    verifyNoMoreInteractions(mockFieldState);
-  }
-
   @Test
   public void testHasNestedChildFields() throws IOException {
     FieldAndFacetState mockFieldState = mock(FieldAndFacetState.class);