Split index-time and search-time sort definitions (#137880)

At search time, we can now always use our custom comparators that
allow pruning on secondary sorts. At index time, we need to always use
Lucene SortFields, which is fine as pruning isn't used at this point in any case.

Author: romseygeek

server/src/main/java/org/elasticsearch/index/IndexSortConfig.java

@@ -375,7 +375,7 @@ public Sort buildIndexSort(
             if (fieldData == null) {
                 throw new IllegalArgumentException("docvalues not found for index sort field:[" + sortSpec.field + "]");
             }
-            sortFields[i] = fieldData.sortField(this.indexCreatedVersion, sortSpec.missingValue, mode, null, reverse);
+            sortFields[i] = fieldData.indexSort(this.indexCreatedVersion, sortSpec.missingValue, mode, reverse);
             validateIndexSortField(sortFields[i]);
         }
         return new Sort(sortFields);
@@ -430,6 +430,8 @@ public static SortField.Type getSortFieldType(SortField sortField) {
             return SortField.Type.STRING;
         } else if (sortField instanceof SortedNumericSortField) {
             return ((SortedNumericSortField) sortField).getNumericType();
+        } else if (sortField.getComparatorSource() instanceof IndexFieldData.XFieldComparatorSource fcs) {
+            return fcs.sortType();
         } else {
             return sortField.getType();
         }

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

@@ -69,13 +69,20 @@ public interface IndexFieldData<FD extends LeafFieldData> {
      * Returns the {@link SortField} to use for sorting depending on the version of the index.
      */
     default SortField sortField(
+        boolean indexSort,
         IndexVersion indexCreatedVersion,
         @Nullable Object missingValue,
         MultiValueMode sortMode,
         Nested nested,
         boolean reverse
     ) {
-        return sortField(missingValue, sortMode, nested, reverse);
+        return indexSort
+            ? indexSort(indexCreatedVersion, missingValue, sortMode, reverse)
+            : sortField(missingValue, sortMode, nested, reverse);
+    }
+
+    default SortField indexSort(IndexVersion indexCreatedVersion, @Nullable Object missingValue, MultiValueMode sortMode, boolean reverse) {
+        return sortField(missingValue, sortMode, null, reverse);
     }
 
     /**
@@ -231,6 +238,10 @@ public Object missingObject(Object missingValue, boolean reversed) {
 
         public abstract SortField.Type reducedType();
 
+        public SortField.Type sortType() {
+            return reducedType();
+        }
+
         /**
          * Return a missing value that is understandable by {@link SortField#setMissingValue(Object)}.
          * Most implementations return null because they already replace the value at the fielddata level.

server/src/main/java/org/elasticsearch/index/fielddata/IndexNumericFieldData.java

@@ -86,6 +86,7 @@ public final ValuesSourceType getValuesSourceType() {
      * match the field's <code>numericType</code>.
      */
     public final SortField sortField(
+        boolean indexSort,
         NumericType targetNumericType,
         Object missingValue,
         MultiValueMode sortMode,
@@ -104,7 +105,7 @@ public final SortField sortField(
         boolean requiresCustomComparator = nested != null
             || (sortMode != MultiValueMode.MAX && sortMode != MultiValueMode.MIN)
             || targetNumericType != getNumericType();
-        if (sortRequiresCustomComparator() || requiresCustomComparator) {
+        if (indexSort == false || sortRequiresCustomComparator()) {
             SortField sortField = new SortField(getFieldName(), source, reverse);
             sortField.setOptimizeSortWithPoints(requiresCustomComparator == false && canUseOptimizedSort(indexType()));
             return sortField;
@@ -138,18 +139,24 @@ private static boolean canUseOptimizedSort(IndexType indexType) {
 
     @Override
     public final SortField sortField(Object missingValue, MultiValueMode sortMode, Nested nested, boolean reverse) {
-        return sortField(getNumericType(), missingValue, sortMode, nested, reverse);
+        return sortField(false, getNumericType(), missingValue, sortMode, nested, reverse);
+    }
+
+    @Override
+    public SortField indexSort(IndexVersion indexCreatedVersion, Object missingValue, MultiValueMode sortMode, boolean reverse) {
+        return sortField(true, indexCreatedVersion, missingValue, sortMode, null, reverse);
     }
 
     @Override
     public SortField sortField(
+        boolean indexSort,
         IndexVersion indexCreatedVersion,
         Object missingValue,
         MultiValueMode sortMode,
         Nested nested,
         boolean reverse
     ) {
-        SortField sortField = sortField(missingValue, sortMode, nested, reverse);
+        SortField sortField = sortField(indexSort, getNumericType(), missingValue, sortMode, nested, reverse);
         if (getNumericType() == NumericType.DATE_NANOSECONDS
             && indexCreatedVersion.before(IndexVersions.V_7_14_0)
             && missingValue.equals("_last")

server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java

@@ -78,9 +78,7 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
 
         final double dMissingValue = (Double) missingObject(missingValue, reversed);
-        // NOTE: it's important to pass null as a missing value in the constructor so that
-        // the comparator doesn't check docsWithField since we replace missing values in select()
-        return new DoubleComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        return new DoubleComparator(numHits, fieldname, dMissingValue, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new DoubleLeafComparator(context) {

server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java

@@ -73,7 +73,7 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         final float fMissingValue = (Float) missingObject(missingValue, reversed);
         // NOTE: it's important to pass null as a missing value in the constructor so that
         // the comparator doesn't check docsWithField since we replace missing values in select()
-        return new FloatComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        return new FloatComparator(numHits, fieldname, fMissingValue, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new FloatLeafComparator(context) {

server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/HalfFloatValuesComparatorSource.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.LeafFieldComparator;
 import org.apache.lucene.search.Pruning;
+import org.apache.lucene.search.SortField;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.fielddata.DenseDoubleValues;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
@@ -33,14 +34,17 @@ public HalfFloatValuesComparatorSource(
         super(indexFieldData, missingValue, sortMode, nested);
     }
 
+    @Override
+    public SortField.Type sortType() {
+        return SortField.Type.CUSTOM;
+    }
+
     @Override
     public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning enableSkipping, boolean reversed) {
         assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
 
         final float fMissingValue = (Float) missingObject(missingValue, reversed);
-        // NOTE: it's important to pass null as a missing value in the constructor so that
-        // the comparator doesn't check docsWithField since we replace missing values in select()
-        return new HalfFloatComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        return new HalfFloatComparator(numHits, fieldname, fMissingValue, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new HalfFloatLeafComparator(context) {

server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/IntValuesComparatorSource.java

@@ -52,9 +52,7 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
 
         final int iMissingValue = (Integer) missingObject(missingValue, reversed);
-        // NOTE: it's important to pass null as a missing value in the constructor so that
-        // the comparator doesn't check docsWithField since we replace missing values in select()
-        return new IntComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        return new IntComparator(numHits, fieldname, iMissingValue, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new IntLeafComparator(context) {

server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/LongValuesComparatorSource.java

@@ -105,9 +105,7 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         assert indexFieldData == null || fieldname.equals(indexFieldData.getFieldName());
 
         final long lMissingValue = (Long) missingObject(missingValue, reversed);
-        // NOTE: it's important to pass null as a missing value in the constructor so that
-        // the comparator doesn't check docsWithField since we replace missing values in select()
-        return new XLongComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        return new XLongComparator(numHits, fieldname, lMissingValue, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 final int maxDoc = context.reader().maxDoc();

server/src/main/java/org/elasticsearch/lucene/comparators/XUpdateableDocIdSetIterator.java

@@ -69,4 +69,8 @@ public int docIDRunEnd() throws IOException {
             return super.docIDRunEnd();
         }
     }
+
+    public DocIdSetIterator getDelegate() {
+        return in;
+    }
 }

server/src/main/java/org/elasticsearch/search/sort/FieldSortBuilder.java

@@ -355,10 +355,10 @@ public SortFieldAndFormat build(SearchExecutionContext context) throws IOExcepti
             }
             IndexNumericFieldData numericFieldData = (IndexNumericFieldData) fieldData;
             NumericType resolvedType = resolveNumericType(numericType);
-            field = numericFieldData.sortField(resolvedType, missing, localSortMode(), nested, reverse);
+            field = numericFieldData.sortField(false, resolvedType, missing, localSortMode(), nested, reverse);
             isNanosecond = resolvedType == NumericType.DATE_NANOSECONDS;
         } else {
-            field = fieldData.sortField(context.indexVersionCreated(), missing, localSortMode(), nested, reverse);
+            field = fieldData.sortField(false, context.indexVersionCreated(), missing, localSortMode(), nested, reverse);
             if (fieldData instanceof IndexNumericFieldData) {
                 isNanosecond = ((IndexNumericFieldData) fieldData).getNumericType() == NumericType.DATE_NANOSECONDS;
             }