Speed up sorts on secondary sort fields

docs/changelog/137533.yaml

@@ -0,0 +1,5 @@
+pr: 137533
+summary: Speed up sorts on secondary sort fields
+area: Search
+type: enhancement
+issues: []

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

@@ -634,12 +634,8 @@ private static SortField rewriteMergeSortField(SortField sortField) {
             SortField newSortField = new SortField(sortField.getField(), SortField.Type.STRING, sortField.getReverse());
             newSortField.setMissingValue(sortField.getMissingValue());
             return newSortField;
-        } else if (sortField.getClass() == SortedNumericSortField.class) {
-            SortField newSortField = new SortField(
-                sortField.getField(),
-                ((SortedNumericSortField) sortField).getNumericType(),
-                sortField.getReverse()
-            );
+        } else if (sortField instanceof SortedNumericSortField snsf) {
+            SortField newSortField = new SortField(sortField.getField(), snsf.getNumericType(), sortField.getReverse());
             newSortField.setMissingValue(sortField.getMissingValue());
             return newSortField;
         } else if (sortField.getClass() == ShardDocSortField.class) {
@@ -651,9 +647,6 @@ private static SortField rewriteMergeSortField(SortField sortField) {
 
     static void writeSortField(StreamOutput out, SortField sortField) throws IOException {
         sortField = rewriteMergeSortField(sortField);
-        if (sortField.getClass() != SortField.class) {
-            throw new IllegalArgumentException("Cannot serialize SortField impl [" + sortField + "]");
-        }
         out.writeOptionalString(sortField.getField());
         if (sortField.getComparatorSource() != null) {
             IndexFieldData.XFieldComparatorSource comparatorSource = (IndexFieldData.XFieldComparatorSource) sortField

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

@@ -80,7 +80,8 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         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, null, null, reversed, Pruning.NONE) {
+        // TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
+        return new DoubleComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
             @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,8 @@ 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, null, null, reversed, Pruning.NONE) {
+        // TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
+        return new FloatComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new FloatLeafComparator(context) {

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

@@ -39,7 +39,8 @@ 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 HalfFloatComparator(numHits, fieldname, null, reversed, enableSkipping) {
+        // TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
+        return new HalfFloatComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
                 return new HalfFloatLeafComparator(context) {

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

@@ -60,7 +60,7 @@ public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws I
                 return new IntLeafComparator(context) {
                     @Override
                     protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
-                        return wrap(getLongValues(context, iMissingValue));
+                        return wrap(getLongValues(context, iMissingValue), context.reader().maxDoc());
                     }
                 };
             }

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

@@ -8,15 +8,16 @@
  */
 package org.elasticsearch.index.fielddata.fieldcomparator;
 
+import org.apache.lucene.index.DocValuesSkipper;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.LeafFieldComparator;
 import org.apache.lucene.search.LongValues;
 import org.apache.lucene.search.Pruning;
+import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
-import org.apache.lucene.search.comparators.LongComparator;
 import org.apache.lucene.util.BitSet;
 import org.elasticsearch.common.time.DateUtils;
 import org.elasticsearch.common.util.BigArrays;
@@ -28,6 +29,8 @@
 import org.elasticsearch.index.fielddata.LeafNumericFieldData;
 import org.elasticsearch.index.fielddata.SortedNumericLongValues;
 import org.elasticsearch.index.fielddata.plain.SortedNumericIndexFieldData;
+import org.elasticsearch.lucene.comparators.XLongComparator;
+import org.elasticsearch.lucene.comparators.XNumericComparator;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.MultiValueMode;
 import org.elasticsearch.search.sort.BucketedSort;
@@ -103,13 +106,48 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
         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 LongComparator(numHits, null, null, reversed, Pruning.NONE) {
+        return new XLongComparator(numHits, fieldname, null, reversed, enableSkipping) {
             @Override
             public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
+                final int maxDoc = context.reader().maxDoc();
                 return new LongLeafComparator(context) {
                     @Override
                     protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
-                        return wrap(getLongValues(context, lMissingValue));
+                        return wrap(getLongValues(context, lMissingValue), maxDoc);
+                    }
+
+                    @Override
+                    protected XNumericComparator<Long>.CompetitiveDISIBuilder buildCompetitiveDISIBuilder(LeafReaderContext context)
+                        throws IOException {
+                        Sort indexSort = context.reader().getMetaData().sort();
+                        if (indexSort == null) {
+                            return super.buildCompetitiveDISIBuilder(context);
+                        }
+                        SortField[] sortFields = indexSort.getSort();
+                        if (sortFields.length != 2) {
+                            return super.buildCompetitiveDISIBuilder(context);
+                        }
+                        if (sortFields[1].getField().equals(field) == false) {
+                            return super.buildCompetitiveDISIBuilder(context);
+                        }
+                        DocValuesSkipper skipper = context.reader().getDocValuesSkipper(field);
+                        DocValuesSkipper primaryFieldSkipper = context.reader().getDocValuesSkipper(sortFields[0].getField());
+                        if (primaryFieldSkipper == null || skipper.docCount() != maxDoc || primaryFieldSkipper.docCount() != maxDoc) {
+                            return super.buildCompetitiveDISIBuilder(context);
+                        }
+                        return new CompetitiveDISIBuilder(this) {
+                            @Override
+                            protected int docCount() {
+                                return skipper.docCount();
+                            }
+
+                            @Override
+                            protected void doUpdateCompetitiveIterator() {
+                                competitiveIterator.update(
+                                    new SecondarySortIterator(docValues, skipper, primaryFieldSkipper, minValueAsLong, maxValueAsLong)
+                                );
+                            }
+                        };
                     }
                 };
             }
@@ -163,31 +201,43 @@ public Object missingObject(Object missingValue, boolean reversed) {
         return super.missingObject(missingValue, reversed);
     }
 
-    protected static NumericDocValues wrap(LongValues longValues) {
+    protected static NumericDocValues wrap(LongValues longValues, int maxDoc) {
         return new NumericDocValues() {
+
+            int doc = -1;
+
             @Override
             public long longValue() throws IOException {
                 return longValues.longValue();
             }
 
             @Override
             public boolean advanceExact(int target) throws IOException {
+                doc = target;
                 return longValues.advanceExact(target);
             }
 
             @Override
             public int docID() {
-                throw new UnsupportedOperationException();
+                return doc;
             }
 
             @Override
             public int nextDoc() throws IOException {
-                throw new UnsupportedOperationException();
+                return advance(doc + 1);
             }
 
             @Override
             public int advance(int target) throws IOException {
-                throw new UnsupportedOperationException();
+                if (target >= maxDoc) {
+                    return doc = NO_MORE_DOCS;
+                }
+                // All documents are guaranteed to have a value, as all invocations of getLongValues
+                // always return `true` from `advanceExact()`
+                boolean hasValue = longValues.advanceExact(target);
+                assert hasValue : "LongValuesComparatorSource#wrap called with a LongValues that has missing values";
+                doc = target;
+                return target;
             }
 
             @Override

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

@@ -0,0 +1,168 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.fielddata.fieldcomparator;
+
+import org.apache.lucene.index.DocValuesSkipper;
+import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.search.DocIdSetIterator;
+
+import java.io.IOException;
+
+/**
+ * A competitive DocIdSetIterator that examines the values of a secondary
+ * sort field and tries to exclude documents with values outside a given
+ * range, using DocValueSkippers on the primary sort field to advance rapidly
+ * to the next block of values.
+ */
+class SecondarySortIterator extends DocIdSetIterator {
+
+    final NumericDocValues values;
+
+    final DocValuesSkipper valueSkipper;
+    final DocValuesSkipper primaryFieldSkipper;
+    final long minValue;
+    final long maxValue;
+
+    int docID = -1;
+    boolean skipperMatch;
+    int primaryFieldUpTo = -1;
+    int valueFieldUpTo = -1;
+
+    SecondarySortIterator(
+        NumericDocValues values,
+        DocValuesSkipper valueSkipper,
+        DocValuesSkipper primaryFieldSkipper,
+        long minValue,
+        long maxValue
+    ) {
+        this.values = values;
+        this.valueSkipper = valueSkipper;
+        this.primaryFieldSkipper = primaryFieldSkipper;
+        this.minValue = minValue;
+        this.maxValue = maxValue;
+
+        valueFieldUpTo = valueSkipper.maxDocID(0);
+        primaryFieldUpTo = primaryFieldSkipper.maxDocID(0);
+    }
+
+    @Override
+    public int docID() {
+        return docID;
+    }
+
+    @Override
+    public int nextDoc() throws IOException {
+        return advance(docID + 1);
+    }
+
+    @Override
+    public int advance(int target) throws IOException {
+        skipperMatch = false;
+        target = values.advance(target);
+        if (target == DocIdSetIterator.NO_MORE_DOCS) {
+            return docID = target;
+        }
+        while (true) {
+            if (target > valueFieldUpTo) {
+                valueSkipper.advance(target);
+                valueFieldUpTo = valueSkipper.maxDocID(0);
+                long minValue = valueSkipper.minValue(0);
+                long maxValue = valueSkipper.maxValue(0);
+                if (minValue > this.maxValue || maxValue < this.minValue) {
+                    // outside the desired range, skip forward
+                    for (int level = 1; level < valueSkipper.numLevels(); level++) {
+                        minValue = valueSkipper.minValue(level);
+                        maxValue = valueSkipper.maxValue(level);
+                        if (minValue > this.maxValue || maxValue < this.minValue) {
+                            valueFieldUpTo = valueSkipper.maxDocID(level);
+                        } else {
+                            break;
+                        }
+                    }
+
+                    int upTo = valueFieldUpTo;
+                    if (maxValue < this.minValue) {
+                        // We've moved past the end of the valid values in the secondary sort field
+                        // for this primary value. Advance the primary skipper to find the starting point
+                        // for the next primary value, where the secondary field values will have reset
+                        primaryFieldSkipper.advance(target);
+                        primaryFieldUpTo = primaryFieldSkipper.maxDocID(0);
+                        if (primaryFieldSkipper.minValue(0) == primaryFieldSkipper.maxValue(0)) {
+                            for (int level = 1; level < primaryFieldSkipper.numLevels(); level++) {
+                                if (primaryFieldSkipper.minValue(level) == primaryFieldSkipper.maxValue(level)) {
+                                    primaryFieldUpTo = primaryFieldSkipper.maxDocID(level);
+                                } else {
+                                    break;
+                                }
+                            }
+                        }
+                        if (primaryFieldUpTo > upTo) {
+                            upTo = primaryFieldUpTo;
+                        }
+                    }
+
+                    target = values.advance(upTo + 1);
+                    if (target == DocIdSetIterator.NO_MORE_DOCS) {
+                        return docID = target;
+                    }
+                } else if (minValue >= this.minValue && maxValue <= this.maxValue) {
+                    assert valueSkipper.docCount(0) == valueSkipper.maxDocID(0) - valueSkipper.minDocID(0) + 1;
+                    skipperMatch = true;
+                    return docID = target;
+                }
+            }
+
+            long value = values.longValue();
+            if (value < minValue && target > primaryFieldUpTo) {
+                primaryFieldSkipper.advance(target);
+                primaryFieldUpTo = primaryFieldSkipper.maxDocID(0);
+                if (primaryFieldSkipper.minValue(0) == primaryFieldSkipper.maxValue(0)) {
+                    for (int level = 1; level < primaryFieldSkipper.numLevels(); level++) {
+                        if (primaryFieldSkipper.minValue(level) == primaryFieldSkipper.maxValue(level)) {
+                            primaryFieldUpTo = primaryFieldSkipper.maxDocID(level);
+                        } else {
+                            break;
+                        }
+                    }
+                    target = values.advance(primaryFieldUpTo + 1);
+                    if (target == DocIdSetIterator.NO_MORE_DOCS) {
+                        return docID = target;
+                    }
+                } else {
+                    target = values.nextDoc();
+                    if (target == DocIdSetIterator.NO_MORE_DOCS) {
+                        return docID = target;
+                    }
+                }
+            } else if (value >= minValue && value <= maxValue) {
+                return docID = target;
+            } else {
+                target = values.nextDoc();
+                if (target == DocIdSetIterator.NO_MORE_DOCS) {
+                    return docID = target;
+                }
+            }
+        }
+    }
+
+    @Override
+    public int docIDRunEnd() throws IOException {
+        if (skipperMatch) {
+            return valueFieldUpTo + 1;
+        }
+        return super.docIDRunEnd();
+    }
+
+    @Override
+    public long cost() {
+        return values.cost();
+    }
+
+}