Skip to content

Commit 53a9125

Browse files
carlosdelestcarlosdelest
committed
Fixing, adding dimensions
1 parent acbdea2 commit 53a9125

28 files changed

+423
-118
lines changed

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

Lines changed: 25 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -569,7 +569,7 @@ private abstract static class DenseVectorValuesBlockReader<T extends KnnVectorVa
569569
@Override
570570
public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
571571
// Doubles from doc values ensures that the values are in order
572-
try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count() - offset, dimensions)) {
572+
try (BlockLoader.DenseVectorBuilder builder = factory.denseVectors(docs.count() - offset, dimensions)) {
573573
for (int i = offset; i < docs.count(); i++) {
574574
read(docs.get(i), builder);
575575
}
@@ -579,24 +579,22 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boole
579579

580580
@Override
581581
public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
582-
read(docId, (BlockLoader.FloatBuilder) builder);
582+
read(docId, (BlockLoader.DenseVectorBuilder) builder);
583583
}
584584

585-
private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
585+
private void read(int doc, BlockLoader.DenseVectorBuilder builder) throws IOException {
586586
assertDimensions();
587587

588588
if (iterator.docID() > doc) {
589589
builder.appendNull();
590590
} else if (iterator.docID() == doc || iterator.advance(doc) == doc) {
591-
builder.beginPositionEntry();
592591
appendDoc(builder);
593-
builder.endPositionEntry();
594592
} else {
595593
builder.appendNull();
596594
}
597595
}
598596

599-
protected abstract void appendDoc(BlockLoader.FloatBuilder builder) throws IOException;
597+
protected abstract void appendDoc(BlockLoader.DenseVectorBuilder builder) throws IOException;
600598

601599
@Override
602600
public int docId() {
@@ -615,11 +613,9 @@ private static class FloatDenseVectorValuesBlockReader extends DenseVectorValues
615613
super(floatVectorValues, dimensions);
616614
}
617615

618-
protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
616+
protected void appendDoc(BlockLoader.DenseVectorBuilder builder) throws IOException {
619617
float[] floats = vectorValues.vectorValue(iterator.index());
620-
for (float aFloat : floats) {
621-
builder.appendFloat(aFloat);
622-
}
618+
builder.appendDenseVector(floats);
623619
}
624620

625621
@Override
@@ -641,17 +637,15 @@ private static class FloatDenseVectorNormalizedValuesBlockReader extends DenseVe
641637
}
642638

643639
@Override
644-
protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
640+
protected void appendDoc(BlockLoader.DenseVectorBuilder builder) throws IOException {
645641
float magnitude = 1.0f;
646642
// If all vectors are normalized, no doc values will be present. The vector may be normalized already, so we may not have a
647643
// stored magnitude for all docs
648644
if ((magnitudeDocValues != null) && magnitudeDocValues.advanceExact(iterator.docID())) {
649645
magnitude = Float.intBitsToFloat((int) magnitudeDocValues.longValue());
650646
}
651647
float[] floats = vectorValues.vectorValue(iterator.index());
652-
for (float aFloat : floats) {
653-
builder.appendFloat(aFloat * magnitude);
654-
}
648+
builder.appendDenseVector(floats);
655649
}
656650

657651
@Override
@@ -665,11 +659,14 @@ private static class ByteDenseVectorValuesBlockReader extends DenseVectorValuesB
665659
super(floatVectorValues, dimensions);
666660
}
667661

668-
protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
662+
protected void appendDoc(BlockLoader.DenseVectorBuilder builder) throws IOException {
669663
byte[] bytes = vectorValues.vectorValue(iterator.index());
670-
for (byte aFloat : bytes) {
671-
builder.appendFloat(aFloat);
664+
float[] floats = new float[bytes.length];
665+
for (int i = 0; i < bytes.length; i++) {
666+
floats[i] = bytes[i];
672667
}
668+
669+
builder.appendDenseVector(floats);
673670
}
674671

675672
@Override
@@ -1068,12 +1065,12 @@ public int docId() {
10681065

10691066
@Override
10701067
public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
1071-
read(docId, (BlockLoader.FloatBuilder) builder);
1068+
read(docId, (BlockLoader.DenseVectorBuilder) builder);
10721069
}
10731070

10741071
@Override
10751072
public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
1076-
try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count() - offset, dimensions)) {
1073+
try (BlockLoader.DenseVectorBuilder builder = factory.denseVectors(docs.count() - offset, dimensions)) {
10771074
for (int i = offset; i < docs.count(); i++) {
10781075
int doc = docs.get(i);
10791076
read(doc, builder);
@@ -1082,23 +1079,20 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boole
10821079
}
10831080
}
10841081

1085-
private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
1082+
private void read(int doc, BlockLoader.DenseVectorBuilder builder) throws IOException {
10861083
if (docValues.advanceExact(doc) == false) {
10871084
builder.appendNull();
10881085
return;
10891086
}
10901087
BytesRef bytesRef = docValues.binaryValue();
10911088
assert bytesRef.length > 0;
10921089
decodeDenseVector(bytesRef, scratch);
1093-
1094-
builder.beginPositionEntry();
10951090
writeScratchToBuilder(scratch, builder);
1096-
builder.endPositionEntry();
10971091
}
10981092

10991093
protected abstract void decodeDenseVector(BytesRef bytesRef, T scratch);
11001094

1101-
protected abstract void writeScratchToBuilder(T scratch, BlockLoader.FloatBuilder builder);
1095+
protected abstract void writeScratchToBuilder(T scratch, BlockLoader.DenseVectorBuilder builder);
11021096
}
11031097

11041098
private static class FloatDenseVectorFromBinary extends AbstractDenseVectorFromBinary<float[]> {
@@ -1107,10 +1101,8 @@ private static class FloatDenseVectorFromBinary extends AbstractDenseVectorFromB
11071101
}
11081102

11091103
@Override
1110-
protected void writeScratchToBuilder(float[] scratch, BlockLoader.FloatBuilder builder) {
1111-
for (float value : scratch) {
1112-
builder.appendFloat(value);
1113-
}
1104+
protected void writeScratchToBuilder(float[] scratch, BlockLoader.DenseVectorBuilder builder) {
1105+
builder.appendDenseVector(scratch.clone());
11141106
}
11151107

11161108
@Override
@@ -1138,10 +1130,12 @@ public String toString() {
11381130
return "ByteDenseVectorFromBinary.Bytes";
11391131
}
11401132

1141-
protected void writeScratchToBuilder(byte[] scratch, BlockLoader.FloatBuilder builder) {
1142-
for (byte value : scratch) {
1143-
builder.appendFloat(value);
1133+
protected void writeScratchToBuilder(byte[] scratch, BlockLoader.DenseVectorBuilder builder) {
1134+
float[] vector = new float[scratch.length];
1135+
for (int i = 0; i < scratch.length; i++) {
1136+
vector[i] = scratch[i];
11441137
}
1138+
builder.appendDenseVector(vector);
11451139
}
11461140

11471141
protected void decodeDenseVector(BytesRef bytesRef, byte[] scratch) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ interface BlockFactory {
410410
/**
411411
* Build a builder to load dense vectors without any loading constraints.
412412
*/
413-
FloatBuilder denseVectors(int expectedVectorsCount, int dimensions);
413+
DenseVectorBuilder denseVectors(int expectedVectorsCount, int dimensions);
414414

415415
/**
416416
* Build a builder to load ints as loaded from doc values.

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

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -121,40 +121,43 @@ public DoublesBuilder appendDouble(double value) {
121121
}
122122

123123
@Override
124-
public BlockLoader.FloatBuilder denseVectors(int expectedCount, int dimensions) {
125-
class FloatsBuilder extends TestBlock.Builder implements BlockLoader.FloatBuilder {
124+
public BlockLoader.DenseVectorBuilder denseVectors(int expectedCount, int dimensions) {
125+
class DenseBuilder extends TestBlock.Builder implements BlockLoader.DenseVectorBuilder {
126126
int numElements = 0;
127127

128-
private FloatsBuilder() {
128+
private DenseBuilder() {
129129
super(expectedCount);
130130
}
131131

132-
@Override
133-
public BlockLoader.FloatBuilder appendFloat(float value) {
134-
add(value);
135-
numElements++;
136-
return this;
137-
}
138-
139132
@Override
140133
public Builder appendNull() {
141134
throw new IllegalArgumentException("dense vectors should not have null values");
142135
}
143136

144137
@Override
145138
public Builder endPositionEntry() {
146-
assert numElements == dimensions : "expected " + dimensions + " dimensions, but got " + numElements;
147-
numElements = 0;
148-
return super.endPositionEntry();
139+
throw new UnsupportedOperationException();
140+
}
141+
142+
@Override
143+
public Builder beginPositionEntry() {
144+
throw new UnsupportedOperationException();
149145
}
150146

151147
@Override
152148
public TestBlock build() {
153149
assert numElements == 0 : "endPositionEntry() was not called for the last entry";
154150
return super.build();
155151
}
152+
153+
@Override
154+
public BlockLoader.DenseVectorBuilder appendDenseVector(float[] value) {
155+
add(value);
156+
numElements++;
157+
return this;
158+
}
156159
}
157-
return new FloatsBuilder();
160+
return new DenseBuilder();
158161
}
159162

160163
@Override

‎x-pack/plugin/esql/compute/src/main/java/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
requires org.elasticsearch.geo;
2222
requires org.elasticsearch.xcore;
2323
requires hppc;
24+
requires unboundid.ldapsdk;
2425

2526
exports org.elasticsearch.compute;
2627
exports org.elasticsearch.compute.aggregation;

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -492,15 +492,8 @@ public DenseVectorBlock.Builder newDenseVectorBlockBuilder(int estimatedSize) {
492492
return new DenseVectorBlockBuilder(estimatedSize, this);
493493
}
494494

495-
public final DenseVectorBlock newDenseVectorArrayBlock(
496-
float[][] values,
497-
int dimensions,
498-
int positionCount,
499-
int[] firstValueIndexes,
500-
BitSet nulls,
501-
MvOrdering mvOrdering
502-
) {
503-
return newDenseVectorArrayBlock(values, dimensions, positionCount, firstValueIndexes, nulls, mvOrdering, 0L);
495+
public DenseVectorBlock.Builder newDenseVectorBlockBuilder(int estimatedSize, int dimensions) {
496+
return new DenseVectorBlockBuilder(estimatedSize, this);
504497
}
505498

506499
public DenseVectorBlock newDenseVectorArrayBlock(
@@ -512,20 +505,20 @@ public DenseVectorBlock newDenseVectorArrayBlock(
512505
MvOrdering mvOrdering,
513506
long preAdjustedBytes
514507
) {
515-
var b = new DenseVectorArrayBlock(values, pc, fvi, nulls, mvOrdering, this);
508+
var b = new DenseVectorArrayBlock(values, pc, dimensions, fvi, nulls, mvOrdering, this);
516509
adjustBreaker(b.ramBytesUsed() - preAdjustedBytes);
517510
return b;
518511
}
519512

520513
public DenseVectorVector.Builder newDenseVectorVectorBuilder(int estimatedSize, int dimensions) {
521-
return new DenseVectorVectorBuilder(dimensions, estimatedSize, this);
514+
return new DenseVectorVectorBuilder(estimatedSize, dimensions, this);
522515
}
523516

524517
/**
525518
* Build a {@link DenseVectorVector.FixedBuilder} that never grows.
526519
*/
527520
public DenseVectorVector.FixedBuilder newDenseVectorVectorFixedBuilder(int dimensions, int size) {
528-
return new DenseVectorVectorFixedBuilder(dimensions, size, this);
521+
return new DenseVectorVectorFixedBuilder(size, dimensions, this);
529522
}
530523

531524
/**
@@ -539,8 +532,8 @@ public DenseVectorVector.FixedBuilder newDenseVectorVectorFixedBuilder(int dimen
539532
* }
540533
* var vector = blockFactory.newDenseVectorArrayVector(values, positionCount, preAdjustedBytes);
541534
*/
542-
public DenseVectorVector newDenseVectorArrayVector(float[][] values, int positionCount, long preAdjustedBytes) {
543-
var b = new DenseVectorArrayVector(values, positionCount, this);
535+
public DenseVectorVector newDenseVectorArrayVector(float[][] values, int positionCount, int dimensions, long preAdjustedBytes) {
536+
var b = new DenseVectorArrayVector(values, positionCount, dimensions,this);
544537
adjustBreaker(b.ramBytesUsed() - preAdjustedBytes);
545538
return b;
546539
}
@@ -556,9 +549,8 @@ public DenseVectorBlock newConstantDenseVectorBlockWith(float[] value, int posit
556549
}
557550

558551
public DenseVectorVector newConstantDenseVectorVector(float[] value, int positions) {
559-
adjustBreaker(ConstantDenseVectorVector.RAM_BYTES_USED);
560552
var v = new ConstantDenseVectorVector(value, positions, this);
561-
assert v.ramBytesUsed() == ConstantDenseVectorVector.RAM_BYTES_USED;
553+
adjustBreaker(v.ramBytesUsed());
562554
return v;
563555
}
564556
}

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantDenseVectorVector.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,6 @@
1919
*/
2020
final class ConstantDenseVectorVector extends AbstractVector implements DenseVectorVector {
2121

22-
static final long RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(ConstantDenseVectorVector.class);
23-
2422
private final float[] value;
2523

2624
ConstantDenseVectorVector(float[] value, int positionCount, BlockFactory blockFactory) {
@@ -56,7 +54,7 @@ public DenseVectorBlock keepMask(BooleanVector mask) {
5654
}
5755
return (DenseVectorBlock) blockFactory().newConstantNullBlock(getPositionCount());
5856
}
59-
try (DenseVectorBlock.Builder builder = blockFactory().newDenseVectorBlockBuilder(getPositionCount())) {
57+
try (DenseVectorBlock.Builder builder = blockFactory().newDenseVectorBlockBuilder(getPositionCount(), value.length)) {
6058
// TODO if X-ArrayBlock used BooleanVector for it's null mask then we could shuffle references here.
6159
for (int p = 0; p < getPositionCount(); p++) {
6260
if (mask.getBoolean(p)) {
@@ -110,7 +108,7 @@ public DenseVectorVector deepCopy(BlockFactory blockFactory) {
110108

111109
@Override
112110
public long ramBytesUsed() {
113-
return RAM_BYTES_USED;
111+
return RamUsageEstimator.sizeOf(value);
114112
}
115113

116114
@Override

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,12 @@ public float[] getDenseVector(int valueIndex) {
299299
throw new UnsupportedOperationException("null block");
300300
}
301301

302+
@Override
303+
public int dimensions() {
304+
assert false : "null block";
305+
throw new UnsupportedOperationException("null block");
306+
}
307+
302308
@Override
303309
public int getTotalValueCount() {
304310
return 0;

0 commit comments

Comments
 (0)