Make searcher available through DocLookup (#924)

Author: aprudhomme

src/main/java/com/yelp/nrtsearch/server/doc/DocLookup.java

@@ -19,6 +19,7 @@
 import java.util.Collection;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager;
 import org.apache.lucene.index.LeafReaderContext;
 
 /**
@@ -28,6 +29,7 @@
 public class DocLookup {
   private final Function<String, FieldDef> fieldDefLookup;
   private final Supplier<Collection<String>> allFieldNamesSupplier;
+  private final SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy;
 
   /**
    * Constructor.
@@ -37,9 +39,11 @@ public class DocLookup {
    */
   public DocLookup(
       Function<String, FieldDef> fieldDefLookup,
-      Supplier<Collection<String>> allFieldNamesSupplier) {
+      Supplier<Collection<String>> allFieldNamesSupplier,
+      SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy) {
     this.fieldDefLookup = fieldDefLookup;
     this.allFieldNamesSupplier = allFieldNamesSupplier;
+    this.searcherAndTaxonomy = searcherAndTaxonomy;
   }
 
   /**
@@ -82,4 +86,12 @@ public FieldDef getFieldDefOrThrow(String fieldName) {
   public Collection<String> getAllFieldNames() {
     return allFieldNamesSupplier.get();
   }
+
+  /**
+   * Get the searcher for this request, or null if unavailable. Currently, this is only available
+   * for query supplied virtual or runtime fields.
+   */
+  public SearcherTaxonomyManager.SearcherAndTaxonomy getSearcherAndTaxonomy() {
+    return searcherAndTaxonomy;
+  }
 }

src/main/java/com/yelp/nrtsearch/server/index/FieldUpdateUtils.java

@@ -170,7 +170,7 @@ public static void parseVirtualField(Field field, FieldAndFacetState.Builder fie
         factory.newFactory(
             params,
             new DocLookup(
-                fieldStateBuilder.getFields()::get, fieldStateBuilder.getFields()::keySet));
+                fieldStateBuilder.getFields()::get, fieldStateBuilder.getFields()::keySet, null));
 
     FieldDef virtualFieldDef = new VirtualFieldDef(field.getName(), values);
     fieldStateBuilder.addField(virtualFieldDef, field);

src/main/java/com/yelp/nrtsearch/server/index/IndexState.java

@@ -101,7 +101,8 @@ public abstract class IndexState implements Closeable {
   /**
    * Index level doc values lookup. Generates {@link SegmentDocLookup} for a given lucene segment.
    */
-  public final DocLookup docLookup = new DocLookup(this::getField, () -> getAllFields().keySet());
+  public final DocLookup docLookup =
+      new DocLookup(this::getField, () -> getAllFields().keySet(), null);
 
   /**
    * Holds the configuration for parallel fetch operations.

src/main/java/com/yelp/nrtsearch/server/search/SearchRequestProcessor.java

@@ -135,8 +135,17 @@ public static SearchContext buildContextForRequest(
         .setExplain(searchRequest.getExplain())
         .setWarming(warming);
 
-    Map<String, FieldDef> queryVirtualFields = getVirtualFields(indexState, searchRequest);
-    Map<String, FieldDef> queryRuntimeFields = getRuntimeFields(indexState, searchRequest);
+    // We do this to make the searcher available to scripting engines, by adding it to the
+    // DocLookup. It would be better to change the script factory interface to take a
+    // context object containing these objects separately. We should consider making
+    // these changes in a future minor version.
+    DocLookup indexLookupWithSearcher =
+        new DocLookup(
+            indexState::getField, () -> indexState.getAllFields().keySet(), searcherAndTaxonomy);
+    Map<String, FieldDef> queryVirtualFields =
+        getVirtualFields(indexLookupWithSearcher, searchRequest);
+    Map<String, FieldDef> queryRuntimeFields =
+        getRuntimeFields(indexLookupWithSearcher, searchRequest);
 
     Map<String, FieldDef> queryFields = new HashMap<>(queryVirtualFields);
 
@@ -148,7 +157,7 @@ public static SearchContext buildContextForRequest(
         getRetrieveFields(searchRequest.getRetrieveFieldsList(), queryFields);
     contextBuilder.setRetrieveFields(Collections.unmodifiableMap(retrieveFields));
 
-    DocLookup docLookup = new DocLookup(queryFields::get, queryFields::keySet);
+    DocLookup docLookup = new DocLookup(queryFields::get, queryFields::keySet, searcherAndTaxonomy);
     contextBuilder.setDocLookup(docLookup);
 
     String rootQueryNestedPath =
@@ -363,7 +372,7 @@ private static void setVectorTotalHits(
    * @throws IllegalArgumentException if there are multiple virtual fields with the same name
    */
   private static Map<String, FieldDef> getVirtualFields(
-      IndexState indexState, SearchRequest searchRequest) {
+      DocLookup docLookup, SearchRequest searchRequest) {
     if (searchRequest.getVirtualFieldsList().isEmpty()) {
       return new HashMap<>();
     }
@@ -378,7 +387,7 @@ private static Map<String, FieldDef> getVirtualFields(
           ScriptService.getInstance().compile(vf.getScript(), ScoreScript.CONTEXT);
       Map<String, Object> params = ScriptParamsUtils.decodeParams(vf.getScript().getParamsMap());
       FieldDef virtualField =
-          new VirtualFieldDef(vf.getName(), factory.newFactory(params, indexState.docLookup));
+          new VirtualFieldDef(vf.getName(), factory.newFactory(params, docLookup));
       virtualFields.put(vf.getName(), virtualField);
     }
     return virtualFields;
@@ -390,7 +399,7 @@ private static Map<String, FieldDef> getVirtualFields(
    * @throws IllegalArgumentException if there are multiple runtime fields with the same name
    */
   private static Map<String, FieldDef> getRuntimeFields(
-      IndexState indexState, SearchRequest searchRequest) {
+      DocLookup docLookup, SearchRequest searchRequest) {
     if (searchRequest.getRuntimeFieldsList().isEmpty()) {
       return Map.of();
     }
@@ -404,8 +413,7 @@ private static Map<String, FieldDef> getRuntimeFields(
       RuntimeScript.Factory factory =
           ScriptService.getInstance().compile(vf.getScript(), RuntimeScript.CONTEXT);
       Map<String, Object> params = ScriptParamsUtils.decodeParams(vf.getScript().getParamsMap());
-      RuntimeScript.SegmentFactory segmentFactory =
-          factory.newFactory(params, indexState.docLookup);
+      RuntimeScript.SegmentFactory segmentFactory = factory.newFactory(params, docLookup);
       FieldDef runtimeField = new RuntimeFieldDef(vf.getName(), segmentFactory);
       runtimeFields.put(vf.getName(), runtimeField);
     }

src/test/java/com/yelp/nrtsearch/server/doc/DocLookupTest.java

@@ -25,6 +25,7 @@
 import java.util.HashSet;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import org.apache.lucene.facet.taxonomy.SearcherTaxonomyManager;
 import org.apache.lucene.index.LeafReaderContext;
 import org.junit.Before;
 import org.junit.Test;
@@ -33,6 +34,7 @@ public class DocLookupTest {
 
   private Function<String, FieldDef> mockFieldDefLookup;
   private Supplier<Collection<String>> mockAllFieldNamesSupplier;
+  private SearcherTaxonomyManager.SearcherAndTaxonomy mockSearcherAndTaxonomy;
   private FieldDef mockFieldDef1;
   private FieldDef mockFieldDef2;
   private DocLookup docLookup;
@@ -42,6 +44,7 @@ public class DocLookupTest {
   public void setUp() {
     mockFieldDefLookup = mock(Function.class);
     mockAllFieldNamesSupplier = mock(Supplier.class);
+    mockSearcherAndTaxonomy = mock(SearcherTaxonomyManager.SearcherAndTaxonomy.class);
     mockFieldDef1 = mock(FieldDef.class);
     mockFieldDef2 = mock(FieldDef.class);
 
@@ -53,7 +56,8 @@ public void setUp() {
     when(mockFieldDef2.getType()).thenReturn("INT");
     when(mockFieldDef2.getFacetValueType()).thenReturn(IndexableFieldDef.FacetValueType.NO_FACETS);
 
-    docLookup = new DocLookup(mockFieldDefLookup, mockAllFieldNamesSupplier);
+    docLookup =
+        new DocLookup(mockFieldDefLookup, mockAllFieldNamesSupplier, mockSearcherAndTaxonomy);
   }
 
   @Test
@@ -217,7 +221,7 @@ public void testGetSegmentLookup() {
   public void testConstructor_NullFieldDefLookup() {
     // Act & Assert
     try {
-      new DocLookup(null, mockAllFieldNamesSupplier);
+      new DocLookup(null, mockAllFieldNamesSupplier, null);
       // Constructor doesn't validate null parameters, so this should not throw
     } catch (Exception e) {
       fail("Constructor should not throw for null fieldDefLookup: " + e.getMessage());
@@ -228,10 +232,68 @@ public void testConstructor_NullFieldDefLookup() {
   public void testConstructor_NullAllFieldNamesSupplier() {
     // Act & Assert
     try {
-      new DocLookup(mockFieldDefLookup, null);
+      new DocLookup(mockFieldDefLookup, null, null);
       // Constructor doesn't validate null parameters, so this should not throw
     } catch (Exception e) {
       fail("Constructor should not throw for null allFieldNamesSupplier: " + e.getMessage());
     }
   }
+
+  @Test
+  public void testGetSearcherAndTaxonomy_ValidSearcherAndTaxonomy() {
+    // Act
+    SearcherTaxonomyManager.SearcherAndTaxonomy result = docLookup.getSearcherAndTaxonomy();
+
+    // Assert
+    assertEquals(mockSearcherAndTaxonomy, result);
+  }
+
+  @Test
+  public void testGetSearcherAndTaxonomy_NullSearcherAndTaxonomy() {
+    // Arrange
+    DocLookup docLookupWithNull =
+        new DocLookup(mockFieldDefLookup, mockAllFieldNamesSupplier, null);
+
+    // Act
+    SearcherTaxonomyManager.SearcherAndTaxonomy result = docLookupWithNull.getSearcherAndTaxonomy();
+
+    // Assert
+    assertNull(result);
+  }
+
+  @Test
+  public void testGetSearcherAndTaxonomy_MultipleCalls() {
+    // Act
+    SearcherTaxonomyManager.SearcherAndTaxonomy result1 = docLookup.getSearcherAndTaxonomy();
+    SearcherTaxonomyManager.SearcherAndTaxonomy result2 = docLookup.getSearcherAndTaxonomy();
+
+    // Assert
+    assertEquals(mockSearcherAndTaxonomy, result1);
+    assertEquals(mockSearcherAndTaxonomy, result2);
+    assertSame(result1, result2); // Should return the same instance
+  }
+
+  @Test
+  public void testConstructor_NullSearcherAndTaxonomy() {
+    // Act & Assert
+    try {
+      new DocLookup(mockFieldDefLookup, mockAllFieldNamesSupplier, null);
+      // Constructor doesn't validate null parameters, so this should not throw
+    } catch (Exception e) {
+      fail("Constructor should not throw for null searcherAndTaxonomy: " + e.getMessage());
+    }
+  }
+
+  @Test
+  public void testConstructor_AllParameters() {
+    // Act & Assert
+    try {
+      DocLookup testDocLookup =
+          new DocLookup(mockFieldDefLookup, mockAllFieldNamesSupplier, mockSearcherAndTaxonomy);
+      assertNotNull(testDocLookup);
+      assertEquals(mockSearcherAndTaxonomy, testDocLookup.getSearcherAndTaxonomy());
+    } catch (Exception e) {
+      fail("Constructor should not throw for valid parameters: " + e.getMessage());
+    }
+  }
 }