tests: added cosine and dot_product HNSW Index to Model #30

Author: dan-obx

tests/common.py

@@ -52,9 +52,9 @@ def create_test_objectbox(db_name: Optional[str] = None, clear_db: bool = True)
     model.entity(TestEntity, last_property_id=IdUid(27, 1027))
     model.entity(TestEntityDatetime, last_property_id=IdUid(4, 2004))
     model.entity(TestEntityFlex, last_property_id=IdUid(2, 3002))
-    model.entity(VectorEntity, last_property_id=IdUid(3, 4003))
+    model.entity(VectorEntity, last_property_id=IdUid(5, 4005))
     model.last_entity_id = IdUid(4, 4)
-    model.last_index_id = IdUid(3, 40001)
+    model.last_index_id = IdUid(5, 40003)
 
     return objectbox.Builder().model(model).directory(db_path).build()
 

tests/model.py

@@ -57,7 +57,7 @@ class TestEntityFlex:
 class VectorEntity:
     id = Id(id=1, uid=4001)
     name = Property(str, type=PropertyType.string, id=2, uid=4002)
-    vector = Property(np.ndarray, type=PropertyType.floatVector, id=3, uid=4003,
+    vector_euclidean = Property(np.ndarray, type=PropertyType.floatVector, id=3, uid=4003,
                       index=HnswIndex(
                           id=3, uid=40001,
                           dimensions=2, distance_type=HnswDistanceType.EUCLIDEAN)

tests/test_hnsw.py

@@ -15,9 +15,11 @@ def _find_expected_nn(points: np.ndarray, query: np.ndarray, n: int):
     return np.argsort(d)[:n]
 
 
-def _test_random_points(num_points: int, num_query_points: int, seed: Optional[int] = None):
+def _test_random_points(num_points: int, num_query_points: int, seed: Optional[int] = None, distance_type: HnswDistanceType = HnswDistanceType.EUCLIDEAN, min_score: float = 0.5):
     """ Generates random points in a 2d plane; checks the queried NN against the expected. """
 
+    vector_field_name = "vector_"+distance_type.name.lower()
+
     print(f"Test random points; Points: {num_points}, Query points: {num_query_points}, Seed: {seed}")
 
     k = 10
@@ -37,7 +39,7 @@ def _test_random_points(num_points: int, num_query_points: int, seed: Optional[i
     for i in range(points.shape[0]):
         object_ = VectorEntity()
         object_.name = f"point_{i}"
-        object_.vector = points[i]
+        setattr(object_, vector_field_name, points[i])
         objects.append(object_)
     box.put(*objects)
     print(f"DB seeded with {box.count()} random points!")
@@ -58,50 +60,62 @@ def _test_random_points(num_points: int, num_query_points: int, seed: Optional[i
 
         # Run ANN with OBX
         qb = box.query()
-        qb.nearest_neighbors_f32("vector", query_point, k)
+        qb.nearest_neighbors_f32(vector_field_name, query_point, k)
         query = qb.build()
         obx_result = [id_ for id_, score in query.find_ids_with_scores()]  # Ignore score
         assert len(obx_result) == k
 
         # We would like at least half of the expected results, to be returned by the search (in any order)
         # Remember: it's an approximate search!
         search_score = len(np.intersect1d(expected_result, obx_result)) / k
-        assert search_score >= 0.5  # TODO likely could be increased
+        assert search_score >= min_score  # TODO likely could be increased
 
     print(f"Done!")
 
 
 def test_random_points():
-    _test_random_points(num_points=100, num_query_points=10, seed=10)
-    _test_random_points(num_points=100, num_query_points=10, seed=11)
-    _test_random_points(num_points=100, num_query_points=10, seed=12)
-    _test_random_points(num_points=100, num_query_points=10, seed=13)
-    _test_random_points(num_points=100, num_query_points=10, seed=14)
-    _test_random_points(num_points=100, num_query_points=10, seed=15)
-
-
-def test_combined_nn_search():
-    """ Tests NN search combined with regular query conditions, offset and limit. """
-
+        
+    min_score = 0.5
+    distance_type = HnswDistanceType.EUCLIDEAN  
+    _test_random_points(num_points=100, num_query_points=10, seed=10, distance_type=distance_type, min_score=min_score)
+    _test_random_points(num_points=100, num_query_points=10, seed=11, distance_type=distance_type, min_score=min_score)
+    _test_random_points(num_points=100, num_query_points=10, seed=12, distance_type=distance_type, min_score=min_score)
+    _test_random_points(num_points=100, num_query_points=10, seed=13, distance_type=distance_type, min_score=min_score)
+    _test_random_points(num_points=100, num_query_points=10, seed=14, distance_type=distance_type, min_score=min_score)
+    _test_random_points(num_points=100, num_query_points=10, seed=15, distance_type=distance_type, min_score=min_score)
+
+    # TODO: Cosine and Dot Product may result in 0 score
+
+def _test_combined_nn_search(distance_type: HnswDistanceType = HnswDistanceType.EUCLIDEAN):
+    
     db = create_test_objectbox()
 
     box = objectbox.Box(db, VectorEntity)
 
-    box.put(VectorEntity(name="Power of red", vector=[1, 1]))
-    box.put(VectorEntity(name="Blueberry", vector=[2, 2]))
-    box.put(VectorEntity(name="Red", vector=[3, 3]))
-    box.put(VectorEntity(name="Blue sea", vector=[4, 4]))
-    box.put(VectorEntity(name="Lightblue", vector=[5, 5]))
-    box.put(VectorEntity(name="Red apple", vector=[6, 6]))
-    box.put(VectorEntity(name="Hundred", vector=[7, 7]))
-    box.put(VectorEntity(name="Tired", vector=[8, 8]))
-    box.put(VectorEntity(name="Power of blue", vector=[9, 9]))
-
+    vector_field_name = "vector_"+distance_type.name.lower()
+   
+    values = [ 
+        ("Power of red", [1, 1]),
+        ("Blueberry", [2, 2]),
+        ("Red", [3, 3]),
+        ("Blue sea", [4, 4]),
+        ("Lightblue", [5, 5]),
+        ("Red apple", [6, 6]),
+        ("Hundred", [7, 7]),
+        ("Tired", [8, 8]),
+        ("Power of blue", [9, 9])
+    ]
+    for value in values:
+        entity = VectorEntity()
+        setattr(entity, "name", value[0])
+        setattr(entity, vector_field_name, value[1])
+        box.put(entity)
+        
     assert box.count() == 9
 
     # Test condition + NN search
     qb = box.query()
-    qb.nearest_neighbors_f32("vector", [4.1, 4.2], 6)
+    qb.nearest_neighbors_f32(vector_field_name, [4.1, 4.2], 6)
     qb.contains_string("name", "red", case_sensitive=False)
     query = qb.build()
     # 4, 5, 3, 6, 2, 7
@@ -121,7 +135,7 @@ def test_combined_nn_search():
 
     # Regular condition + NN search
     qb = box.query()
-    qb.nearest_neighbors_f32("vector", [9.2, 8.9], 7)
+    qb.nearest_neighbors_f32(vector_field_name, [9.2, 8.9], 7)
     qb.starts_with_string("name", "Blue", case_sensitive=True)
     query = qb.build()
 
@@ -131,7 +145,7 @@ def test_combined_nn_search():
 
     # Regular condition + NN search
     qb = box.query()
-    qb.nearest_neighbors_f32("vector", [7.7, 7.7], 8)
+    qb.nearest_neighbors_f32(vector_field_name, [7.7, 7.7], 8)
     qb.contains_string("name", "blue", case_sensitive=False)
     query = qb.build()
     # 8, 7, 9, 6, 5, 4, 3, 2
@@ -157,3 +171,10 @@ def test_combined_nn_search():
     assert len(search_results) == 2
     assert search_results[0] == 4
     assert search_results[1] == 5
+
+
+def test_combined_nn_search():
+    """ Tests NN search combined with regular query conditions, offset and limit. """
+    distance_type = HnswDistanceType.EUCLIDEAN
+    _test_combined_nn_search(distance_type)
+    # TODO: Cosine, DotProduct  diverges see below

tests/test_query.py

@@ -16,9 +16,9 @@ def test_basics():
     box_test_entity.put(TestEntity(str="bar", int64=456))
 
     box_vector_entity = objectbox.Box(ob, VectorEntity)
-    box_vector_entity.put(VectorEntity(name="Object 1", vector=[1, 1]))
-    box_vector_entity.put(VectorEntity(name="Object 2", vector=[2, 2]))
-    box_vector_entity.put(VectorEntity(name="Object 3", vector=[3, 3]))
+    box_vector_entity.put(VectorEntity(name="Object 1", vector_euclidean=[1, 1]))
+    box_vector_entity.put(VectorEntity(name="Object 2", vector_euclidean=[2, 2]))
+    box_vector_entity.put(VectorEntity(name="Object 3", vector_euclidean=[3, 3]))
 
     # String query
     str_prop: Property = TestEntity.get_property("str")
@@ -98,7 +98,7 @@ def test_basics():
     assert query.remove() == 1
 
     # NN query
-    vector_prop: Property = VectorEntity.get_property("vector")
+    vector_prop: Property = VectorEntity.get_property("vector_euclidean")
 
     query = box_vector_entity.query(vector_prop.nearest_neighbor([2.1, 2.1], 2)).build()
     assert query.count() == 2
@@ -258,11 +258,11 @@ def test_set_parameter():
     box_test_entity.put(TestEntity(str="Barrakuda", int64=4, int32=386, int8=60))
 
     box_vector_entity = objectbox.Box(db, VectorEntity)
-    box_vector_entity.put(VectorEntity(name="Object 1", vector=[1, 1]))
-    box_vector_entity.put(VectorEntity(name="Object 2", vector=[2, 2]))
-    box_vector_entity.put(VectorEntity(name="Object 3", vector=[3, 3]))
-    box_vector_entity.put(VectorEntity(name="Object 4", vector=[4, 4]))
-    box_vector_entity.put(VectorEntity(name="Object 5", vector=[5, 5]))
+    box_vector_entity.put(VectorEntity(name="Object 1", vector_euclidean=[1, 1]))
+    box_vector_entity.put(VectorEntity(name="Object 2", vector_euclidean=[2, 2]))
+    box_vector_entity.put(VectorEntity(name="Object 3", vector_euclidean=[3, 3]))
+    box_vector_entity.put(VectorEntity(name="Object 4", vector_euclidean=[4, 4]))
+    box_vector_entity.put(VectorEntity(name="Object 5", vector_euclidean=[5, 5]))
 
     qb = box_test_entity.query()
     qb.starts_with_string("str", "fo", case_sensitive=False)
@@ -280,22 +280,22 @@ def test_set_parameter():
     assert query.find_ids() == [3]
 
     qb = box_vector_entity.query()
-    qb.nearest_neighbors_f32("vector", [3.4, 3.4], 3)
+    qb.nearest_neighbors_f32("vector_euclidean", [3.4, 3.4], 3)
     query = qb.build()
     assert query.find_ids() == sorted([3, 4, 2])
 
     # set_parameter_vector_f32
     # set_parameter_int (NN count)
-    query.set_parameter_vector_f32("vector", [4.9, 4.9])
+    query.set_parameter_vector_f32("vector_euclidean", [4.9, 4.9])
     assert query.find_ids() == sorted([5, 4, 3])
 
-    query.set_parameter_vector_f32("vector", [0, 0])
+    query.set_parameter_vector_f32("vector_euclidean", [0, 0])
     assert query.find_ids() == sorted([1, 2, 3])
 
-    query.set_parameter_vector_f32("vector", [2.5, 2.1])
+    query.set_parameter_vector_f32("vector_euclidean", [2.5, 2.1])
     assert query.find_ids() == sorted([2, 3, 1])
 
-    query.set_parameter_int("vector", 2)
+    query.set_parameter_int("vector_euclidean", 2)
     assert query.find_ids() == sorted([2, 3])
 
 
@@ -307,11 +307,11 @@ def test_set_parameter_alias():
     box.put(TestEntity(str="FooBar", int64=10, int32=49, int8=45))
 
     box_vector = objectbox.Box(db, VectorEntity)
-    box_vector.put(VectorEntity(name="Object 1", vector=[1, 1]))
-    box_vector.put(VectorEntity(name="Object 2", vector=[2, 2]))
-    box_vector.put(VectorEntity(name="Object 3", vector=[3, 3]))
-    box_vector.put(VectorEntity(name="Object 4", vector=[4, 4]))
-    box_vector.put(VectorEntity(name="Object 5", vector=[5, 5]))
+    box_vector.put(VectorEntity(name="Object 1", vector_euclidean=[1, 1]))
+    box_vector.put(VectorEntity(name="Object 2", vector_euclidean=[2, 2]))
+    box_vector.put(VectorEntity(name="Object 3", vector_euclidean=[3, 3]))
+    box_vector.put(VectorEntity(name="Object 4", vector_euclidean=[4, 4]))
+    box_vector.put(VectorEntity(name="Object 5", vector_euclidean=[5, 5]))
 
     str_prop: Property = TestEntity.get_property("str")
     int32_prop: Property = TestEntity.get_property("int32")
@@ -354,7 +354,7 @@ def test_set_parameter_alias():
     assert query.find()[0].str == "FooBar"
 
     # Test set parameter alias on vector
-    vector_prop: Property = VectorEntity.get_property("vector")
+    vector_prop: Property = VectorEntity.get_property("vector_euclidean")
 
     query = box_vector.query(vector_prop.nearest_neighbor([3.4, 3.4], 3).alias("nearest_neighbour_filter")).build()
     assert query.count() == 3