diff --git a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto index 527b28754..69bfbd4ec 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto +++ b/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 response = 1; // Custom response sent by the plugin -} \ No newline at end of file +} diff --git a/clientlib/src/main/proto/yelp/nrtsearch/search.proto b/clientlib/src/main/proto/yelp/nrtsearch/search.proto index 28a85f004..536879f25 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/search.proto +++ b/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 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 diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java index f2b503696..c4e18b0cc 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java +++ b/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 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 sliceHits, diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/RuntimeFieldDef.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/RuntimeFieldDef.java new file mode 100644 index 000000000..070913115 --- /dev/null +++ b/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; + } +} diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java index 5c1bb5885..6f4141373 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java +++ b/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 nonVirtualFields = new ArrayList<>(); List virtualFields = new ArrayList<>(); + List 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 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); + } } diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java new file mode 100644 index 000000000..1699ce179 --- /dev/null +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java @@ -0,0 +1,185 @@ +/* + * 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.script; + +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 java.io.IOException; +import java.util.Map; +import java.util.Objects; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.queries.function.FunctionValues; +import org.apache.lucene.queries.function.ValueSource; + +/** + * Script to produce a value for a given document. Implementations must have an execute function. + * This class conforms with the script compile contract, see {@link ScriptContext}. The script has + * 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 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[] {}; + + /** + * ObjectScript constructor. + * + * @param params script parameters from {@link com.yelp.nrtsearch.server.grpc.Script} + * @param docLookup index level doc values lookup + * @param leafContext lucene segment context + */ + public RuntimeScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + this.params = params; + this.segmentDocLookup = docLookup.getSegmentLookup(leafContext); + } + + /** + * Main script function. + * + * @return object value computed for document + */ + public abstract Object execute(); + + /** Set segment level id for the next document to score. * */ + public void setDocId(int docId) { + segmentDocLookup.setDocId(docId); + } + + /** Get the script parameters provided in the request. */ + public Map getParams() { + return params; + } + + /** Get doc values for the current document. */ + public Map> 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. + * + *

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 params; + private final DocLookup docLookup; + + public SegmentFactory(Map params, DocLookup docLookup) { + this.params = params; + this.docLookup = docLookup; + } + + public Map getParams() { + return params; + } + + public DocLookup getDocLookup() { + return docLookup; + } + + /** + * Create a {@link FunctionValues} instance for the given lucene segment. + * + * @param ctx Context map for + * @param context segment context + * @return script to produce values for the given segment + */ + 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; + } + } + + /** + * Factory required from the compilation of a ScoreScript. Used to produce request level {@link + * ValueSource}. See script compile contract {@link ScriptContext}. + */ + public interface Factory { + /** + * Create request level {@link SegmentFactory}. + * + * @param params parameters from script request + * @param docLookup index level doc value lookup provider + * @return {@link ValueSource} to evaluate script + */ + ValueSource newFactory(Map 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 + public static final ScriptContext CONTEXT = + new ScriptContext<>( + "runtime", Factory.class, RuntimeScript.SegmentFactory.class, RuntimeScript.class); +} diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/ScriptService.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/ScriptService.java index 01b3b4a9e..7d1914288 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/ScriptService.java +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/ScriptService.java @@ -47,6 +47,7 @@ public class ScriptService { ImmutableList.>builder() .add(FacetScript.CONTEXT) .add(ScoreScript.CONTEXT) + .add(RuntimeScript.CONTEXT) .build(); private final Map scriptEngineMap = new HashMap<>(); diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/js/JsScriptEngine.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/js/JsScriptEngine.java index 64dbcac57..881855163 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/js/JsScriptEngine.java +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/js/JsScriptEngine.java @@ -15,6 +15,7 @@ */ package com.yelp.nrtsearch.server.luceneserver.script.js; +import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript; import com.yelp.nrtsearch.server.luceneserver.script.ScoreScript; import com.yelp.nrtsearch.server.luceneserver.script.ScriptContext; import com.yelp.nrtsearch.server.luceneserver.script.ScriptEngine; @@ -24,6 +25,7 @@ import org.apache.lucene.expressions.Bindings; import org.apache.lucene.expressions.Expression; import org.apache.lucene.expressions.js.JavascriptCompiler; +import org.apache.lucene.queries.function.ValueSource; /** * Script engine that provides a language based on javascript expressions. Expressions are compiled @@ -58,7 +60,7 @@ public String getLang() { */ @Override public T compile(String source, ScriptContext context) { - if (!context.equals(ScoreScript.CONTEXT)) { + if (!(context.equals(ScoreScript.CONTEXT) || context.equals(RuntimeScript.CONTEXT))) { throw new IllegalArgumentException("Unsupported script context: " + context.name); } Expression expr; @@ -69,25 +71,54 @@ public T compile(String source, ScriptContext context) { throw new IllegalArgumentException( String.format("could not parse expression: %s", source), pe); } - ScoreScript.Factory factory = - ((params, docLookup) -> { - Map scriptParams; - Bindings fieldBindings; - Object bindingsParam = params.get("bindings"); - if (bindingsParam instanceof Bindings) { - fieldBindings = (Bindings) bindingsParam; + if (context.equals(ScoreScript.CONTEXT)) { + ScoreScript.Factory factory = + ((params, docLookup) -> { + Map 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)); - }); - return context.factoryClazz.cast(factory); + // 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)); + }); + + return context.factoryClazz.cast(factory); + } + + if (context.equals(RuntimeScript.CONTEXT)) { + RuntimeScript.Factory runtimeFactory = + ((params, docLookup) -> { + Map 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); + } + return null; } } diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java index a0cce8aa7..caa52ad1f 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java @@ -21,6 +21,7 @@ import com.yelp.nrtsearch.server.grpc.PluginRescorer; import com.yelp.nrtsearch.server.grpc.ProfileResult; import com.yelp.nrtsearch.server.grpc.QueryRescorer; +import com.yelp.nrtsearch.server.grpc.RuntimeField; import com.yelp.nrtsearch.server.grpc.SearchRequest; import com.yelp.nrtsearch.server.grpc.SearchResponse; import com.yelp.nrtsearch.server.grpc.VirtualField; @@ -30,6 +31,7 @@ import com.yelp.nrtsearch.server.luceneserver.doc.DefaultSharedDocContext; import com.yelp.nrtsearch.server.luceneserver.field.FieldDef; 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.highlights.HighlightFetchTask; import com.yelp.nrtsearch.server.luceneserver.highlights.HighlighterService; @@ -40,6 +42,7 @@ import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreOperation; import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreTask; import com.yelp.nrtsearch.server.luceneserver.rescore.RescorerCreator; +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.search.collectors.AdditionalCollectorManager; @@ -123,8 +126,12 @@ public static SearchContext buildContextForRequest( .setExplain(searchRequest.getExplain()); Map queryVirtualFields = getVirtualFields(indexState, searchRequest); + Map queryRuntimeFields = getRuntimeFields(indexState, searchRequest); Map queryFields = new HashMap<>(queryVirtualFields); + for (String key : queryRuntimeFields.keySet()) { + queryFields.put(key, queryRuntimeFields.get(key)); + } addIndexFields(indexState, queryFields); contextBuilder.setQueryFields(Collections.unmodifiableMap(queryFields)); @@ -232,6 +239,33 @@ private static Map getVirtualFields( return virtualFields; } + /** + * Parses any runtimeFields, which define dynamic (expression) fields for this one request. + * + * @throws IllegalArgumentException if there are multiple runtime fields with the same name + */ + private static Map getRuntimeFields( + IndexState indexState, SearchRequest searchRequest) { + if (searchRequest.getRuntimeFieldsList().isEmpty()) { + return new HashMap<>(); + } + + Map runtimeFields = new HashMap<>(); + for (RuntimeField vf : searchRequest.getRuntimeFieldsList()) { + if (runtimeFields.containsKey(vf.getName())) { + throw new IllegalArgumentException( + "Multiple definitions of runtime field: " + vf.getName()); + } + RuntimeScript.Factory factory = + ScriptService.getInstance().compile(vf.getScript(), RuntimeScript.CONTEXT); + Map params = ScriptParamsUtils.decodeParams(vf.getScript().getParamsMap()); + FieldDef runtimeField = + new RuntimeFieldDef(vf.getName(), factory.newFactory(params, indexState.docLookup)); + runtimeFields.put(vf.getName(), runtimeField); + } + return runtimeFields; + } + /** * Get map of fields that need to be retrieved for the given request. * @@ -287,7 +321,7 @@ private static void addIndexFields(IndexState indexState, Map /** If a field's value can be retrieved */ private static boolean isRetrievable(FieldDef fieldDef) { - if (fieldDef instanceof VirtualFieldDef) { + if (fieldDef instanceof VirtualFieldDef || fieldDef instanceof RuntimeFieldDef) { return true; } if (fieldDef instanceof IndexableFieldDef) {