diff --git a/docs/changelog/128848.yaml b/docs/changelog/128848.yaml new file mode 100644 index 0000000000000..a928206d5568d --- /dev/null +++ b/docs/changelog/128848.yaml @@ -0,0 +1,5 @@ +pr: 128848 +summary: Add `bucketedSort` based on int +area: Search +type: enhancement +issues: [] diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/IntValuesComparatorSource.java b/server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/IntValuesComparatorSource.java index 7c2b3888f7a47..184288d0437b8 100644 --- a/server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/IntValuesComparatorSource.java +++ b/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; + } + }; + } + }; + } } diff --git a/server/src/main/java/org/elasticsearch/search/sort/BucketedSort.java b/server/src/main/java/org/elasticsearch/search/sort/BucketedSort.java index 0696974ae2b53..3b47881511c59 100644 --- a/server/src/main/java/org/elasticsearch/search/sort/BucketedSort.java +++ b/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; + } + } + } } diff --git a/server/src/test/java/org/elasticsearch/search/sort/BucketedSortForIntsTests.java b/server/src/test/java/org/elasticsearch/search/sort/BucketedSortForIntsTests.java new file mode 100644 index 0000000000000..4516dd25cfe3b --- /dev/null +++ b/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 { + @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(); + } +} diff --git a/server/src/test/java/org/elasticsearch/search/sort/BucketedSortForLongsTests.java b/server/src/test/java/org/elasticsearch/search/sort/BucketedSortForLongsTests.java index eca336b5293af..cead6264f576f 100644 --- a/server/src/test/java/org/elasticsearch/search/sort/BucketedSortForLongsTests.java +++ b/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)); } }