dhruva_RP-13223-Parent_field_access_from_nested_query

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

@@ -25,9 +25,15 @@
  */
 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;
   }
 
   /**
@@ -37,7 +43,7 @@ public DocLookup(Function<String, FieldDef> fieldDefLookup) {
    * @return lookup accessor for given segment context
    */
   public SegmentDocLookup getSegmentLookup(LeafReaderContext context) {
-    return new SegmentDocLookup(fieldDefLookup, context);
+    return new SegmentDocLookup(fieldDefLookup, context, queryNestedPath);
   }
 
   /**

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

@@ -39,13 +39,22 @@ 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 SegmentDocLookup parentLookup = null; // Lazy initialized
 
   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;
   }
 
   /**
@@ -74,8 +83,6 @@ 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.
    *
-   * <p>For "_PARENT." notation, this checks if the underlying parent field can have doc values.
-   *
    * @param key field name
    * @return if this field may have stored doc values
    */
@@ -86,10 +93,6 @@ 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();
@@ -107,9 +110,8 @@ public boolean containsValue(Object value) {
    * Get the {@link LoadedDocValues} for a given field. Creates a new instance or uses one from the
    * cache. The data is loaded for the current set document id.
    *
-   * <p>Clients can explicitly access parent fields using the "_PARENT." notation. For example,
-   * "_PARENT.biz_feature_a" will directly access the "biz_feature_a" field from the parent
-   * document. Multiple levels are supported like "_PARENT._PARENT.field_name".
+   * <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.
    *
    * @param key field name
    * @return {@link LoadedDocValues} implementation for the given field
@@ -122,16 +124,27 @@ public LoadedDocValues<?> get(Object key) {
     Objects.requireNonNull(key);
     String fieldName = key.toString();
 
-    if (fieldName.startsWith("_PARENT.")) {
-      String remainingFieldName = fieldName.substring("_PARENT.".length());
-      SegmentDocLookup parentLookup = getParentLookup();
-      if (parentLookup == null) {
-        throw new IllegalArgumentException(
-            "Could not access parent field: "
-                + remainingFieldName
-                + " (document may not be nested or parent field may not exist)");
+    if (queryNestedPath != null && !queryNestedPath.isEmpty() && !"_root".equals(queryNestedPath)) {
+      int parentLevels = resolveParentLevels(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)");
+        }
+
+        for (int i = 1; i < parentLevels; i++) {
+          currentLookup = currentLookup.getParentLookup();
+          if (currentLookup == null) {
+            throw new IllegalArgumentException(
+                "Could not access field: " + fieldName + " (required parent level not accessible)");
+          }
+        }
+
+        return currentLookup.get(fieldName);
       }
-      return parentLookup.get(remainingFieldName);
     }
 
     LoadedDocValues<?> docValues = loaderCache.get(fieldName);
@@ -213,6 +226,47 @@ private int getParentDocId() {
     return docId + offset;
   }
 
+  /**
+   * Utility method to resolve the relative path from current nested location to target field.
+   * Returns the number of parent levels to traverse to access the field.
+   *
+   * @param currentNestedPath current nested document path (e.g., "reviews.generation")
+   * @param targetFieldPath absolute field path (e.g., "biz_name" or "reviews.rating")
+   * @return number of parent levels to traverse, or -1 if field is in current or child level
+   */
+  private static int resolveParentLevels(String currentNestedPath, String targetFieldPath) {
+    if (currentNestedPath == null
+        || currentNestedPath.isEmpty()
+        || "_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;
+
+    // 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,32 +134,36 @@ public static SearchContext buildContextForRequest(
         .setExplain(searchRequest.getExplain())
         .setWarming(warming);
 
-    Map<String, FieldDef> queryVirtualFields = getVirtualFields(indexState, searchRequest);
-    Map<String, FieldDef> queryRuntimeFields = getRuntimeFields(indexState, searchRequest);
+    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> queryFields = new HashMap<>(queryVirtualFields);
+    // 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);
 
+    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));
 
-    DocLookup docLookup = new DocLookup(queryFields::get);
-    contextBuilder.setDocLookup(docLookup);
-
-    String rootQueryNestedPath =
-        indexState.resolveQueryNestedPath(searchRequest.getQueryNestedPath());
-    contextBuilder.setQueryNestedPath(rootQueryNestedPath);
+    String queryText = searchRequest.getQueryText();
     Query query =
         extractQuery(
-            indexState,
-            searchRequest.getQueryText(),
-            searchRequest.getQuery(),
-            rootQueryNestedPath,
-            docLookup);
+            indexState, queryText, searchRequest.getQuery(), rootQueryNestedPath, docLookup);
     if (profileResult != null) {
       profileResult.setParsedQuery(query.toString());
     }
@@ -352,7 +356,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) {
+      IndexState indexState, SearchRequest searchRequest, DocLookup docLookup) {
     if (searchRequest.getVirtualFieldsList().isEmpty()) {
       return new HashMap<>();
     }
@@ -367,7 +371,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;
@@ -379,7 +383,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) {
+      IndexState indexState, SearchRequest searchRequest, DocLookup docLookup) {
     if (searchRequest.getRuntimeFieldsList().isEmpty()) {
       return Map.of();
     }
@@ -393,8 +397,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);
     }