Skip to content

Commit 018af55

Browse files
committed
tests: add another hnsw unit test objectbox#24
1 parent ff4c72a commit 018af55

File tree

1 file changed

+80
-1
lines changed

1 file changed

+80
-1
lines changed

tests/test_hnsw.py

+80-1
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def _test_random_points(num_points: int, num_query_points: int, seed: Optional[i
5858

5959
# Run ANN with OBX
6060
query_builder = QueryBuilder(db, box)
61-
query_builder.nearest_neighbors_f32(VectorEntity.get_property("vector")._id, query_point, k)
61+
query_builder.nearest_neighbors_f32("vector", query_point, k)
6262
query = query_builder.build()
6363
obx_result = [id_ for id_, score in query.find_ids_with_scores()] # Ignore score
6464
assert len(obx_result) == k
@@ -78,3 +78,82 @@ def test_random_points():
7878
_test_random_points(num_points=100, num_query_points=10, seed=13)
7979
_test_random_points(num_points=100, num_query_points=10, seed=14)
8080
_test_random_points(num_points=100, num_query_points=10, seed=15)
81+
82+
83+
def test_combined_nn_search():
84+
""" Tests NN search combined with regular query conditions, offset and limit. """
85+
86+
db = create_test_objectbox()
87+
88+
box = objectbox.Box(db, VectorEntity)
89+
90+
box.put(VectorEntity(name="Power of red", vector=[1, 1]))
91+
box.put(VectorEntity(name="Blueberry", vector=[2, 2]))
92+
box.put(VectorEntity(name="Red", vector=[3, 3]))
93+
box.put(VectorEntity(name="Blue sea", vector=[4, 4]))
94+
box.put(VectorEntity(name="Lightblue", vector=[5, 5]))
95+
box.put(VectorEntity(name="Red apple", vector=[6, 6]))
96+
box.put(VectorEntity(name="Hundred", vector=[7, 7]))
97+
box.put(VectorEntity(name="Tired", vector=[8, 8]))
98+
box.put(VectorEntity(name="Power of blue", vector=[9, 9]))
99+
100+
assert box.count() == 9
101+
102+
# Test condition + NN search
103+
query = box.query() \
104+
.nearest_neighbors_f32("vector", [4.1, 4.2], 6) \
105+
.contains_string("name", "red", case_sensitive=False) \
106+
.build()
107+
# 4, 5, 3, 6, 2, 7
108+
# Filtered: 3, 6, 7
109+
search_results = query.find_with_scores()
110+
assert len(search_results) == 3
111+
assert search_results[0][0].name == "Red"
112+
assert search_results[1][0].name == "Red apple"
113+
assert search_results[2][0].name == "Hundred"
114+
115+
# Test offset/limit on find_with_scores (result is ordered by score desc)
116+
query.offset(1)
117+
query.limit(1)
118+
search_results = query.find_with_scores()
119+
assert len(search_results) == 1
120+
assert search_results[0][0].name == "Red apple"
121+
122+
# Regular condition + NN search
123+
query = box.query() \
124+
.nearest_neighbors_f32("vector", [9.2, 8.9], 7) \
125+
.starts_with_string("name", "Blue", case_sensitive=True) \
126+
.build()
127+
128+
search_results = query.find_with_scores()
129+
assert len(search_results) == 1
130+
assert search_results[0][0].name == "Blue sea"
131+
132+
# Regular condition + NN search
133+
query = box.query() \
134+
.nearest_neighbors_f32("vector", [7.7, 7.7], 8) \
135+
.contains_string("name", "blue", case_sensitive=False) \
136+
.build()
137+
# 8, 7, 9, 6, 5, 4, 3, 2
138+
# Filtered: 9, 5, 4, 2
139+
search_results = query.find_ids_with_scores()
140+
assert len(search_results) == 4
141+
assert search_results[0][0] == 9
142+
assert search_results[1][0] == 5
143+
assert search_results[2][0] == 4
144+
assert search_results[3][0] == 2
145+
146+
search_results = query.find_ids()
147+
assert len(search_results) == 4
148+
assert search_results[0] == 2
149+
assert search_results[1] == 4
150+
assert search_results[2] == 5
151+
assert search_results[3] == 9
152+
153+
# Test offset/limit on find_ids (result is ordered by ID asc)
154+
query.offset(1)
155+
query.limit(2)
156+
search_results = query.find_ids()
157+
assert len(search_results) == 2
158+
assert search_results[0] == 4
159+
assert search_results[1] == 5

0 commit comments

Comments
 (0)