Skip to content

Commit 5f1d8bf

Browse files
romseygeekromseygeek
authored
Add DenseDoubleValues and DenseLongValues (#137817)
These are extensions of DoubleValues and LongValues that are guaranteed to have values for every document, which means they can also be used for iteration. This allows us to re-enable pruning on sorts that use them as sources for comparators.
1 parent 97f96b4 commit 5f1d8bf

File tree

13 files changed

+173
-55
lines changed

13 files changed

+173
-55
lines changed

‎server/src/main/java/org/elasticsearch/index/fielddata/DenseDoubleValues.java‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.fielddata;
11+
12+
import org.apache.lucene.index.NumericDocValues;
13+
import org.apache.lucene.search.DoubleValues;
14+
15+
import java.io.IOException;
16+
import java.util.function.DoubleToLongFunction;
17+
18+
/**
19+
* DoubleValues implementation that is guaranteed to have a value
20+
* for every document in a reader
21+
*/
22+
public abstract class DenseDoubleValues extends DoubleValues {
23+
24+
@Override
25+
public final boolean advanceExact(int doc) throws IOException {
26+
doAdvanceExact(doc);
27+
return true;
28+
}
29+
30+
protected abstract void doAdvanceExact(int doc) throws IOException;
31+
32+
/**
33+
* Represent a DenseDoubleValues as a NumericDocValues instance
34+
* @param in the DenseDoubleValues to wrap
35+
* @param maxDoc the maxDoc of the current reader
36+
* @param converter a function to convert the double-valued output of DenseDoubleValues to a long
37+
*/
38+
public static NumericDocValues asNumericDocValues(DenseDoubleValues in, int maxDoc, DoubleToLongFunction converter) {
39+
return new NumericDocValues() {
40+
41+
int doc = -1;
42+
43+
@Override
44+
public long longValue() throws IOException {
45+
return converter.applyAsLong(in.doubleValue());
46+
}
47+
48+
@Override
49+
public boolean advanceExact(int target) throws IOException {
50+
doc = target;
51+
return in.advanceExact(target);
52+
}
53+
54+
@Override
55+
public int docID() {
56+
return doc;
57+
}
58+
59+
@Override
60+
public int nextDoc() throws IOException {
61+
return advance(doc + 1);
62+
}
63+
64+
@Override
65+
public int advance(int target) throws IOException {
66+
if (target >= maxDoc) {
67+
return doc = NO_MORE_DOCS;
68+
}
69+
in.advanceExact(target);
70+
return doc = target;
71+
}
72+
73+
@Override
74+
public long cost() {
75+
return maxDoc;
76+
}
77+
};
78+
}
79+
}

‎server/src/main/java/org/elasticsearch/index/fielddata/DenseLongValues.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.fielddata;
11+
12+
import org.apache.lucene.search.LongValues;
13+
14+
import java.io.IOException;
15+
16+
/**
17+
* LongValues implementation that is guaranteed to have a value
18+
* for every document in a reader
19+
*/
20+
public abstract class DenseLongValues extends LongValues {
21+
22+
@Override
23+
public final boolean advanceExact(int doc) throws IOException {
24+
doAdvanceExact(doc);
25+
return true;
26+
}
27+
28+
protected abstract void doAdvanceExact(int doc) throws IOException;
29+
}

‎server/src/main/java/org/elasticsearch/index/fielddata/FieldData.java‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -614,19 +614,18 @@ public boolean advanceExact(int doc) throws IOException {
614614
* document, returns the same value as {@code values} if there is a value
615615
* for the current document and {@code missing} otherwise.
616616
*/
617-
public static LongValues replaceMissing(LongValues values, long missing) {
618-
return new LongValues() {
617+
public static DenseLongValues replaceMissing(LongValues values, long missing) {
618+
return new DenseLongValues() {
619619

620620
private long value;
621621

622622
@Override
623-
public boolean advanceExact(int target) throws IOException {
623+
public void doAdvanceExact(int target) throws IOException {
624624
if (values.advanceExact(target)) {
625625
value = values.longValue();
626626
} else {
627627
value = missing;
628628
}
629-
return true;
630629
}
631630

632631
@Override
@@ -641,19 +640,18 @@ public long longValue() {
641640
* document, returns the same value as {@code values} if there is a value
642641
* for the current document and {@code missing} otherwise.
643642
*/
644-
public static NumericDoubleValues replaceMissing(NumericDoubleValues values, double missing) {
645-
return new NumericDoubleValues() {
643+
public static DenseDoubleValues replaceMissing(NumericDoubleValues values, double missing) {
644+
return new DenseDoubleValues() {
646645

647646
private double value;
648647

649648
@Override
650-
public boolean advanceExact(int target) throws IOException {
649+
public void doAdvanceExact(int target) throws IOException {
651650
if (values.advanceExact(target)) {
652651
value = values.doubleValue();
653652
} else {
654653
value = missing;
655654
}
656-
return true;
657655
}
658656

659657
@Override

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@
2121
import org.apache.lucene.util.BitSet;
2222
import org.elasticsearch.common.util.BigArrays;
2323
import org.elasticsearch.core.Nullable;
24+
import org.elasticsearch.index.fielddata.DenseDoubleValues;
2425
import org.elasticsearch.index.fielddata.FieldData;
2526
import org.elasticsearch.index.fielddata.IndexFieldData;
2627
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
27-
import org.elasticsearch.index.fielddata.NumericDoubleValues;
2828
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
2929
import org.elasticsearch.search.DocValueFormat;
3030
import org.elasticsearch.search.MultiValueMode;
@@ -59,7 +59,7 @@ protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws
5959
return indexFieldData.load(context).getDoubleValues();
6060
}
6161

62-
private NumericDoubleValues getNumericDocValues(LeafReaderContext context, double missingValue) throws IOException {
62+
private DenseDoubleValues getDenseDoubleValues(LeafReaderContext context, double missingValue) throws IOException {
6363
final SortedNumericDoubleValues values = getValues(context);
6464
if (nested == null) {
6565
return FieldData.replaceMissing(sortMode.select(values), missingValue);
@@ -80,14 +80,17 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
8080
final double dMissingValue = (Double) missingObject(missingValue, reversed);
8181
// NOTE: it's important to pass null as a missing value in the constructor so that
8282
// the comparator doesn't check docsWithField since we replace missing values in select()
83-
// TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
84-
return new DoubleComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
83+
return new DoubleComparator(numHits, fieldname, null, reversed, enableSkipping) {
8584
@Override
8685
public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
8786
return new DoubleLeafComparator(context) {
8887
@Override
8988
protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
90-
return DoubleValuesComparatorSource.this.getNumericDocValues(context, dMissingValue).getRawDoubleValues();
89+
return DenseDoubleValues.asNumericDocValues(
90+
getDenseDoubleValues(context, dMissingValue),
91+
context.reader().maxDoc(),
92+
Double::doubleToRawLongBits
93+
);
9194
}
9295

9396
@Override
@@ -113,7 +116,7 @@ public BucketedSort newBucketedSort(
113116
@Override
114117
public Leaf forLeaf(LeafReaderContext ctx) throws IOException {
115118
return new Leaf(ctx) {
116-
private final NumericDoubleValues docValues = getNumericDocValues(ctx, dMissingValue);
119+
private final DenseDoubleValues docValues = getDenseDoubleValues(ctx, dMissingValue);
117120
private double docValue;
118121

119122
@Override

‎server/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@
2020
import org.apache.lucene.util.BitSet;
2121
import org.elasticsearch.common.util.BigArrays;
2222
import org.elasticsearch.core.Nullable;
23+
import org.elasticsearch.index.fielddata.DenseDoubleValues;
2324
import org.elasticsearch.index.fielddata.FieldData;
2425
import org.elasticsearch.index.fielddata.IndexFieldData;
2526
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
26-
import org.elasticsearch.index.fielddata.NumericDoubleValues;
2727
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
2828
import org.elasticsearch.search.DocValueFormat;
2929
import org.elasticsearch.search.MultiValueMode;
@@ -54,7 +54,7 @@ public SortField.Type reducedType() {
5454
return SortField.Type.FLOAT;
5555
}
5656

57-
NumericDoubleValues getNumericDocValues(LeafReaderContext context, double missingValue) throws IOException {
57+
DenseDoubleValues getDenseDoubleValues(LeafReaderContext context, double missingValue) throws IOException {
5858
final SortedNumericDoubleValues values = indexFieldData.load(context).getDoubleValues();
5959
if (nested == null) {
6060
return FieldData.replaceMissing(sortMode.select(values), missingValue);
@@ -73,14 +73,17 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
7373
final float fMissingValue = (Float) missingObject(missingValue, reversed);
7474
// NOTE: it's important to pass null as a missing value in the constructor so that
7575
// the comparator doesn't check docsWithField since we replace missing values in select()
76-
// TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
77-
return new FloatComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
76+
return new FloatComparator(numHits, fieldname, null, reversed, enableSkipping) {
7877
@Override
7978
public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
8079
return new FloatLeafComparator(context) {
8180
@Override
8281
protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
83-
return FloatValuesComparatorSource.this.getNumericDocValues(context, fMissingValue).getRawFloatValues();
82+
return DenseDoubleValues.asNumericDocValues(
83+
getDenseDoubleValues(context, fMissingValue),
84+
context.reader().maxDoc(),
85+
v -> Float.floatToRawIntBits((float) v)
86+
);
8487
}
8588
};
8689
}
@@ -106,7 +109,7 @@ public boolean needsScores() {
106109
@Override
107110
public Leaf forLeaf(LeafReaderContext ctx) throws IOException {
108111
return new Leaf(ctx) {
109-
private final NumericDoubleValues docValues = getNumericDocValues(ctx, dMissingValue);
112+
private final DenseDoubleValues docValues = getDenseDoubleValues(ctx, dMissingValue);
110113
private float docValue;
111114

112115
@Override

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.apache.lucene.search.LeafFieldComparator;
1515
import org.apache.lucene.search.Pruning;
1616
import org.elasticsearch.core.Nullable;
17+
import org.elasticsearch.index.fielddata.DenseDoubleValues;
1718
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
1819
import org.elasticsearch.search.MultiValueMode;
1920

@@ -39,14 +40,17 @@ public FieldComparator<?> newComparator(String fieldname, int numHits, Pruning e
3940
final float fMissingValue = (Float) missingObject(missingValue, reversed);
4041
// NOTE: it's important to pass null as a missing value in the constructor so that
4142
// the comparator doesn't check docsWithField since we replace missing values in select()
42-
// TODO we can re-enable pruning here if we allow NumericDoubleValues to expose an iterator
43-
return new HalfFloatComparator(numHits, fieldname, null, reversed, Pruning.NONE) {
43+
return new HalfFloatComparator(numHits, fieldname, null, reversed, enableSkipping) {
4444
@Override
4545
public LeafFieldComparator getLeafComparator(LeafReaderContext context) throws IOException {
4646
return new HalfFloatLeafComparator(context) {
4747
@Override
4848
protected NumericDocValues getNumericDocValues(LeafReaderContext context, String field) throws IOException {
49-
return HalfFloatValuesComparatorSource.this.getNumericDocValues(context, fMissingValue).getRawFloatValues();
49+
return DenseDoubleValues.asNumericDocValues(
50+
getDenseDoubleValues(context, fMissingValue),
51+
context.reader().maxDoc(),
52+
v -> Float.floatToRawIntBits((float) v)
53+
);
5054
}
5155
};
5256
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import org.elasticsearch.common.time.DateUtils;
2323
import org.elasticsearch.common.util.BigArrays;
2424
import org.elasticsearch.core.Nullable;
25+
import org.elasticsearch.index.fielddata.DenseLongValues;
2526
import org.elasticsearch.index.fielddata.FieldData;
2627
import org.elasticsearch.index.fielddata.IndexFieldData;
2728
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
@@ -88,7 +89,7 @@ private SortedNumericLongValues loadDocValues(LeafReaderContext context) {
8889
return converter != null ? converter.apply(values) : values;
8990
}
9091

91-
LongValues getLongValues(LeafReaderContext context, long missingValue) throws IOException {
92+
DenseLongValues getLongValues(LeafReaderContext context, long missingValue) throws IOException {
9293
final SortedNumericLongValues values = loadDocValues(context);
9394
if (nested == null) {
9495
return FieldData.replaceMissing(sortMode.select(values), missingValue);
@@ -201,7 +202,7 @@ public Object missingObject(Object missingValue, boolean reversed) {
201202
return super.missingObject(missingValue, reversed);
202203
}
203204

204-
protected static NumericDocValues wrap(LongValues longValues, int maxDoc) {
205+
protected static NumericDocValues wrap(DenseLongValues longValues, int maxDoc) {
205206
return new NumericDocValues() {
206207

207208
int doc = -1;

‎server/src/main/java/org/elasticsearch/index/query/functionscore/DecayFunctionBuilder.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.index.query.functionscore;
1111

1212
import org.apache.lucene.index.LeafReaderContext;
13+
import org.apache.lucene.search.DoubleValues;
1314
import org.apache.lucene.search.Explanation;
1415
import org.elasticsearch.ElasticsearchParseException;
1516
import org.elasticsearch.common.ParsingException;
@@ -31,7 +32,6 @@
3132
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
3233
import org.elasticsearch.index.fielddata.IndexNumericFieldData;
3334
import org.elasticsearch.index.fielddata.MultiGeoPointValues;
34-
import org.elasticsearch.index.fielddata.NumericDoubleValues;
3535
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
3636
import org.elasticsearch.index.fielddata.SortingNumericDoubleValues;
3737
import org.elasticsearch.index.mapper.DateFieldMapper;
@@ -391,7 +391,7 @@ public boolean needsScores() {
391391
}
392392

393393
@Override
394-
protected NumericDoubleValues distance(LeafReaderContext context) {
394+
protected DoubleValues distance(LeafReaderContext context) {
395395
final MultiGeoPointValues geoPointValues = fieldData.load(context).getPointValues();
396396
return FieldData.replaceMissing(mode.select(new SortingNumericDoubleValues() {
397397
@Override
@@ -484,7 +484,7 @@ public boolean needsScores() {
484484
}
485485

486486
@Override
487-
protected NumericDoubleValues distance(LeafReaderContext context) {
487+
protected DoubleValues distance(LeafReaderContext context) {
488488
final SortedNumericDoubleValues doubleValues = fieldData.load(context).getDoubleValues();
489489
return FieldData.replaceMissing(mode.select(new SortingNumericDoubleValues() {
490490
@Override
@@ -585,11 +585,11 @@ public AbstractDistanceScoreFunction(
585585
* guaranteed that the value actually exists. If it does not, we assume
586586
* the user handles this case in the query and return 0.
587587
* */
588-
protected abstract NumericDoubleValues distance(LeafReaderContext context);
588+
protected abstract DoubleValues distance(LeafReaderContext context);
589589

590590
@Override
591591
public final LeafScoreFunction getLeafScoreFunction(final LeafReaderContext ctx) {
592-
final NumericDoubleValues distance = distance(ctx);
592+
final DoubleValues distance = distance(ctx);
593593
return new LeafScoreFunction() {
594594

595595
@Override

0 commit comments

Comments
 (0)