Skip to content

Commit 987e7af

Browse files
martijnvgmartijnvg
authored
Support bulk reading doc values for byte, short and integer field types. (#134753)
Otherwise byte, short and integer field types can't bulk load like long field types. And these field types are being used.
1 parent c8c635b commit 987e7af

File tree

14 files changed

+521
-41
lines changed

14 files changed

+521
-41
lines changed

‎server/src/main/java/org/elasticsearch/index/codec/tsdb/es819/ES819TSDBDocValuesProducer.java‎

Lines changed: 64 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -390,7 +390,8 @@ public BlockLoader.Block tryRead(
390390
BlockLoader.Docs docs,
391391
int offset,
392392
boolean nullsFiltered,
393-
BlockDocValuesReader.ToDouble toDouble
393+
BlockDocValuesReader.ToDouble toDouble,
394+
boolean toInt
394395
) throws IOException {
395396
assert toDouble == null;
396397
if (ords instanceof BaseDenseNumericValues denseOrds) {
@@ -471,7 +472,8 @@ public BlockLoader.Block tryRead(
471472
BlockLoader.Docs docs,
472473
int offset,
473474
boolean nullsFiltered,
474-
BlockDocValuesReader.ToDouble toDouble
475+
BlockDocValuesReader.ToDouble toDouble,
476+
boolean toInt
475477
) throws IOException {
476478
return null;
477479
}
@@ -524,7 +526,8 @@ public BlockLoader.Block tryRead(
524526
BlockLoader.Docs docs,
525527
int offset,
526528
boolean nullsFiltered,
527-
BlockDocValuesReader.ToDouble toDouble
529+
BlockDocValuesReader.ToDouble toDouble,
530+
boolean toInt
528531
) throws IOException {
529532
return null;
530533
}
@@ -574,7 +577,8 @@ public BlockLoader.Block tryRead(
574577
BlockLoader.Docs docs,
575578
int offset,
576579
boolean nullsFiltered,
577-
BlockDocValuesReader.ToDouble toDouble
580+
BlockDocValuesReader.ToDouble toDouble,
581+
boolean toInt
578582
) throws IOException {
579583
return null;
580584
}
@@ -1402,9 +1406,10 @@ public BlockLoader.Block tryRead(
14021406
BlockLoader.Docs docs,
14031407
int offset,
14041408
boolean nullsFiltered,
1405-
BlockDocValuesReader.ToDouble toDouble
1409+
BlockDocValuesReader.ToDouble toDouble,
1410+
boolean toInt
14061411
) throws IOException {
1407-
try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, docs.count() - offset)) {
1412+
try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, docs.count() - offset, toInt)) {
14081413
return tryRead(singletonLongBuilder, docs, offset);
14091414
}
14101415
}
@@ -1525,7 +1530,8 @@ public BlockLoader.Block tryRead(
15251530
BlockLoader.Docs docs,
15261531
int offset,
15271532
boolean nullsFiltered,
1528-
BlockDocValuesReader.ToDouble toDouble
1533+
BlockDocValuesReader.ToDouble toDouble,
1534+
boolean toInt
15291535
) throws IOException {
15301536
if (nullsFiltered == false) {
15311537
return null;
@@ -1566,7 +1572,7 @@ public BlockLoader.Block tryRead(
15661572
assert disi.index() == firstIndex + i : "unexpected disi index " + (firstIndex + i) + "!=" + disi.index();
15671573
}
15681574
}
1569-
try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, valueCount)) {
1575+
try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, valueCount, toInt)) {
15701576
for (int i = 0; i < valueCount;) {
15711577
final int index = firstIndex + i;
15721578
final int blockIndex = index >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT;
@@ -1884,10 +1890,15 @@ public void close() {}
18841890
static BlockLoader.SingletonLongBuilder singletonLongBuilder(
18851891
BlockLoader.BlockFactory factory,
18861892
BlockDocValuesReader.ToDouble toDouble,
1887-
int valueCount
1893+
int valueCount,
1894+
boolean toInt
18881895
) {
1896+
assert (toInt && toDouble != null) == false;
1897+
18891898
if (toDouble != null) {
18901899
return new SingletonLongToDoubleDelegate(factory.singletonDoubles(valueCount), toDouble);
1900+
} else if (toInt) {
1901+
return new SingletonLongtoIntDelegate(factory.singletonInts(valueCount));
18911902
} else {
18921903
return factory.singletonLongs(valueCount);
18931904
}
@@ -1941,4 +1952,48 @@ public void close() {
19411952
}
19421953
}
19431954

1955+
static final class SingletonLongtoIntDelegate implements BlockLoader.SingletonLongBuilder {
1956+
private final BlockLoader.SingletonIntBuilder builder;
1957+
1958+
SingletonLongtoIntDelegate(BlockLoader.SingletonIntBuilder builder) {
1959+
this.builder = builder;
1960+
}
1961+
1962+
@Override
1963+
public BlockLoader.SingletonLongBuilder appendLong(long value) {
1964+
throw new UnsupportedOperationException();
1965+
}
1966+
1967+
@Override
1968+
public BlockLoader.SingletonLongBuilder appendLongs(long[] values, int from, int length) {
1969+
builder.appendLongs(values, from, length);
1970+
return this;
1971+
}
1972+
1973+
@Override
1974+
public BlockLoader.Block build() {
1975+
return builder.build();
1976+
}
1977+
1978+
@Override
1979+
public BlockLoader.Builder appendNull() {
1980+
throw new UnsupportedOperationException();
1981+
}
1982+
1983+
@Override
1984+
public BlockLoader.Builder beginPositionEntry() {
1985+
throw new UnsupportedOperationException();
1986+
}
1987+
1988+
@Override
1989+
public BlockLoader.Builder endPositionEntry() {
1990+
throw new UnsupportedOperationException();
1991+
}
1992+
1993+
@Override
1994+
public void close() {
1995+
builder.close();
1996+
}
1997+
}
1998+
19441999
}

‎server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ static class SingletonLongs extends BlockDocValuesReader implements NumericDocVa
137137
@Override
138138
public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
139139
if (numericDocValues instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
140-
BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, null);
140+
BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, null, false);
141141
if (result != null) {
142142
return result;
143143
}
@@ -262,7 +262,7 @@ public AllReader reader(LeafReaderContext context) throws IOException {
262262
}
263263
}
264264

265-
private static class SingletonInts extends BlockDocValuesReader {
265+
static class SingletonInts extends BlockDocValuesReader implements NumericDocValuesAccessor {
266266
private final NumericDocValues numericDocValues;
267267

268268
SingletonInts(NumericDocValues numericDocValues) {
@@ -271,6 +271,12 @@ private static class SingletonInts extends BlockDocValuesReader {
271271

272272
@Override
273273
public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
274+
if (numericDocValues instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
275+
BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, null, true);
276+
if (result != null) {
277+
return result;
278+
}
279+
}
274280
try (BlockLoader.IntBuilder builder = factory.intsFromDocValues(docs.count() - offset)) {
275281
for (int i = offset; i < docs.count(); i++) {
276282
int doc = docs.get(i);
@@ -303,9 +309,14 @@ public int docId() {
303309
public String toString() {
304310
return "BlockDocValuesReader.SingletonInts";
305311
}
312+
313+
@Override
314+
public NumericDocValues numericDocValues() {
315+
return numericDocValues;
316+
}
306317
}
307318

308-
private static class Ints extends BlockDocValuesReader {
319+
static class Ints extends BlockDocValuesReader {
309320
private final SortedNumericDocValues numericDocValues;
310321

311322
Ints(SortedNumericDocValues numericDocValues) {
@@ -409,7 +420,7 @@ static class SingletonDoubles extends BlockDocValuesReader implements NumericDoc
409420
@Override
410421
public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
411422
if (docValues instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
412-
BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, toDouble);
423+
BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, toDouble, false);
413424
if (result != null) {
414425
return result;
415426
}
@@ -764,7 +775,7 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boole
764775
return readSingleDoc(factory, docs.get(offset));
765776
}
766777
if (ordinals instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
767-
BlockLoader.Block block = direct.tryRead(factory, docs, offset, nullsFiltered, null);
778+
BlockLoader.Block block = direct.tryRead(factory, docs, offset, nullsFiltered, null, false);
768779
if (block != null) {
769780
return block;
770781
}

‎server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,17 @@ interface OptionalColumnAtATimeReader {
6868
*
6969
* @param nullsFiltered if {@code true}, then target docs are guaranteed to have a value for the field.
7070
* see {@link ColumnAtATimeReader#read(BlockFactory, Docs, int, boolean)}
71-
* @param toDouble a function to convert long values to double, or null if no conversion is needed/supported
71+
* @param toDouble a function to convert long values to double, or null if no conversion is needed/supported
72+
* @param toInt whether to convert to int in case int block / vector is needed
7273
*/
7374
@Nullable
7475
BlockLoader.Block tryRead(
7576
BlockFactory factory,
7677
Docs docs,
7778
int offset,
7879
boolean nullsFiltered,
79-
BlockDocValuesReader.ToDouble toDouble
80+
BlockDocValuesReader.ToDouble toDouble,
81+
boolean toInt
8082
) throws IOException;
8183
}
8284

@@ -445,6 +447,17 @@ interface BlockFactory {
445447
*/
446448
SingletonLongBuilder singletonLongs(int expectedCount);
447449

450+
/**
451+
* Build a specialized builder for singleton dense int based fields with the following constraints:
452+
* <ul>
453+
* <li>Only one value per document can be collected</li>
454+
* <li>No more than expectedCount values can be collected</li>
455+
* </ul>
456+
*
457+
* @param expectedCount The maximum number of values to be collected.
458+
*/
459+
SingletonIntBuilder singletonInts(int expectedCount);
460+
448461
/**
449462
* Build a specialized builder for singleton dense double based fields with the following constraints:
450463
* <ul>
@@ -570,6 +583,13 @@ interface SingletonDoubleBuilder extends Builder {
570583
SingletonDoubleBuilder appendLongs(BlockDocValuesReader.ToDouble toDouble, long[] values, int from, int length);
571584
}
572585

586+
/**
587+
* Specialized builder for collecting dense arrays of double values.
588+
*/
589+
interface SingletonIntBuilder extends Builder {
590+
SingletonIntBuilder appendLongs(long[] values, int from, int length);
591+
}
592+
573593
interface LongBuilder extends Builder {
574594
/**
575595
* Appends a long to the current entry.

0 commit comments

Comments
 (0)