diff --git a/clientlib/src/main/java/com/yelp/nrtsearch/server/utils/StructJsonUtils.java b/clientlib/src/main/java/com/yelp/nrtsearch/server/utils/StructJsonUtils.java index 820f52b6c..79b5ab0b4 100644 --- a/clientlib/src/main/java/com/yelp/nrtsearch/server/utils/StructJsonUtils.java +++ b/clientlib/src/main/java/com/yelp/nrtsearch/server/utils/StructJsonUtils.java @@ -71,6 +71,22 @@ public static Struct convertMapToStruct(Map map, boolean longAsS return builder.build(); } + /** + * Convert an Iterable of native java types into a protobuf ListValue. This iterable may contain + * null, Boolean, String, Number, Iterable, or Map (String key). + * + * @param iterable iterable of java native types + * @param longAsString if Long values should be encoded as String + * @return protobuf list value for iterable + */ + public static ListValue convertIterableToListValue(Iterable iterable, boolean longAsString) { + ListValue.Builder listValueBuilder = ListValue.newBuilder(); + for (Object e : iterable) { + listValueBuilder.addValues(convertObjectToValue(e, longAsString)); + } + return listValueBuilder.build(); + } + /** * Convert an Iterable of native java types into a protobuf Value. This iterable may contain null, * Boolean, String, Number, Iterable, or Map (String key). 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..4d21600b8 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/search.proto +++ b/clientlib/src/main/proto/yelp/nrtsearch/search.proto @@ -452,6 +452,8 @@ message SearchRequest { bool explain = 25; // Search nested object fields for each hit map inner_hits = 26; + // Defines runtime fields for this query + repeated RuntimeField runtimeFields = 27; } /* Inner Hit search request */ @@ -478,6 +480,14 @@ 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 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 { string lang = 1; // script language string source = 2; // script source @@ -598,12 +608,12 @@ message SearchResponse { google.protobuf.Struct structValue = 8; // Value for structured data // Value for VECTOR FieldType Vector vectorValue = 9; + google.protobuf.ListValue listValue = 10; } message Vector { repeated float value = 1; } - } message CompositeFieldValue { 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 a8b8d9666..0dd9783c1 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java @@ -37,15 +37,18 @@ 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; +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; import com.yelp.nrtsearch.server.luceneserver.search.SearchRequestProcessor; import com.yelp.nrtsearch.server.luceneserver.search.SearcherResult; import com.yelp.nrtsearch.server.monitoring.SearchResponseCollector; +import com.yelp.nrtsearch.server.utils.ObjectToCompositeFieldTransformer; import java.io.IOException; import java.util.ArrayList; import java.util.Comparator; @@ -763,7 +766,6 @@ private CompositeFieldValue getFieldForHit( // We detect invalid field above: assert fd != null; - if (fd instanceof VirtualFieldDef) { VirtualFieldDef virtualFieldDef = (VirtualFieldDef) fd; @@ -786,6 +788,18 @@ 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; + RuntimeScript.SegmentFactory segmentFactory = runtimeFieldDef.getSegmentFactory(); + RuntimeScript values = segmentFactory.newInstance(leaf); + int docID = hit.getLuceneDocId() - leaf.docBase; + // Check if the value is available for the current document + if (values != null) { + values.setDocId(docID); + Object obj = values.execute(); + ObjectToCompositeFieldTransformer.enrichCompositeField(obj, compositeFieldValue); + } + } 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 @@ -794,6 +808,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()) { @@ -896,6 +911,12 @@ private static void fetchSlice( sliceSegment, fieldDefEntry.getKey(), (VirtualFieldDef) fieldDefEntry.getValue()); + } else if (fieldDefEntry.getValue() instanceof RuntimeFieldDef) { + fetchRuntimeFromSegmentFactory( + sliceHits, + sliceSegment, + fieldDefEntry.getKey(), + (RuntimeFieldDef) fieldDefEntry.getValue()); } else if (fieldDefEntry.getValue() instanceof IndexableFieldDef) { IndexableFieldDef indexableFieldDef = (IndexableFieldDef) fieldDefEntry.getValue(); if (indexableFieldDef.hasDocValues()) { @@ -951,6 +972,30 @@ private static void fetchFromValueSource( } } + /** Fetch field value from runtime field's Object. */ + private static void fetchRuntimeFromSegmentFactory( + List sliceHits, + LeafReaderContext sliceSegment, + String name, + RuntimeFieldDef runtimeFieldDef) + throws IOException { + 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) { + values.setDocId(docID); + Object obj = values.execute(); + SearchResponse.Hit.CompositeFieldValue.Builder compositeFieldValue = + SearchResponse.Hit.CompositeFieldValue.newBuilder(); + ObjectToCompositeFieldTransformer.enrichCompositeField(obj, compositeFieldValue); + 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/facet/DrillSidewaysImpl.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/facet/DrillSidewaysImpl.java index ca24f9eb0..09859924b 100644 --- a/src/main/java/com/yelp/nrtsearch/server/luceneserver/facet/DrillSidewaysImpl.java +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/facet/DrillSidewaysImpl.java @@ -28,6 +28,7 @@ import com.yelp.nrtsearch.server.luceneserver.field.IndexableFieldDef; import com.yelp.nrtsearch.server.luceneserver.field.IntFieldDef; import com.yelp.nrtsearch.server.luceneserver.field.LongFieldDef; +import com.yelp.nrtsearch.server.luceneserver.field.RuntimeFieldDef; import com.yelp.nrtsearch.server.luceneserver.field.VirtualFieldDef; import com.yelp.nrtsearch.server.luceneserver.script.FacetScript; import com.yelp.nrtsearch.server.luceneserver.script.ScriptService; @@ -324,10 +325,11 @@ private static com.yelp.nrtsearch.server.grpc.FacetResult getFieldFacetResult( } FacetResult facetResult; - if (!(fieldDef instanceof IndexableFieldDef) && !(fieldDef instanceof VirtualFieldDef)) { + if (!(fieldDef instanceof IndexableFieldDef) + && !(fieldDef instanceof VirtualFieldDef || fieldDef instanceof RuntimeFieldDef)) { throw new IllegalArgumentException( String.format( - "field %s is neither a virtual field nor registered as an indexable field. Facets are supported only for these types", + "field %s is neither a virtual/runtime field nor registered as an indexable field. Facets are supported only for these types", fieldName)); } if (!facet.getNumericRangeList().isEmpty()) { 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 new file mode 100644 index 000000000..1aae0326e --- /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 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 RuntimeScript.SegmentFactory segmentFactory; + private final IndexableFieldDef.FacetValueType facetValueType; + + /** + * Field constructor. + * + * @param name name of field + * @param segmentFactory lucene value source used to produce field value from documents + */ + public RuntimeFieldDef(String name, RuntimeScript.SegmentFactory segmentFactory) { + super(name); + this.segmentFactory = segmentFactory; + this.facetValueType = FacetValueType.NO_FACETS; + } + + /** + * Get value source for this field. + * + * @return Segment factory to create the expression. + */ + public RuntimeScript.SegmentFactory getSegmentFactory() { + return segmentFactory; + } + + @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..0cb314575 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 @@ -119,6 +119,7 @@ public static UpdatedFieldInfo updateFields( parseVirtualField(field, fieldStateBuilder); newFields.put(field.getName(), field); } + return new UpdatedFieldInfo(newFields, fieldStateBuilder.build()); } 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..cb2ae7746 --- /dev/null +++ b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java @@ -0,0 +1,104 @@ +/* + * 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 com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript.SegmentFactory; +import java.io.IOException; +import java.util.Map; +import org.apache.lucene.index.LeafReaderContext; + +/** + * 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 final Map params; + private final SegmentDocLookup segmentDocLookup; + + // 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; + } + + /** Factory interface for creating a RuntimeScript bound to a lucene segment. */ + public interface SegmentFactory { + + /** + * Create a RuntimeScript instance for a lucene segment. + * + * @param context lucene segment context + * @return segment level RuntimeScript + * @throws IOException + */ + RuntimeScript newInstance(LeafReaderContext context) throws IOException; + } + + /** + * 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 RuntimeScript.SegmentFactory}. + * + * @param params parameters from script request + * @param docLookup index level doc value lookup provider + * @return {@link RuntimeScript.SegmentFactory} to evaluate script + */ + SegmentFactory newFactory(Map params, DocLookup docLookup); + } + + // compile context for the RuntimeScript, contains script type info + public static final ScriptContext CONTEXT = + new ScriptContext<>("runtime", Factory.class, 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/search/SearchRequestProcessor.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java index a0cce8aa7..9a91290ea 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,11 @@ public static SearchContext buildContextForRequest( .setExplain(searchRequest.getExplain()); Map queryVirtualFields = getVirtualFields(indexState, searchRequest); + Map queryRuntimeFields = getRuntimeFields(indexState, searchRequest); Map queryFields = new HashMap<>(queryVirtualFields); + + addToQueryFields(queryFields, queryRuntimeFields); addIndexFields(indexState, queryFields); contextBuilder.setQueryFields(Collections.unmodifiableMap(queryFields)); @@ -232,6 +238,34 @@ 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 Map.of(); + } + + 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()); + RuntimeScript.SegmentFactory segmentFactory = + factory.newFactory(params, indexState.docLookup); + FieldDef runtimeField = new RuntimeFieldDef(vf.getName(), segmentFactory); + runtimeFields.put(vf.getName(), runtimeField); + } + return runtimeFields; + } + /** * Get map of fields that need to be retrieved for the given request. * @@ -268,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. * @@ -287,7 +338,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) { @@ -333,7 +384,8 @@ private static Query extractQuery( /** If field is non-null it overrides any specified defaultField. */ private static QueryBuilder createQueryParser(IndexState state, String field) { - // TODO: Support "queryParser" field provided by user e.g. MultiFieldQueryParser, + // TODO: Support "queryParser" field provided by user e.g. + // MultiFieldQueryParser, // SimpleQueryParser, classic List fields; if (field != null) { @@ -458,7 +510,8 @@ private static InnerHitContext buildInnerHitContext( String innerHitName, InnerHit innerHit, boolean explain) { - // Do not apply nestedPath here. This is query is used to create a shared weight. + // Do not apply nestedPath here. This is query is used to create a shared + // weight. Query childQuery = extractQuery(indexState, "", innerHit.getInnerQuery(), null); return InnerHitContextBuilder.Builder() .withInnerHitName(innerHitName) diff --git a/src/main/java/com/yelp/nrtsearch/server/utils/ObjectToCompositeFieldTransformer.java b/src/main/java/com/yelp/nrtsearch/server/utils/ObjectToCompositeFieldTransformer.java new file mode 100644 index 000000000..1491f9ad8 --- /dev/null +++ b/src/main/java/com/yelp/nrtsearch/server/utils/ObjectToCompositeFieldTransformer.java @@ -0,0 +1,53 @@ +/* + * Copyright 2024 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.utils; + +import com.yelp.nrtsearch.server.grpc.SearchResponse; +import java.util.Map; + +public class ObjectToCompositeFieldTransformer { + + public static void enrichCompositeField( + Object obj, SearchResponse.Hit.CompositeFieldValue.Builder compositeFieldValue) { + 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))); + } 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)); + } else if (obj instanceof Map) { + compositeFieldValue.addFieldValue( + SearchResponse.Hit.FieldValue.newBuilder() + .setStructValue(StructJsonUtils.convertMapToStruct((Map) obj))); + } else if (obj instanceof Boolean) { + compositeFieldValue.addFieldValue( + SearchResponse.Hit.FieldValue.newBuilder().setBooleanValue((Boolean) obj)); + } else if (obj instanceof Iterable) { + compositeFieldValue.addFieldValue( + SearchResponse.Hit.FieldValue.newBuilder() + .setListValue(StructJsonUtils.convertIterableToListValue((Iterable) obj, false))); + } + } +} diff --git a/src/test/java/com/yelp/nrtsearch/server/luceneserver/facet/RuntimeScriptFacetsTest.java b/src/test/java/com/yelp/nrtsearch/server/luceneserver/facet/RuntimeScriptFacetsTest.java new file mode 100644 index 000000000..ece01679f --- /dev/null +++ b/src/test/java/com/yelp/nrtsearch/server/luceneserver/facet/RuntimeScriptFacetsTest.java @@ -0,0 +1,390 @@ +/* + * 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.facet; + +import static org.junit.Assert.assertEquals; + +import com.google.protobuf.Value; +import com.yelp.nrtsearch.server.config.LuceneServerConfiguration; +import com.yelp.nrtsearch.server.grpc.*; +import com.yelp.nrtsearch.server.luceneserver.ServerTestCase; +import com.yelp.nrtsearch.server.luceneserver.doc.DocLookup; +import com.yelp.nrtsearch.server.luceneserver.doc.LoadedDocValues; +import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript; +import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript.SegmentFactory; +import com.yelp.nrtsearch.server.luceneserver.script.ScriptContext; +import com.yelp.nrtsearch.server.luceneserver.script.ScriptEngine; +import com.yelp.nrtsearch.server.plugins.Plugin; +import com.yelp.nrtsearch.server.plugins.ScriptPlugin; +import io.grpc.testing.GrpcCleanupRule; +import java.io.IOException; +import java.util.*; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NoMergePolicy; +import org.junit.ClassRule; +import org.junit.Test; + +public class RuntimeScriptFacetsTest extends ServerTestCase { + private static final int NUM_DOCS = 100; + private static final int TOP_HITS = 10; + + @ClassRule public static final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule(); + + public static class TestRuntimeScriptPlugin extends Plugin implements ScriptPlugin { + + public Iterable getScriptEngines(List> contexts) { + return Collections.singletonList(new TestRuntimeScriptEngine()); + } + + public static class TestRuntimeScriptEngine implements ScriptEngine { + + @Override + public String getLang() { + return "test_p"; + } + + @Override + public T compile(String source, ScriptContext context) { + return context.factoryClazz.cast(new TestFactory(source)); + } + + public static class TestFactory implements RuntimeScript.Factory { + private final String source; + + public TestFactory(String source) { + this.source = source; + } + + @Override + public SegmentFactory newFactory(Map params, DocLookup docLookup) { + return new TestSegmentFactory(params, docLookup, source); + } + } + + public static class TestSegmentFactory implements RuntimeScript.SegmentFactory { + private final Map params; + private final DocLookup docLookup; + private final String source; + + public TestSegmentFactory(Map params, DocLookup docLookup, String source) { + this.params = params; + this.docLookup = docLookup; + this.source = source; + } + + @Override + public RuntimeScript newInstance(LeafReaderContext context) { + switch (source) { + case "int": + return new IntScript(params, docLookup, context); + case "string": + return new StringScript(params, docLookup, context); + case "map": + return new MapScript(params, docLookup, context); + case "list": + return new ListScript(params, docLookup, context); + case "docValue": + return new DocValueScript(params, docLookup, context); + } + throw new IllegalStateException("Unsupported script source: " + source); + } + + public static class IntScript extends RuntimeScript { + + public IntScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + super(params, docLookup, leafContext); + } + + @Override + public Object execute() { + return 2; + } + } + + public static class StringScript extends RuntimeScript { + + public StringScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + super(params, docLookup, leafContext); + } + + @Override + public Object execute() { + return "2"; + } + } + + public static class MapScript extends RuntimeScript { + + public MapScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + super(params, docLookup, leafContext); + } + + @Override + public Object execute() { + return Collections.singletonMap("key", 2.0); + } + } + + public static class ListScript extends RuntimeScript { + + public ListScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + super(params, docLookup, leafContext); + } + + @Override + public Object execute() { + List nums = new ArrayList(); + nums.add("1"); + nums.add("2"); + return nums; + } + } + + public static class DocValueScript extends RuntimeScript { + + public DocValueScript( + Map params, DocLookup docLookup, LeafReaderContext leafContext) { + super(params, docLookup, leafContext); + } + + @Override + public Object execute() { + Map> doc = getDoc(); + return doc.get("atom_1").get(0) + "_" + doc.get("atom_2").get(0); + } + } + } + } + } + + protected List getIndices() { + return Collections.singletonList(DEFAULT_TEST_INDEX); + } + + protected FieldDefRequest getIndexDef(String name) throws IOException { + return getFieldsFromResourceFile("/script/runtime_field_script.json"); + } + + protected void initIndex(String name) throws Exception { + IndexWriter writer = getGlobalState().getIndex(name).getShard(0).writer; + // don't want any merges for these tests + writer.getConfig().setMergePolicy(NoMergePolicy.INSTANCE); + + // add documents one chunk at a time to ensure multiple index segments + List requestChunk = new ArrayList<>(); + for (int id = 0; id < NUM_DOCS; ++id) { + requestChunk.add( + AddDocumentRequest.newBuilder() + .setIndexName(name) + .putFields( + "doc_id", + AddDocumentRequest.MultiValuedField.newBuilder() + .addValue(String.valueOf(id)) + .build()) + .putFields( + "int_field", + AddDocumentRequest.MultiValuedField.newBuilder() + .addValue(String.valueOf(id)) + .build()) + .putFields( + "atom_1", + AddDocumentRequest.MultiValuedField.newBuilder() + .addValue(String.valueOf(id % 3)) + .build()) + .putFields( + "atom_2", + AddDocumentRequest.MultiValuedField.newBuilder() + .addValue(String.valueOf(id % 2)) + .build()) + .build()); + + if (requestChunk.size() == TOP_HITS) { + addDocuments(requestChunk.stream()); + requestChunk.clear(); + writer.commit(); + } + } + } + + @Override + protected List getPlugins(LuceneServerConfiguration configuration) { + return Collections.singletonList(new TestRuntimeScriptPlugin()); + } + + @Test + public void RuntimeScriptForInt() { + + RuntimeField runtimeField = + RuntimeField.newBuilder() + .setScript(Script.newBuilder().setLang("test_p").setSource("int").build()) + .setName("runtime_field") + .build(); + + List expectedValues = new ArrayList<>(); + for (int id = 0; id < TOP_HITS; ++id) { + expectedValues.add(2); + } + SearchResponse response = doQuery(runtimeField); + assertEquals(TOP_HITS, response.getHitsCount()); + for (int id = 0; id < TOP_HITS; ++id) { + assertEquals( + response.getHits(id).getFieldsMap().get("runtime_field").getFieldValue(0).getIntValue(), + expectedValues.get(id)); + } + } + + @Test + public void RuntimeScriptForString() { + + RuntimeField runtimeField = + RuntimeField.newBuilder() + .setScript(Script.newBuilder().setLang("test_p").setSource("string").build()) + .setName("runtime_field") + .build(); + + List expectedValues = new ArrayList<>(); + for (int id = 0; id < TOP_HITS; ++id) { + expectedValues.add("2"); + } + SearchResponse response = doQuery(runtimeField); + assertEquals(TOP_HITS, response.getHitsCount()); + for (int id = 0; id < TOP_HITS; ++id) { + assertEquals( + response.getHits(id).getFieldsMap().get("runtime_field").getFieldValue(0).getTextValue(), + expectedValues.get(id)); + } + } + + @Test + public void RuntimeScriptForMap() { + + RuntimeField runtimeField = + RuntimeField.newBuilder() + .setScript(Script.newBuilder().setLang("test_p").setSource("map").build()) + .setName("runtime_field") + .build(); + + List expectedValues = new ArrayList<>(); + for (int id = 0; id < TOP_HITS; ++id) { + expectedValues.add(2.0); + } + SearchResponse response = doQuery(runtimeField); + assertEquals(TOP_HITS, response.getHitsCount()); + for (int id = 0; id < TOP_HITS; ++id) { + assertEquals( + response + .getHits(id) + .getFieldsMap() + .get("runtime_field") + .getFieldValue(0) + .getStructValue() + .getFieldsMap() + .get("key") + .getNumberValue(), + expectedValues.get(id)); + } + } + + @Test + public void RuntimeScriptForList() { + + RuntimeField runtimeField = + RuntimeField.newBuilder() + .setScript(Script.newBuilder().setLang("test_p").setSource("list").build()) + .setName("runtime_field") + .build(); + + List> expectedValues = new ArrayList<>(); + for (int id = 0; id < TOP_HITS; ++id) { + List nums = new ArrayList(); + nums.add(Value.newBuilder().setStringValue("1")); + nums.add(Value.newBuilder().setStringValue("2")); + List values = Collections.unmodifiableList(nums); + expectedValues.add(values); + } + SearchResponse response = doQuery(runtimeField); + assertEquals(TOP_HITS, response.getHitsCount()); + for (int id = 0; id < TOP_HITS; ++id) { + String respOne = + response + .getHits(id) + .getFieldsMap() + .get("runtime_field") + .getFieldValueList() + .get(0) + .getListValue() + .getValuesList() + .get(0) + .getStringValue(); + String respTwo = + response + .getHits(id) + .getFieldsMap() + .get("runtime_field") + .getFieldValueList() + .get(0) + .getListValue() + .getValuesList() + .get(1) + .getStringValue(); + assertEquals(respOne, expectedValues.get(id).get(0).getStringValue()); + assertEquals(respTwo, expectedValues.get(id).get(1).getStringValue()); + } + } + + @Test + public void RuntimeScriptForDocValue() { + + RuntimeField runtimeField = + RuntimeField.newBuilder() + .setScript(Script.newBuilder().setLang("test_p").setSource("docValue").build()) + .setName("runtime_field") + .build(); + + List expectedValues = new ArrayList<>(); + for (int id = 0; id < TOP_HITS; ++id) { + expectedValues.add(String.valueOf(id % 3) + "_" + String.valueOf(id % 2)); + } + SearchResponse response = doQuery(runtimeField); + assertEquals(TOP_HITS, response.getHitsCount()); + for (int id = 0; id < TOP_HITS; id++) { + assertEquals( + response.getHits(id).getFieldsMap().get("runtime_field").getFieldValue(0).getTextValue(), + expectedValues.get(id)); + } + } + + private SearchResponse doQuery(RuntimeField runtimeField) { + return getGrpcServer() + .getBlockingStub() + .search( + SearchRequest.newBuilder() + .setIndexName(DEFAULT_TEST_INDEX) + .setStartHit(0) + .setTopHits(10) + .addRetrieveFields("doc_id") + .addRetrieveFields("atom_1") + .addRetrieveFields("atom_2") + .addRetrieveFields("runtime_field") + .addRuntimeFields(runtimeField) + .build()); + } +} diff --git a/src/test/resources/script/runtime_field_script.json b/src/test/resources/script/runtime_field_script.json new file mode 100644 index 000000000..c637ce153 --- /dev/null +++ b/src/test/resources/script/runtime_field_script.json @@ -0,0 +1,26 @@ +{ + "indexName": "test_index", + "field": [ + { + "name": "doc_id", + "type": "ATOM", + "storeDocValues": true + }, + { + "name": "int_field", + "type": "INT", + "storeDocValues": true, + "search": true + }, + { + "name": "atom_1", + "type": "ATOM", + "storeDocValues": true + }, + { + "name": "atom_2", + "type": "ATOM", + "storeDocValues": true + } + ] +}