Skip to content

Commit 41b08c9

Browse files
dnhatndnhatn
committed
Speed up reading dimension fields in TS
1 parent f7a7f0b commit 41b08c9

File tree

5 files changed

+88
-16
lines changed

5 files changed

+88
-16
lines changed

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/ShardContext.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,4 +53,10 @@ public interface ShardContext {
5353
* Returns something to load values from this field into a {@link Block}.
5454
*/
5555
BlockLoader blockLoader(String name, boolean asUnsupportedSource, MappedFieldType.FieldExtractPreference fieldExtractPreference);
56+
57+
/**
58+
* Returns the {@link MappedFieldType} for the given field name.
59+
* By default, this delegate to {@link org.elasticsearch.index.query.SearchExecutionContext#getFieldType(String)}
60+
*/
61+
MappedFieldType fieldType(String name);
5662
}

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperator.java‎

Lines changed: 63 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -133,11 +133,12 @@ public Page getCheckedOutput() throws IOException {
133133
if (docCollector != null) {
134134
blocks[blockIndex++] = docCollector.build().asBlock();
135135
}
136-
blocks[blockIndex++] = tsHashesBuilder.build().asBlock();
136+
OrdinalBytesRefVector tsidVector = tsHashesBuilder.build();
137+
blocks[blockIndex++] = tsidVector.asBlock();
137138
tsHashesBuilder = new TsidBuilder(blockFactory, Math.min(remainingDocs, maxPageSize));
138139
blocks[blockIndex++] = timestampsBuilder.build().asBlock();
139140
timestampsBuilder = blockFactory.newLongVectorBuilder(Math.min(remainingDocs, maxPageSize));
140-
System.arraycopy(fieldsReader.buildBlocks(), 0, blocks, blockIndex, fieldsToExtracts.size());
141+
System.arraycopy(fieldsReader.buildBlocks(tsidVector.getOrdinalsVector()), 0, blocks, blockIndex, fieldsToExtracts.size());
141142
page = new Page(currentPagePos, blocks);
142143
currentPagePos = 0;
143144
}
@@ -217,6 +218,7 @@ void readDocsForNextPage() throws IOException {
217218
}
218219

219220
private boolean readValuesForOneTsid(PriorityQueue<LeafIterator> sub) throws IOException {
221+
boolean first = true;
220222
do {
221223
LeafIterator top = sub.top();
222224
currentPagePos++;
@@ -226,7 +228,8 @@ private boolean readValuesForOneTsid(PriorityQueue<LeafIterator> sub) throws IOE
226228
}
227229
tsHashesBuilder.appendOrdinal();
228230
timestampsBuilder.appendLong(top.timestamp);
229-
fieldsReader.readValues(top.segmentOrd, top.docID);
231+
fieldsReader.readValues(top.segmentOrd, top.docID, first == false);
232+
first = false;
230233
if (top.nextDoc()) {
231234
sub.updateTop();
232235
} else if (top.docID == DocIdSetIterator.NO_MORE_DOCS) {
@@ -350,6 +353,7 @@ static final class ShardLevelFieldsReader implements Releasable {
350353
private final BlockLoaderFactory blockFactory;
351354
private final SegmentLevelFieldsReader[] segments;
352355
private final BlockLoader[] loaders;
356+
private final boolean[] dimensions;
353357
private final Block.Builder[] builders;
354358
private final StoredFieldsSpec storedFieldsSpec;
355359
private final SourceLoader sourceLoader;
@@ -377,10 +381,14 @@ static final class ShardLevelFieldsReader implements Releasable {
377381
sourceLoader = null;
378382
}
379383
this.storedFieldsSpec = storedFieldsSpec;
384+
this.dimensions = new boolean[fields.size()];
385+
for (int i = 0; i < fields.size(); i++) {
386+
dimensions[i] = shardContext.fieldType(fields.get(i).name()).isDimension();
387+
}
380388
}
381389

382-
void readValues(int segment, int docID) throws IOException {
383-
segments[segment].read(docID, builders);
390+
void readValues(int segment, int docID, boolean nonDimensionFieldsOnly) throws IOException {
391+
segments[segment].read(docID, builders, nonDimensionFieldsOnly, dimensions);
384392
}
385393

386394
void prepareForReading(int estimatedSize) throws IOException {
@@ -396,12 +404,46 @@ void prepareForReading(int estimatedSize) throws IOException {
396404
}
397405
}
398406

399-
Block[] buildBlocks() {
400-
Block[] blocks = Block.Builder.buildAll(builders);
401-
Arrays.fill(builders, null);
407+
Block[] buildBlocks(IntVector tsidOrdinals) {
408+
final Block[] blocks = new Block[loaders.length];
409+
try {
410+
for (int i = 0; i < builders.length; i++) {
411+
if (dimensions[i]) {
412+
blocks[i] = buildBlockForDimensionField(builders[i], tsidOrdinals);
413+
} else {
414+
blocks[i] = builders[i].build();
415+
}
416+
}
417+
Arrays.fill(builders, null);
418+
} finally {
419+
if (blocks.length > 0 && blocks[blocks.length - 1] == null) {
420+
Releasables.close(blocks);
421+
}
422+
}
402423
return blocks;
403424
}
404425

426+
private Block buildBlockForDimensionField(Block.Builder builder, IntVector tsidOrdinals) {
427+
try (var values = builder.build()) {
428+
if (values.asVector() instanceof BytesRefVector bytes) {
429+
tsidOrdinals.incRef();
430+
values.incRef();
431+
return new OrdinalBytesRefVector(tsidOrdinals, bytes).asBlock();
432+
} else if (values.areAllValuesNull()) {
433+
return blockFactory.factory.newConstantNullBlock(tsidOrdinals.getPositionCount());
434+
} else {
435+
final int positionCount = tsidOrdinals.getPositionCount();
436+
try (var newBuilder = values.elementType().newBlockBuilder(positionCount, blockFactory.factory)) {
437+
for (int p = 0; p < positionCount; p++) {
438+
int pos = tsidOrdinals.getInt(p);
439+
newBuilder.copyFrom(values, pos, pos + 1);
440+
}
441+
return newBuilder.build();
442+
}
443+
}
444+
}
445+
}
446+
405447
@Override
406448
public void close() {
407449
Releasables.close(builders);
@@ -435,10 +477,18 @@ private void reinitializeIfNeeded(SourceLoader sourceLoader, StoredFieldsSpec st
435477
}
436478
}
437479

438-
void read(int docId, Block.Builder[] builder) throws IOException {
480+
void read(int docId, Block.Builder[] builder, boolean nonDimensionFieldsOnly, boolean[] dimensions) throws IOException {
439481
storedFields.advanceTo(docId);
440-
for (int i = 0; i < rowStride.length; i++) {
441-
rowStride[i].read(docId, storedFields, builder[i]);
482+
if (nonDimensionFieldsOnly) {
483+
for (int i = 0; i < rowStride.length; i++) {
484+
if (dimensions[i] == false) {
485+
rowStride[i].read(docId, storedFields, builder[i]);
486+
}
487+
}
488+
} else {
489+
for (int i = 0; i < rowStride.length; i++) {
490+
rowStride[i].read(docId, storedFields, builder[i]);
491+
}
442492
}
443493
}
444494
}
@@ -480,9 +530,9 @@ public void close() {
480530
Releasables.close(dictBuilder, ordinalsBuilder);
481531
}
482532

483-
BytesRefVector build() throws IOException {
533+
OrdinalBytesRefVector build() throws IOException {
484534
BytesRefVector dict = null;
485-
BytesRefVector result = null;
535+
OrdinalBytesRefVector result = null;
486536
IntVector ordinals = null;
487537
try {
488538
dict = dictBuilder.build();

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,5 +324,10 @@ public Optional<SortAndFormats> buildSort(List<SortBuilder<?>> sorts) {
324324
public String shardIdentifier() {
325325
return "test";
326326
}
327+
328+
@Override
329+
public MappedFieldType fieldType(String name) {
330+
throw new UnsupportedOperationException();
331+
}
327332
}
328333
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperatorTests.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,17 @@ public static TimeSeriesSourceOperatorFactory createTimeSeriesSourceOperator(
438438
} catch (IOException e) {
439439
throw new UncheckedIOException(e);
440440
}
441-
var ctx = new LuceneSourceOperatorTests.MockShardContext(reader, 0);
441+
var ctx = new LuceneSourceOperatorTests.MockShardContext(reader, 0) {
442+
@Override
443+
public MappedFieldType fieldType(String name) {
444+
for (ExtractField e : extractFields) {
445+
if (e.ft.name().equals(name)) {
446+
return e.ft;
447+
}
448+
}
449+
throw new IllegalArgumentException("Unknown field [" + name + "]");
450+
}
451+
};
442452
Function<ShardContext, Query> queryFunction = c -> new MatchAllDocsQuery();
443453

444454
var fieldInfos = extractFields.stream()

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ private static class DefaultShardContextForUnmappedField extends DefaultShardCon
177177
}
178178

179179
@Override
180-
protected @Nullable MappedFieldType fieldType(String name) {
180+
public @Nullable MappedFieldType fieldType(String name) {
181181
var superResult = super.fieldType(name);
182182
return superResult == null && name.equals(unmappedEsField.getName())
183183
? new KeywordFieldMapper.KeywordFieldType(name, false /* isIndexed */, false /* hasDocValues */, Map.of() /* meta */)
@@ -459,7 +459,8 @@ public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() {
459459
return loader;
460460
}
461461

462-
protected @Nullable MappedFieldType fieldType(String name) {
462+
@Override
463+
public @Nullable MappedFieldType fieldType(String name) {
463464
return ctx.getFieldType(name);
464465
}
465466

0 commit comments

Comments
 (0)