utils: bubble up vector_distance_f32 and vector_distance_to_relevance #47

Author: loryruta

objectbox/c.py

@@ -360,6 +360,12 @@ def c_array_pointer(py_list: Union[List[Any], np.ndarray], c_type):
     return ctypes.cast(c_array(py_list, c_type), ctypes.POINTER(c_type))
 
 
+# OBX_C_API float obx_vector_distance_float32(OBXVectorDistanceType type, const float* vector1, const float* vector2, size_t dimension);
+obx_vector_distance_float32 = c_fn("obx_vector_distance_float32", ctypes.c_float, [OBXVectorDistanceType, ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), ctypes.c_size_t])
+
+# OBX_C_API float obx_vector_distance_to_relevance(OBXVectorDistanceType type, float distance);
+obx_vector_distance_to_relevance = c_fn("obx_vector_distance_to_relevance", ctypes.c_float, [OBXVectorDistanceType, ctypes.c_float])
+
 # OBX_model* (void);
 obx_model = c_fn('obx_model', OBX_model_p, [])
 

objectbox/query_builder.py

@@ -2,10 +2,11 @@
 import numpy as np
 from typing import *
 
+from objectbox.c import *
 from objectbox.model.properties import Property
 from objectbox.objectbox import ObjectBox
 from objectbox.query import Query
-from objectbox.c import *
+from objectbox.utils import check_float_vector
 
 
 class QueryBuilder:
@@ -108,10 +109,11 @@ def between_2ints(self, prop: Union[int, str, Property], value_a: int, value_b:
         cond = obx_qb_between_2ints(self._c_builder, prop_id, value_a, value_b)
         return cond
 
-    def nearest_neighbors_f32(self, prop: Union[int, str, Property], query_vector: Union[np.ndarray, List[float]],
+    def nearest_neighbors_f32(self,
+                              prop: Union[int, str, Property],
+                              query_vector: Union[np.ndarray, List[float]],
                               element_count: int):
-        if isinstance(query_vector, np.ndarray) and query_vector.dtype != np.float32:
-            raise Exception(f"query_vector dtype is expected to be np.float32, got: {query_vector.dtype}")
+        check_float_vector(query_vector, "query_vector")
         prop_id = self._entity.get_property_id(prop)
         c_query_vector = c_array(query_vector, ctypes.c_float)
         cond = obx_qb_nearest_neighbors_f32(self._c_builder, prop_id, c_query_vector, element_count)

objectbox/utils.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+from objectbox.c import *
+from objectbox.model.properties import VectorDistanceType
+
+
+def check_float_vector(vector: Union[np.ndarray, List[float]], vector_name: str):
+    """ Checks that the given vector is a float vector (either np.ndarray or Python's list). """
+    if isinstance(vector, np.ndarray) and vector.dtype != np.float32:
+        raise Exception(f"{vector_name} dtype is expected to be np.float32, got: {vector.dtype}")
+    elif isinstance(vector, list) and len(vector) > 0 and (type(vector[0]) is not float):
+        raise Exception(f"{vector_name} is expected to be a float list, got vector[0]: {type(vector[0])}")
+
+
+def vector_distance_f32(distance_type: VectorDistanceType,
+                        vector1: Union[np.ndarray, List[float]],
+                        vector2: Union[np.ndarray, List[float]],
+                        dimension: int) -> float:
+    """ Utility function to calculate the distance of two vectors. """
+    check_float_vector(vector1, "vector1")
+    check_float_vector(vector2, "vector2")
+    return obx_vector_distance_float32(distance_type,
+                                       c_array(vector1, ctypes.c_float),
+                                       c_array(vector2, ctypes.c_float),
+                                       dimension)
+
+
+def vector_distance_to_relevance(distance_type: VectorDistanceType, distance: float) -> float:
+    """ Converts the given distance to a relevance score in range [0.0, 1.0], according to its type. """
+    return obx_vector_distance_to_relevance(distance_type, distance)