Skip to content

Commit d196be0

Browse files
carlosdelestcarlosdelest
committed
Add index = false support
1 parent 2754697 commit d196be0

File tree

2 files changed

+75
-42
lines changed

2 files changed

+75
-42
lines changed

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

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,15 +53,18 @@
5353
import org.elasticsearch.index.mapper.ArraySourceValueFetcher;
5454
import org.elasticsearch.index.mapper.BlockDocValuesReader;
5555
import org.elasticsearch.index.mapper.BlockLoader;
56+
import org.elasticsearch.index.mapper.BlockSourceReader;
5657
import org.elasticsearch.index.mapper.DocumentParserContext;
5758
import org.elasticsearch.index.mapper.FieldMapper;
5859
import org.elasticsearch.index.mapper.MappedFieldType;
5960
import org.elasticsearch.index.mapper.Mapper;
6061
import org.elasticsearch.index.mapper.MapperBuilderContext;
6162
import org.elasticsearch.index.mapper.MapperParsingException;
6263
import org.elasticsearch.index.mapper.MappingParser;
64+
import org.elasticsearch.index.mapper.NumberFieldMapper;
6365
import org.elasticsearch.index.mapper.SimpleMappedFieldType;
6466
import org.elasticsearch.index.mapper.SourceLoader;
67+
import org.elasticsearch.index.mapper.SourceValueFetcher;
6568
import org.elasticsearch.index.mapper.TextSearchInfo;
6669
import org.elasticsearch.index.mapper.ValueFetcher;
6770
import org.elasticsearch.index.query.SearchExecutionContext;
@@ -91,6 +94,7 @@
9194
import java.util.Map;
9295
import java.util.Objects;
9396
import java.util.Optional;
97+
import java.util.Set;
9498
import java.util.function.Function;
9599
import java.util.function.Supplier;
96100
import java.util.stream.Stream;
@@ -2079,6 +2083,18 @@ protected Object parseSourceValue(Object value) {
20792083
};
20802084
}
20812085

2086+
private SourceValueFetcher sourceValueFetcher(Set<String> sourcePaths) {
2087+
return new SourceValueFetcher(sourcePaths, null) {
2088+
@Override
2089+
protected Object parseSourceValue(Object value) {
2090+
if (value.equals("")) {
2091+
return null;
2092+
}
2093+
return NumberFieldMapper.NumberType.FLOAT.parse(value, false);
2094+
}
2095+
};
2096+
}
2097+
20822098
@Override
20832099
public DocValueFormat docValueFormat(String format, ZoneId timeZone) {
20842100
return DocValueFormat.DENSE_VECTOR;
@@ -2316,7 +2332,12 @@ ElementType getElementType() {
23162332

23172333
@Override
23182334
public BlockLoader blockLoader(MappedFieldType.BlockLoaderContext blContext) {
2319-
return new BlockDocValuesReader.DenseVectorBlockLoader(name());
2335+
if (indexed) {
2336+
return new BlockDocValuesReader.DenseVectorBlockLoader(name());
2337+
}
2338+
2339+
BlockSourceReader.LeafIteratorLookup lookup = BlockSourceReader.lookupMatchingAll();
2340+
return new BlockSourceReader.DoublesBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup);
23202341
}
23212342
}
23222343

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java‎

Lines changed: 53 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,15 @@
1313
import org.elasticsearch.action.index.IndexRequestBuilder;
1414
import org.elasticsearch.cluster.metadata.IndexMetadata;
1515
import org.elasticsearch.common.settings.Settings;
16+
import org.elasticsearch.xcontent.XContentBuilder;
17+
import org.elasticsearch.xcontent.XContentFactory;
1618
import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
1719
import org.junit.Before;
1820

21+
import java.io.IOException;
1922
import java.util.ArrayList;
2023
import java.util.HashMap;
2124
import java.util.List;
22-
import java.util.Locale;
2325
import java.util.Map;
2426
import java.util.Set;
2527

@@ -39,19 +41,26 @@ public class DenseVectorFieldTypeIT extends AbstractEsqlIntegTestCase {
3941
);
4042

4143
private final String indexType;
44+
private final boolean index;
4245
private int numDims;
4346
private int numDocs;
4447

4548
@ParametersFactory
4649
public static Iterable<Object[]> parameters() throws Exception {
47-
return DENSE_VECTOR_INDEX_TYPES.stream().map(type -> new Object[] { type }).toList();
50+
List<Object[]> params = new ArrayList<>();
51+
for (String indexType : DENSE_VECTOR_INDEX_TYPES) {
52+
params.add(new Object[] { indexType, true });
53+
}
54+
params.add(new Object[] {null, false});
55+
return params;
4856
}
4957

50-
public DenseVectorFieldTypeIT(@Name("indexType") String indexType) {
58+
public DenseVectorFieldTypeIT(@Name("indexType") String indexType, @Name("index") boolean index) {
5159
this.indexType = indexType;
60+
this.index = index;
5261
}
5362

54-
private Map<Integer, List<Float>> indexedDocs = new HashMap<>();
63+
private final Map<Integer, List<Float>> indexedVectors = new HashMap<>();
5564

5665
public void testRetrieveFieldType() {
5766
var query = """
@@ -65,18 +74,20 @@ public void testRetrieveFieldType() {
6574
}
6675

6776
@SuppressWarnings("unchecked")
68-
public void testRetrieveOrderedDenseVectorFieldData() {
77+
public void testRetrieveTopNDenseVectorFieldData() {
6978
var query = """
7079
FROM test
80+
| KEEP id, vector
7181
| SORT id ASC
72-
""";
82+
""";
7383

7484
try (var resp = run(query)) {
7585
List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
76-
indexedDocs.forEach((id, vector) -> {
86+
indexedVectors.forEach((id, vector) -> {
7787
var values = valuesList.get(id);
7888
assertEquals(id, values.get(0));
7989
List<Double> vectors = (List<Double>) values.get(1);
90+
assertNotNull(vectors);
8091
assertEquals(vector.size(), vectors.size());
8192
for (int i = 0; i < vector.size(); i++) {
8293
assertEquals(vector.get(i), vectors.get(i).floatValue(), 0F);
@@ -86,18 +97,22 @@ public void testRetrieveOrderedDenseVectorFieldData() {
8697
}
8798

8899
@SuppressWarnings("unchecked")
89-
public void testRetrieveUnOrderedDenseVectorFieldData() {
90-
var query = "FROM test";
100+
public void testRetrieveDenseVectorFieldData() {
101+
var query = """
102+
FROM test
103+
| KEEP id, vector
104+
""";
91105

92106
try (var resp = run(query)) {
93107
List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
94-
assertEquals(valuesList.size(), indexedDocs.size());
108+
assertEquals(valuesList.size(), indexedVectors.size());
95109
valuesList.forEach(value -> {;
96110
assertEquals(2, value.size());
97111
Integer id = (Integer) value.get(0);
98112
List<Double> vector = (List<Double>) value.get(1);
99-
100-
List<Float> expectedVector = indexedDocs.get(id);
113+
assertNotNull(vector);
114+
List<Float> expectedVector = indexedVectors.get(id);
115+
assertNotNull(expectedVector);
101116
for (int i = 0; i < vector.size(); i++) {
102117
assertEquals(expectedVector.get(i), vector.get(i).floatValue(), 0F);
103118
}
@@ -106,36 +121,25 @@ public void testRetrieveUnOrderedDenseVectorFieldData() {
106121
}
107122

108123
@Before
109-
public void setup() {
110-
numDims = randomIntBetween(64, 256);
111-
numDocs = randomIntBetween(10, 100);
112-
for (int i = 0; i < numDocs; i++) {
113-
List<Float> vector = new ArrayList<>(numDims);
114-
for (int j = 0; j < numDims; j++) {
115-
// vector.add(randomFloat());
116-
vector.add(1.0f);
117-
}
118-
indexedDocs.put(i, vector);
119-
}
120-
124+
public void setup() throws IOException {
121125
var indexName = "test";
122126
var client = client().admin().indices();
123-
var mapping = String.format(Locale.ROOT, """
124-
{
125-
"properties": {
126-
"id": {
127-
"type": "integer"
128-
},
129-
"vector": {
130-
"type": "dense_vector",
131-
"similarity": "l2_norm",
132-
"index_options": {
133-
"type": "%s"
134-
}
135-
}
136-
}
137-
}
138-
""", indexType);
127+
XContentBuilder mapping = XContentFactory.jsonBuilder()
128+
.startObject()
129+
.startObject("properties")
130+
.startObject("id")
131+
.field("type", "integer")
132+
.endObject()
133+
.startObject("vector")
134+
.field("type", "dense_vector")
135+
.field("index", index);
136+
if (index) {
137+
mapping.field("similarity", "l2_norm");
138+
}
139+
if (indexType != null) {
140+
mapping.startObject("index_options").field("type", indexType).endObject();
141+
}
142+
mapping.endObject().endObject().endObject();
139143
Settings settings = Settings.builder()
140144
.put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
141145
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, randomIntBetween(1, 5))
@@ -146,10 +150,18 @@ public void setup() {
146150
.setSettings(settings);
147151
assertAcked(CreateRequest);
148152

153+
numDims = randomIntBetween(32, 64) * 2; // min 64, even number
154+
numDocs = randomIntBetween(10, 100);
149155
IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
150156
for (int i = 0; i < numDocs; i++) {
151-
docs[i] = prepareIndex("test").setId("" + i).setSource("id", i, "vector", indexedDocs.get(i));
157+
List<Float> vector = new ArrayList<>(numDims);
158+
for (int j = 0; j < numDims; j++) {
159+
vector.add(randomFloat());
160+
}
161+
docs[i] = prepareIndex("test").setId("" + i).setSource("id", String.valueOf(i), "vector", vector);
162+
indexedVectors.put(i, vector);
152163
}
164+
153165
indexRandom(true, docs);
154166
}
155167
}

0 commit comments

Comments
 (0)