Skip to content

Commit c951ee7

Browse files
carlosdelestcarlosdelest
committed
Use dense vector specific methods for creating builders, and check dimensions
1 parent f8741ca commit c951ee7

File tree

6 files changed

+100
-36
lines changed

6 files changed

+100
-36
lines changed

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

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -510,39 +510,43 @@ public String toString() {
510510

511511
public static class DenseVectorBlockLoader extends DocValuesBlockLoader {
512512
private final String fieldName;
513+
private final int dimensions;
513514

514-
public DenseVectorBlockLoader(String fieldName) {
515+
public DenseVectorBlockLoader(String fieldName, int dimensions) {
515516
this.fieldName = fieldName;
517+
this.dimensions = dimensions;
516518
}
517519

518520
@Override
519521
public Builder builder(BlockFactory factory, int expectedCount) {
520-
return factory.floats(expectedCount);
522+
return factory.denseVectors(expectedCount, dimensions);
521523
}
522524

523525
@Override
524526
public AllReader reader(LeafReaderContext context) throws IOException {
525527
FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
526528
if (floatVectorValues != null) {
527-
return new FloatVectorValuesBlockReader(floatVectorValues);
529+
return new DenseVectorValuesBlockReader(floatVectorValues, dimensions);
528530
}
529531
return new ConstantNullsReader();
530532
}
531533
}
532534

533-
private static class FloatVectorValuesBlockReader extends BlockDocValuesReader {
535+
private static class DenseVectorValuesBlockReader extends BlockDocValuesReader {
534536
private final FloatVectorValues floatVectorValues;
535537
private final KnnVectorValues.DocIndexIterator iterator;
538+
private final int dimensions;
536539

537-
FloatVectorValuesBlockReader(FloatVectorValues floatVectorValues) {
540+
DenseVectorValuesBlockReader(FloatVectorValues floatVectorValues, int dimensions) {
538541
this.floatVectorValues = floatVectorValues;
539542
iterator = floatVectorValues.iterator();
543+
this.dimensions = dimensions;
540544
}
541545

542546
@Override
543547
public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
544548
// Doubles from doc values ensures that the values are in order
545-
try (BlockLoader.FloatBuilder builder = factory.floats(docs.count())) {
549+
try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count(), dimensions)) {
546550
for (int i = 0; i < docs.count(); i++) {
547551
int doc = docs.get(i);
548552
if (doc < iterator.docID()) {
@@ -563,6 +567,8 @@ private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException
563567
if (iterator.advance(doc) == doc) {
564568
builder.beginPositionEntry();
565569
float[] floats = floatVectorValues.vectorValue(iterator.index());
570+
assert floats.length == dimensions
571+
: "unexpected dimensions for vector value; expected " + dimensions + " but got " + floats.length;
566572
for (float aFloat : floats) {
567573
builder.appendFloat(aFloat);
568574
}
@@ -844,7 +850,7 @@ public DenseVectorFromBinaryBlockLoader(String fieldName, int dims, IndexVersion
844850

845851
@Override
846852
public Builder builder(BlockFactory factory, int expectedCount) {
847-
return factory.floats(expectedCount);
853+
return factory.denseVectors(expectedCount, dims);
848854
}
849855

850856
@Override
@@ -860,6 +866,7 @@ public AllReader reader(LeafReaderContext context) throws IOException {
860866
private static class DenseVectorFromBinary extends BlockDocValuesReader {
861867
private final BinaryDocValues docValues;
862868
private final IndexVersion indexVersion;
869+
private final int dimensions;
863870
private final float[] scratch;
864871

865872
private int docID = -1;
@@ -868,11 +875,12 @@ private static class DenseVectorFromBinary extends BlockDocValuesReader {
868875
this.docValues = docValues;
869876
this.scratch = new float[dims];
870877
this.indexVersion = indexVersion;
878+
this.dimensions = dims;
871879
}
872880

873881
@Override
874882
public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
875-
try (BlockLoader.FloatBuilder builder = factory.floats(docs.count())) {
883+
try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count(), dimensions)) {
876884
for (int i = 0; i < docs.count(); i++) {
877885
int doc = docs.get(i);
878886
if (doc < docID) {

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -374,15 +374,9 @@ interface BlockFactory {
374374
DoubleBuilder doubles(int expectedCount);
375375

376376
/**
377-
* Build a builder to load floats as loaded from doc values.
378-
* Doc values load floats in sorted order.
377+
* Build a builder to load dense vectors without any loading constraints.
379378
*/
380-
FloatBuilder floatsFromDocValues(int expectedCount);
381-
382-
/**
383-
* Build a builder to load floats without any loading constraints.
384-
*/
385-
FloatBuilder floats(int expectedCount);
379+
FloatBuilder denseVectors(int expectedVectorsCount, int dimensions);
386380

387381
/**
388382
* Build a builder to load ints as loaded from doc values.

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -304,19 +304,22 @@ public String toString() {
304304
/**
305305
* Load {@code float}s from {@code _source}.
306306
*/
307-
public static class FloatsBlockLoader extends SourceBlockLoader {
308-
public FloatsBlockLoader(ValueFetcher fetcher, LeafIteratorLookup lookup) {
307+
public static class DenseVectorBlockLoader extends SourceBlockLoader {
308+
private final int dimensions;
309+
310+
public DenseVectorBlockLoader(ValueFetcher fetcher, LeafIteratorLookup lookup, int dimensions) {
309311
super(fetcher, lookup);
312+
this.dimensions = dimensions;
310313
}
311314

312315
@Override
313316
public Builder builder(BlockFactory factory, int expectedCount) {
314-
return factory.floats(expectedCount);
317+
return factory.denseVectors(expectedCount, dimensions);
315318
}
316319

317320
@Override
318321
public RowStrideReader rowStrideReader(LeafReaderContext context, DocIdSetIterator iter) {
319-
return new Floats(fetcher, iter);
322+
return new DenseVectors(fetcher, iter);
320323
}
321324

322325
@Override
@@ -325,8 +328,8 @@ protected String name() {
325328
}
326329
}
327330

328-
private static class Floats extends BlockSourceReader {
329-
Floats(ValueFetcher fetcher, DocIdSetIterator iter) {
331+
private static class DenseVectors extends BlockSourceReader {
332+
DenseVectors(ValueFetcher fetcher, DocIdSetIterator iter) {
330333
super(fetcher, iter);
331334
}
332335

‎server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
import org.elasticsearch.index.query.SearchExecutionContext;
7171
import org.elasticsearch.search.DocValueFormat;
7272
import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
73+
import org.elasticsearch.search.lookup.Source;
7374
import org.elasticsearch.search.vectors.DenseVectorQuery;
7475
import org.elasticsearch.search.vectors.ESDiversifyingChildrenByteKnnVectorQuery;
7576
import org.elasticsearch.search.vectors.ESDiversifyingChildrenFloatKnnVectorQuery;
@@ -90,6 +91,7 @@
9091
import java.time.ZoneId;
9192
import java.util.Arrays;
9293
import java.util.HexFormat;
94+
import java.util.List;
9395
import java.util.Locale;
9496
import java.util.Map;
9597
import java.util.Objects;
@@ -2330,15 +2332,15 @@ public BlockLoader blockLoader(MappedFieldType.BlockLoaderContext blContext) {
23302332
}
23312333

23322334
if (indexed) {
2333-
return new BlockDocValuesReader.DenseVectorBlockLoader(name());
2335+
return new BlockDocValuesReader.DenseVectorBlockLoader(name(), dims);
23342336
}
23352337

23362338
if (hasDocValues() && (blContext.fieldExtractPreference() != FieldExtractPreference.STORED || isSyntheticSource)) {
23372339
return new BlockDocValuesReader.DenseVectorFromBinaryBlockLoader(name(), dims, indexVersionCreated);
23382340
}
23392341

23402342
BlockSourceReader.LeafIteratorLookup lookup = BlockSourceReader.lookupMatchingAll();
2341-
return new BlockSourceReader.FloatsBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup);
2343+
return new BlockSourceReader.DenseVectorBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup, dims);
23422344
}
23432345

23442346
private SourceValueFetcher sourceValueFetcher(Set<String> sourcePaths) {
@@ -2350,6 +2352,13 @@ protected Object parseSourceValue(Object value) {
23502352
}
23512353
return NumberFieldMapper.NumberType.FLOAT.parse(value, false);
23522354
}
2355+
2356+
@Override
2357+
public List<Object> fetchValues(Source source, int doc, List<Object> ignoredValues) {
2358+
List<Object> result = super.fetchValues(source, doc, ignoredValues);
2359+
assert result.size() == dims : "Unexpected number of dimensions; got " + result.size() + " but expected " + dims;
2360+
return result;
2361+
}
23532362
};
23542363
}
23552364
}

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

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -76,18 +76,34 @@ public DoublesBuilder appendDouble(double value) {
7676
}
7777

7878
@Override
79-
public BlockLoader.FloatBuilder floatsFromDocValues(int expectedCount) {
80-
return floats(expectedCount);
81-
}
82-
83-
@Override
84-
public BlockLoader.FloatBuilder floats(int expectedCount) {
79+
public BlockLoader.FloatBuilder denseVectors(int expectedCount, int dimensions) {
8580
class FloatsBuilder extends TestBlock.Builder implements BlockLoader.FloatBuilder {
81+
int numElements = 0;
82+
8683
@Override
8784
public BlockLoader.FloatBuilder appendFloat(float value) {
8885
add(value);
86+
numElements++;
8987
return this;
9088
}
89+
90+
@Override
91+
public Builder appendNull() {
92+
throw new IllegalArgumentException("dense vectors should not have null values");
93+
}
94+
95+
@Override
96+
public Builder endPositionEntry() {
97+
assert numElements == dimensions : "expected " + dimensions + " dimensions, but got " + numElements;
98+
numElements = 0;
99+
return super.endPositionEntry();
100+
}
101+
102+
@Override
103+
public TestBlock build() {
104+
assert numElements == 0 : "endPositionEntry() was not called for the last entry";
105+
return super.build();
106+
}
91107
}
92108
return new FloatsBuilder();
93109
}

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

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.elasticsearch.compute.data.DocBlock;
2222
import org.elasticsearch.compute.data.DocVector;
2323
import org.elasticsearch.compute.data.ElementType;
24+
import org.elasticsearch.compute.data.FloatBlock;
2425
import org.elasticsearch.compute.data.IntVector;
2526
import org.elasticsearch.compute.data.Page;
2627
import org.elasticsearch.compute.data.SingletonOrdinalsBuilder;
@@ -655,13 +656,46 @@ public BlockLoader.DoubleBuilder doubles(int expectedCount) {
655656
}
656657

657658
@Override
658-
public BlockLoader.FloatBuilder floatsFromDocValues(int expectedCount) {
659-
return factory.newFloatBlockBuilder(expectedCount).mvOrdering(Block.MvOrdering.SORTED_ASCENDING);
660-
}
659+
public BlockLoader.FloatBuilder denseVectors(int expectedVectorsCount, int dimensions) {
660+
FloatBlock.Builder builder = factory.newFloatBlockBuilder(expectedVectorsCount * dimensions);
661661

662-
@Override
663-
public BlockLoader.FloatBuilder floats(int expectedCount) {
664-
return factory.newFloatBlockBuilder(expectedCount);
662+
return new BlockLoader.FloatBuilder() {
663+
int numElements = 0;
664+
665+
@Override
666+
public BlockLoader.FloatBuilder appendFloat(float value) {
667+
numElements++;
668+
return builder.appendFloat(value);
669+
}
670+
671+
@Override
672+
public BlockLoader.Block build() {
673+
assert numElements == 0 : "endPositionEntry() was not called before build()";
674+
return builder.build();
675+
}
676+
677+
@Override
678+
public BlockLoader.Builder appendNull() {
679+
throw new IllegalArgumentException("null values are not expected for dense vectors");
680+
}
681+
682+
@Override
683+
public BlockLoader.Builder beginPositionEntry() {
684+
return builder.beginPositionEntry();
685+
}
686+
687+
@Override
688+
public BlockLoader.Builder endPositionEntry() {
689+
assert numElements == dimensions : "expected " + dimensions + " dimensions, but got " + numElements + " elements";
690+
numElements = 0;
691+
return builder.endPositionEntry();
692+
}
693+
694+
@Override
695+
public void close() {
696+
builder.close();
697+
}
698+
};
665699
}
666700

667701
@Override

0 commit comments

Comments
 (0)