search_after query optimization with missing value comparison

server/src/main/java/org/opensearch/search/SearchService.java

@@ -36,7 +36,9 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.util.BytesRef;
 import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionRunnable;
 import org.opensearch.action.IndicesRequest;
@@ -79,6 +81,7 @@
 import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.IndexSettings;
+import org.opensearch.index.IndexSortConfig;
 import org.opensearch.index.engine.Engine;
 import org.opensearch.index.mapper.DerivedFieldResolver;
 import org.opensearch.index.mapper.DerivedFieldResolverFactory;
@@ -136,10 +139,10 @@
 import org.opensearch.search.rescore.RescorerBuilder;
 import org.opensearch.search.searchafter.SearchAfterBuilder;
 import org.opensearch.search.sort.FieldSortBuilder;
+import org.opensearch.search.sort.FieldStats;
 import org.opensearch.search.sort.MinAndMax;
 import org.opensearch.search.sort.SortAndFormats;
 import org.opensearch.search.sort.SortBuilder;
-import org.opensearch.search.sort.SortOrder;
 import org.opensearch.search.startree.StarTreeQueryContext;
 import org.opensearch.search.startree.StarTreeQueryHelper;
 import org.opensearch.search.suggest.Suggest;
@@ -1692,8 +1695,13 @@ private CanMatchResponse canMatch(ShardSearchRequest request, boolean checkRefre
                 );
                 Rewriteable.rewrite(request.getRewriteable(), context, false);
                 final boolean aliasFilterCanMatch = request.getAliasFilter().getQueryBuilder() instanceof MatchNoneQueryBuilder == false;
-                FieldSortBuilder sortBuilder = FieldSortBuilder.getPrimaryFieldSortOrNull(request.source());
-                MinAndMax<?> minMax = sortBuilder != null ? FieldSortBuilder.getMinMaxOrNull(context, sortBuilder) : null;
+                final FieldSortBuilder sortBuilder = FieldSortBuilder.getPrimaryFieldSortOrNull(request.source());
+                final SortAndFormats primarySort = sortBuilder != null
+                    ? SortBuilder.buildSort(Collections.singletonList(sortBuilder), context).get()
+                    : null;
+                final FieldStats fieldStats = sortBuilder != null
+                    ? FieldSortBuilder.getFieldStatsOrNullForShard(context, sortBuilder)
+                    : null;
                 boolean canMatch;
                 if (canRewriteToMatchNone(request.source())) {
                     QueryBuilder queryBuilder = request.source().query();
@@ -1704,44 +1712,71 @@ private CanMatchResponse canMatch(ShardSearchRequest request, boolean checkRefre
                 }
                 final FieldDoc searchAfterFieldDoc = getSearchAfterFieldDoc(request, context);
                 final Integer trackTotalHitsUpto = request.source() == null ? null : request.source().trackTotalHitsUpTo();
-                canMatch = canMatch && canMatchSearchAfter(searchAfterFieldDoc, minMax, sortBuilder, trackTotalHitsUpto);
+                canMatch = canMatch
+                    && canMatchSearchAfter(
+                        searchAfterFieldDoc,
+                        fieldStats,
+                        primarySort,
+                        trackTotalHitsUpto,
+                        FieldSortBuilder.isSingleSort(request.source())
+                    );
 
-                return new CanMatchResponse(canMatch || hasRefreshPending, minMax);
+                return new CanMatchResponse(canMatch || hasRefreshPending, fieldStats != null ? fieldStats.minAndMax() : null);
             }
         }
     }
 
     public static boolean canMatchSearchAfter(
         FieldDoc searchAfter,
-        MinAndMax<?> minMax,
-        FieldSortBuilder primarySortField,
-        Integer trackTotalHitsUpto
+        FieldStats fieldStats,
+        SortAndFormats primarySort,
+        Integer trackTotalHitsUpto,
+        boolean singleSort
     ) {
-        // Check for sort.missing == null, since in case of missing values sort queries, if segment/shard's min/max
-        // is out of search_after range, it still should be printed and hence we should not skip segment/shard.
         // Skipping search on shard/segment entirely can cause mismatch on total_tracking_hits, hence skip only if
         // track_total_hits is false.
         if (searchAfter != null
-            && minMax != null
-            && primarySortField != null
-            && primarySortField.missing() == null
+            && searchAfter.fields[0] != null
+            && fieldStats != null
+            && primarySort != null
             && Objects.equals(trackTotalHitsUpto, TRACK_TOTAL_HITS_DISABLED)) {
             final Object searchAfterPrimary = searchAfter.fields[0];
-            if (primarySortField.order() == SortOrder.DESC) {
-                if (minMax.compareMin(searchAfterPrimary) > 0) {
+            final SortField primarySortField = primarySort.sort.getSort()[0];
+            if (primarySortField.getReverse()) {
+                final boolean nonMatch = fieldStats.minAndMax().compareMin(searchAfterPrimary) > (singleSort ? -1 : 0);
+                if (nonMatch) {
                     // In Desc order, if segment/shard minimum is gt search_after, the segment/shard won't be competitive
-                    return false;
+                    return fieldStats.allDocsHaveValue() == false && canMatchMissingValue(primarySortField, searchAfterPrimary, singleSort);
                 }
             } else {
-                if (minMax.compareMax(searchAfterPrimary) < 0) {
+                final boolean nonMatch = fieldStats.minAndMax().compareMax(searchAfterPrimary) < (singleSort ? 1 : 0);
+                if (nonMatch) {
                     // In ASC order, if segment/shard maximum is lt search_after, the segment/shard won't be competitive
-                    return false;
+                    return fieldStats.allDocsHaveValue() == false && canMatchMissingValue(primarySortField, searchAfterPrimary, singleSort);
                 }
             }
         }
         return true;
     }
 
+    private static boolean canMatchMissingValue(SortField primarySortField, Object primarySearchAfter, boolean singleSort) {
+        final Object missingValue = primarySortField.getMissingValue();
+        if (primarySortField.getReverse()) {
+            // the missing value of Type.STRING can only be SortField.STRING_FIRS or SortField.STRING_LAST
+            if (primarySearchAfter instanceof BytesRef) {
+                assert IndexSortConfig.getSortFieldType(primarySortField) == SortField.Type.STRING;
+                return missingValue == SortField.STRING_FIRST;
+            }
+            return MinAndMax.compare(primarySearchAfter, missingValue) > (singleSort ? 0 : -1);
+        } else {
+            if (primarySearchAfter instanceof BytesRef) {
+                assert IndexSortConfig.getSortFieldType(primarySortField) == SortField.Type.STRING;
+                return missingValue == SortField.STRING_LAST;
+            }
+            return MinAndMax.compare(primarySearchAfter, missingValue) < (singleSort ? 0 : 1);
+        }
+    }
+
     private static FieldDoc getSearchAfterFieldDoc(ShardSearchRequest request, QueryShardContext context) throws IOException {
         if (context != null && request != null && request.source() != null && request.source().sorts() != null) {
             final List<SortBuilder<?>> sorts = request.source().sorts();

server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java

@@ -81,7 +81,7 @@
 import org.opensearch.search.query.QueryPhase;
 import org.opensearch.search.query.QuerySearchResult;
 import org.opensearch.search.sort.FieldSortBuilder;
-import org.opensearch.search.sort.MinAndMax;
+import org.opensearch.search.sort.FieldStats;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -585,17 +585,18 @@
             // Only applied on primary sort field and primary search_after.
             FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(searchContext.request().source());
             if (primarySortField != null) {
-                MinAndMax<?> minMax = FieldSortBuilder.getMinMaxOrNullForSegment(
+                final FieldStats fieldStats = FieldSortBuilder.getFieldStatsOrNullForSegment(
                     this.searchContext.getQueryShardContext(),
                     ctx,
                     primarySortField,
                     searchContext.sort()
                 );
                 return SearchService.canMatchSearchAfter(
                     searchContext.searchAfter(),
-                    minMax,
-                    primarySortField,
-                    searchContext.trackTotalHitsUpTo()
+                    fieldStats,
+                    searchContext.sort(),
+                    searchContext.trackTotalHitsUpTo(),
+                    FieldSortBuilder.isSingleSort(searchContext.request().source())
                 );
             }
         }

server/src/main/java/org/opensearch/search/searchafter/SearchAfterBuilder.java

@@ -138,6 +138,10 @@
             if (values[i] != null) {
                 fieldValues[i] = convertValueFromSortField(values[i], sortField, format);
             } else {
+                SortField.Type sortType = extractSortType(sortField);
+                if (sortType != SortField.Type.STRING && sortType != SortField.Type.STRING_VAL) {
+                    throw new IllegalArgumentException("search after value of type [" + sortType + "] cannot be null");
+                }
                 fieldValues[i] = null;
             }
         }

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

@@ -72,7 +72,6 @@
 import org.opensearch.search.builder.SearchSourceBuilder;
 
 import java.io.IOException;
-import java.math.BigInteger;
 import java.util.Collections;
 import java.util.Locale;
 import java.util.Objects;
@@ -613,30 +612,39 @@
     }
 
     /**
-     * Return the {@link MinAndMax} indexed value for shard from the provided {@link FieldSortBuilder} or <code>null</code> if unknown.
+     * Indicates whether the sort is based on a single sort field or not.
+     *
+     * @return {@code true} if the sort is based on a single sort field, {@code false} otherwise
+     */
+    public static boolean isSingleSort(SearchSourceBuilder source) {
+        return source != null && source.sorts() != null && source.sorts().size() == 1;
+    }
+
+    /**
+     * Return the {@link FieldStats} indexed value for shard from the provided {@link FieldSortBuilder} or {@code null} if unknown.
      * The value can be extracted on non-nested indexed mapped fields of type keyword, numeric or date, other fields
-     * and configurations return <code>null</code>.
+     * and configurations return {@code null}.
      */
-    public static MinAndMax<?> getMinMaxOrNull(QueryShardContext context, FieldSortBuilder sortBuilder) throws IOException {
+    public static FieldStats getFieldStatsOrNullForShard(QueryShardContext context, FieldSortBuilder sortBuilder) throws IOException {
         final SortAndFormats sort = SortBuilder.buildSort(Collections.singletonList(sortBuilder), context).get();
-        return getMinMaxOrNullInternal(context.getIndexReader(), context, sortBuilder, sort);
+        return getFieldStatsOrNullInternal(context.getIndexReader(), context, sortBuilder, sort);
     }
 
     /**
-     * Return the {@link MinAndMax} indexed value for segment from the provided {@link FieldSortBuilder} or <code>null</code> if unknown.
+     * Return the {@link FieldStats} indexed value for segment from the provided {@link FieldSortBuilder} or {@code null} if unknown.
      * The value can be extracted on non-nested indexed mapped fields of type keyword, numeric or date, other fields
-     * and configurations return <code>null</code>.
+     * and configurations return {@code null}.
      */
-    public static MinAndMax<?> getMinMaxOrNullForSegment(
+    public static FieldStats getFieldStatsOrNullForSegment(
         QueryShardContext context,
         LeafReaderContext ctx,
         FieldSortBuilder sortBuilder,
         SortAndFormats sort
     ) throws IOException {
-        return getMinMaxOrNullInternal(ctx.reader(), context, sortBuilder, sort);
+        return getFieldStatsOrNullInternal(ctx.reader(), context, sortBuilder, sort);
     }
 
-    private static MinAndMax<?> getMinMaxOrNullInternal(
+    private static FieldStats getFieldStatsOrNullInternal(
         IndexReader reader,
         QueryShardContext context,
         FieldSortBuilder sortBuilder,
@@ -655,60 +663,63 @@
             case INT:
             case DOUBLE:
             case FLOAT:
-                return extractNumericMinAndMax(reader, sortField, fieldType, sortBuilder);
+                return extractNumericFieldStats(reader, sortField, fieldType, sortBuilder);
             case STRING:
             case STRING_VAL:
                 if (fieldType.unwrap() instanceof KeywordFieldMapper.KeywordFieldType) {
                     Terms terms = MultiTerms.getTerms(reader, fieldType.name());
                     if (terms == null) {
                         return null;
                     }
-                    return terms.getMin() != null ? new MinAndMax<>(terms.getMin(), terms.getMax()) : null;
+                    MinAndMax<?> minAndMax = terms.getMin() != null ? new MinAndMax<>(terms.getMin(), terms.getMax()) : null;
+                    return new FieldStats(minAndMax, terms.getDocCount() == reader.maxDoc());
                 }
                 break;
         }
         return null;
     }
 
-    private static MinAndMax<?> extractNumericMinAndMax(
+    private static FieldStats extractNumericFieldStats(
         IndexReader reader,
         SortField sortField,
         MappedFieldType fieldType,
         FieldSortBuilder sortBuilder
     ) throws IOException {
         String fieldName = fieldType.name();
-        if (PointValues.size(reader, fieldName) == 0) {
+        final int docCount = PointValues.getDocCount(reader, fieldName);
+        // TODO: should we deal with the case that all docs have no value?
+        if (docCount == 0) {
             return null;
         }
-        if (fieldType.unwrap() instanceof NumberFieldType) {
-            NumberFieldType numberFieldType = (NumberFieldType) fieldType;
+        final boolean allDocsHaveValue = docCount == reader.maxDoc();
+        MinAndMax<?> minAndMax = null;
+        if (fieldType.unwrap() instanceof NumberFieldType numberFieldType) {
             Number minPoint = numberFieldType.parsePoint(PointValues.getMinPackedValue(reader, fieldName));
             Number maxPoint = numberFieldType.parsePoint(PointValues.getMaxPackedValue(reader, fieldName));
+            // TODO: deal with half float and unsigned long
             switch (IndexSortConfig.getSortFieldType(sortField)) {
                 case LONG:
-                    if (numberFieldType.numericType() == NumericType.UNSIGNED_LONG) {
-                        // The min and max are expected to be BigInteger numbers
-                        return new MinAndMax<>((BigInteger) minPoint, (BigInteger) maxPoint);
-                    } else {
-                        return new MinAndMax<>(minPoint.longValue(), maxPoint.longValue());
-                    }
+                    minAndMax = new MinAndMax<>(minPoint.longValue(), maxPoint.longValue());
+                    break;
                 case INT:
-                    return new MinAndMax<>(minPoint.intValue(), maxPoint.intValue());
+                    minAndMax = new MinAndMax<>(minPoint.intValue(), maxPoint.intValue());
+                    break;
                 case DOUBLE:
-                    return new MinAndMax<>(minPoint.doubleValue(), maxPoint.doubleValue());
+                    minAndMax = new MinAndMax<>(minPoint.doubleValue(), maxPoint.doubleValue());
+                    break;
                 case FLOAT:
-                    return new MinAndMax<>(minPoint.floatValue(), maxPoint.floatValue());
+                    minAndMax = new MinAndMax<>(minPoint.floatValue(), maxPoint.floatValue());
+                    break;
                 default:
                     return null;
             }
-        } else if (fieldType.unwrap() instanceof DateFieldType) {
-            DateFieldType dateFieldType = (DateFieldType) fieldType;
+        } else if (fieldType.unwrap() instanceof DateFieldType dateFieldType) {
             Function<byte[], Long> dateConverter = createDateConverter(sortBuilder, dateFieldType);
             Long min = dateConverter.apply(PointValues.getMinPackedValue(reader, fieldName));
             Long max = dateConverter.apply(PointValues.getMaxPackedValue(reader, fieldName));
-            return new MinAndMax<>(min, max);
+            minAndMax = new MinAndMax<>(min, max);
         }
-        return null;
+        return new FieldStats(minAndMax, allDocsHaveValue);
     }
 
     private static Function<byte[], Long> createDateConverter(FieldSortBuilder sortBuilder, DateFieldType dateFieldType) {

server/src/main/java/org/opensearch/search/sort/FieldStats.java

@@ -0,0 +1,24 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.sort;
+
+import java.util.Objects;
+
+/**
+ * A class that encapsulates some stats for a given field.
+ *
+ * @param minAndMax the minimum and maximum value of the given field
+ * @param allDocsHaveValue whether all docs have value for the given field
+ * @opensearch.internal
+ */
+public record FieldStats(MinAndMax<?> minAndMax, boolean allDocsHaveValue) {
+    public FieldStats {
+        Objects.requireNonNull(minAndMax);
+    }
+}

server/src/main/java/org/opensearch/search/sort/MinAndMax.java

@@ -110,21 +110,15 @@ public int compareMax(Object object) {
         return compare(getMax(), object);
     }
 
-    private int compare(T one, Object two) {
-        if (one instanceof Long) {
-            return Long.compare((Long) one, (Long) two);
-        } else if (one instanceof Integer) {
-            return Integer.compare((Integer) one, (Integer) two);
-        } else if (one instanceof Float) {
-            return Float.compare((Float) one, (Float) two);
-        } else if (one instanceof Double) {
-            return Double.compare((Double) one, (Double) two);
-        } else if (one instanceof BigInteger) {
-            return ((BigInteger) one).compareTo((BigInteger) two);
-        } else if (one instanceof BytesRef) {
-            return ((BytesRef) one).compareTo((BytesRef) two);
-        } else {
-            throw new UnsupportedOperationException("compare type not supported : " + one.getClass());
-        }
+    public static int compare(Object one, Object two) {
+        return switch (one) {
+            case Long v -> Long.compare(v, (Long) two);
+            case Integer v -> Integer.compare(v, (Integer) two);
+            case Float v -> Float.compare(v, (Float) two);
+            case Double v -> Double.compare(v, (Double) two);
+            case BigInteger bigInteger -> bigInteger.compareTo((BigInteger) two);
+            case BytesRef bytesRef -> bytesRef.compareTo((BytesRef) two);
+            default -> throw new UnsupportedOperationException("compare type not supported : " + one.getClass());
+        };
     }
 }