Add tests for doc values

Yelp · Mar 21, 2024 · 8a04251 · 8a04251
1 parent 86eda24
commit 8a04251
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 1 deletion.
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java
@@ -976,6 +976,7 @@ private static void fetchRuntimeFromSegmentFactory(
         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();

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
@@ -70,7 +70,7 @@ public Map<String, LoadedDocValues<?>> getDoc() {
     return segmentDocLookup;
   }
 
-  /** Factory interface for creating a RuntimeScript bound to a lucene segment. */
+  /** Factory interface for creating a RuntimeSsrc/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java cript bound to a lucene segment. */
   public interface SegmentFactory {
 
     /**

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
@@ -21,6 +21,7 @@
 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;
@@ -95,6 +96,8 @@ public RuntimeScript newInstance(LeafReaderContext context) {
               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);
         }
@@ -153,6 +156,20 @@ public Object execute() {
             return nums;
           }
         }
+
+        public static class DocValueScript extends RuntimeScript {
+
+          public DocValueScript(
+              Map<String, Object> params, DocLookup docLookup, LeafReaderContext leafContext) {
+            super(params, docLookup, leafContext);
+          }
+
+          @Override
+          public Object execute() {
+            Map<String, LoadedDocValues<?>> doc = getDoc();
+            return doc.get("atom_1").get(0) + "_" + doc.get("atom_2").get(0);
+          }
+        }
       }
     }
   }
@@ -317,6 +334,28 @@ public void RuntimeScriptForList() {
     }
   }
 
+  @Test
+  public void RuntimeScriptForDocValue() {
+
+    RuntimeField runtimeField =
+        RuntimeField.newBuilder()
+            .setScript(Script.newBuilder().setLang("painless").setSource("docValue").build())
+            .setName("runtime_field")
+            .build();
+
+    List expectedValues = new ArrayList<>();
+    for (int id = 0; id < SEGMENT_CHUNK; ++id) {
+      expectedValues.add(String.valueOf(id % 3) + "_" + String.valueOf(id % 2));
+    }
+    SearchResponse response = doQuery(runtimeField);
+    assertEquals(SEGMENT_CHUNK, response.getHitsCount());
+    for (int id = 0; id < SEGMENT_CHUNK; id++) {
+      assertEquals(
+          response.getHits(id).getFieldsMap().get("runtime_field").getFieldValue(0).getTextValue(),
+          expectedValues.get(id));
+    }
+  }
+
   private SearchResponse doQuery(RuntimeField runtimeField) {
     return getGrpcServer()
         .getBlockingStub()
@@ -326,6 +365,8 @@ private SearchResponse doQuery(RuntimeField runtimeField) {
                 .setStartHit(0)
                 .setTopHits(10)
                 .addRetrieveFields("doc_id")
+                .addRetrieveFields("atom_1")
+                .addRetrieveFields("atom_2")
                 .addRetrieveFields("runtime_field")
                 .addRuntimeFields(runtimeField)
                 .build());