sort knn results by dist (#5505)

* sort knn results by dist * fix doctest GitOrigin-RevId: 88c01da923dc4e7ab6ab2306d585687b63016abd
pathwaycom · Jan 25, 2024 · 38f5e7d · 38f5e7d
1 parent 9793c6d
commit 38f5e7d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/python/pathway/stdlib/ml/classifiers/_knn_lsh.py b/python/pathway/stdlib/ml/classifiers/_knn_lsh.py
@@ -262,8 +262,13 @@ def knns(self) -> list[tuple[pw.Pointer, float]]:
                         :neighs
                     ]  # neighs - 1 in argpartition, because of 0-based indexing
                     k_distances = distances[knn_ids]
-                    ret = np.array(ids_filtered)[knn_ids]
-                    return list(zip(ret, k_distances))
+                    final_ids = np.array(ids_filtered)[knn_ids]
+                    if len(final_ids) == 0:
+                        return []
+                    sorted_dists, sorted_ids_by_dist = zip(
+                        *sorted(zip(k_distances, final_ids))
+                    )
+                    return list(zip(sorted_ids_by_dist, sorted_dists))
 
         knn_result: pw.Table = compute_knns_transformer(  # type: ignore
             training_data=data, flattened=flattened

diff --git a/python/pathway/stdlib/ml/index.py b/python/pathway/stdlib/ml/index.py
@@ -166,7 +166,7 @@ def get_nearest_items(
         coords  | nn                | __time__ | __diff__
         (1, 1)  | ((0, 0), (2, 3))  | 4        | 1
         (1, 1)  | ((0, 0), (2, 3))  | 6        | -1
-        (1, 1)  | ((2, 2), (0, 0))  | 6        | 1
+        (1, 1)  | ((0, 0), (2, 2))  | 6        | 1
         (-3, 1) | ((0, 0), (2, 2))  | 8        | 1
         (-3, 1) | ((0, 0), (2, 2))  | 10       | -1
         (-3, 1) | ((-3, 3), (0, 0)) | 10       | 1