Approximation framework to support numeric search_after queries

CHANGELOG.md

@@ -52,6 +52,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Star-Tree] Add search support for ip field type ([#18671](https://github.com/opensearch-project/OpenSearch/pull/18671))
 - [Derived Source] Add integration of derived source feature across various paths like get/search/recovery ([#18565](https://github.com/opensearch-project/OpenSearch/pull/18565))
 - Supporting Scripted Metric Aggregation when reducing aggregations in InternalValueCount and InternalAvg ([18411](https://github.com/opensearch-project/OpenSearch/pull18411)))
+- Support `search_after` numeric queries with Approximation Framework ([#18896](https://github.com/opensearch-project/OpenSearch/pull/18896))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

modules/mapper-extras/src/main/java/org/opensearch/index/mapper/ScaledFloatFieldMapper.java

@@ -217,6 +217,22 @@ public byte[] encodePoint(Number value) {
             return point;
         }
 
+        @Override
+        public byte[] encodePoint(Object value, boolean roundUp) {
+            double doubleValue = parse(value);
+            long scaledValue = Math.round(scale(doubleValue));
+            if (roundUp) {
+                if (scaledValue < Long.MAX_VALUE) {
+                    scaledValue = scaledValue + 1;
+                }
+            } else {
+                if (scaledValue > Long.MIN_VALUE) {
+                    scaledValue = scaledValue - 1;
+                }
+            }
+            return encodePoint(scaledValue);
+        }
+
         public double getScalingFactor() {
             return scalingFactor;
         }

server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java

@@ -619,6 +619,23 @@ public byte[] encodePoint(Number value) {
             return point;
         }
 
+        @Override
+        public byte[] encodePoint(Object value, boolean roundUp) {
+            // Always parse with roundUp=false to get consistent date math
+            // In this method the parseToLong is only used for date math rounding operations
+            long timestamp = parseToLong(value, false, null, null, null);
+            if (roundUp) {
+                if (timestamp < Long.MAX_VALUE) {
+                    timestamp = timestamp + 1;
+                }
+            } else {
+                if (timestamp > Long.MIN_VALUE) {
+                    timestamp = timestamp - 1;
+                }
+            }
+            return encodePoint(timestamp);
+        }
+
         @Override
         public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();

server/src/main/java/org/opensearch/index/mapper/NumberFieldMapper.java

@@ -283,6 +283,17 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Float numericValue = parse(value, true);
+                if (roundUp) {
+                    numericValue = HalfFloatPoint.nextUp(numericValue);
+                } else {
+                    numericValue = HalfFloatPoint.nextDown(numericValue);
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return HalfFloatPoint.sortableShortToHalfFloat((short) value);
@@ -459,6 +470,17 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Float numericValue = parse(value, true);
+                if (roundUp) {
+                    numericValue = FloatPoint.nextUp(numericValue);
+                } else {
+                    numericValue = FloatPoint.nextDown(numericValue);
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return NumericUtils.sortableIntToFloat((int) value);
@@ -626,6 +648,17 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Double numericValue = parse(value, true);
+                if (roundUp) {
+                    numericValue = DoublePoint.nextUp(numericValue);
+                } else {
+                    numericValue = DoublePoint.nextDown(numericValue);
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return NumericUtils.sortableLongToDouble(value);
@@ -789,6 +822,23 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Byte numericValue = parse(value, true);
+                if (roundUp) {
+                    // ASC: exclusive lower bound
+                    if (numericValue < Byte.MAX_VALUE) {
+                        numericValue = (byte) (numericValue + 1);
+                    }
+                } else {
+                    // DESC: exclusive upper bound
+                    if (numericValue > Byte.MIN_VALUE) {
+                        numericValue = (byte) (numericValue - 1);
+                    }
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return objectToDouble(value);
@@ -873,6 +923,22 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Short numericValue = parse(value, true);
+                if (roundUp) {
+                    // ASC: exclusive lower bound
+                    if (numericValue < Short.MAX_VALUE) {
+                        numericValue = (short) (numericValue + 1);
+                    }
+                } else {
+                    if (numericValue > Short.MIN_VALUE) {
+                        numericValue = (short) (numericValue - 1);
+                    }
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return (double) value;
@@ -953,6 +1019,23 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Integer numericValue = parse(value, true);
+                // Always apply exclusivity
+                if (roundUp) {
+                    if (numericValue < Integer.MAX_VALUE) {
+                        numericValue = numericValue + 1;
+                    }
+                } else {
+                    if (numericValue > Integer.MIN_VALUE) {
+                        numericValue = numericValue - 1;
+                    }
+                }
+
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return (double) value;
@@ -1139,6 +1222,23 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                Long numericValue = parse(value, true);
+                if (roundUp) {
+                    // ASC: exclusive lower bound
+                    if (numericValue < Long.MAX_VALUE) {
+                        numericValue = numericValue + 1;
+                    }
+                } else {
+                    // DESC: exclusive upper bound
+                    if (numericValue > Long.MIN_VALUE) {
+                        numericValue = numericValue - 1;
+                    }
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return (double) value;
@@ -1281,6 +1381,22 @@ public byte[] encodePoint(Number value) {
                 return point;
             }
 
+            @Override
+            public byte[] encodePoint(Object value, boolean roundUp) {
+                BigInteger numericValue = parse(value, true);
+                if (roundUp) {
+                    if (numericValue.compareTo(Numbers.MAX_UNSIGNED_LONG_VALUE) < 0) {
+                        numericValue = numericValue.add(BigInteger.ONE);
+                    }
+                } else {
+                    // DESC: exclusive upper bound
+                    if (numericValue.compareTo(Numbers.MIN_UNSIGNED_LONG_VALUE) > 0) {
+                        numericValue = numericValue.subtract(BigInteger.ONE);
+                    }
+                }
+                return encodePoint(numericValue);
+            }
+
             @Override
             public double toDoubleValue(long value) {
                 return Numbers.unsignedLongToDouble(value);
@@ -1851,6 +1967,11 @@ public byte[] encodePoint(Number value) {
             return type.encodePoint(value);
         }
 
+        @Override
+        public byte[] encodePoint(Object value, boolean roundUp) {
+            return type.encodePoint(value, roundUp);
+        }
+
         @Override
         public double toDoubleValue(long value) {
             return type.toDoubleValue(value);

server/src/main/java/org/opensearch/index/mapper/NumericPointEncoder.java

@@ -13,4 +13,12 @@
  */
 public interface NumericPointEncoder {
     byte[] encodePoint(Number value);
+
+    /**
+     * Encodes an Object value to byte array for Approximation Framework search_after optimization.
+     * @param value the search_after value as Object
+     * @param roundUp whether to round up (for lower bounds) or down (for upper bounds)
+     * @return encoded byte array
+     */
+    byte[] encodePoint(Object value, boolean roundUp);
 }

server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

@@ -31,6 +31,8 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.DocIdSetBuilder;
 import org.apache.lucene.util.IntsRef;
+import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.NumericPointEncoder;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.SortOrder;
@@ -52,10 +54,9 @@ public class ApproximatePointRangeQuery extends ApproximateQuery {
     public static final Function<byte[], String> UNSIGNED_LONG_FORMAT = bytes -> BigIntegerPoint.decodeDimension(bytes, 0).toString();
 
     private int size;
-
     private SortOrder sortOrder;
-
-    public final PointRangeQuery pointRangeQuery;
+    public PointRangeQuery pointRangeQuery;
+    private final Function<byte[], String> valueToString;
 
     public ApproximatePointRangeQuery(
         String field,
@@ -78,6 +79,7 @@ protected ApproximatePointRangeQuery(
     ) {
         this.size = size;
         this.sortOrder = sortOrder;
+        this.valueToString = valueToString;
         this.pointRangeQuery = new PointRangeQuery(field, lowerPoint, upperPoint, numDims) {
             @Override
             protected String toString(int dimension, byte[] value) {
@@ -114,12 +116,12 @@ public void visit(QueryVisitor visitor) {
 
     @Override
     public final ConstantScoreWeight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+        final ArrayUtil.ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(pointRangeQuery.getBytesPerDim());
+
         Weight pointRangeQueryWeight = pointRangeQuery.createWeight(searcher, scoreMode, boost);
 
         return new ConstantScoreWeight(this, boost) {
 
-            private final ArrayUtil.ByteArrayComparator comparator = ArrayUtil.getUnsignedComparator(pointRangeQuery.getBytesPerDim());
-
             // we pull this from PointRangeQuery since it is final
             private boolean matches(byte[] packedValue) {
                 for (int dim = 0; dim < pointRangeQuery.getNumDims(); dim++) {
@@ -138,7 +140,6 @@ private boolean matches(byte[] packedValue) {
 
             // we pull this from PointRangeQuery since it is final
             private PointValues.Relation relate(byte[] minPackedValue, byte[] maxPackedValue) {
-
                 boolean crosses = false;
 
                 for (int dim = 0; dim < pointRangeQuery.getNumDims(); dim++) {
@@ -352,6 +353,7 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
                 if (checkValidPointValues(values) == false) {
                     return null;
                 }
+                // values.size(): total points indexed, In most cases: values.size() ≈ number of documents (assuming single-valued fields)
                 if (size > values.size()) {
                     return pointRangeQueryWeight.scorerSupplier(context);
                 } else {
@@ -423,6 +425,19 @@ public boolean isCacheable(LeafReaderContext ctx) {
         };
     }
 
+    private byte[] computeEffectiveBound(SearchContext context, boolean isLowerBound) {
+        byte[] originalBound = isLowerBound ? pointRangeQuery.getLowerPoint() : pointRangeQuery.getUpperPoint();
+        boolean isAscending = sortOrder == null || sortOrder.equals(SortOrder.ASC);
+        if ((isLowerBound && isAscending) || (isLowerBound == false && isAscending == false)) {
+            Object searchAfterValue = context.request().source().searchAfter()[0];
+            MappedFieldType fieldType = context.getQueryShardContext().fieldMapper(pointRangeQuery.getField());
+            if (fieldType instanceof NumericPointEncoder encoder) {
+                return encoder.encodePoint(searchAfterValue, isLowerBound);
+            }
+        }
+        return originalBound;
+    }
+
     @Override
     public boolean canApproximate(SearchContext context) {
         if (context == null) {
@@ -435,7 +450,6 @@ public boolean canApproximate(SearchContext context) {
         if (context.trackTotalHitsUpTo() == SearchContext.TRACK_TOTAL_HITS_ACCURATE) {
             return false;
         }
-
         // size 0 could be set for caching
         if (context.from() + context.size() == 0) {
             this.setSize(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
@@ -459,12 +473,24 @@ public boolean canApproximate(SearchContext context) {
                     // Cannot sort documents missing this field.
                     return false;
                 }
+                this.setSortOrder(primarySortField.order());
                 if (context.request().source().searchAfter() != null) {
-                    // TODO: We *could* optimize searchAfter, especially when this is the only sort field, but existing pruning is pretty
-                    // good.
-                    return false;
+                    byte[] lower;
+                    byte[] upper;
+                    if (sortOrder == SortOrder.ASC) {
+                        lower = computeEffectiveBound(context, true);
+                        upper = pointRangeQuery.getUpperPoint();
+                    } else {
+                        lower = pointRangeQuery.getLowerPoint();
+                        upper = computeEffectiveBound(context, false);
+                    }
+                    this.pointRangeQuery = new PointRangeQuery(pointRangeQuery.getField(), lower, upper, pointRangeQuery.getNumDims()) {
+                        @Override
+                        protected String toString(int dimension, byte[] value) {
+                            return valueToString.apply(value);
+                        }
+                    };
                 }
-                this.setSortOrder(primarySortField.order());
             }
             return context.request().source().terminateAfter() == SearchContext.DEFAULT_TERMINATE_AFTER;
         }