fix ATOM sort and searchAfter (#732)

Yelp · Sep 27, 2024 · bef28b4 · bef28b4
1 parent 6dd6085
commit bef28b4
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 4 deletions.
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/AtomFieldDef.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/field/AtomFieldDef.java
@@ -26,6 +26,7 @@
 import org.apache.lucene.index.IndexOptions;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.SortedSetSortField;
+import org.apache.lucene.util.BytesRef;
 
 /** Field class for 'ATOM' field type. Uses {@link KeywordAnalyzer} for text analysis. */
 public class AtomFieldDef extends TextBaseFieldDef implements Sortable {
@@ -50,7 +51,7 @@ public String getType() {
 
   @Override
   public Object parseLastValue(String value) {
-    return value;
+    return new BytesRef(value);
   }
 
   @Override

diff --git a/...ain/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/SortFieldCollector.java b/...ain/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/SortFieldCollector.java
@@ -30,6 +30,7 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopFieldCollector;
 import org.apache.lucene.search.TopFieldDocs;
+import org.apache.lucene.util.BytesRef;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -89,8 +90,14 @@ public void fillLastHit(SearchResponse.SearchState.Builder stateBuilder, ScoreDo
     LastHitInfo.Builder lastHitBuilder = LastHitInfo.newBuilder();
     lastHitBuilder.setLastDocId(lastHit.doc);
     for (Object fv : fd.fields) {
-      stateBuilder.addLastFieldValues(fv.toString());
-      lastHitBuilder.addLastFieldValues(fv.toString());
+      String fvstr;
+      if (fv instanceof BytesRef) {
+        fvstr = ((BytesRef) fv).utf8ToString();
+      } else {
+        fvstr = fv.toString();
+      }
+      stateBuilder.addLastFieldValues(fvstr);
+      lastHitBuilder.addLastFieldValues(fvstr);
     }
     stateBuilder.setLastHitInfo(lastHitBuilder.build());
   }

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/sort/SortParser.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/sort/SortParser.java
@@ -31,6 +31,7 @@
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
+import org.apache.lucene.util.BytesRef;
 
 /**
  * Class to handle creation of a {@link Sort} used to sort documents by field values for queries.
@@ -195,7 +196,11 @@ private static SearchResponse.Hit.CompositeFieldValue getValueForSortField(
         break;
       case STRING:
       case STRING_VAL:
-        fieldValue.setTextValue((String) sortValue);
+        if (sortValue instanceof BytesRef) {
+          fieldValue.setTextValue(((BytesRef) sortValue).utf8ToString());
+        } else {
+          fieldValue.setTextValue((String) sortValue);
+        }
         break;
       case CUSTOM:
         // could be anything, try to determine from value class

diff --git a/src/test/java/com/yelp/nrtsearch/server/grpc/SortFieldTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/SortFieldTest.java
@@ -478,6 +478,38 @@ public void testReverseSortQueryVirtualField() {
     }
   }
 
+  @Test
+  public void testSortAtomDocId() {
+    QuerySortField querySortField =
+        QuerySortField.newBuilder()
+            .setFields(
+                SortFields.newBuilder()
+                    .addSortedFields(SortType.newBuilder().setFieldName("doc_id").build())
+                    .build())
+            .build();
+    SearchResponse searchResponse = doSortQuery(querySortField);
+    assertEquals(5, searchResponse.getHitsCount());
+    List<String> expectedIds = Arrays.asList("0", "1", "10", "11", "12");
+    assertFields(expectedIds, searchResponse.getHitsList());
+  }
+
+  @Test
+  public void testSortAtomDocIdSearchAfter() {
+    QuerySortField querySortField =
+        QuerySortField.newBuilder()
+            .setFields(
+                SortFields.newBuilder()
+                    .addSortedFields(SortType.newBuilder().setFieldName("doc_id").build())
+                    .build())
+            .build();
+    LastHitInfo searchAfter = LastHitInfo.newBuilder().addLastFieldValues("1").build();
+    SearchResponse searchResponse = dosSortQuerySearchAfter(querySortField, searchAfter);
+    assertEquals(5, searchResponse.getHitsCount());
+
+    List<String> expectedIds = Arrays.asList("1", "10", "11", "12", "13");
+    assertFields(expectedIds, searchResponse.getHitsList());
+  }
+
   @Test
   public void testSortDocId() {
     QuerySortField querySortField =
@@ -783,6 +815,35 @@ public void testSortLanLonDistanceInInnerHit() {
     }
   }
 
+  @Test
+  public void testSortIntFieldWithSearchAfter() {
+    QuerySortField querySortField =
+        QuerySortField.newBuilder()
+            .setFields(
+                SortFields.newBuilder()
+                    .addSortedFields(SortType.newBuilder().setFieldName("int_field").build())
+                    .build())
+            .build();
+    LastHitInfo searchAfter = LastHitInfo.newBuilder().addLastFieldValues("1").build();
+    SearchResponse searchResponse = dosSortQuerySearchAfter(querySortField, searchAfter);
+    assertEquals(5, searchResponse.getHitsCount());
+
+    List<String> expectedIds = Arrays.asList("91", "92", "93", "94", "95");
+    assertFields(expectedIds, searchResponse.getHitsList());
+
+    List<Integer> expectedSort = Arrays.asList(1, 2, 3, 4, 5);
+    for (int i = 0; i < searchResponse.getHitsCount(); ++i) {
+      var hit = searchResponse.getHits(i);
+      assertEquals(1, hit.getSortedFieldsCount());
+      assertEquals(
+          expectedSort.get(i).intValue(),
+          hit.getSortedFieldsOrThrow("int_field").getFieldValue(0).getIntValue());
+
+      assertEquals(0.0, hit.getScore(), 0);
+      assertEquals(6, hit.getFieldsCount());
+    }
+  }
+
   @Test
   public void testSortDocIdWithSearchAfter() {
     QuerySortField querySortField =