Replace custom sort field with SortedSetSortField and SortedNumericSortField when possible (#23827)

Currently for field sorting we always use a custom sort field and a custom comparator source.
Though for numeric fields this custom sort field could be replaced with a standard SortedNumericSortField unless
the field is nested especially since we removed the FieldData for numerics.
We can also use a SortedSetSortField for string sort based on doc_values when the field is not nested.

This change replaces IndexFieldData#comparatorSource with IndexFieldData#sortField that returns a Sorted{Set,Numeric}SortField when possible or a custom
 sort field when the field sort spec is not handled by the SortedSortFields.

Author: jimczi

core/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -57,6 +57,8 @@
 import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.search.grouping.CollapseTopFieldDocs;
+import org.apache.lucene.search.SortedNumericSortField;
+import org.apache.lucene.search.SortedSetSortField;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
@@ -552,7 +554,22 @@ public static void writeSortField(StreamOutput out, SortField sortField) throws
             SortField newSortField = new SortField(sortField.getField(), SortField.Type.DOUBLE);
             newSortField.setMissingValue(sortField.getMissingValue());
             sortField = newSortField;
+        } else if (sortField.getClass() == SortedSetSortField.class) {
+            // for multi-valued sort field, we replace the SortedSetSortField with a simple SortField.
+            // It works because the sort field is only used to merge results from different shards.
+            SortField newSortField = new SortField(sortField.getField(), SortField.Type.STRING, sortField.getReverse());
+            newSortField.setMissingValue(sortField.getMissingValue());
+            sortField = newSortField;
+        } else if (sortField.getClass() == SortedNumericSortField.class) {
+            // for multi-valued sort field, we replace the SortedSetSortField with a simple SortField.
+            // It works because the sort field is only used to merge results from different shards.
+            SortField newSortField = new SortField(sortField.getField(),
+                ((SortedNumericSortField) sortField).getNumericType(),
+                sortField.getReverse());
+            newSortField.setMissingValue(sortField.getMissingValue());
+            sortField = newSortField;
         }
+
         if (sortField.getClass() != SortField.class) {
             throw new IllegalArgumentException("Cannot serialize SortField impl [" + sortField + "]");
         }

core/src/main/java/org/elasticsearch/index/fielddata/IndexFieldData.java

@@ -85,9 +85,9 @@ public static MemoryStorageFormat fromString(String string) {
     FD loadDirect(LeafReaderContext context) throws Exception;
 
     /**
-     * Comparator used for sorting.
+     * Returns the {@link SortField} to used for sorting.
      */
-    XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested);
+    SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse);
 
     /**
      * Clears any resources associated with this field data.
@@ -136,17 +136,17 @@ public DocIdSetIterator innerDocs(LeafReaderContext ctx) throws IOException {
         }
 
         /** Whether missing values should be sorted first. */
-        protected final boolean sortMissingFirst(Object missingValue) {
+        public final boolean sortMissingFirst(Object missingValue) {
             return "_first".equals(missingValue);
         }
 
         /** Whether missing values should be sorted last, this is the default. */
-        protected final boolean sortMissingLast(Object missingValue) {
+        public final boolean sortMissingLast(Object missingValue) {
             return missingValue == null || "_last".equals(missingValue);
         }
 
         /** Return the missing object value according to the reduced type of the comparator. */
-        protected final Object missingObject(Object missingValue, boolean reversed) {
+        public final Object missingObject(Object missingValue, boolean reversed) {
             if (sortMissingFirst(missingValue) || sortMissingLast(missingValue)) {
                 final boolean min = sortMissingFirst(missingValue) ^ reversed;
                 switch (reducedType()) {

core/src/main/java/org/elasticsearch/index/fielddata/ordinals/GlobalOrdinalsIndexFieldData.java

@@ -20,6 +20,7 @@
 
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
 import org.apache.lucene.util.Accountable;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.AbstractIndexComponent;
@@ -68,7 +69,7 @@ public String getFieldName() {
     }
 
     @Override
-    public XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested) {
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse) {
         throw new UnsupportedOperationException("no global ordinals sorting yet");
     }
 

core/src/main/java/org/elasticsearch/index/fielddata/plain/AbstractGeoPointDVIndexFieldData.java

@@ -21,6 +21,7 @@
 
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
@@ -43,7 +44,7 @@ public abstract class AbstractGeoPointDVIndexFieldData extends DocValuesIndexFie
     }
 
     @Override
-    public final XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested) {
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse) {
         throw new IllegalArgumentException("can't sort on geo_point field without using specific sorting feature, like geo_distance");
     }
 

core/src/main/java/org/elasticsearch/index/fielddata/plain/AbstractIndexGeoPointFieldData.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.index.fielddata.plain;
 
+import org.apache.lucene.search.SortField;
 import org.apache.lucene.spatial.geopoint.document.GeoPointField;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefIterator;
@@ -28,6 +29,7 @@
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.fielddata.AtomicGeoPointFieldData;
+import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexFieldData.XFieldComparatorSource.Nested;
 import org.elasticsearch.index.fielddata.IndexFieldDataCache;
 import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
@@ -104,7 +106,7 @@ public GeoPoint next() throws IOException {
     }
 
     @Override
-    public final XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested) {
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse) {
         throw new IllegalArgumentException("can't sort on geo_point field without using specific sorting feature, like geo_distance");
     }
 

core/src/main/java/org/elasticsearch/index/fielddata/plain/AbstractIndexOrdinalsFieldData.java

@@ -55,11 +55,6 @@ protected AbstractIndexOrdinalsFieldData(IndexSettings indexSettings, String fie
         this.minSegmentSize = minSegmentSize;
     }
 
-    @Override
-    public XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested) {
-        return new BytesRefFieldComparatorSource(this, missingValue, sortMode, nested);
-    }
-
     @Override
     public IndexOrdinalsFieldData loadGlobal(DirectoryReader indexReader) {
         if (indexReader.leaves().size() <= 1) {

core/src/main/java/org/elasticsearch/index/fielddata/plain/AbstractLatLonPointDVIndexFieldData.java

@@ -24,6 +24,7 @@
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.Index;
@@ -46,8 +47,7 @@ public abstract class AbstractLatLonPointDVIndexFieldData extends DocValuesIndex
     }
 
     @Override
-    public final XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode,
-                                                         XFieldComparatorSource.Nested nested) {
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, XFieldComparatorSource.Nested nested, boolean reverse) {
         throw new IllegalArgumentException("can't sort on geo_point field without using specific sorting feature, like geo_distance");
     }
 

core/src/main/java/org/elasticsearch/index/fielddata/plain/BinaryDVIndexFieldData.java

@@ -20,9 +20,12 @@
 package org.elasticsearch.index.fielddata.plain;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
+import org.apache.lucene.search.SortedSetSortField;
+import org.apache.lucene.search.SortedSetSelector;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.fielddata.IndexFieldData;
-import org.elasticsearch.index.fielddata.IndexFieldData.XFieldComparatorSource.Nested;
 import org.elasticsearch.index.fielddata.fieldcomparator.BytesRefFieldComparatorSource;
 import org.elasticsearch.search.MultiValueMode;
 
@@ -43,7 +46,21 @@ public BinaryDVAtomicFieldData loadDirect(LeafReaderContext context) throws Exce
     }
 
     @Override
-    public org.elasticsearch.index.fielddata.IndexFieldData.XFieldComparatorSource comparatorSource(Object missingValue, MultiValueMode sortMode, Nested nested) {
-        return new BytesRefFieldComparatorSource(this, missingValue, sortMode, nested);
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, XFieldComparatorSource.Nested nested, boolean reverse) {
+        XFieldComparatorSource source = new BytesRefFieldComparatorSource(this, missingValue, sortMode, nested);
+        /**
+         * Check if we can use a simple {@link SortedSetSortField} compatible with index sorting and
+         * returns a custom sort field otherwise.
+         */
+        if (nested != null ||
+                (sortMode != MultiValueMode.MAX && sortMode != MultiValueMode.MIN) ||
+                (source.sortMissingFirst(missingValue) == false && source.sortMissingLast(missingValue) == false)) {
+            return new SortField(getFieldName(), source, reverse);
+        }
+        SortField sortField = new SortedSetSortField(fieldName, reverse,
+            sortMode == MultiValueMode.MAX ? SortedSetSelector.Type.MAX : SortedSetSelector.Type.MIN);
+        sortField.setMissingValue(source.sortMissingLast(missingValue) ^ reverse ?
+            SortedSetSortField.STRING_LAST : SortedSetSortField.STRING_FIRST);
+        return sortField;
     }
 }

core/src/main/java/org/elasticsearch/index/fielddata/plain/BytesBinaryDVIndexFieldData.java

@@ -21,6 +21,7 @@
 
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
@@ -41,7 +42,7 @@ public BytesBinaryDVIndexFieldData(Index index, String fieldName) {
     }
 
     @Override
-    public final XFieldComparatorSource comparatorSource(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested) {
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse) {
         throw new IllegalArgumentException("can't sort on binary field");
     }
 

core/src/main/java/org/elasticsearch/index/fielddata/plain/IndexIndexFieldData.java

@@ -24,17 +24,21 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.RandomAccessOrds;
 import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.search.SortField;
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.fielddata.AtomicOrdinalsFieldData;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexFieldDataCache;
 import org.elasticsearch.index.fielddata.IndexOrdinalsFieldData;
+import org.elasticsearch.index.fielddata.fieldcomparator.BytesRefFieldComparatorSource;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.TextFieldMapper;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.search.MultiValueMode;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -124,6 +128,12 @@ public AtomicOrdinalsFieldData loadDirect(LeafReaderContext context)
         return atomicFieldData;
     }
 
+    @Override
+    public SortField sortField(@Nullable Object missingValue, MultiValueMode sortMode, XFieldComparatorSource.Nested nested, boolean reverse) {
+        final XFieldComparatorSource source = new BytesRefFieldComparatorSource(this, missingValue, sortMode, nested);
+        return new SortField(getFieldName(), source, reverse);
+    }
+
     @Override
     public IndexOrdinalsFieldData loadGlobal(DirectoryReader indexReader) {
         return this;