FIX: Use intro sort for simultaneous_sort (2 way partitioning) (scikit-learn#33252)

Author: cakedev0

doc/whats_new/upcoming_changes/sklearn.metrics/33252.fix.rst

@@ -0,0 +1,5 @@
+- Fixed :func:`metrics.pairwise_distances_argmin` and
+  :func:`metrics.pairwise_distances_argmin_min` to avoid a quadratic-time path
+  when many distances are identical, which could lead to severe slowdowns or
+  even a stack overflow (segmentation fault) on large inputs.
+  By :user:`Arthur Lacote <cakedev0>`.

doc/whats_new/upcoming_changes/sklearn.neighbors/33252.fix.rst

@@ -0,0 +1,6 @@
+- Fixed a quadratic-time path in the internal ``simultaneous_sort`` used by
+  :class:`neighbors.BallTree` and :class:`neighbors.KDTree` queries when many
+  distances are identical, which could lead to severe slowdowns or even a stack
+  overflow (segmentation fault) on large inputs. Neighbor searches with tied
+  distances no longer degrade badly in runtime.
+  By :user:`Arthur Lacote <cakedev0>`.

sklearn/neighbors/_binary_tree.pxi.tp

@@ -583,9 +583,9 @@ cdef class NeighborsHeap{{name_suffix}}:
         cdef intp_t row
         for row in range(self.distances.shape[0]):
             _simultaneous_sort(
-                dist=&self.distances[row, 0],
-                idx=&self.indices[row, 0],
-                size=self.distances.shape[1],
+                values=&self.distances[row, 0],
+                indices=&self.indices[row, 0],
+                n=self.distances.shape[1],
             )
         return 0
 

sklearn/tree/_criterion.pyx

@@ -13,7 +13,7 @@ from scipy.special.cython_special cimport xlogy
 
 from sklearn.tree._utils cimport log
 from sklearn.tree._utils cimport WeightedFenwickTree
-from sklearn.utils._sorting cimport sort
+from sklearn.utils._sorting cimport simultaneous_sort
 
 # EPSILON is used in the Poisson criterion
 cdef float64_t EPSILON = 10 * np.finfo('double').eps
@@ -1138,7 +1138,7 @@ cdef inline void compute_ranks(
     cdef intp_t i
     for i in range(n):
         sorted_indices[i] = i
-    sort(sorted_y, sorted_indices, n)
+    simultaneous_sort(sorted_y, sorted_indices, n, use_three_way_partition=True)
     for i in range(n):
         ranks[sorted_indices[i]] = i
 

sklearn/tree/_partitioner.pyx

@@ -21,7 +21,7 @@ cnp.import_array()
 from scipy.sparse import issparse
 
 from sklearn.tree._splitter cimport SplitRecord
-from sklearn.utils._sorting cimport sort
+from sklearn.utils._sorting cimport simultaneous_sort
 
 # Constant to switch between algorithm non zero value extract algorithm
 # in SparsePartitioner
@@ -100,7 +100,12 @@ cdef class DensePartitioner:
             for i in range(self.start, self.end):
                 feature_values[i] = X[samples[i], current_feature]
 
-        sort(&feature_values[self.start], &samples[self.start], self.end - self.start - n_missing)
+        simultaneous_sort(
+            &feature_values[self.start],
+            &samples[self.start],
+            self.end - self.start - n_missing,
+            use_three_way_partition=True,
+        )
         self.n_missing = n_missing
 
     cdef void shift_missing_to_the_left(self) noexcept nogil:
@@ -331,12 +336,18 @@ cdef class SparsePartitioner:
 
         self.extract_nnz(current_feature)
         # Sort the positive and negative parts of `feature_values`
-        sort(&feature_values[self.start], &samples[self.start], self.end_negative - self.start)
+        simultaneous_sort(
+            &feature_values[self.start],
+            &samples[self.start],
+            self.end_negative - self.start,
+            use_three_way_partition=True,
+        )
         if self.start_positive < self.end:
-            sort(
+            simultaneous_sort(
                 &feature_values[self.start_positive],
                 &samples[self.start_positive],
-                self.end - self.start_positive
+                self.end - self.start_positive,
+                use_three_way_partition=True,
             )
 
         # Update index_to_samples to take into account the sort

sklearn/utils/_sorting.pxd

@@ -2,14 +2,9 @@ from sklearn.utils._typedefs cimport intp_t
 
 from cython cimport floating
 
-cdef int simultaneous_sort(
-    floating *dist,
-    intp_t *idx,
-    intp_t size,
-) noexcept nogil
-
-cdef void sort(
-    floating* feature_values,
-    intp_t* samples,
+cdef void simultaneous_sort(
+    floating* values,
+    intp_t* indices,
     intp_t n,
+    bint use_three_way_partition=*,
 ) noexcept nogil