Skip to content

Commit 04ce871

Browse files
benwtrentbenwtrent
committed
adding tests
1 parent 2f22dd9 commit 04ce871

File tree

3 files changed

+129
-19
lines changed

3 files changed

+129
-19
lines changed

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/135_knn_query_nested_search_ivf.yml‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ setup:
4040
name: cow.jpg
4141
nested:
4242
- paragraph_id: 0
43-
vector: [230.0, 300.33, -34.8988, 15.555, -200.0]
43+
vector: [230, 300.33, -34.8988, 15.555, -200]
4444
- paragraph_id: 1
45-
vector: [240.0, 300, -3, 1, -20]
45+
vector: [240, 300, -3, 1, -20]
4646

4747
- do:
4848
index:
@@ -52,11 +52,11 @@ setup:
5252
name: moose.jpg
5353
nested:
5454
- paragraph_id: 0
55-
vector: [-0.5, 100.0, -13, 14.8, -156.0]
55+
vector: [-0.5, 100, -13, 14.8, -156]
5656
- paragraph_id: 2
57-
vector: [0, 100.0, 0, 14.8, -156.0]
57+
vector: [0, 100, 0, 14.8, -156]
5858
- paragraph_id: 3
59-
vector: [0, 1.0, 0, 1.8, -15.0]
59+
vector: [0, 1, 0, 1.8, -15]
6060

6161
- do:
6262
index:
@@ -66,7 +66,7 @@ setup:
6666
name: rabbit.jpg
6767
nested:
6868
- paragraph_id: 0
69-
vector: [0.5, 111.3, -13.0, 14.8, -156.0]
69+
vector: [0.5, 111.3, -13, 14.8, -156]
7070
- do:
7171
indices.forcemerge:
7272
index: test
@@ -88,8 +88,9 @@ setup:
8888
query:
8989
knn:
9090
field: nested.vector
91-
query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
91+
query_vector: [-0.5, 90, -10, 14.8, -156]
9292
num_candidates: 3
93+
- match: {hits.total.value: 3}
9394

9495
- match: {hits.hits.0._id: "2"}
9596
- match: {hits.hits.0.fields.name.0: "moose.jpg"}
@@ -108,7 +109,7 @@ setup:
108109
query:
109110
knn:
110111
field: nested.vector
111-
query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
112+
query_vector: [ -0.5, 90, -10, 14.8, -156 ]
112113
num_candidates: 3
113114
inner_hits: { size: 1, "fields": [ "nested.paragraph_id" ], _source: false }
114115

‎server/src/internalClusterTest/java/org/elasticsearch/search/nested/VectorNestedIT.java‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@
99

1010
package org.elasticsearch.search.nested;
1111

12+
import org.apache.lucene.search.join.ScoreMode;
1213
import org.elasticsearch.action.get.GetResponse;
1314
import org.elasticsearch.cluster.health.ClusterHealthStatus;
1415
import org.elasticsearch.common.settings.Settings;
1516
import org.elasticsearch.index.query.InnerHitBuilder;
17+
import org.elasticsearch.index.query.QueryBuilders;
1618
import org.elasticsearch.search.vectors.KnnSearchBuilder;
19+
import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
1720
import org.elasticsearch.test.ESIntegTestCase;
1821

1922
import java.util.List;
@@ -75,4 +78,106 @@ public void testSimpleNested() throws Exception {
7578
);
7679
}
7780

81+
public void testSimpleIVFNested() throws Exception {
82+
assertAcked(
83+
prepareCreate("test").setMapping(
84+
jsonBuilder().startObject()
85+
.startObject("properties")
86+
.startObject("name")
87+
.field("type", "keyword")
88+
.endObject()
89+
.startObject("nested")
90+
.field("type", "nested")
91+
.startObject("properties")
92+
.startObject("paragraph_id")
93+
.field("type", "keyword")
94+
.endObject()
95+
.startObject("vector")
96+
.field("type", "dense_vector")
97+
.field("index", true)
98+
.field("dims", 5)
99+
.field("similarity", "l2_norm")
100+
.startObject("index_options")
101+
.field("type", "bbq_ivf")
102+
.endObject()
103+
.endObject()
104+
.endObject()
105+
.endObject()
106+
.endObject()
107+
.endObject()
108+
).setSettings(Settings.builder().put(indexSettings()).put("index.number_of_shards", 1))
109+
);
110+
ensureGreen();
111+
112+
prepareIndex("test").setId("1")
113+
.setSource(
114+
jsonBuilder().startObject()
115+
.field("name", "cow.jpg")
116+
.startArray("nested")
117+
.startObject()
118+
.field("paragraph_id", "0")
119+
.field("vector", new float[] { 230, 300.33f, -34.8988f, 15.555f, -200 })
120+
.endObject()
121+
.startObject()
122+
.field("paragraph_id", "1")
123+
.field("vector", new float[] { 240, 300, -3, 1, -20 })
124+
.endObject()
125+
.endArray()
126+
.endObject()
127+
)
128+
.get();
129+
prepareIndex("test").setId("2")
130+
.setSource(
131+
jsonBuilder().startObject()
132+
.field("name", "moose.jpg")
133+
.startArray("nested")
134+
.startObject()
135+
.field("paragraph_id", "0")
136+
.field("vector", new float[] { -0.5f, 100, -13, 14.8f, -156 })
137+
.endObject()
138+
.startObject()
139+
.field("paragraph_id", "1")
140+
.field("vector", new float[] { 0, 100, 0, 14.8f, -156 })
141+
.endObject()
142+
.startObject()
143+
.field("paragraph_id", "2")
144+
.field("vector", new float[] { 0, 1, 0, 1.8f, -15 })
145+
.endObject()
146+
.endArray()
147+
.endObject()
148+
)
149+
.get();
150+
prepareIndex("test").setId("3")
151+
.setSource(
152+
jsonBuilder().startObject()
153+
.field("name", "rabbit.jpg")
154+
.startArray("nested")
155+
.startObject()
156+
.field("paragraph_id", "0")
157+
.field("vector", new float[] { 0.5f, 111.3f, -13f, 14.8f, -156f })
158+
.endObject()
159+
.endArray()
160+
.endObject()
161+
)
162+
.get();
163+
client().admin().indices().prepareForceMerge("test").setMaxNumSegments(1).get();
164+
refresh();
165+
166+
waitForRelocation(ClusterHealthStatus.GREEN);
167+
168+
assertResponse(
169+
prepareSearch("test").setQuery(
170+
QueryBuilders.nestedQuery(
171+
"nested",
172+
new KnnVectorQueryBuilder("nested.vector", new float[] { -0.5f, 90, -10, 14.8f, -156 }, null, 3, null, null),
173+
ScoreMode.Max
174+
)
175+
).setAllowPartialSearchResults(false),
176+
response -> {
177+
assertThat(response.getHits().getHits().length, greaterThan(0));
178+
assertThat(response.getHits().getHits()[0].getId(), equalTo("2"));
179+
}
180+
);
181+
}
182+
78183
}

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

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@
7272
import org.elasticsearch.index.mapper.TextSearchInfo;
7373
import org.elasticsearch.index.mapper.ValueFetcher;
7474
import org.elasticsearch.index.query.SearchExecutionContext;
75+
import org.elasticsearch.logging.LogManager;
76+
import org.elasticsearch.logging.Logger;
7577
import org.elasticsearch.search.DocValueFormat;
7678
import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
7779
import org.elasticsearch.search.lookup.Source;
@@ -117,6 +119,7 @@
117119
* A {@link FieldMapper} for indexing a dense vector of floats.
118120
*/
119121
public class DenseVectorFieldMapper extends FieldMapper {
122+
private static final Logger logger = LogManager.getLogger(DenseVectorFieldMapper.class);
120123
public static final String COSINE_MAGNITUDE_FIELD_SUFFIX = "._magnitude";
121124
private static final float EPS = 1e-3f;
122125
public static final int BBQ_MIN_DIMS = 64;
@@ -2518,17 +2521,17 @@ && isNotUnitVector(squaredMagnitude)) {
25182521
}
25192522
Query knnQuery;
25202523
if (indexOptions instanceof BBQIVFIndexOptions bbqIndexOptions) {
2521-
knnQuery = parentFilter != null ?
2522-
new DiversifyingChildrenIVFKnnFloatVectorQuery(
2523-
name(),
2524-
queryVector,
2525-
adjustedK,
2526-
numCands,
2527-
filter,
2528-
parentFilter,
2529-
bbqIndexOptions.defaultNProbe
2530-
) :
2531-
new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe);
2524+
knnQuery = parentFilter != null
2525+
? new DiversifyingChildrenIVFKnnFloatVectorQuery(
2526+
name(),
2527+
queryVector,
2528+
adjustedK,
2529+
numCands,
2530+
filter,
2531+
parentFilter,
2532+
bbqIndexOptions.defaultNProbe
2533+
)
2534+
: new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe);
25322535
} else {
25332536
knnQuery = parentFilter != null
25342537
? new ESDiversifyingChildrenFloatKnnVectorQuery(
@@ -2558,6 +2561,7 @@ && isNotUnitVector(squaredMagnitude)) {
25582561
similarity.score(similarityThreshold, elementType, dims)
25592562
);
25602563
}
2564+
logger.warn("Executing knn query [{}]", knnQuery);
25612565
return knnQuery;
25622566
}
25632567

0 commit comments

Comments
 (0)