Skip to content

Commit 743d0c8

Browse files
martijnvgmartijnvg
committed
add TSIDOrdinalsBuilder
1 parent abd9117 commit 743d0c8

File tree

8 files changed

+357
-3
lines changed

8 files changed

+357
-3
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1325,7 +1325,7 @@ public void loadBlock(BlockLoader.SingletonLongBuilder builder, BlockLoader.Docs
13251325
}
13261326

13271327
@Override
1328-
public void loadBlock(BlockLoader.SingletonOrdinalsBuilder builder, BlockLoader.Docs docs, int offset)
1328+
public void loadBlock(BlockLoader.TSIDOrdinalsBuilder builder, BlockLoader.Docs docs, int offset)
13291329
throws IOException {
13301330
assert maxOrd >= 0;
13311331
for (int i = offset; i < docs.count(); i++) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public interface SingletonDocValuesBlockLoader {
1717

1818
void loadBlock(BlockLoader.SingletonLongBuilder builder, BlockLoader.Docs docs, int offset) throws IOException;
1919

20-
void loadBlock(BlockLoader.SingletonOrdinalsBuilder builder, BlockLoader.Docs docs, int offset) throws IOException;
20+
void loadBlock(BlockLoader.TSIDOrdinalsBuilder builder, BlockLoader.Docs docs, int offset) throws IOException;
2121

2222
int docID();
2323

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,8 @@ interface BlockFactory {
423423
*/
424424
SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count);
425425

426+
TSIDOrdinalsBuilder tsidOrdinalsBuilder(SortedDocValues ordinals, int count);
427+
426428
/**
427429
* Build a reader for reading {@link SortedSetDocValues}
428430
*/
@@ -521,6 +523,14 @@ interface SingletonOrdinalsBuilder extends Builder {
521523
SingletonOrdinalsBuilder appendOrd(int value);
522524
}
523525

526+
interface TSIDOrdinalsBuilder extends Builder {
527+
528+
TSIDOrdinalsBuilder appendOrd(long value);
529+
530+
TSIDOrdinalsBuilder appendOrds(long[] values, int from, int length);
531+
532+
}
533+
524534
interface SortedSetOrdinalsBuilder extends Builder {
525535
/**
526536
* Appends an ordinal to the builder.

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public static final class TSIDs implements ColumnAtATimeReader {
6868

6969
@Override
7070
public Block read(BlockFactory factory, Docs docs, int offset) throws IOException {
71-
try (SingletonOrdinalsBuilder builder = factory.singletonOrdinalsBuilder(sorted, docs.count() - offset)) {
71+
try (TSIDOrdinalsBuilder builder = factory.tsidOrdinalsBuilder(sorted, docs.count() - offset)) {
7272
blockLoader.loadBlock(builder, docs, offset);
7373
return builder.build();
7474
}

‎test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,61 @@ public BlockLoader.Block constantBytes(BytesRef value, int count) {
271271
return builder.build();
272272
}
273273

274+
@Override
275+
public BlockLoader.TSIDOrdinalsBuilder tsidOrdinalsBuilder(SortedDocValues ordinals, int expectedCount) {
276+
final long[] ords = new long[expectedCount];
277+
return new BlockLoader.TSIDOrdinalsBuilder() {
278+
279+
int count;
280+
281+
@Override
282+
public BlockLoader.TSIDOrdinalsBuilder appendOrd(long value) {
283+
ords[count++] = value;
284+
return this;
285+
}
286+
287+
@Override
288+
public BlockLoader.TSIDOrdinalsBuilder appendOrds(long[] values, int from, int length) {
289+
try {
290+
System.arraycopy(values, from, ords, expectedCount, length);
291+
} catch (ArrayIndexOutOfBoundsException e) {
292+
throw e;
293+
}
294+
count += length;
295+
return this;
296+
}
297+
298+
@Override
299+
public BlockLoader.Block build() {
300+
return new TestBlock(Arrays.stream(ords).mapToInt(Math::toIntExact).mapToObj(ord -> {
301+
try {
302+
return BytesRef.deepCopyOf(ordinals.lookupOrd(ord));
303+
} catch (IOException e) {
304+
throw new UncheckedIOException(e);
305+
}
306+
}).collect(Collectors.toUnmodifiableList()));
307+
}
308+
309+
@Override
310+
public BlockLoader.Builder appendNull() {
311+
throw new UnsupportedOperationException();
312+
}
313+
314+
@Override
315+
public BlockLoader.Builder beginPositionEntry() {
316+
throw new UnsupportedOperationException();
317+
}
318+
319+
@Override
320+
public BlockLoader.Builder endPositionEntry() {
321+
throw new UnsupportedOperationException();
322+
}
323+
324+
@Override
325+
public void close() {}
326+
};
327+
}
328+
274329
@Override
275330
public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int expectedCount) {
276331
class SingletonOrdsBuilder extends TestBlock.Builder implements BlockLoader.SingletonOrdinalsBuilder {

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,11 @@ public BlockLoader.Builder nulls(int expectedCount) {
8686
return ElementType.NULL.newBlockBuilder(expectedCount, factory);
8787
}
8888

89+
@Override
90+
public BlockLoader.TSIDOrdinalsBuilder tsidOrdinalsBuilder(SortedDocValues ordinals, int count) {
91+
return new TSIDOrdinalsBuilder(factory, ordinals, count);
92+
}
93+
8994
@Override
9095
public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count) {
9196
return new SingletonOrdinalsBuilder(factory, ordinals, count);

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

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
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; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.compute.lucene.read;
9+
10+
import org.apache.lucene.index.SortedDocValues;
11+
import org.apache.lucene.util.RamUsageEstimator;
12+
import org.elasticsearch.compute.data.Block;
13+
import org.elasticsearch.compute.data.BlockFactory;
14+
import org.elasticsearch.compute.data.BytesRefBlock;
15+
import org.elasticsearch.compute.data.BytesRefVector;
16+
import org.elasticsearch.compute.data.IntBlock;
17+
import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
18+
import org.elasticsearch.core.Releasable;
19+
import org.elasticsearch.core.Releasables;
20+
import org.elasticsearch.index.mapper.BlockLoader;
21+
22+
import java.io.IOException;
23+
import java.io.UncheckedIOException;
24+
25+
/**
26+
* Fork of {@link SingletonOrdinalsBuilder} but specialized for _tsid field.
27+
* The _tsid field is dense and used as primary sort field, and therefore we optimize a little more.
28+
* Additionally, this implementation can collect complete value blocks.
29+
*/
30+
public class TSIDOrdinalsBuilder implements BlockLoader.TSIDOrdinalsBuilder, Releasable, Block.Builder {
31+
private final BlockFactory blockFactory;
32+
private final SortedDocValues docValues;
33+
private final long[] ords;
34+
private int count;
35+
36+
public TSIDOrdinalsBuilder(BlockFactory blockFactory, SortedDocValues docValues, int initialSize) {
37+
this.blockFactory = blockFactory;
38+
this.docValues = docValues;
39+
blockFactory.adjustBreaker(nativeOrdsSize(initialSize));
40+
// tsdb codec uses long array natively. Use this array the capture native encoded ordinals and then later in build() convert:
41+
this.ords = new long[initialSize];
42+
}
43+
44+
@Override
45+
public TSIDOrdinalsBuilder appendOrd(long ord) {
46+
ords[count++] = ord;
47+
return this;
48+
}
49+
50+
@Override
51+
public BlockLoader.TSIDOrdinalsBuilder appendOrds(long[] values, int from, int length) {
52+
System.arraycopy(values, from, ords, count, length);
53+
count += length;
54+
return this;
55+
}
56+
57+
@Override
58+
public long estimatedBytes() {
59+
return (long) ords.length * Long.BYTES;
60+
}
61+
62+
@Override
63+
public BytesRefBlock build() {
64+
// TODO: detect when constant block can be used like is done in: https://github.com/elastic/elasticsearch/pull/132456
65+
return buildOrdinal();
66+
}
67+
68+
BytesRefBlock buildOrdinal() {
69+
assert ords.length == count;
70+
71+
int minOrd = Math.toIntExact(ords[0]);
72+
int maxOrd = Math.toIntExact(ords[count - 1]);
73+
int valueCount = maxOrd - minOrd + 1;
74+
75+
long breakerSize = ordsSize(count);
76+
blockFactory.adjustBreaker(breakerSize);
77+
78+
BytesRefVector bytesVector = null;
79+
IntBlock ordinalBlock = null;
80+
try {
81+
// Convert back to int[] and remap ordinals
82+
int[] newOrds = new int[count];
83+
for (int i = 0; i < count; i++) {
84+
newOrds[i] = Math.toIntExact(ords[i]) - minOrd;
85+
}
86+
try (BytesRefVector.Builder bytesBuilder = blockFactory.newBytesRefVectorBuilder(valueCount)) {
87+
for (int ord = minOrd; ord <= maxOrd; ord++) {
88+
bytesBuilder.appendBytesRef(docValues.lookupOrd(ord));
89+
}
90+
bytesVector = bytesBuilder.build();
91+
} catch (IOException e) {
92+
throw new UncheckedIOException("error resolving tsid ordinals", e);
93+
}
94+
ordinalBlock = blockFactory.newIntArrayVector(newOrds, newOrds.length).asBlock();
95+
final OrdinalBytesRefBlock result = new OrdinalBytesRefBlock(ordinalBlock, bytesVector);
96+
assert ords.length == result.getPositionCount();
97+
98+
bytesVector = null;
99+
ordinalBlock = null;
100+
return result;
101+
} finally {
102+
Releasables.close(() -> blockFactory.adjustBreaker(-breakerSize), ordinalBlock, bytesVector);
103+
}
104+
}
105+
106+
@Override
107+
public void close() {
108+
blockFactory.adjustBreaker(-nativeOrdsSize(ords.length));
109+
}
110+
111+
@Override
112+
public TSIDOrdinalsBuilder appendNull() {
113+
throw new UnsupportedOperationException();
114+
}
115+
116+
@Override
117+
public TSIDOrdinalsBuilder beginPositionEntry() {
118+
throw new UnsupportedOperationException("should only have one value per doc");
119+
}
120+
121+
@Override
122+
public TSIDOrdinalsBuilder endPositionEntry() {
123+
throw new UnsupportedOperationException("should only have one value per doc");
124+
}
125+
126+
@Override
127+
public Block.Builder copyFrom(Block block, int beginInclusive, int endExclusive) {
128+
throw new UnsupportedOperationException();
129+
}
130+
131+
@Override
132+
public Block.Builder mvOrdering(Block.MvOrdering mvOrdering) {
133+
throw new UnsupportedOperationException();
134+
}
135+
136+
private static long ordsSize(int ordsCount) {
137+
return RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) ordsCount * Integer.BYTES;
138+
}
139+
140+
private static long nativeOrdsSize(int ordsCount) {
141+
return RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + (long) ordsCount * Long.BYTES;
142+
}
143+
144+
}

0 commit comments

Comments
 (0)