Change ValueSource to segment factory.

clientlib/src/main/proto/yelp/nrtsearch/search.proto

@@ -452,7 +452,8 @@ message SearchRequest {
     bool explain = 25;
     // Search nested object fields for each hit
     map<string, InnerHit> inner_hits = 26;
-    repeated RuntimeField runtimeFields = 27; //Defines runtime fields for this query.
+    // Defines runtime fields for this query
+    repeated RuntimeField runtimeFields = 27;
 }
 
 /* Inner Hit search request */
@@ -481,8 +482,10 @@ message VirtualField {
 
 /* Runtime field used during search */
 message RuntimeField {
-    Script script = 1; // Script defining this field's values.
-    string name = 2; // Runtime field's name. Must be different from registered fields and any other runtime fields.
+    // Script defining this field's values.
+    Script script = 1;
+    // Runtime field's name. Must be different from registered fields and any other runtime fields.
+    string name = 2;
 }
 
 message Script {

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

@@ -41,6 +41,7 @@
 import com.yelp.nrtsearch.server.luceneserver.field.VirtualFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitFetchTask;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreTask;
+import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript;
 import com.yelp.nrtsearch.server.luceneserver.search.FieldFetchContext;
 import com.yelp.nrtsearch.server.luceneserver.search.SearchContext;
 import com.yelp.nrtsearch.server.luceneserver.search.SearchCutoffWrapper.CollectionTimeoutException;
@@ -49,7 +50,6 @@
 import com.yelp.nrtsearch.server.monitoring.VerboseIndexCollector;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.Date;
 import java.util.HashMap;
@@ -70,8 +70,6 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
-import org.apache.lucene.queries.function.FunctionValues;
-import org.apache.lucene.queries.function.ValueSource;
 import org.apache.lucene.search.DoubleValues;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.ReferenceManager;
@@ -787,13 +785,6 @@ public boolean advanceExact(int doc) throws IOException {
         doubleValues.advanceExact(docID);
         compositeFieldValue.addFieldValue(
             FieldValue.newBuilder().setDoubleValue(doubleValues.doubleValue()));
-      } else if (fd instanceof RuntimeFieldDef) {
-        RuntimeFieldDef runtimeFieldDef = (RuntimeFieldDef) fd;
-
-        assert !Double.isNaN(hit.getScore()) || runtimeFieldDef.getValuesSource() != null;
-
-        Object obj = runtimeFieldDef.getValuesSource();
-        compositeFieldValue.addFieldValue(FieldValue.newBuilder().setStructValue((Struct) obj));
       } else if (fd instanceof IndexableFieldDef && ((IndexableFieldDef) fd).hasDocValues()) {
         int docID = hit.getLuceneDocId() - leaf.docBase;
         // it may be possible to cache this if there are multiple hits in the same segment
@@ -906,7 +897,7 @@ private static void fetchSlice(
               fieldDefEntry.getKey(),
               (VirtualFieldDef) fieldDefEntry.getValue());
         } else if (fieldDefEntry.getValue() instanceof RuntimeFieldDef) {
-          fetchRuntimeFromValueSource(
+          fetchRuntimeFromSegmentFactory(
               sliceHits,
               sliceSegment,
               fieldDefEntry.getKey(),
@@ -967,21 +958,20 @@ private static void fetchFromValueSource(
     }
 
     /** Fetch field value from runtime field's Object. */
-    private static void fetchRuntimeFromValueSource(
+    private static void fetchRuntimeFromSegmentFactory(
         List<SearchResponse.Hit.Builder> sliceHits,
         LeafReaderContext sliceSegment,
         String name,
         RuntimeFieldDef runtimeFieldDef)
         throws IOException {
-      ValueSource valueSource = runtimeFieldDef.getValuesSource();
-      Map context = Collections.emptyMap();
-      FunctionValues values = valueSource.getValues(context, sliceSegment);
+      RuntimeScript.SegmentFactory segmentFactory = runtimeFieldDef.getSegmentFactory();
+      RuntimeScript values = segmentFactory.newInstance(sliceSegment);
 
       for (SearchResponse.Hit.Builder hit : sliceHits) {
         int docID = hit.getLuceneDocId() - sliceSegment.docBase;
         // Check if the value is available for the current document
         if (values != null) {
-          Object obj = values.objectVal(docID);
+          Object obj = values.execute();
           SearchResponse.Hit.CompositeFieldValue.Builder compositeFieldValue =
               SearchResponse.Hit.CompositeFieldValue.newBuilder();
           if (obj instanceof Float) {
@@ -990,11 +980,17 @@ private static void fetchRuntimeFromValueSource(
           } else if (obj instanceof String) {
             compositeFieldValue.addFieldValue(
                 SearchResponse.Hit.FieldValue.newBuilder().setTextValue(String.valueOf(obj)));
-          }
-          if (obj instanceof Long) {
+          } else if (obj instanceof Double) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setDoubleValue((Double) obj));
+          } else if (obj instanceof Long) {
             compositeFieldValue.addFieldValue(
                 SearchResponse.Hit.FieldValue.newBuilder().setLongValue((Long) obj));
+          } else if (obj instanceof Integer) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setIntValue((Integer) obj));
           }
+          // TODO: Add support for list and map.
           hit.putFields(name, compositeFieldValue.build());
         }
       }

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

@@ -36,7 +36,12 @@ public DocLookup(IndexState indexState) {
    * @return lookup accessor for given segment context
    */
   public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
-    return new SegmentDocLookup(indexState, context);
+    try {
+      return new SegmentDocLookup(indexState, context);
+    } catch (Exception e) {
+      System.out.println((e.getStackTrace()));
+      return null;
+    }
   }
 
   /**

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

@@ -55,6 +55,11 @@ public FieldDefCreator(LuceneServerConfiguration configuration) {
         (name, field) -> {
           throw new UnsupportedOperationException("Virtual fields should be created directly");
         });
+    register(
+        "RUNTIME",
+        (name, field) -> {
+          throw new UnsupportedOperationException("Runtime fields should be created directly");
+        });
     register("VECTOR", VectorFieldDef::new);
     register("CONTEXT_SUGGEST", ContextSuggestFieldDef::new);
   }

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

@@ -16,25 +16,25 @@
 package com.yelp.nrtsearch.server.luceneserver.field;
 
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef.FacetValueType;
-import org.apache.lucene.queries.function.ValueSource;
+import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript;
 
 /**
  * Field definition for a runtime field. Runtime fields are able to produce a value for each given
  * index document.
  */
 public class RuntimeFieldDef extends FieldDef {
-  private final ValueSource valuesSource;
+  private final RuntimeScript.SegmentFactory segmentFactory;
   private final IndexableFieldDef.FacetValueType facetValueType;
 
   /**
    * Field constructor.
    *
    * @param name name of field
-   * @param valuesSource lucene value source used to produce field value from documents
+   * @param RuntimeScript.SegmentFactory lucene value source used to produce field value from documents
    */
-  public RuntimeFieldDef(String name, ValueSource valuesSource) {
+  public RuntimeFieldDef(String name, RuntimeScript.SegmentFactory segmentFactory) {
     super(name);
-    this.valuesSource = valuesSource;
+    this.segmentFactory = segmentFactory;
     this.facetValueType = FacetValueType.NO_FACETS;
   }
 
@@ -43,8 +43,8 @@ public RuntimeFieldDef(String name, ValueSource valuesSource) {
    *
    * @return lucene value source used to produce field value from documents
    */
-  public ValueSource getValuesSource() {
-    return valuesSource;
+  public RuntimeScript.SegmentFactory getSegmentFactory() {
+    return segmentFactory;
   }
 
   @Override

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

@@ -238,9 +238,9 @@ public static void parseRuntimeField(Field field, FieldAndFacetState.Builder fie
       throw new IllegalArgumentException("Only js lang supported for index runtime fields");
     }
     // js scripts use Bindings instead of DocLookup
-    ValueSource values = factory.newFactory(params, null);
+    RuntimeScript.SegmentFactory segmentFactory = factory.newFactory(params, null);
 
-    FieldDef runtimeFieldDef = new RuntimeFieldDef(field.getName(), values);
+    FieldDef runtimeFieldDef = new RuntimeFieldDef(field.getName(), segmentFactory);
     fieldStateBuilder.addField(runtimeFieldDef, field);
     logger.info("REGISTER: " + runtimeFieldDef.getName() + " -> " + runtimeFieldDef);
   }

src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java

@@ -18,6 +18,7 @@
 import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup;
 import com.yelp.nrtsearch.server.luceneserver.doc.LoadedDocValues;
 import com.yelp.nrtsearch.server.luceneserver.doc.SegmentDocLookup;
+import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript.SegmentFactory;
 import java.io.IOException;
 import java.util.Map;
 import java.util.Objects;
@@ -31,16 +32,9 @@
  * access to the query parameters, the document doc values through {@link SegmentDocLookup}.
  */
 public abstract class RuntimeScript {
-
-  private static final int DOC_UNSET = -1;
   private final Map<String, Object> params;
   private final SegmentDocLookup segmentDocLookup;
 
-  private Object obj;
-
-  private int docId = DOC_UNSET;
-  private int scoreDocId = DOC_UNSET;
-
   // names for parameters to execute
   public static final String[] PARAMETERS = new String[] {};
 
@@ -79,107 +73,37 @@ public Map<String, LoadedDocValues<?>> getDoc() {
     return segmentDocLookup;
   }
 
-  /**
-   * Simple abstract implementation of a {@link ValueSource} this can be extended for engines that
-   * need to implement a custom {@link RuntimeScript}. The newInstance and needs_score must be
-   * implemented. If more state is needed, the equals/hashCode should be redefined appropriately.
-   *
-   * <p>This class conforms with the script compile contract, see {@link ScriptContext}. However,
-   * Engines are also free to create there own {@link ValueSource} implementations instead.
-   */
-  public abstract static class SegmentFactory extends ValueSource {
-    private final Map<String, Object> params;
-    private final DocLookup docLookup;
-
-    public SegmentFactory(Map<String, Object> params, DocLookup docLookup) {
-      this.params = params;
-      this.docLookup = docLookup;
-    }
-
-    public Map<String, Object> getParams() {
-      return params;
-    }
-
-    public DocLookup getDocLookup() {
-      return docLookup;
-    }
+  /** Factory interface for creating a RuntimeScript bound to a lucene segment. */
+  public interface SegmentFactory {
 
     /**
-     * Create a {@link FunctionValues} instance for the given lucene segment.
+     * Create a RuntimeScript instance for a lucene segment.
      *
-     * @param ctx Context map for
-     * @param context segment context
-     * @return script to produce values for the given segment
+     * @param context lucene segment context
+     * @return segment level RuntimeScript
+     * @throws IOException
      */
-    public abstract FunctionValues newInstance(Map ctx, LeafReaderContext context);
-
-    /**
-     * Get if this script will need access to the document score.
-     *
-     * @return if this script uses the document score.
-     */
-    public abstract boolean needs_score();
-
-    /** Redirect {@link ValueSource} interface to script contract method. */
-    @Override
-    public FunctionValues getValues(Map context, LeafReaderContext ctx) throws IOException {
-      return newInstance(context, ctx);
-    }
-
-    @Override
-    public int hashCode() {
-      return Objects.hash(params, docLookup);
-    }
-
-    @Override
-    public boolean equals(Object obj) {
-      if (obj == null) {
-        return false;
-      }
-      if (obj.getClass() != this.getClass()) {
-        return false;
-      }
-      RuntimeScript.SegmentFactory factory = (RuntimeScript.SegmentFactory) obj;
-      return Objects.equals(factory.params, this.params)
-          && Objects.equals(factory.docLookup, this.docLookup);
-    }
-
-    @Override
-    public String toString() {
-      return "RuntimeScriptValuesSource: params: " + params + ", docLookup: " + docLookup;
-    }
+    RuntimeScript newInstance(LeafReaderContext context) throws IOException;
   }
 
   /**
-   * Factory required from the compilation of a ScoreScript. Used to produce request level {@link
-   * ValueSource}. See script compile contract {@link ScriptContext}.
+   * Factory required for the compilation of a RuntimeScript. Used to produce request level {@link
+   * SegmentFactory}. See script compile contract {@link ScriptContext}.
    */
   public interface Factory {
     /**
-     * Create request level {@link SegmentFactory}.
+     * Create request level {@link RuntimeScript.SegmentFactory}.
      *
      * @param params parameters from script request
      * @param docLookup index level doc value lookup provider
-     * @return {@link ValueSource} to evaluate script
+     * @return {@link RuntimeScript.SegmentFactory} to evaluate script
      */
-    ValueSource newFactory(Map<String, Object> params, DocLookup docLookup);
+    SegmentFactory newFactory(Map<String, Object> params, DocLookup docLookup);
   }
 
-  /**
-   * Advance script to a given segment document.
-   *
-   * @param doc segment doc id
-   * @return if there is data for the given id, this should always be the case
-   */
-  public boolean advanceExact(int doc) {
-    segmentDocLookup.setDocId(doc);
-    docId = doc;
-    scoreDocId = DOC_UNSET;
-    return true;
-  }
 
-  // compile context for the ScoreScript, contains script type info
+  // compile context for the RuntimeScript, contains script type info
   public static final ScriptContext<Factory> CONTEXT =
       new ScriptContext<>(
-          "runtime", Factory.class, RuntimeScript.SegmentFactory.class, RuntimeScript.class);
+          "runtime", Factory.class, SegmentFactory.class, RuntimeScript.class);
 }

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

@@ -95,30 +95,31 @@ public <T> T compile(String source, ScriptContext<T> context) {
       return context.factoryClazz.cast(factory);
     }
 
-    if (context.equals(RuntimeScript.CONTEXT)) {
-      RuntimeScript.Factory runtimeFactory =
-          ((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;
-            }
-            // TODO: Support returning other typees.
-            return ValueSource.fromDoubleValuesSource(
-                expr.getDoubleValuesSource(new JsScriptBindings(fieldBindings, scriptParams)));
-          });
-      return context.factoryClazz.cast(runtimeFactory);
-    }
+//    if (context.equals(RuntimeScript.CONTEXT)) {
+//      RuntimeScript.Factory runtimeFactory =
+//          ((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;
+//            }
+//            // TODO: Support returning other types.
+////            return ValueSource.fromDoubleValuesSource(
+////                expr.getDoubleValuesSource(new JsScriptBindings(fieldBindings, scriptParams)));
+//            return expr.getDoubleValuesSource(new JsScriptBindings(fieldBindings, scriptParams).getDoubleValuesSource());
+//          });
+//      return context.factoryClazz.cast(runtimeFactory);
+//    }
     return null;
   }
 }