Addressing PR comments, adding parameter to the tryGetFromParentDocument

Author: dhruva-r

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

@@ -111,8 +111,9 @@ public LoadedDocValues<?> get(Object key) {
     Objects.requireNonNull(key);
     String fieldName = key.toString();
     LoadedDocValues<?> docValues = loaderCache.get(fieldName);
+    FieldDef fieldDef = null;
     if (docValues == null) {
-      FieldDef fieldDef = fieldDefLookup.apply(fieldName);
+      fieldDef = fieldDefLookup.apply(fieldName);
       if (fieldDef == null) {
         throw new IllegalArgumentException("Field does not exist: " + fieldName);
       }
@@ -133,7 +134,10 @@ public LoadedDocValues<?> get(Object key) {
           "Could not set doc: " + docId + ", field: " + fieldName, e);
     }
     if (docValues.isEmpty()) {
-      LoadedDocValues<?> parentDocValues = tryGetFromParentDocument(fieldName);
+      if (fieldDef == null) {
+        fieldDef = fieldDefLookup.apply(fieldName);
+      }
+      LoadedDocValues<?> parentDocValues = tryGetFromParentDocument(fieldName, fieldDef);
       if (parentDocValues != null) {
         return parentDocValues;
       }
@@ -145,11 +149,12 @@ public LoadedDocValues<?> get(Object key) {
    * Attempt to retrieve the field from the parent document using NESTED_DOCUMENT_OFFSET.
    *
    * @param fieldName the name of the field to retrieve
+   * @param fieldDef the definition of the field to retrieve
    * @return LoadedDocValues from parent document, or null if not found or not a nested document
    * @throws IllegalArgumentException if there are issues accessing the offset field or parent
    *     document
    */
-  private LoadedDocValues<?> tryGetFromParentDocument(String fieldName) {
+  private LoadedDocValues<?> tryGetFromParentDocument(String fieldName, FieldDef fieldDef) {
     FieldDef offsetFieldDef;
     try {
       offsetFieldDef = IndexState.getMetaField(IndexState.NESTED_DOCUMENT_OFFSET);
@@ -178,6 +183,7 @@ private LoadedDocValues<?> tryGetFromParentDocument(String fieldName) {
           "Could not set doc: " + docId + " for NESTED_DOCUMENT_OFFSET field", e);
     }
 
+    // If there's no offset value, this is not a nested document and therefore we should terminate
     if (offsetDocValues.isEmpty()) {
       return null;
     }
@@ -187,11 +193,6 @@ private LoadedDocValues<?> tryGetFromParentDocument(String fieldName) {
 
     // The offset represents the exact number of documents to jump forward to reach the parent
     int parentDocId = docId + offset;
-
-    FieldDef fieldDef = fieldDefLookup.apply(fieldName);
-    if (fieldDef == null) {
-      throw new IllegalArgumentException("Field does not exist: " + fieldName);
-    }
     if (!(fieldDef instanceof IndexableFieldDef<?> indexableFieldDef)) {
       throw new IllegalArgumentException("Field cannot have doc values: " + fieldName);
     }

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

@@ -22,6 +22,7 @@
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.yelp.nrtsearch.server.ServerTestCase;
+import com.yelp.nrtsearch.server.field.FieldDef;
 import com.yelp.nrtsearch.server.grpc.AddDocumentRequest;
 import com.yelp.nrtsearch.server.grpc.FieldDefRequest;
 import com.yelp.nrtsearch.server.index.IndexState;
@@ -102,13 +103,15 @@ public void testParentAccessWithValidNestedDocs() throws Exception {
       SegmentDocLookup docLookup = new SegmentDocLookup(indexState::getField, leafContext);
 
       Method tryGetMethod =
-          SegmentDocLookup.class.getDeclaredMethod("tryGetFromParentDocument", String.class);
+          SegmentDocLookup.class.getDeclaredMethod(
+              "tryGetFromParentDocument", String.class, FieldDef.class);
       tryGetMethod.setAccessible(true);
 
       // Test 1: Child document should be able to access parent field
       docLookup.setDocId(0);
+      FieldDef docIdFieldDef = indexState.getField("doc_id");
       LoadedDocValues<?> parentField =
-          (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id");
+          (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id", docIdFieldDef);
 
       if (parentField != null && !parentField.isEmpty()) {
         assertNotNull("Should find parent field from child document", parentField);
@@ -119,8 +122,10 @@ public void testParentAccessWithValidNestedDocs() throws Exception {
       // Test 2: Non-existent field should return null
       LoadedDocValues<?> nonExistentField = null;
       try {
+        FieldDef nonExistentFieldDef = indexState.getField("non_existent_field");
         nonExistentField =
-            (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "non_existent_field");
+            (LoadedDocValues<?>)
+                tryGetMethod.invoke(docLookup, "non_existent_field", nonExistentFieldDef);
       } catch (Exception e) {
         if (e.getCause() instanceof IllegalArgumentException) {
           nonExistentField = null;
@@ -133,7 +138,7 @@ public void testParentAccessWithValidNestedDocs() throws Exception {
       // Test 3: Parent document (no offset field) should return null
       docLookup.setDocId(2);
       LoadedDocValues<?> fromParentDoc =
-          (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id");
+          (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id", docIdFieldDef);
       assertNull("Parent document should return null (no offset field)", fromParentDoc);
     } finally {
       if (s != null) {
@@ -155,7 +160,8 @@ public void testParentAccessWithDifferentFieldTypes() throws Exception {
       SegmentDocLookup docLookup = new SegmentDocLookup(indexState::getField, leafContext);
 
       Method tryGetMethod =
-          SegmentDocLookup.class.getDeclaredMethod("tryGetFromParentDocument", String.class);
+          SegmentDocLookup.class.getDeclaredMethod(
+              "tryGetFromParentDocument", String.class, FieldDef.class);
       tryGetMethod.setAccessible(true);
 
       docLookup.setDocId(0);
@@ -164,8 +170,9 @@ public void testParentAccessWithDifferentFieldTypes() throws Exception {
 
       for (String fieldName : fieldsToTest) {
         try {
+          FieldDef fieldDef = indexState.getField(fieldName);
           LoadedDocValues<?> result =
-              (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, fieldName);
+              (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, fieldName, fieldDef);
           assertTrue("Field access should not throw unexpected exceptions: " + fieldName, true);
         } catch (Exception e) {
           assertTrue(
@@ -193,14 +200,17 @@ public void testParentAccessWithoutNestedDocuments() throws Exception {
       SegmentDocLookup docLookup = new SegmentDocLookup(indexState::getField, leafContext);
 
       Method tryGetMethod =
-          SegmentDocLookup.class.getDeclaredMethod("tryGetFromParentDocument", String.class);
+          SegmentDocLookup.class.getDeclaredMethod(
+              "tryGetFromParentDocument", String.class, FieldDef.class);
       tryGetMethod.setAccessible(true);
 
       int maxDoc = leafContext.reader().maxDoc();
       if (maxDoc > 1) {
         docLookup.setDocId(maxDoc - 1);
 
-        LoadedDocValues<?> result = (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id");
+        FieldDef docIdFieldDef = indexState.getField("doc_id");
+        LoadedDocValues<?> result =
+            (LoadedDocValues<?>) tryGetMethod.invoke(docLookup, "doc_id", docIdFieldDef);
         assertNull("Document without nested structure should return null", result);
       }
     } finally {