Add bucketedSort based on int (elastic#128848)

Add bucketedSort on Int

Follow up on elastic#127968

Author: mayya-sharipova

docs/changelog/128848.yaml

@@ -0,0 +1,5 @@
+pr: 128848
+summary: Add `bucketedSort` based on int
+area: Search
+type: enhancement
+issues: []

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

@@ -15,10 +15,14 @@
 import org.apache.lucene.search.Pruning;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.comparators.IntComparator;
+import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData.NumericType;
+import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.MultiValueMode;
+import org.elasticsearch.search.sort.BucketedSort;
+import org.elasticsearch.search.sort.SortOrder;
 
 import java.io.IOException;
 
@@ -62,6 +66,39 @@ protected NumericDocValues getNumericDocValues(LeafReaderContext context, String
         };
     }
 
-    // TODO: add newBucketedSort based on integer values
+    @Override
+    public BucketedSort newBucketedSort(
+        BigArrays bigArrays,
+        SortOrder sortOrder,
+        DocValueFormat format,
+        int bucketSize,
+        BucketedSort.ExtraData extra
+    ) {
+        return new BucketedSort.ForInts(bigArrays, sortOrder, format, bucketSize, extra) {
+            private final int iMissingValue = (Integer) missingObject(missingValue, sortOrder == SortOrder.DESC);
+
+            @Override
+            public Leaf forLeaf(LeafReaderContext ctx) throws IOException {
+                return new Leaf(ctx) {
+                    private final NumericDocValues docValues = getNumericDocValues(ctx, iMissingValue);
+                    private int docValue;
+
+                    @Override
+                    protected boolean advanceExact(int doc) throws IOException {
+                        if (docValues.advanceExact(doc)) {
+                            docValue = (int) docValues.longValue();
+                            return true;
+                        }
+                        return false;
+                    }
+
+                    @Override
+                    protected int docValue() {
+                        return docValue;
+                    }
+                };
+            }
+        };
+    }
 
 }

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.common.util.DoubleArray;
 import org.elasticsearch.common.util.FloatArray;
+import org.elasticsearch.common.util.IntArray;
 import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -714,4 +715,93 @@ protected final boolean docBetterThan(long index) {
             }
         }
     }
+
+    public abstract static class ForInts extends BucketedSort {
+        private IntArray values;
+
+        @SuppressWarnings("this-escape")
+        public ForInts(BigArrays bigArrays, SortOrder sortOrder, DocValueFormat format, int bucketSize, ExtraData extra) {
+            super(bigArrays, sortOrder, format, bucketSize, extra);
+            boolean success = false;
+            try {
+                values = bigArrays.newIntArray(1, false);
+                success = true;
+            } finally {
+                if (success == false) {
+                    close();
+                }
+            }
+            initGatherOffsets();
+        }
+
+        @Override
+        public final boolean needsScores() {
+            return false;
+        }
+
+        @Override
+        protected final BigArray values() {
+            return values;
+        }
+
+        @Override
+        protected final void growValues(long minSize) {
+            values = bigArrays.grow(values, minSize);
+        }
+
+        @Override
+        protected final int getNextGatherOffset(long rootIndex) {
+            return values.get(rootIndex);
+        }
+
+        @Override
+        protected final void setNextGatherOffset(long rootIndex, int offset) {
+            values.set(rootIndex, offset);
+        }
+
+        @Override
+        protected final SortValue getValue(long index) {
+            return SortValue.from(values.get(index));
+        }
+
+        @Override
+        protected final boolean betterThan(long lhs, long rhs) {
+            return getOrder().reverseMul() * Integer.compare(values.get(lhs), values.get(rhs)) < 0;
+        }
+
+        @Override
+        protected final void swap(long lhs, long rhs) {
+            int tmp = values.get(lhs);
+            values.set(lhs, values.get(rhs));
+            values.set(rhs, tmp);
+        }
+
+        protected abstract class Leaf extends BucketedSort.Leaf {
+            protected Leaf(LeafReaderContext ctx) {
+                super(ctx);
+            }
+
+            /**
+             * Return the value for of this sort for the document to which
+             * we just {@link #advanceExact(int) moved}. This should be fast
+             * because it is called twice per competitive hit when in heap
+             * mode, once for {@link #docBetterThan(long)} and once
+             * for {@link #setIndexToDocValue(long)}.
+             */
+            protected abstract int docValue();
+
+            @Override
+            public final void setScorer(Scorable scorer) {}
+
+            @Override
+            protected final void setIndexToDocValue(long index) {
+                values.set(index, docValue());
+            }
+
+            @Override
+            protected final boolean docBetterThan(long index) {
+                return getOrder().reverseMul() * Integer.compare(docValue(), values.get(index)) < 0;
+            }
+        }
+    }
 }

server/src/test/java/org/elasticsearch/search/sort/BucketedSortForIntsTests.java

@@ -0,0 +1,54 @@
+/*
+ * 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.search.sort;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.search.DocValueFormat;
+
+public class BucketedSortForIntsTests extends BucketedSortTestCase<BucketedSort.ForInts> {
+    @Override
+    public BucketedSort.ForInts build(
+        SortOrder sortOrder,
+        DocValueFormat format,
+        int bucketSize,
+        BucketedSort.ExtraData extra,
+        double[] values
+    ) {
+        return new BucketedSort.ForInts(bigArrays(), sortOrder, format, bucketSize, extra) {
+            @Override
+            public Leaf forLeaf(LeafReaderContext ctx) {
+                return new Leaf(ctx) {
+                    int index = -1;
+
+                    @Override
+                    protected boolean advanceExact(int doc) {
+                        index = doc;
+                        return doc < values.length;
+                    }
+
+                    @Override
+                    protected int docValue() {
+                        return (int) values[index];
+                    }
+                };
+            }
+        };
+    }
+
+    @Override
+    protected SortValue expectedSortValue(double v) {
+        return SortValue.from((int) v);
+    }
+
+    @Override
+    protected double randomValue() {
+        return randomInt();
+    }
+}

server/src/test/java/org/elasticsearch/search/sort/BucketedSortForLongsTests.java

@@ -50,6 +50,6 @@ protected SortValue expectedSortValue(double v) {
     @Override
     protected double randomValue() {
         // 2L^50 fits in the mantisa of a double which the test sort of needs.
-        return randomLongBetween(-2L ^ 50, 2L ^ 50);
+        return randomLongBetween(-(1L << 50), (1L << 50));
     }
 }