Reverting design to use _PARENT notation

Author: dhruva-r

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

@@ -25,15 +25,9 @@
  */
 public class DocLookup {
   private final Function<String, FieldDef> fieldDefLookup;
-  private final String queryNestedPath;
 
   public DocLookup(Function<String, FieldDef> fieldDefLookup) {
-    this(fieldDefLookup, null);
-  }
-
-  public DocLookup(Function<String, FieldDef> fieldDefLookup, String queryNestedPath) {
     this.fieldDefLookup = fieldDefLookup;
-    this.queryNestedPath = queryNestedPath;
   }
 
   /**
@@ -43,7 +37,7 @@ public DocLookup(Function<String, FieldDef> fieldDefLookup, String queryNestedPa
    * @return lookup accessor for given segment context
    */
   public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
-    return new SegmentDocLookup(fieldDefLookup, context, queryNestedPath);
+    return new SegmentDocLookup(fieldDefLookup, context);
   }
 
   /**

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

@@ -39,23 +39,14 @@ public class SegmentDocLookup implements Map<String, LoadedDocValues<?>> {
   private final Function<String, FieldDef> fieldDefLookup;
   private final LeafReaderContext context;
   private final Map<String, LoadedDocValues<?>> loaderCache = new HashMap<>();
-  private final String queryNestedPath;
 
   private int docId = -1;
   private int parentDocId = -1;
   private SegmentDocLookup parentLookup = null;
 
   public SegmentDocLookup(Function<String, FieldDef> fieldDefLookup, LeafReaderContext context) {
-    this(fieldDefLookup, context, null);
-  }
-
-  public SegmentDocLookup(
-      Function<String, FieldDef> fieldDefLookup,
-      LeafReaderContext context,
-      String queryNestedPath) {
     this.fieldDefLookup = fieldDefLookup;
     this.context = context;
-    this.queryNestedPath = queryNestedPath;
   }
 
   /**
@@ -82,7 +73,7 @@ public boolean isEmpty() {
 
   /**
    * Check if a given field name is capable of having doc values. This does not mean there is data
-   * present, just that there can be.
+   * present, just that there can be. Handles "_PARENT." prefix for parent field access.
    *
    * @param key field name
    * @return if this field may have stored doc values
@@ -94,6 +85,10 @@ public boolean containsKey(Object key) {
     }
     String fieldName = key.toString();
 
+    if (fieldName.startsWith("_PARENT.")) {
+      fieldName = fieldName.substring("_PARENT.".length());
+    }
+
     try {
       FieldDef field = fieldDefLookup.apply(fieldName);
       return field instanceof IndexableFieldDef && ((IndexableFieldDef<?>) field).hasDocValues();
@@ -112,7 +107,7 @@ public boolean containsValue(Object value) {
    * cache. The data is loaded for the current set document id.
    *
    * <p>The system automatically determines if a field requires parent document access based on the
-   * current nested path. Fields are resolved automatically without requiring explicit notation.
+   * "_PARENT." prefix in the field name. Fields with this prefix will access the parent document.
    *
    * @param key field name
    * @return {@link LoadedDocValues} implementation for the given field
@@ -125,18 +120,16 @@ public LoadedDocValues<?> get(Object key) {
     Objects.requireNonNull(key);
     String fieldName = key.toString();
 
-    if ((!IndexState.ROOT.equals(queryNestedPath))) {
-      int parentLevels = computeParentLevels(queryNestedPath, fieldName);
-      if (parentLevels > 0) {
-        SegmentDocLookup currentLookup = getParentLookup();
-        if (currentLookup == null) {
-          throw new IllegalArgumentException(
-              "Could not access parent field: "
-                  + fieldName
-                  + " (document may not be nested or parent field may not exist)");
-        }
-        return currentLookup.get(fieldName);
+    if (fieldName.startsWith("_PARENT.")) {
+      String actualFieldName = fieldName.substring("_PARENT.".length());
+      SegmentDocLookup parentLookup = getParentLookup();
+      if (parentLookup == null) {
+        throw new IllegalArgumentException(
+            "Could not access parent field: "
+                + fieldName
+                + " (document may not be nested or parent field may not exist)");
       }
+      return parentLookup.get(actualFieldName);
     }
 
     LoadedDocValues<?> docValues = loaderCache.get(fieldName);
@@ -222,52 +215,6 @@ private int getParentDocId() {
     return docId + offset;
   }
 
-  /** Computes the number of parent levels to traverse to access a field. */
-  private static int computeParentLevels(String currentNestedPath, String targetFieldPath) {
-    if (currentNestedPath == null
-        || currentNestedPath.isEmpty()
-        || IndexState.ROOT.equals(currentNestedPath)) {
-      return -1; // Field is at current level or below
-    }
-
-    if (targetFieldPath.startsWith(currentNestedPath + ".")
-        || targetFieldPath.equals(currentNestedPath)) {
-      return -1; // Field is at current level or below
-    }
-
-    String[] currentPathParts = currentNestedPath.split("\\.");
-    String[] targetPathParts = targetFieldPath.split("\\.");
-
-    // Find common prefix
-    int commonPrefixLength = 0;
-    int minLength = Math.min(currentPathParts.length, targetPathParts.length);
-    for (int i = 0; i < minLength; i++) {
-      if (currentPathParts[i].equals(targetPathParts[i])) {
-        commonPrefixLength++;
-      } else {
-        break;
-      }
-    }
-
-    int levelsUp = currentPathParts.length - commonPrefixLength;
-
-    // TODO - Enable support for more than one level of nesting
-    if (levelsUp > 1) {
-      throw new IllegalArgumentException(
-          "Field access requires "
-              + levelsUp
-              + " parent levels, but only 1 level of nesting is supported. "
-              + "Field: "
-              + targetFieldPath
-              + ", Current path: "
-              + currentNestedPath);
-    }
-
-    // If we need to go up to access the field, return the number of levels
-    // If levelsUp is 0, it means the field is at the same level or below
-    return levelsUp > 0 ? levelsUp : -1;
-  }
-
   @Override
   public LoadedDocValues<?> put(String key, LoadedDocValues<?> value) {
     throw new UnsupportedOperationException();

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

@@ -134,36 +134,32 @@ public static SearchContext buildContextForRequest(
         .setExplain(searchRequest.getExplain())
         .setWarming(warming);
 
-    String rootQueryNestedPath =
-        indexState.resolveQueryNestedPath(searchRequest.getQueryNestedPath());
-    contextBuilder.setQueryNestedPath(rootQueryNestedPath);
-
-    // First, create a basic queryFields map with index fields only
-    Map<String, FieldDef> queryFields = new HashMap<>();
-    addIndexFields(indexState, queryFields);
-
-    // Create DocLookup with nested path for use by virtual/runtime fields
-    DocLookup docLookup = new DocLookup(queryFields::get, rootQueryNestedPath);
-    contextBuilder.setDocLookup(docLookup);
+    Map<String, FieldDef> queryVirtualFields = getVirtualFields(indexState, searchRequest);
+    Map<String, FieldDef> queryRuntimeFields = getRuntimeFields(indexState, searchRequest);
 
-    // Now add virtual and runtime fields that can use the nested-path-aware DocLookup
-    Map<String, FieldDef> queryVirtualFields =
-        getVirtualFields(indexState, searchRequest, docLookup);
-    Map<String, FieldDef> queryRuntimeFields =
-        getRuntimeFields(indexState, searchRequest, docLookup);
+    Map<String, FieldDef> queryFields = new HashMap<>(queryVirtualFields);
 
-    addToQueryFields(queryFields, queryVirtualFields);
     addToQueryFields(queryFields, queryRuntimeFields);
+    addIndexFields(indexState, queryFields);
     contextBuilder.setQueryFields(Collections.unmodifiableMap(queryFields));
 
     Map<String, FieldDef> retrieveFields =
         getRetrieveFields(searchRequest.getRetrieveFieldsList(), queryFields);
     contextBuilder.setRetrieveFields(Collections.unmodifiableMap(retrieveFields));
 
-    String queryText = searchRequest.getQueryText();
+    DocLookup docLookup = new DocLookup(queryFields::get);
+    contextBuilder.setDocLookup(docLookup);
+
+    String rootQueryNestedPath =
+        indexState.resolveQueryNestedPath(searchRequest.getQueryNestedPath());
+    contextBuilder.setQueryNestedPath(rootQueryNestedPath);
     Query query =
         extractQuery(
-            indexState, queryText, searchRequest.getQuery(), rootQueryNestedPath, docLookup);
+            indexState,
+            searchRequest.getQueryText(),
+            searchRequest.getQuery(),
+            rootQueryNestedPath,
+            docLookup);
     if (profileResult != null) {
       profileResult.setParsedQuery(query.toString());
     }
@@ -356,7 +352,7 @@ private static Query resolveKnnQueryAndBoost(
    * @throws IllegalArgumentException if there are multiple virtual fields with the same name
    */
   private static Map<String, FieldDef> getVirtualFields(
-      IndexState indexState, SearchRequest searchRequest, DocLookup docLookup) {
+      IndexState indexState, SearchRequest searchRequest) {
     if (searchRequest.getVirtualFieldsList().isEmpty()) {
       return new HashMap<>();
     }
@@ -371,7 +367,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, docLookup));
+          new VirtualFieldDef(vf.getName(), factory.newFactory(params, indexState.docLookup));
       virtualFields.put(vf.getName(), virtualField);
     }
     return virtualFields;
@@ -383,7 +379,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) {
+      IndexState indexState, SearchRequest searchRequest) {
     if (searchRequest.getRuntimeFieldsList().isEmpty()) {
       return Map.of();
     }
@@ -397,7 +393,8 @@ 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, docLookup);
+      RuntimeScript.SegmentFactory segmentFactory =
+          factory.newFactory(params, indexState.docLookup);
       FieldDef runtimeField = new RuntimeFieldDef(vf.getName(), segmentFactory);
       runtimeFields.put(vf.getName(), runtimeField);
     }