Fix search_after with geo distance sorting (#26891)

Support for search_after and geo distance sorting is broken when the optimized LatLonDocValuesField.distanceSort is used.
This commit fixes the parsing of the search_after value for this case.

Author: jimczi

core/src/main/java/org/elasticsearch/search/searchafter/SearchAfterBuilder.java

@@ -30,7 +30,6 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.text.Text;
-import org.elasticsearch.common.xcontent.ToXContent.Params;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
@@ -131,21 +130,25 @@ public static FieldDoc buildFieldDoc(SortAndFormats sort, Object[] values) {
         return new FieldDoc(Integer.MAX_VALUE, 0, fieldValues);
     }
 
-    private static SortField.Type extractSortType(SortField sortField) {
-        if (sortField instanceof SortedSetSortField) {
+    /**
+     * Returns the inner {@link SortField.Type} expected for this sort field.
+     */
+    static SortField.Type extractSortType(SortField sortField) {
+        if (sortField.getComparatorSource() instanceof IndexFieldData.XFieldComparatorSource) {
+            return ((IndexFieldData.XFieldComparatorSource) sortField.getComparatorSource()).reducedType();
+        } else if (sortField instanceof SortedSetSortField) {
             return SortField.Type.STRING;
         } else if (sortField instanceof SortedNumericSortField) {
             return ((SortedNumericSortField) sortField).getNumericType();
+        } else if ("LatLonPointSortField".equals(sortField.getClass().getSimpleName())) {
+            // for geo distance sorting
+            return SortField.Type.DOUBLE;
         } else {
             return sortField.getType();
         }
     }
 
-    private static Object convertValueFromSortField(Object value, SortField sortField, DocValueFormat format) {
-        if (sortField.getComparatorSource() instanceof IndexFieldData.XFieldComparatorSource) {
-            IndexFieldData.XFieldComparatorSource cmpSource = (IndexFieldData.XFieldComparatorSource) sortField.getComparatorSource();
-            return convertValueFromSortType(sortField.getField(), cmpSource.reducedType(), value, format);
-        }
+    static Object convertValueFromSortField(Object value, SortField sortField, DocValueFormat format) {
         SortField.Type sortType = extractSortType(sortField);
         return convertValueFromSortType(sortField.getField(), sortType, value, format);
     }

core/src/test/java/org/elasticsearch/search/searchafter/SearchAfterBuilderTests.java

@@ -19,6 +19,11 @@
 
 package org.elasticsearch.search.searchafter;
 
+import org.apache.lucene.document.LatLonDocValuesField;
+import org.apache.lucene.search.FieldComparator;
+import org.apache.lucene.search.SortField;
+import org.apache.lucene.search.SortedNumericSortField;
+import org.apache.lucene.search.SortedSetSortField;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.text.Text;
@@ -27,13 +32,16 @@
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
+import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.search.MultiValueMode;
 import org.elasticsearch.test.ESTestCase;
-import org.hamcrest.Matchers;
 
 import java.io.IOException;
 import java.util.Collections;
 
+import static org.elasticsearch.search.searchafter.SearchAfterBuilder.extractSortType;
 import static org.elasticsearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;
+import static org.hamcrest.Matchers.equalTo;
 
 public class SearchAfterBuilderTests extends ESTestCase {
     private static final int NUMBER_OF_TESTBUILDERS = 20;
@@ -182,7 +190,7 @@ public void testWithNullArray() throws Exception {
             builder.setSortValues(null);
             fail("Should fail on null array.");
         } catch (NullPointerException e) {
-            assertThat(e.getMessage(), Matchers.equalTo("Values cannot be null."));
+            assertThat(e.getMessage(), equalTo("Values cannot be null."));
         }
     }
 
@@ -192,7 +200,7 @@ public void testWithEmptyArray() throws Exception {
             builder.setSortValues(new Object[0]);
             fail("Should fail on empty array.");
         } catch (IllegalArgumentException e) {
-            assertThat(e.getMessage(), Matchers.equalTo("Values must contains at least one value."));
+            assertThat(e.getMessage(), equalTo("Values must contains at least one value."));
         }
     }
 
@@ -215,4 +223,29 @@ private static void randomSearchFromBuilderWithSortValueThrows(Object containing
         Exception e = expectThrows(IllegalArgumentException.class, () -> builder.setSortValues(values));
         assertEquals(e.getMessage(), "Can't handle search_after field value of type [" + containing.getClass() + "]");
     }
+
+    public void testExtractSortType() throws Exception {
+        SortField.Type type = extractSortType(LatLonDocValuesField.newDistanceSort("field", 0.0, 180.0));
+        assertThat(type, equalTo(SortField.Type.DOUBLE));
+        IndexFieldData.XFieldComparatorSource source = new IndexFieldData.XFieldComparatorSource(null, MultiValueMode.MIN, null) {
+            @Override
+            public SortField.Type reducedType() {
+                return SortField.Type.STRING;
+            }
+
+            @Override
+            public FieldComparator<?> newComparator(String fieldname, int numHits, int sortPos, boolean reversed) {
+                return null;
+            }
+        };
+
+        type = extractSortType(new SortField("field", source));
+        assertThat(type, equalTo(SortField.Type.STRING));
+
+        type = extractSortType(new SortedNumericSortField("field", SortField.Type.DOUBLE));
+        assertThat(type, equalTo(SortField.Type.DOUBLE));
+
+        type = extractSortType(new SortedSetSortField("field", false));
+        assertThat(type, equalTo(SortField.Type.STRING));
+    }
 }