Query: add find_ids_by_score()

Author: greenrobot

objectbox/c.py

@@ -881,6 +881,9 @@ def c_array_pointer(py_list: Union[List[Any], np.ndarray], c_type):
 # OBX_C_API OBX_id_score_array* obx_query_find_ids_with_scores(OBX_query* query);
 obx_query_find_ids_with_scores = c_fn('obx_query_find_ids_with_scores', OBX_id_score_array_p, [OBX_query_p])
 
+# OBX_C_API OBX_id_array* obx_query_find_ids_by_score(OBX_query* query);
+obx_query_find_ids_by_score = c_fn('obx_query_find_ids_by_score', OBX_id_array_p, [OBX_query_p])
+
 # OBX_C_API obx_err obx_query_count(OBX_query* query, uint64_t* out_count);
 obx_query_count = c_fn_rc('obx_query_count', [OBX_query_p, ctypes.POINTER(ctypes.c_uint64)])
 

objectbox/query.py

@@ -88,6 +88,20 @@ def find_ids_with_scores(self) -> List[Tuple[int, float]]:
         finally:
             obx_id_score_array_free(c_id_score_array_p)
 
+    def find_ids_by_score(self) -> List[int]:
+        """ 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
+            if c_id_array.count == 0:
+                return []
+            ids = ctypes.cast(c_id_array.ids, ctypes.POINTER(obx_id * c_id_array.count))
+            return list(ids.contents)
+        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

@@ -163,6 +163,13 @@ def _test_combined_nn_search(distance_type: VectorDistanceType = VectorDistanceT
     assert search_results[2][0] == 4
     assert search_results[3][0] == 2
 
+    search_results = query.find_ids_by_score()
+    assert len(search_results) == 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