Fix some issues

Signed-off-by: Binlong Gao <[email protected]>

Author: gaobinlong

server/src/main/java/org/apache/lucene/search/grouping/CollapsingTopDocsCollector.java

@@ -85,7 +85,10 @@ public final class CollapsingTopDocsCollector<T> extends FirstPassGroupingCollec
         this.after = after;
 
         if (after != null) {
-            // we have only one sort field which is the collapse field
+            // we should have only one sort field which is the collapse field
+            if (sort.getSort().length != 1 || !sort.getSort()[0].getField().equals(collapseField)) {
+                throw new IllegalArgumentException("The after parameter can only be used when the sort is based on the collapse field");
+            }
             SortField field = sort.getSort()[0];
             afterComparator = field.getComparator(1, Pruning.NONE);
 

server/src/main/java/org/opensearch/search/collapse/CollapseContext.java

@@ -109,7 +109,7 @@ public CollapsingTopDocsCollector<?> createTopDocs(Sort sort, int topN, FieldDoc
         } else if (fieldType != null && fieldType.unwrap() instanceof NumberFieldMapper.NumberFieldType) {
             return CollapsingTopDocsCollector.createNumeric(fieldName, fieldType, sort, topN, searchAfter);
         } else {
-            throw new IllegalStateException("unknown type for collapse field " + fieldName + ", only keywords and numbers are accepted");
+            throw new IllegalStateException("unsupported type for collapse field " + fieldName + ", only keywords and numbers are accepted");
         }
     }
 }

server/src/test/java/org/opensearch/lucene/grouping/CollapsingTopDocsCollectorTests.java

@@ -494,6 +494,47 @@ public void testInconsistentShardIndicesException() {
         assertEquals("Inconsistent order of shard indices", exception.getMessage());
     }
 
+    public void testSearchAfterValidation() {
+        MappedFieldType fieldType = new MockFieldMapper.FakeFieldType("category");
+
+        // Test multiple sort fields - should fail
+        Sort multiSort = new Sort(new SortField("category", SortField.Type.INT), new SortField("score", SortField.Type.FLOAT));
+        FieldDoc multiAfter = new FieldDoc(0, Float.NaN, new Object[] { 1, 1.0f });
+        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> {
+            CollapsingTopDocsCollector.createNumeric("category", fieldType, multiSort, 10, multiAfter);
+        });
+        assertEquals("The after parameter can only be used when the sort is based on the collapse field", exception.getMessage());
+
+        // Test wrong sort field - should fail
+        Sort wrongSort = new Sort(new SortField("different_field", SortField.Type.INT));
+        FieldDoc wrongAfter = new FieldDoc(0, Float.NaN, new Object[] { 1 });
+        exception = expectThrows(IllegalArgumentException.class, () -> {
+            CollapsingTopDocsCollector.createNumeric("category", fieldType, wrongSort, 10, wrongAfter);
+        });
+        assertEquals("The after parameter can only be used when the sort is based on the collapse field", exception.getMessage());
+
+        // Test correct sort field - should succeed
+        Sort correctSort = new Sort(new SortField("category", SortField.Type.INT));
+        FieldDoc correctAfter = new FieldDoc(0, Float.NaN, new Object[] { 1 });
+        CollapsingTopDocsCollector<?> collector = CollapsingTopDocsCollector.createNumeric(
+            "category",
+            fieldType,
+            correctSort,
+            10,
+            correctAfter
+        );
+        assertNotNull(collector);
+
+        // Test keyword field with multiple sorts - should fail
+        MappedFieldType keywordFieldType = new MockFieldMapper.FakeFieldType("tag");
+        Sort keywordMultiSort = new Sort(new SortField("tag", SortField.Type.STRING), new SortField("score", SortField.Type.FLOAT));
+        FieldDoc keywordAfter = new FieldDoc(0, Float.NaN, new Object[] { "A", 1.0f });
+        exception = expectThrows(IllegalArgumentException.class, () -> {
+            CollapsingTopDocsCollector.createKeyword("tag", keywordFieldType, keywordMultiSort, 10, keywordAfter);
+        });
+        assertEquals("The after parameter can only be used when the sort is based on the collapse field", exception.getMessage());
+    }
+
     public void testSearchAfterWithNumericCollapse() throws IOException {
         testSearchAfterCollapse(new NumericDVProducer(), true);
     }
@@ -509,20 +550,19 @@ private <T extends Comparable<T>> void testSearchAfterCollapse(CollapsingDocValu
         final Directory dir = newDirectory();
         final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
 
-        // Create test documents
         for (int i = 0; i < numDocs; i++) {
             Document doc = new Document();
             T groupValue = dvProducer.randomGroup(maxGroup);
             dvProducer.add(doc, groupValue, false);
-            doc.add(new NumericDocValuesField("sort1", i)); // Ascending sort values
             w.addDocument(doc);
         }
 
         final IndexReader reader = w.getReader();
         final IndexSearcher searcher = newSearcher(reader);
 
         SortField collapseField = dvProducer.sortField(false);
-        Sort sort = new Sort(new SortField("sort1", SortField.Type.INT), collapseField);
+        // Use collapse field as sort field to satisfy validation
+        Sort sort = new Sort(collapseField);
         MappedFieldType fieldType = new MockFieldMapper.FakeFieldType(collapseField.getField());
 
         // First search without search_after
@@ -546,13 +586,6 @@ private <T extends Comparable<T>> void testSearchAfterCollapse(CollapsingDocValu
         searcher.search(new MatchAllDocsQuery(), collector2);
         CollapseTopFieldDocs results2 = collector2.getTopDocs();
 
-        // Verify search_after works correctly
-        if (results2.scoreDocs.length > 0) {
-            FieldDoc firstAfterDoc = (FieldDoc) results2.scoreDocs[0];
-            // First doc in second page should be after the last doc in first page
-            assertTrue("Search after should skip previous results", compareFieldDocs(after, firstAfterDoc, sort) < 0);
-        }
-
         // Verify no overlap between pages
         Set<Integer> firstPageDocs = new HashSet<>();
         for (ScoreDoc doc : results1.scoreDocs) {
@@ -572,16 +605,18 @@ public void testSearchAfterWithEmptyResults() throws IOException {
         final Directory dir = newDirectory();
         final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
 
-        // Add a single document
-        Document doc = new Document();
-        doc.add(new NumericDocValuesField("group", 1));
-        doc.add(new NumericDocValuesField("sort1", 100));
-        w.addDocument(doc);
+        final int numDocs = 100;
+        for (int i = 0; i < numDocs; i++) {
+            Document doc = new Document();
+            doc.add(new NumericDocValuesField("group", i));
+            w.addDocument(doc);
+        }
 
         final IndexReader reader = w.getReader();
         final IndexSearcher searcher = newSearcher(reader);
 
-        Sort sort = new Sort(new SortField("sort1", SortField.Type.INT));
+        // Use collapse field as sort field to satisfy validation
+        Sort sort = new Sort(new SortField("group", SortField.Type.INT));
         MappedFieldType fieldType = new MockFieldMapper.FakeFieldType("group");
 
         // Create search_after that's beyond all documents
@@ -593,28 +628,13 @@ public void testSearchAfterWithEmptyResults() throws IOException {
         CollapseTopFieldDocs results = collector.getTopDocs();
 
         assertEquals("Should have no results after last document", 0, results.scoreDocs.length);
-        assertEquals("Total hits should reflect all documents processed with search_after", 1, results.totalHits.value());
+        assertEquals("Total hits should reflect all documents processed with search_after", numDocs, results.totalHits.value());
 
         w.close();
         reader.close();
         dir.close();
     }
 
-    private int compareFieldDocs(FieldDoc doc1, FieldDoc doc2, Sort sort) {
-        SortField[] fields = sort.getSort();
-        for (int i = 0; i < fields.length && i < doc1.fields.length && i < doc2.fields.length; i++) {
-            @SuppressWarnings("unchecked")
-            Comparable<Object> val1 = (Comparable<Object>) doc1.fields[i];
-            Object val2 = doc2.fields[i];
-
-            int cmp = val1.compareTo(val2);
-            if (cmp != 0) {
-                return fields[i].getReverse() ? -cmp : cmp;
-            }
-        }
-        return Integer.compare(doc1.doc, doc2.doc);
-    }
-
     // Helper classes for test data
     private static class NumericDVProducer implements CollapsingDocValuesProducer<Long> {
         @Override