diff --git a/clientlib/src/main/proto/yelp/nrtsearch/search.proto b/clientlib/src/main/proto/yelp/nrtsearch/search.proto index 536879f25..9480421ea 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/search.proto +++ b/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 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 { 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 c4e18b0cc..cf8837e36 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java +++ b/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 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()); } } diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/DocLookup.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/DocLookup.java index 97b71cd30..096b19117 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/doc/DocLookup.java +++ b/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; + } } /** diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/FieldDefCreator.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/FieldDefCreator.java index 4c840b0ba..fb788544e 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/FieldDefCreator.java +++ b/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); } 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 index 070913115..019f78fdb 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/RuntimeFieldDef.java +++ b/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 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 6f4141373..984e23635 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 @@ -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); } 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 index 1699ce179..3c6c318b8 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java +++ b/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 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> 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; - } + /** 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 params, DocLookup docLookup); + SegmentFactory 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 + // compile context for the RuntimeScript, contains script type info public static final ScriptContext CONTEXT = new ScriptContext<>( - "runtime", Factory.class, RuntimeScript.SegmentFactory.class, RuntimeScript.class); + "runtime", Factory.class, SegmentFactory.class, RuntimeScript.class); } 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 881855163..c37f18e17 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 @@ -95,30 +95,31 @@ public T compile(String source, ScriptContext context) { 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); - } +// 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 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; } } 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 caa52ad1f..d0619526e 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 @@ -129,9 +129,8 @@ public static SearchContext buildContextForRequest( Map queryRuntimeFields = getRuntimeFields(indexState, searchRequest); Map queryFields = new HashMap<>(queryVirtualFields); - for (String key : queryRuntimeFields.keySet()) { - queryFields.put(key, queryRuntimeFields.get(key)); - } + + addToQueryFields(queryFields, queryRuntimeFields); addIndexFields(indexState, queryFields); contextBuilder.setQueryFields(Collections.unmodifiableMap(queryFields)); @@ -259,8 +258,9 @@ private static Map getRuntimeFields( RuntimeScript.Factory factory = ScriptService.getInstance().compile(vf.getScript(), RuntimeScript.CONTEXT); Map params = ScriptParamsUtils.decodeParams(vf.getScript().getParamsMap()); + RuntimeScript.SegmentFactory segmentFactory = factory.newFactory(params, indexState.docLookup); FieldDef runtimeField = - new RuntimeFieldDef(vf.getName(), factory.newFactory(params, indexState.docLookup)); + new RuntimeFieldDef(vf.getName(), segmentFactory); runtimeFields.put(vf.getName(), runtimeField); } return runtimeFields; @@ -302,6 +302,23 @@ private static Map getRetrieveFields( return retrieveFields; } + /** + * Add index fields to given query fields map. + * + * @param indexState state for query index + * @param queryFields mutable current map of query fields + * @throws IllegalArgumentException if any index field already exists + */ + private static void addToQueryFields(Map queryFields, Map otherFields) { + for (String key : otherFields.keySet()) { + FieldDef current = queryFields.put(key, otherFields.get(key)); + if (current != null) { + throw new IllegalArgumentException( + "QueryFields: " + key + " specified multiple times"); + } + } + } + /** * Add index fields to given query fields map. *