Create runtime fields

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

@@ -516,6 +516,8 @@ enum FieldType {
     // List of float values
     VECTOR = 15;
     CONTEXT_SUGGEST = 16; //Field used for contextual suggest fields
+    // Runtime fields
+    RUNTIME = 17;
 }
 
 //How the tokens should be indexed.
@@ -1132,4 +1134,4 @@ message CustomRequest {
 
 message CustomResponse {
     map<string, string> response = 1; // Custom response sent by the plugin
-}
+}

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

@@ -452,6 +452,7 @@ 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.
 }
 
 /* Inner Hit search request */
@@ -478,6 +479,12 @@ message VirtualField {
     string name = 2; // Virtual field's name. Must be different from registered fields and any other virtual fields.
 }
 
+/* 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.
+}
+
 message Script {
     string lang = 1; // script language
     string source = 2; // script source

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

@@ -37,6 +37,7 @@
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.ObjectFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.PolygonfieldDef;
+import com.yelp.nrtsearch.server.luceneserver.field.RuntimeFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.VirtualFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitFetchTask;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreTask;
@@ -48,6 +49,7 @@
 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;
@@ -68,6 +70,8 @@
 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;
@@ -761,7 +765,6 @@ private CompositeFieldValue getFieldForHit(
 
       // We detect invalid field above:
       assert fd != null;
-
       if (fd instanceof VirtualFieldDef) {
         VirtualFieldDef virtualFieldDef = (VirtualFieldDef) fd;
 
@@ -784,6 +787,13 @@ 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
@@ -792,6 +802,7 @@ public boolean advanceExact(int doc) throws IOException {
         for (int i = 0; i < docValues.size(); ++i) {
           compositeFieldValue.addFieldValue(docValues.toFieldValue(i));
         }
+
       }
       // retrieve stored fields
       else if (fd instanceof IndexableFieldDef && ((IndexableFieldDef) fd).isStored()) {
@@ -894,6 +905,12 @@ private static void fetchSlice(
               sliceSegment,
               fieldDefEntry.getKey(),
               (VirtualFieldDef) fieldDefEntry.getValue());
+        } else if (fieldDefEntry.getValue() instanceof RuntimeFieldDef) {
+          fetchRuntimeFromValueSource(
+              sliceHits,
+              sliceSegment,
+              fieldDefEntry.getKey(),
+              (RuntimeFieldDef) fieldDefEntry.getValue());
         } else if (fieldDefEntry.getValue() instanceof IndexableFieldDef) {
           IndexableFieldDef indexableFieldDef = (IndexableFieldDef) fieldDefEntry.getValue();
           if (indexableFieldDef.hasDocValues()) {
@@ -949,6 +966,40 @@ private static void fetchFromValueSource(
       }
     }
 
+    /** Fetch field value from runtime field's Object. */
+    private static void fetchRuntimeFromValueSource(
+        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);
+
+      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);
+          SearchResponse.Hit.CompositeFieldValue.Builder compositeFieldValue =
+              SearchResponse.Hit.CompositeFieldValue.newBuilder();
+          if (obj instanceof Float) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setFloatValue((Float) obj));
+          } else if (obj instanceof String) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setTextValue(String.valueOf(obj)));
+          }
+          if (obj instanceof Long) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setLongValue((Long) obj));
+          }
+          hit.putFields(name, compositeFieldValue.build());
+        }
+      }
+    }
+
     /** Fetch field value from its doc value */
     private static void fetchFromDocVales(
         List<SearchResponse.Hit.Builder> sliceHits,

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2020 Yelp Inc.
+ *
+ * Licensed 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 com.yelp.nrtsearch.server.luceneserver.field;
+
+import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef.FacetValueType;
+import org.apache.lucene.queries.function.ValueSource;
+
+/**
+ * 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 IndexableFieldDef.FacetValueType facetValueType;
+
+  /**
+   * Field constructor.
+   *
+   * @param name name of field
+   * @param valuesSource lucene value source used to produce field value from documents
+   */
+  public RuntimeFieldDef(String name, ValueSource valuesSource) {
+    super(name);
+    this.valuesSource = valuesSource;
+    this.facetValueType = FacetValueType.NO_FACETS;
+  }
+
+  /**
+   * Get value source for this field.
+   *
+   * @return lucene value source used to produce field value from documents
+   */
+  public ValueSource getValuesSource() {
+    return valuesSource;
+  }
+
+  @Override
+  public String getType() {
+    return "RUNTIME";
+  }
+
+  /**
+   * Get the facet value type for this field.
+   *
+   * @return field facet value type
+   */
+  @Override
+  public IndexableFieldDef.FacetValueType getFacetValueType() {
+    return facetValueType;
+  }
+}

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

@@ -23,9 +23,11 @@
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.FieldDefCreator;
 import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef;
+import com.yelp.nrtsearch.server.luceneserver.field.RuntimeFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.field.VirtualFieldDef;
 import com.yelp.nrtsearch.server.luceneserver.index.FieldAndFacetState;
 import com.yelp.nrtsearch.server.luceneserver.index.IndexStateManager;
+import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript;
 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;
@@ -35,6 +37,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.apache.lucene.queries.function.ValueSource;
 import org.apache.lucene.search.DoubleValuesSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -91,11 +94,14 @@ public static UpdatedFieldInfo updateFields(
     FieldAndFacetState.Builder fieldStateBuilder = currentState.toBuilder();
     List<Field> nonVirtualFields = new ArrayList<>();
     List<Field> virtualFields = new ArrayList<>();
+    List<Field> runtimeFields = new ArrayList<>();
 
     for (Field field : updateFields) {
       checkFieldName(field.getName());
       if (FieldType.VIRTUAL.equals(field.getType())) {
         virtualFields.add(field);
+      } else if (FieldType.RUNTIME.equals(field.getType())) {
+        runtimeFields.add(field);
       } else {
         nonVirtualFields.add(field);
       }
@@ -119,6 +125,14 @@ public static UpdatedFieldInfo updateFields(
       parseVirtualField(field, fieldStateBuilder);
       newFields.put(field.getName(), field);
     }
+
+    for (Field field : runtimeFields) {
+      if (newFields.containsKey(field.getName())) {
+        throw new IllegalArgumentException("Duplicate field registration: " + field.getName());
+      }
+      parseRuntimeField(field, fieldStateBuilder);
+      newFields.put(field.getName(), field);
+    }
     return new UpdatedFieldInfo(newFields, fieldStateBuilder.build());
   }
 
@@ -198,4 +212,36 @@ public static void parseVirtualField(Field field, FieldAndFacetState.Builder fie
     fieldStateBuilder.addField(virtualFieldDef, field);
     logger.info("REGISTER: " + virtualFieldDef.getName() + " -> " + virtualFieldDef);
   }
+
+  /**
+   * Parse a runtime {@link Field} message and apply changes to the {@link
+   * FieldAndFacetState.Builder}.
+   *
+   * @param field runtime field specification
+   * @param fieldStateBuilder builder for new field state
+   */
+  public static void parseRuntimeField(Field field, FieldAndFacetState.Builder fieldStateBuilder) {
+    RuntimeScript.Factory factory =
+        ScriptService.getInstance().compile(field.getScript(), RuntimeScript.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 runtime fields");
+    }
+    // js scripts use Bindings instead of DocLookup
+    ValueSource values = factory.newFactory(params, null);
+
+    FieldDef runtimeFieldDef = new RuntimeFieldDef(field.getName(), values);
+    fieldStateBuilder.addField(runtimeFieldDef, field);
+    logger.info("REGISTER: " + runtimeFieldDef.getName() + " -> " + runtimeFieldDef);
+  }
 }