Skip to content

Commit 2f22dd9

Browse files
benwtrentbenwtrent
committed
adding tests
1 parent 12367e0 commit 2f22dd9

File tree

3 files changed

+210
-20
lines changed

3 files changed

+210
-20
lines changed

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

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
setup:
2+
- requires:
3+
cluster_features: "mapper.ivf_nested_support"
4+
reason: 'ivf nested support required'
5+
- do:
6+
indices.create:
7+
index: test
8+
body:
9+
settings:
10+
index:
11+
number_of_shards: 1
12+
mappings:
13+
properties:
14+
name:
15+
type: keyword
16+
nested:
17+
type: nested
18+
properties:
19+
paragraph_id:
20+
type: keyword
21+
vector:
22+
type: dense_vector
23+
dims: 5
24+
index: true
25+
similarity: l2_norm
26+
index_options:
27+
type: bbq_ivf
28+
29+
aliases:
30+
my_alias:
31+
filter:
32+
term:
33+
name: "rabbit.jpg"
34+
35+
- do:
36+
index:
37+
index: test
38+
id: "1"
39+
body:
40+
name: cow.jpg
41+
nested:
42+
- paragraph_id: 0
43+
vector: [230.0, 300.33, -34.8988, 15.555, -200.0]
44+
- paragraph_id: 1
45+
vector: [240.0, 300, -3, 1, -20]
46+
47+
- do:
48+
index:
49+
index: test
50+
id: "2"
51+
body:
52+
name: moose.jpg
53+
nested:
54+
- paragraph_id: 0
55+
vector: [-0.5, 100.0, -13, 14.8, -156.0]
56+
- paragraph_id: 2
57+
vector: [0, 100.0, 0, 14.8, -156.0]
58+
- paragraph_id: 3
59+
vector: [0, 1.0, 0, 1.8, -15.0]
60+
61+
- do:
62+
index:
63+
index: test
64+
id: "3"
65+
body:
66+
name: rabbit.jpg
67+
nested:
68+
- paragraph_id: 0
69+
vector: [0.5, 111.3, -13.0, 14.8, -156.0]
70+
- do:
71+
indices.forcemerge:
72+
index: test
73+
max_num_segments: 1
74+
75+
- do:
76+
indices.refresh: {}
77+
78+
---
79+
"nested kNN search that returns diverse parents docs":
80+
- do:
81+
search:
82+
index: test
83+
body:
84+
fields: [ "name" ]
85+
query:
86+
nested:
87+
path: nested
88+
query:
89+
knn:
90+
field: nested.vector
91+
query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
92+
num_candidates: 3
93+
94+
- match: {hits.hits.0._id: "2"}
95+
- match: {hits.hits.0.fields.name.0: "moose.jpg"}
96+
97+
- match: {hits.hits.1._id: "3"}
98+
- match: {hits.hits.1.fields.name.0: "rabbit.jpg"}
99+
100+
- do:
101+
search:
102+
index: test
103+
body:
104+
fields: [ "name" ]
105+
query:
106+
nested:
107+
path: nested
108+
query:
109+
knn:
110+
field: nested.vector
111+
query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
112+
num_candidates: 3
113+
inner_hits: { size: 1, "fields": [ "nested.paragraph_id" ], _source: false }
114+
115+
- match: {hits.total.value: 3}
116+
117+
- match: { hits.hits.0._id: "2" }
118+
- match: { hits.hits.0.fields.name.0: "moose.jpg" }
119+
- match: { hits.hits.0.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0" }
120+
121+
- match: { hits.hits.1._id: "3" }
122+
- match: { hits.hits.1.fields.name.0: "rabbit.jpg" }
123+
- match: { hits.hits.1.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0" }
124+
125+
- match: { hits.hits.2._id: "1" }
126+
- match: { hits.hits.2.fields.name.0: "cow.jpg" }
127+
- match: { hits.hits.2.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0" }
128+
129+
---
130+
"nested kNN search pre-filtered on alias with filter on top level fields":
131+
- do:
132+
search:
133+
index: my_alias # filter on name: "rabbit.jpg"
134+
body:
135+
fields: [ "name" ]
136+
query:
137+
nested:
138+
path: nested
139+
query:
140+
knn:
141+
field: nested.vector
142+
query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
143+
num_candidates: 1
144+
inner_hits: { size: 1, "fields": [ "nested.paragraph_id" ], _source: false }
145+
146+
- match: {hits.total.value: 1} # as alias is passed as pre-filter, we get a single result
147+
- match: {hits.hits.0._id: "3"}
148+
- match: {hits.hits.0.fields.name.0: "rabbit.jpg"}
149+
- match: { hits.hits.0.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0" }
150+
151+
---
152+
"nested kNN search post-filtered on top level fields":
153+
- do:
154+
search:
155+
index: test
156+
body:
157+
fields: [ "name" ]
158+
query:
159+
bool:
160+
must:
161+
- term:
162+
name: "rabbit.jpg"
163+
- nested:
164+
path: nested
165+
query:
166+
knn:
167+
field: nested.vector
168+
query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
169+
num_candidates: 1
170+
- match: { hits.total.value: 0 } # no hits because returned single vector did not pass post-filter
171+
172+
- do:
173+
search:
174+
index: test
175+
body:
176+
fields: [ "name" ]
177+
query:
178+
bool:
179+
must:
180+
- term:
181+
name: "rabbit.jpg"
182+
- nested:
183+
path: nested
184+
query:
185+
knn:
186+
field: nested.vector
187+
query_vector: [ -0.5, 90.0, -10, 14.8, -156.0 ]
188+
num_candidates: 3
189+
inner_hits: { size: 1, fields: [ "nested.paragraph_id" ], _source: false }
190+
191+
- match: {hits.total.value: 1}
192+
- match: {hits.hits.0._id: "3"}
193+
- match: {hits.hits.0.fields.name.0: "rabbit.jpg"}
194+
- match: { hits.hits.0.inner_hits.nested.hits.hits.0.fields.nested.0.paragraph_id.0: "0" }

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public class MapperFeatures implements FeatureSpecification {
4242
);
4343
static final NodeFeature NPE_ON_DIMS_UPDATE_FIX = new NodeFeature("mapper.npe_on_dims_update_fix");
4444
static final NodeFeature IVF_FORMAT_CLUSTER_FEATURE = new NodeFeature("mapper.ivf_format_cluster_feature");
45+
static final NodeFeature IVF_NESTED_SUPPORT = new NodeFeature("mapper.ivf_nested_support");
4546

4647
@Override
4748
public Set<NodeFeature> getTestFeatures() {
@@ -70,7 +71,8 @@ public Set<NodeFeature> getTestFeatures() {
7071
NPE_ON_DIMS_UPDATE_FIX,
7172
RESCORE_ZERO_VECTOR_QUANTIZED_VECTOR_MAPPING,
7273
USE_DEFAULT_OVERSAMPLE_VALUE_FOR_BBQ,
73-
IVF_FORMAT_CLUSTER_FEATURE
74+
IVF_FORMAT_CLUSTER_FEATURE,
75+
IVF_NESTED_SUPPORT
7476
);
7577
}
7678
}

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

Lines changed: 13 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@
6464
import org.elasticsearch.index.mapper.Mapper;
6565
import org.elasticsearch.index.mapper.MapperBuilderContext;
6666
import org.elasticsearch.index.mapper.MapperParsingException;
67-
import org.elasticsearch.index.mapper.MappingLookup;
6867
import org.elasticsearch.index.mapper.MappingParser;
6968
import org.elasticsearch.index.mapper.NumberFieldMapper;
7069
import org.elasticsearch.index.mapper.SimpleMappedFieldType;
@@ -77,6 +76,7 @@
7776
import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
7877
import org.elasticsearch.search.lookup.Source;
7978
import org.elasticsearch.search.vectors.DenseVectorQuery;
79+
import org.elasticsearch.search.vectors.DiversifyingChildrenIVFKnnFloatVectorQuery;
8080
import org.elasticsearch.search.vectors.ESDiversifyingChildrenByteKnnVectorQuery;
8181
import org.elasticsearch.search.vectors.ESDiversifyingChildrenFloatKnnVectorQuery;
8282
import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
@@ -1650,7 +1650,7 @@ public boolean supportsElementType(ElementType elementType) {
16501650

16511651
@Override
16521652
public boolean supportsDimension(int dims) {
1653-
return dims >= BBQ_MIN_DIMS;
1653+
return true;
16541654
}
16551655
};
16561656

@@ -2516,12 +2516,19 @@ && isNotUnitVector(squaredMagnitude)) {
25162516
adjustedK = Math.min((int) Math.ceil(k * oversample), OVERSAMPLE_LIMIT);
25172517
numCands = Math.max(adjustedK, numCands);
25182518
}
2519-
if (parentFilter != null && indexOptions instanceof BBQIVFIndexOptions) {
2520-
throw new IllegalArgumentException("IVF index does not support nested queries");
2521-
}
25222519
Query knnQuery;
25232520
if (indexOptions instanceof BBQIVFIndexOptions bbqIndexOptions) {
2524-
knnQuery = new IVFKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter, bbqIndexOptions.defaultNProbe);
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);
25252532
} else {
25262533
knnQuery = parentFilter != null
25272534
? new ESDiversifyingChildrenFloatKnnVectorQuery(
@@ -2764,19 +2771,6 @@ public FieldMapper.Builder getMergeBuilder() {
27642771
return new Builder(leafName(), indexCreatedVersion).init(this);
27652772
}
27662773

2767-
@Override
2768-
public void doValidate(MappingLookup mappers) {
2769-
if (indexOptions instanceof BBQIVFIndexOptions && mappers.nestedLookup().getNestedParent(fullPath()) != null) {
2770-
throw new IllegalArgumentException(
2771-
"["
2772-
+ CONTENT_TYPE
2773-
+ "] fields with index type ["
2774-
+ indexOptions.type
2775-
+ "] cannot be indexed if they're within [nested] mappings"
2776-
);
2777-
}
2778-
}
2779-
27802774
private static IndexOptions parseIndexOptions(String fieldName, Object propNode, IndexVersion indexVersion) {
27812775
@SuppressWarnings("unchecked")
27822776
Map<String, ?> indexOptionsMap = (Map<String, ?>) propNode;

0 commit comments

Comments
 (0)