Skip to content

Commit cd462b8

Browse files
carlosdelestcarlosdelest
committed
Use Float as block type for dense_vector
1 parent a6b0f6c commit cd462b8

File tree

15 files changed

+128
-33
lines changed

15 files changed

+128
-33
lines changed

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -517,7 +517,7 @@ public DenseVectorBlockLoader(String fieldName) {
517517

518518
@Override
519519
public Builder builder(BlockFactory factory, int expectedCount) {
520-
return factory.doubles(expectedCount);
520+
return factory.floats(expectedCount);
521521
}
522522

523523
@Override
@@ -542,7 +542,7 @@ private static class FloatVectorValuesBlockReader extends BlockDocValuesReader {
542542
@Override
543543
public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
544544
// Doubles from doc values ensures that the values are in order
545-
try (BlockLoader.DoubleBuilder builder = factory.doublesFromDocValues(docs.count())) {
545+
try (BlockLoader.FloatBuilder builder = factory.floatsFromDocValues(docs.count())) {
546546
for (int i = 0; i < docs.count(); i++) {
547547
int doc = docs.get(i);
548548
if (doc < iterator.docID()) {
@@ -556,15 +556,15 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOExceptio
556556

557557
@Override
558558
public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
559-
read(docId, (DoubleBuilder) builder);
559+
read(docId, (BlockLoader.FloatBuilder) builder);
560560
}
561561

562-
private void read(int doc, DoubleBuilder builder) throws IOException {
562+
private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
563563
if (iterator.advance(doc) == doc) {
564564
builder.beginPositionEntry();
565565
float[] floats = floatVectorValues.vectorValue(iterator.index());
566566
for (float aFloat : floats) {
567-
builder.appendDouble(aFloat);
567+
builder.appendFloat(aFloat);
568568
}
569569
builder.endPositionEntry();
570570
} else {
@@ -844,7 +844,7 @@ public DenseVectorFromBinaryBlockLoader(String fieldName, int dims, IndexVersion
844844

845845
@Override
846846
public Builder builder(BlockFactory factory, int expectedCount) {
847-
return factory.doubles(expectedCount);
847+
return factory.floats(expectedCount);
848848
}
849849

850850
@Override
@@ -872,7 +872,7 @@ private static class DenseVectorFromBinary extends BlockDocValuesReader {
872872

873873
@Override
874874
public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
875-
try (BlockLoader.DoubleBuilder builder = factory.doubles(docs.count())) {
875+
try (BlockLoader.FloatBuilder builder = factory.floats(docs.count())) {
876876
for (int i = 0; i < docs.count(); i++) {
877877
int doc = docs.get(i);
878878
if (doc < docID) {
@@ -886,10 +886,10 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOExceptio
886886

887887
@Override
888888
public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
889-
read(docId, (DoubleBuilder) builder);
889+
read(docId, (BlockLoader.FloatBuilder) builder);
890890
}
891891

892-
private void read(int doc, DoubleBuilder builder) throws IOException {
892+
private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
893893
this.docID = doc;
894894
if (false == docValues.advanceExact(doc)) {
895895
builder.appendNull();
@@ -901,7 +901,7 @@ private void read(int doc, DoubleBuilder builder) throws IOException {
901901

902902
builder.beginPositionEntry();
903903
for (float value : scratch) {
904-
builder.appendDouble(value);
904+
builder.appendFloat(value);
905905
}
906906
builder.endPositionEntry();
907907
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,17 @@ interface BlockFactory {
373373
*/
374374
DoubleBuilder doubles(int expectedCount);
375375

376+
/**
377+
* Build a builder to load floats as loaded from doc values.
378+
* Doc values load floats in sorted order.
379+
*/
380+
FloatBuilder floatsFromDocValues(int expectedCount);
381+
382+
/**
383+
* Build a builder to load floats without any loading constraints.
384+
*/
385+
FloatBuilder floats(int expectedCount);
386+
376387
/**
377388
* Build a builder to load ints as loaded from doc values.
378389
* Doc values load ints in sorted order.

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,46 @@ public String toString() {
301301
}
302302
}
303303

304+
/**
305+
* Load {@code float}s from {@code _source}.
306+
*/
307+
public static class FloatsBlockLoader extends SourceBlockLoader {
308+
public FloatsBlockLoader(ValueFetcher fetcher, LeafIteratorLookup lookup) {
309+
super(fetcher, lookup);
310+
}
311+
312+
@Override
313+
public Builder builder(BlockFactory factory, int expectedCount) {
314+
return factory.floats(expectedCount);
315+
}
316+
317+
@Override
318+
public RowStrideReader rowStrideReader(LeafReaderContext context, DocIdSetIterator iter) {
319+
return new Floats(fetcher, iter);
320+
}
321+
322+
@Override
323+
protected String name() {
324+
return "Floats";
325+
}
326+
}
327+
328+
private static class Floats extends BlockSourceReader {
329+
Floats(ValueFetcher fetcher, DocIdSetIterator iter) {
330+
super(fetcher, iter);
331+
}
332+
333+
@Override
334+
protected void append(BlockLoader.Builder builder, Object v) {
335+
((BlockLoader.FloatBuilder) builder).appendFloat(((Number) v).floatValue());
336+
}
337+
338+
@Override
339+
public String toString() {
340+
return "BlockSourceReader.Floats";
341+
}
342+
}
343+
304344
/**
305345
* Load {@code int}s from {@code _source}.
306346
*/

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2338,7 +2338,7 @@ public BlockLoader blockLoader(MappedFieldType.BlockLoaderContext blContext) {
23382338
}
23392339

23402340
BlockSourceReader.LeafIteratorLookup lookup = BlockSourceReader.lookupMatchingAll();
2341-
return new BlockSourceReader.DoublesBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup);
2341+
return new BlockSourceReader.FloatsBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup);
23422342
}
23432343

23442344
private SourceValueFetcher sourceValueFetcher(Set<String> sourcePaths) {

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,23 @@ public DoublesBuilder appendDouble(double value) {
7575
return new DoublesBuilder();
7676
}
7777

78+
@Override
79+
public BlockLoader.FloatBuilder floatsFromDocValues(int expectedCount) {
80+
return floats(expectedCount);
81+
}
82+
83+
@Override
84+
public BlockLoader.FloatBuilder floats(int expectedCount) {
85+
class FloatsBuilder extends TestBlock.Builder implements BlockLoader.FloatBuilder {
86+
@Override
87+
public BlockLoader.FloatBuilder appendFloat(float value) {
88+
add(value);
89+
return this;
90+
}
91+
}
92+
return new FloatsBuilder();
93+
}
94+
7895
@Override
7996
public BlockLoader.IntBuilder intsFromDocValues(int expectedCount) {
8097
return ints(expectedCount);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@ private static Block constantBlock(BlockFactory blockFactory, ElementType type,
243243
case BYTES_REF -> blockFactory.newConstantBytesRefBlockWith(toBytesRef(val), size);
244244
case DOUBLE -> blockFactory.newConstantDoubleBlockWith((double) val, size);
245245
case BOOLEAN -> blockFactory.newConstantBooleanBlockWith((boolean) val, size);
246+
case FLOAT -> blockFactory.newConstantFloatBlockWith((float) val, size);
246247
case COMPOSITE -> {
247248
if (val instanceof AggregateMetricDoubleLiteral aggregateMetricDoubleLiteral) {
248249
yield blockFactory.newConstantAggregateMetricDoubleBlock(aggregateMetricDoubleLiteral, size);

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -654,6 +654,16 @@ public BlockLoader.DoubleBuilder doubles(int expectedCount) {
654654
return factory.newDoubleBlockBuilder(expectedCount);
655655
}
656656

657+
@Override
658+
public BlockLoader.FloatBuilder floatsFromDocValues(int expectedCount) {
659+
return factory.newFloatBlockBuilder(expectedCount).mvOrdering(Block.MvOrdering.SORTED_ASCENDING);
660+
}
661+
662+
@Override
663+
public BlockLoader.FloatBuilder floats(int expectedCount) {
664+
return factory.newFloatBlockBuilder(expectedCount);
665+
}
666+
657667
@Override
658668
public BlockLoader.IntBuilder intsFromDocValues(int expectedCount) {
659669
return factory.newIntBlockBuilder(expectedCount).mvOrdering(Block.MvOrdering.SORTED_ASCENDING);

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvAssert.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,8 +146,8 @@ private static void assertMetadata(
146146
// Type.asType translates all bytes references into keywords
147147
continue;
148148
}
149-
if (blockType == Type.DOUBLE && expectedType == DENSE_VECTOR) {
150-
// DENSE_VECTOR is internally represented as a double block
149+
if (blockType == Type.FLOAT && expectedType == DENSE_VECTOR) {
150+
// DENSE_VECTOR is internally represented as a float block
151151
continue;
152152
}
153153
if (blockType == Type.NULL) {

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -486,7 +486,7 @@ public enum Type {
486486
x -> x == null ? null : stringToAggregateMetricDoubleLiteral(x),
487487
AggregateMetricDoubleBlockBuilder.AggregateMetricDoubleLiteral.class
488488
),
489-
DENSE_VECTOR(Double::parseDouble, Double.class),
489+
DENSE_VECTOR(Float::parseFloat, Float.class),
490490
UNSUPPORTED(Type::convertUnsupported, Void.class);
491491

492492
private static Void convertUnsupported(String s) {

‎x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -829,7 +829,7 @@ public static Literal randomLiteral(DataType type) {
829829
throw new UncheckedIOException(e);
830830
}
831831
}
832-
case DENSE_VECTOR -> Arrays.asList(randomArray(10, 10, i -> new Double[10], () -> (double) randomFloat()));
832+
case DENSE_VECTOR -> Arrays.asList(randomArray(10, 10, i -> new Float[10], ESTestCase::randomFloat));
833833
case UNSUPPORTED, OBJECT, DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new IllegalArgumentException(
834834
"can't make random values for [" + type.typeName() + "]"
835835
);

0 commit comments

Comments
 (0)