Query: add find_ids_by_score_numpy()

Author: greenrobot

objectbox/query.py

@@ -102,6 +102,23 @@ def find_ids_by_score(self) -> List[int]:
         finally:
             obx_id_array_free(c_id_array_p)
 
+    def find_ids_by_score_numpy(self) -> np.array:
+        """ Finds object IDs matching the query ordered by their query score (e.g. distance in NN search).
+        The resulting list of IDs is sorted by score in ascending order. """
+        # TODO extract utility function for ID array conversion
+        c_id_array_p = obx_query_find_ids_by_score(self._c_query)
+        try:
+            c_id_array: OBX_id_array = c_id_array_p.contents
+            c_count = c_id_array.count
+            if c_count == 0:
+                return []
+            c_ids = ctypes.cast(c_id_array.ids, ctypes.POINTER(obx_id))
+            numpy_array = np.empty(c_count, dtype=np.uint64)
+            ctypes.memmove(numpy_array.ctypes.data, c_ids, numpy_array.nbytes)
+            return numpy_array
+        finally:
+            obx_id_array_free(c_id_array_p)
+
     def count(self) -> int:
         count = ctypes.c_uint64()
         obx_query_count(self._c_query, ctypes.byref(count))

tests/test_hnsw.py

@@ -170,6 +170,13 @@ def _test_combined_nn_search(distance_type: VectorDistanceType = VectorDistanceT
     assert search_results[2] == 4
     assert search_results[3] == 2
 
+    search_results = query.find_ids_by_score_numpy()
+    assert search_results.size == 4
+    assert search_results[0] == 9
+    assert search_results[1] == 5
+    assert search_results[2] == 4
+    assert search_results[3] == 2
+
     search_results = query.find_ids()
     assert len(search_results) == 4
     assert search_results[0] == 2