Repository: scikit-learn/scikit-learn
Issue category: Performance Issue

Issue:
knn predict unreasonably slow b/c of use of scipy.stats.mode
```python
import numpy as np
from sklearn.datasets import make_blobs
from sklearn.neighbors import KNeighborsClassifier

X, y = make_blobs(centers=2, random_state=4, n_samples=30)
knn = KNeighborsClassifier(algorithm='kd_tree').fit(X, y)

x_min, x_max = X[:, 0].min(), X[:, 0].max()
y_min, y_max = X[:, 1].min(), X[:, 1].max()

xx = np.linspace(x_min, x_max, 1000)
# change 100 to 1000 below and wait a long time                                          
yy = np.linspace(y_min, y_max, 100)                                          

X1, X2 = np.meshgrid(xx, yy)                                                  
X_grid = np.c_[X1.ravel(), X2.ravel()]                                        
decision_values = knn.predict(X_grid)
```

spends all it's time in unique within stats.mode, not within the distance calculation. ``mode`` runs ``unique`` for every row.
I'm pretty sure we can replace the call to mode by some call to making a csr matrix and then argmax.

How much is it worth optimizing this? I feel KNN should be fast in low dimensions and people might actually use this. Having the bottleneck in the wrong place just feels wrong to me ;)
WIP PERF Faster KNeighborsClassifier.predict
Closes https://github.com/scikit-learn/scikit-learn/issues/13783

This makes `KNeighborsClassifier.predict` faster by re-writing `scipy.stats.mode` as an argmax of a sparse array as discussed in the parent issue.

Using the example provided in https://github.com/scikit-learn/scikit-learn/issues/13783#issue-440225690,

**On master**
```
%timeit knn.predict(X_grid)                                                 
2.47 s ± 37.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
**With this PR**
```
%timeit knn.predict(X_grid)                                                 
793 ms ± 39.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

so in this particular case, the `KNeighborsClassifier.predict` is 3.1x faster.

It works in a straightforward way both on weighted an unweighted data making `sklearn.utils.weighted_mode` no longer necessary. 

The downside is that it makes `RadiusNeighborsClassifier.predict` slower by about 30% on the following example,
<details>

```
In [1]: import numpy as np 
   ...: from sklearn.datasets import make_blobs 
   ...: from sklearn.neighbors import RadiusNeighborsClassifier 
   ...:  
   ...: X, y = make_blobs(centers=2, random_state=4, n_samples=30) 
   ...: knn = RadiusNeighborsClassifier(algorithm='kd_tree', radius=2).fit(X, y) 
   ...:  
   ...: x_min, x_max = X[:, 0].min(), X[:, 0].max() 
   ...: y_min, y_max = X[:, 1].min(), X[:, 1].max() 
   ...:  
   ...: xx = np.linspace(x_min, x_max, 100) 
   ...: # change 100 to 1000 below and wait a long time                             
   ...:               
   ...: yy = np.linspace(y_min, y_max, 100)                                         
   ...:   
   ...:  
   ...: X1, X2 = np.meshgrid(xx, yy)                                                
   ...:    
   ...: X_grid = np.c_[X1.ravel(), X2.ravel()]                                      

In [2]: %timeit knn.predict(X_grid)                                                 
1.27 s ± 9.02 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
</details>
The difference is mostly that in `RadiusNeighborsClassifier` we call this function repeatedly on 1D arrays, as opposed to a single 2D array.  In worst case, I could revert back to `scipy.stats.mode` and `sklearn.utils.weighted_mode` for `RadiusNeighborsClassifier` but it's annoying to use two different implementations.

**TODO**
 - [ ] fix the remaining 2 test failures on the comparison between the multi-output and single-output case.
 - [ ] ~investigate RadiusNeighborsClassifier.predict performance regression.~


Hints:
@amueller Do you mean something like this

```
max(top_k, key = list(top_k).count)
```
That isn't going to apply to every row, and involves n_classes passes over
each. Basically because we know the set of class labels, we shouldn't need
to be doing unique.

Yes, we could construct a CSR sparse matrix and sum_duplicates and get the
argmax. Or we could just run bincount and argmax for each row. The question
is if it speeds things up enough to be worth the code.

It might also be possible to use np.add.at to effectively do bincount in
2d...?

Pretty sure doing a CSR construction would speed it up by several orders of magnitude.
@amueller @jnothman Is this something I can try?
You're welcome to submit a pull request!

Cool! Will try it out
Proposed a WIP solution in https://github.com/scikit-learn/scikit-learn/pull/14543
At https://github.com/scikit-learn/scikit-learn/pull/9597#issuecomment-424379575, @TomDLT pointed out that argmax of `predict_proba` is faster than the current `predict` implementation. Any proposal here should compare to using that approach (not yet implemented there) and avoiding mode altogether.
> At #9597 (comment), @TomDLT pointed out that argmax of predict_proba is faster than the current predict implementation. Any proposal here should compare to using that approach (not yet implemented there) and avoiding mode altogether.

Yes, I'm happy  with #9597 and using the argmax as well. Will try to make some benchmarks.
I found that `predict_proba`'s speed is acceptable (at least 10 times faster), maybe we could use a custom function instead before official fix.

```python
def predict(knn, X):
    pro_y = knn.predict_proba(X)
    y = np.argmax(pro_y, axis=1)
    return y
```
Yes, I need to finish  https://github.com/scikit-learn/scikit-learn/pull/14543 to fix it
The radius neighbours input is perfect for transforming into a 2d CSR matrix. Just set indices to the concatenated neighbour indices, and indptr to `array([len(x) for x in neighbor_indices])`.

I suspect you can efficiently use sparse matrices to compute the result for both.
Thanks for the pointer. I'll try to investigate, but first I need to fix the test failure which seems to point to a genuine bug.
@jjerphan In case you are interested to continue it, this should also give an easy performance win for KNN classification :) As a first step maybe just making this happen for `KNeighborsClassifier` would already be a start. 
It might need to re-run benchmarks, merge with main and fix the two remaining test failures.
@rth: thanks for the comment!

I might have a look after https://github.com/scikit-learn/scikit-learn/pull/19950 which might help with improving performance of `neighbours.KNeighborsClassifier.predict` :)


Thanks!

Patch:
diff --git a/.gitignore b/.gitignore
index f6250b4d5f580..89600846100a8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -90,6 +90,7 @@ sklearn/metrics/_dist_metrics.pyx
 sklearn/metrics/_dist_metrics.pxd
 sklearn/metrics/_pairwise_distances_reduction/_argkmin.pxd
 sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx
+sklearn/metrics/_pairwise_distances_reduction/_argkmin_classmode.pyx
 sklearn/metrics/_pairwise_distances_reduction/_base.pxd
 sklearn/metrics/_pairwise_distances_reduction/_base.pyx
 sklearn/metrics/_pairwise_distances_reduction/_datasets_pair.pxd
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
index 1fe4289982993..439b348ce2610 100644
--- a/doc/whats_new/v1.3.rst
+++ b/doc/whats_new/v1.3.rst
@@ -292,6 +292,11 @@ Changelog
   dissimilarity is not a metric and cannot be supported by the BallTree.
   :pr:`25417` by :user:`Guillaume Lemaitre <glemaitre>`.
 
+- |Enhancement| The performance of :meth:`neighbors.KNeighborsClassifier.predict`
+  and of :meth:`neighbors.KNeighborsClassifier.predict_proba` has been improved
+  when `n_neighbors` is large and `algorithm="brute"` with non Euclidean metrics.
+  :pr:`24076` by :user:`Meekail Zain <micky774>`, :user:`Julien Jerphanion <jjerphan>`.
+
 :mod:`sklearn.neural_network`
 .............................
 
diff --git a/setup.cfg b/setup.cfg
index 081e78c92d480..3ed576cedf92f 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -90,6 +90,7 @@ ignore =
     sklearn/metrics/_dist_metrics.pxd
     sklearn/metrics/_pairwise_distances_reduction/_argkmin.pxd
     sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx
+    sklearn/metrics/_pairwise_distances_reduction/_argkmin_classmode.pyx
     sklearn/metrics/_pairwise_distances_reduction/_base.pxd
     sklearn/metrics/_pairwise_distances_reduction/_base.pyx
     sklearn/metrics/_pairwise_distances_reduction/_datasets_pair.pxd
diff --git a/setup.py b/setup.py
index 55c5cc2923527..a083cd1e2c7a0 100755
--- a/setup.py
+++ b/setup.py
@@ -277,6 +277,12 @@ def check_package_status(package, min_version):
             "include_np": True,
             "extra_compile_args": ["-std=c++11"],
         },
+        {
+            "sources": ["_argkmin_classmode.pyx.tp"],
+            "language": "c++",
+            "include_np": True,
+            "extra_compile_args": ["-std=c++11"],
+        },
         {
             "sources": ["_radius_neighbors.pyx.tp", "_radius_neighbors.pxd.tp"],
             "language": "c++",
diff --git a/sklearn/metrics/_pairwise_distances_reduction/__init__.py b/sklearn/metrics/_pairwise_distances_reduction/__init__.py
index 133c854682f0c..baa1c9de03952 100644
--- a/sklearn/metrics/_pairwise_distances_reduction/__init__.py
+++ b/sklearn/metrics/_pairwise_distances_reduction/__init__.py
@@ -90,6 +90,7 @@
     BaseDistancesReductionDispatcher,
     ArgKmin,
     RadiusNeighbors,
+    ArgKminClassMode,
     sqeuclidean_row_norms,
 )
 
@@ -97,5 +98,6 @@
     "BaseDistancesReductionDispatcher",
     "ArgKmin",
     "RadiusNeighbors",
+    "ArgKminClassMode",
     "sqeuclidean_row_norms",
 ]
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx.tp b/sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx.tp
index 3b2bce128dcb6..7caf5e0680dea 100644
--- a/sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx.tp
+++ b/sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx.tp
@@ -178,7 +178,6 @@ cdef class ArgKmin{{name_suffix}}(BaseDistancesReduction{{name_suffix}}):
         self.heaps_r_distances_chunks[thread_num] = &self.argkmin_distances[X_start, 0]
         self.heaps_indices_chunks[thread_num] = &self.argkmin_indices[X_start, 0]
 
-    @final
     cdef void _parallel_on_X_prange_iter_finalize(
         self,
         ITYPE_t thread_num,
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_argkmin_classmode.pyx.tp b/sklearn/metrics/_pairwise_distances_reduction/_argkmin_classmode.pyx.tp
new file mode 100644
index 0000000000000..95577e5754e64
--- /dev/null
+++ b/sklearn/metrics/_pairwise_distances_reduction/_argkmin_classmode.pyx.tp
@@ -0,0 +1,208 @@
+{{py:
+
+implementation_specific_values = [
+    # Values are the following ones:
+    #
+    #       name_suffix, INPUT_DTYPE_t, INPUT_DTYPE
+    #
+    # We also use the float64 dtype and C-type names as defined in
+    # `sklearn.utils._typedefs` to maintain consistency.
+    #
+    ('64', 'DTYPE_t', 'DTYPE'),
+    ('32', 'cnp.float32_t', 'np.float32')
+]
+
+}}
+
+from cython cimport floating, integral
+from cython.parallel cimport parallel, prange
+from libcpp.map cimport map as cpp_map, pair as cpp_pair
+from libc.stdlib cimport free
+
+cimport numpy as cnp
+
+cnp.import_array()
+
+from ...utils._typedefs cimport ITYPE_t, DTYPE_t
+from ...utils._typedefs import ITYPE, DTYPE
+import numpy as np
+from scipy.sparse import issparse
+from sklearn.utils.fixes import threadpool_limits
+
+cpdef enum WeightingStrategy:
+    uniform = 0
+    # TODO: Implement the following options, most likely in
+    # `weighted_histogram_mode`
+    distance = 1
+    callable = 2
+
+{{for name_suffix, INPUT_DTYPE_t, INPUT_DTYPE in implementation_specific_values}}
+from ._argkmin cimport ArgKmin{{name_suffix}}
+from ._datasets_pair cimport DatasetsPair{{name_suffix}}
+
+cdef class ArgKminClassMode{{name_suffix}}(ArgKmin{{name_suffix}}):
+    """
+    {{name_suffix}}bit implementation of ArgKminClassMode.
+    """
+    cdef:
+        const ITYPE_t[:] class_membership,
+        const ITYPE_t[:] unique_labels
+        DTYPE_t[:, :] class_scores
+        cpp_map[ITYPE_t, ITYPE_t] labels_to_index
+        WeightingStrategy weight_type
+
+    @classmethod
+    def compute(
+        cls,
+        X,
+        Y,
+        ITYPE_t k,
+        weights,
+        class_membership,
+        unique_labels,
+        str metric="euclidean",
+        chunk_size=None,
+        dict metric_kwargs=None,
+        str strategy=None,
+    ):
+        """Compute the argkmin reduction with class_membership.
+
+        This classmethod is responsible for introspecting the arguments
+        values to dispatch to the most appropriate implementation of
+        :class:`ArgKminClassMode{{name_suffix}}`.
+
+        This allows decoupling the API entirely from the implementation details
+        whilst maintaining RAII: all temporarily allocated datastructures necessary
+        for the concrete implementation are therefore freed when this classmethod
+        returns.
+
+        No instance _must_ directly be created outside of this class method.
+        """
+        # Use a generic implementation that handles most scipy
+        # metrics by computing the distances between 2 vectors at a time.
+        pda = ArgKminClassMode{{name_suffix}}(
+            datasets_pair=DatasetsPair{{name_suffix}}.get_for(X, Y, metric, metric_kwargs),
+            k=k,
+            chunk_size=chunk_size,
+            strategy=strategy,
+            weights=weights,
+            class_membership=class_membership,
+            unique_labels=unique_labels,
+        )
+
+        # Limit the number of threads in second level of nested parallelism for BLAS
+        # to avoid threads over-subscription (in GEMM for instance).
+        with threadpool_limits(limits=1, user_api="blas"):
+            if pda.execute_in_parallel_on_Y:
+                pda._parallel_on_Y()
+            else:
+                pda._parallel_on_X()
+
+        return pda._finalize_results()
+
+    def __init__(
+        self,
+        DatasetsPair{{name_suffix}} datasets_pair,
+        const ITYPE_t[:] class_membership,
+        const ITYPE_t[:] unique_labels,
+        chunk_size=None,
+        strategy=None,
+        ITYPE_t k=1,
+        weights=None,
+    ):
+        super().__init__(
+            datasets_pair=datasets_pair,
+            chunk_size=chunk_size,
+            strategy=strategy,
+            k=k,
+        )
+
+        if weights == "uniform":
+            self.weight_type = WeightingStrategy.uniform
+        elif weights == "distance":
+            self.weight_type = WeightingStrategy.distance
+        else:
+            self.weight_type = WeightingStrategy.callable
+        self.class_membership = class_membership
+
+        self.unique_labels = unique_labels
+
+        cdef ITYPE_t idx, neighbor_class_idx
+        # Map from set of unique labels to their indices in `class_scores`
+        # Buffer used in building a histogram for one-pass weighted mode
+        self.class_scores = np.zeros(
+            (self.n_samples_X, unique_labels.shape[0]), dtype=DTYPE,
+        )
+
+    def _finalize_results(self):
+        probabilities = np.asarray(self.class_scores)
+        probabilities /= probabilities.sum(axis=1, keepdims=True)
+        return probabilities
+
+    cdef inline void weighted_histogram_mode(
+        self,
+        ITYPE_t sample_index,
+        ITYPE_t* indices,
+        DTYPE_t* distances,
+   ) noexcept nogil:
+        cdef:
+            ITYPE_t neighbor_idx, neighbor_class_idx, label_index, multi_output_index
+            DTYPE_t score_incr = 1
+            # TODO: Implement other WeightingStrategy values
+            bint use_distance_weighting = (
+                self.weight_type == WeightingStrategy.distance
+            )
+
+        # Iterate through the sample k-nearest neighbours
+        for neighbor_rank in range(self.k):
+            # Absolute indice of the neighbor_rank-th Nearest Neighbors
+            # in range [0, n_samples_Y)
+            # TODO: inspect if it worth permuting this condition
+            # and the for-loop above for improved branching.
+            if use_distance_weighting:
+                score_incr = 1 / distances[neighbor_rank]
+            neighbor_idx = indices[neighbor_rank]
+            neighbor_class_idx = self.class_membership[neighbor_idx]
+            self.class_scores[sample_index][neighbor_class_idx] += score_incr
+        return
+
+    cdef void _parallel_on_X_prange_iter_finalize(
+        self,
+        ITYPE_t thread_num,
+        ITYPE_t X_start,
+        ITYPE_t X_end,
+    ) noexcept nogil:
+        cdef:
+            ITYPE_t idx, sample_index
+        for idx in range(X_end - X_start):
+            # One-pass top-one weighted mode
+            # Compute the absolute index in [0, n_samples_X)
+            sample_index = X_start + idx
+            self.weighted_histogram_mode(
+                sample_index,
+                &self.heaps_indices_chunks[thread_num][idx * self.k],
+                &self.heaps_r_distances_chunks[thread_num][idx * self.k],
+            )
+        return
+
+    cdef void _parallel_on_Y_finalize(
+        self,
+    ) noexcept nogil:
+        cdef:
+            ITYPE_t sample_index, thread_num
+
+        with nogil, parallel(num_threads=self.chunks_n_threads):
+            # Deallocating temporary datastructures
+            for thread_num in prange(self.chunks_n_threads, schedule='static'):
+                free(self.heaps_r_distances_chunks[thread_num])
+                free(self.heaps_indices_chunks[thread_num])
+
+            for sample_index in prange(self.n_samples_X, schedule='static'):
+                self.weighted_histogram_mode(
+                    sample_index,
+                    &self.argkmin_indices[sample_index][0],
+                    &self.argkmin_distances[sample_index][0],
+                )
+        return
+
+{{endfor}}
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_dispatcher.py b/sklearn/metrics/_pairwise_distances_reduction/_dispatcher.py
index 576cc64ff5295..73d98f2ebe6b2 100644
--- a/sklearn/metrics/_pairwise_distances_reduction/_dispatcher.py
+++ b/sklearn/metrics/_pairwise_distances_reduction/_dispatcher.py
@@ -4,7 +4,7 @@
 
 from typing import List
 
-from scipy.sparse import isspmatrix_csr
+from scipy.sparse import isspmatrix_csr, issparse
 
 from .._dist_metrics import BOOL_METRICS, METRIC_MAPPING
 
@@ -13,6 +13,12 @@
     ArgKmin64,
     ArgKmin32,
 )
+
+from ._argkmin_classmode import (
+    ArgKminClassMode64,
+    ArgKminClassMode32,
+)
+
 from ._radius_neighbors import (
     RadiusNeighbors64,
     RadiusNeighbors32,
@@ -428,3 +434,187 @@ def compute(
             "Only float64 or float32 datasets pairs are supported at this time, "
             f"got: X.dtype={X.dtype} and Y.dtype={Y.dtype}."
         )
+
+
+class ArgKminClassMode(BaseDistancesReductionDispatcher):
+    """Compute the argkmin of row vectors of X on the ones of Y with labels.
+
+    For each row vector of X, computes the indices of k first the rows
+    vectors of Y with the smallest distances. Computes weighted mode of labels.
+
+    ArgKminClassMode is typically used to perform bruteforce k-nearest neighbors
+    queries when the weighted mode of the labels for the k-nearest neighbors
+    are required, such as in `predict` methods.
+
+    This class is not meant to be instanciated, one should only use
+    its :meth:`compute` classmethod which handles allocation and
+    deallocation consistently.
+    """
+
+    @classmethod
+    def is_usable_for(cls, X, Y, metric) -> bool:
+        """Return True if the dispatcher can be used for the given parameters.
+
+        Parameters
+        ----------
+        X : ndarray of shape (n_samples_X, n_features)
+            The input array to be labelled.
+
+        Y : ndarray of shape (n_samples_Y, n_features)
+            The input array whose labels are provided through the `labels`
+            parameter.
+
+        metric : str, default='euclidean'
+            The distance metric to use. For a list of available metrics, see
+            the documentation of :class:`~sklearn.metrics.DistanceMetric`.
+            Currently does not support `'precomputed'`.
+
+        Returns
+        -------
+        True if the PairwiseDistancesReduction can be used, else False.
+        """
+        return (
+            ArgKmin.is_usable_for(X, Y, metric)
+            # TODO: Support CSR matrices.
+            and not issparse(X)
+            and not issparse(Y)
+            # TODO: implement Euclidean specialization with GEMM.
+            and metric not in ("euclidean", "sqeuclidean")
+        )
+
+    @classmethod
+    def compute(
+        cls,
+        X,
+        Y,
+        k,
+        weights,
+        labels,
+        unique_labels,
+        metric="euclidean",
+        chunk_size=None,
+        metric_kwargs=None,
+        strategy=None,
+    ):
+        """Compute the argkmin reduction.
+
+        Parameters
+        ----------
+        X : ndarray of shape (n_samples_X, n_features)
+            The input array to be labelled.
+
+        Y : ndarray of shape (n_samples_Y, n_features)
+            The input array whose labels are provided through the `labels`
+            parameter.
+
+        k : int
+            The number of nearest neighbors to consider.
+
+        weights : ndarray
+            The weights applied over the `labels` of `Y` when computing the
+            weighted mode of the labels.
+
+        class_membership : ndarray
+            An array containing the index of the class membership of the
+            associated samples in `Y`. This is used in labeling `X`.
+
+        unique_classes : ndarray
+            An array containing all unique indices contained in the
+            corresponding `class_membership` array.
+
+        metric : str, default='euclidean'
+            The distance metric to use. For a list of available metrics, see
+            the documentation of :class:`~sklearn.metrics.DistanceMetric`.
+            Currently does not support `'precomputed'`.
+
+        chunk_size : int, default=None,
+            The number of vectors per chunk. If None (default) looks-up in
+            scikit-learn configuration for `pairwise_dist_chunk_size`,
+            and use 256 if it is not set.
+
+        metric_kwargs : dict, default=None
+            Keyword arguments to pass to specified metric function.
+
+        strategy : str, {'auto', 'parallel_on_X', 'parallel_on_Y'}, default=None
+            The chunking strategy defining which dataset parallelization are made on.
+
+            For both strategies the computations happens with two nested loops,
+            respectively on chunks of X and chunks of Y.
+            Strategies differs on which loop (outer or inner) is made to run
+            in parallel with the Cython `prange` construct:
+
+              - 'parallel_on_X' dispatches chunks of X uniformly on threads.
+                Each thread then iterates on all the chunks of Y. This strategy is
+                embarrassingly parallel and comes with no datastructures
+                synchronisation.
+
+              - 'parallel_on_Y' dispatches chunks of Y uniformly on threads.
+                Each thread processes all the chunks of X in turn. This strategy is
+                a sequence of embarrassingly parallel subtasks (the inner loop on Y
+                chunks) with intermediate datastructures synchronisation at each
+                iteration of the sequential outer loop on X chunks.
+
+              - 'auto' relies on a simple heuristic to choose between
+                'parallel_on_X' and 'parallel_on_Y': when `X.shape[0]` is large enough,
+                'parallel_on_X' is usually the most efficient strategy.
+                When `X.shape[0]` is small but `Y.shape[0]` is large, 'parallel_on_Y'
+                brings more opportunity for parallelism and is therefore more efficient
+                despite the synchronization step at each iteration of the outer loop
+                on chunks of `X`.
+
+              - None (default) looks-up in scikit-learn configuration for
+                `pairwise_dist_parallel_strategy`, and use 'auto' if it is not set.
+
+        Returns
+        -------
+        probabilities : ndarray of shape (n_samples_X, n_classes)
+            An array containing the class probabilities for each sample.
+
+        Notes
+        -----
+        This classmethod is responsible for introspecting the arguments
+        values to dispatch to the most appropriate implementation of
+        :class:`PairwiseDistancesArgKmin`.
+
+        This allows decoupling the API entirely from the implementation details
+        whilst maintaining RAII: all temporarily allocated datastructures necessary
+        for the concrete implementation are therefore freed when this classmethod
+        returns.
+        """
+        if weights not in {"uniform", "distance"}:
+            raise ValueError(
+                "Only the 'uniform' or 'distance' weights options are supported"
+                f" at this time. Got: {weights=}."
+            )
+        if X.dtype == Y.dtype == np.float64:
+            return ArgKminClassMode64.compute(
+                X=X,
+                Y=Y,
+                k=k,
+                weights=weights,
+                class_membership=np.array(labels, dtype=np.intp),
+                unique_labels=np.array(unique_labels, dtype=np.intp),
+                metric=metric,
+                chunk_size=chunk_size,
+                metric_kwargs=metric_kwargs,
+                strategy=strategy,
+            )
+
+        if X.dtype == Y.dtype == np.float32:
+            return ArgKminClassMode32.compute(
+                X=X,
+                Y=Y,
+                k=k,
+                weights=weights,
+                class_membership=np.array(labels, dtype=np.intp),
+                unique_labels=np.array(unique_labels, dtype=np.intp),
+                metric=metric,
+                chunk_size=chunk_size,
+                metric_kwargs=metric_kwargs,
+                strategy=strategy,
+            )
+
+        raise ValueError(
+            "Only float64 or float32 datasets pairs are supported at this time, "
+            f"got: X.dtype={X.dtype} and Y.dtype={Y.dtype}."
+        )
diff --git a/sklearn/neighbors/_base.py b/sklearn/neighbors/_base.py
index 88febbd9a3aea..9e3dfa248f00e 100644
--- a/sklearn/neighbors/_base.py
+++ b/sklearn/neighbors/_base.py
@@ -475,7 +475,10 @@ def _fit(self, X, y=None):
 
                 check_classification_targets(y)
                 self.classes_ = []
-                self._y = np.empty(y.shape, dtype=int)
+                # Using `dtype=np.intp` is necessary since `np.bincount`
+                # (called in _classification.py) fails when dealing
+                # with a float64 array on 32bit systems.
+                self._y = np.empty(y.shape, dtype=np.intp)
                 for k in range(self._y.shape[1]):
                     classes, self._y[:, k] = np.unique(y[:, k], return_inverse=True)
                     self.classes_.append(classes)
diff --git a/sklearn/neighbors/_classification.py b/sklearn/neighbors/_classification.py
index b849d28e131a5..dbc070987d5d0 100644
--- a/sklearn/neighbors/_classification.py
+++ b/sklearn/neighbors/_classification.py
@@ -12,13 +12,24 @@
 import numpy as np
 from ..utils.fixes import _mode
 from ..utils.extmath import weighted_mode
-from ..utils.validation import _is_arraylike, _num_samples
+from ..utils.validation import _is_arraylike, _num_samples, check_is_fitted
 
 import warnings
 from ._base import _get_weights
 from ._base import NeighborsBase, KNeighborsMixin, RadiusNeighborsMixin
 from ..base import ClassifierMixin
+from ..metrics._pairwise_distances_reduction import ArgKminClassMode
 from ..utils._param_validation import StrOptions
+from sklearn.neighbors._base import _check_precomputed
+
+
+def _adjusted_metric(metric, metric_kwargs, p=None):
+    metric_kwargs = metric_kwargs or {}
+    if metric == "minkowski":
+        metric_kwargs["p"] = p
+        if p == 2:
+            metric = "euclidean"
+    return metric, metric_kwargs
 
 
 class KNeighborsClassifier(KNeighborsMixin, ClassifierMixin, NeighborsBase):
@@ -228,7 +239,21 @@ def predict(self, X):
         y : ndarray of shape (n_queries,) or (n_queries, n_outputs)
             Class labels for each data sample.
         """
+        check_is_fitted(self, "_fit_method")
         if self.weights == "uniform":
+            if self._fit_method == "brute" and ArgKminClassMode.is_usable_for(
+                X, self._fit_X, self.metric
+            ):
+                probabilities = self.predict_proba(X)
+                if self.outputs_2d_:
+                    return np.stack(
+                        [
+                            self.classes_[idx][np.argmax(probas, axis=1)]
+                            for idx, probas in enumerate(probabilities)
+                        ],
+                        axis=1,
+                    )
+                return self.classes_[np.argmax(probabilities, axis=1)]
             # In that case, we do not need the distances to perform
             # the weighting so we do not compute them.
             neigh_ind = self.kneighbors(X, return_distance=False)
@@ -277,7 +302,47 @@ def predict_proba(self, X):
             The class probabilities of the input samples. Classes are ordered
             by lexicographic order.
         """
+        check_is_fitted(self, "_fit_method")
         if self.weights == "uniform":
+            # TODO: systematize this mapping of metric for
+            # PairwiseDistancesReductions.
+            metric, metric_kwargs = _adjusted_metric(
+                metric=self.metric, metric_kwargs=self.metric_params, p=self.p
+            )
+            if (
+                self._fit_method == "brute"
+                and ArgKminClassMode.is_usable_for(X, self._fit_X, metric)
+                # TODO: Implement efficient multi-output solution
+                and not self.outputs_2d_
+            ):
+                if self.metric == "precomputed":
+                    X = _check_precomputed(X)
+                else:
+                    X = self._validate_data(
+                        X, accept_sparse="csr", reset=False, order="C"
+                    )
+
+                probabilities = ArgKminClassMode.compute(
+                    X,
+                    self._fit_X,
+                    k=self.n_neighbors,
+                    weights=self.weights,
+                    labels=self._y,
+                    unique_labels=self.classes_,
+                    metric=metric,
+                    metric_kwargs=metric_kwargs,
+                    # `strategy="parallel_on_X"` has in practice be shown
+                    # to be more efficient than `strategy="parallel_on_Y``
+                    # on many combination of datasets.
+                    # Hence, we choose to enforce it here.
+                    # For more information, see:
+                    # https://github.com/scikit-learn/scikit-learn/pull/24076#issuecomment-1445258342  # noqa
+                    # TODO: adapt the heuristic for `strategy="auto"` for
+                    # `ArgKminClassMode` and use `strategy="auto"`.
+                    strategy="parallel_on_X",
+                )
+                return probabilities
+
             # In that case, we do not need the distances to perform
             # the weighting so we do not compute them.
             neigh_ind = self.kneighbors(X, return_distance=False)


Test patch:
diff --git a/sklearn/metrics/tests/test_pairwise_distances_reduction.py b/sklearn/metrics/tests/test_pairwise_distances_reduction.py
index ad0ddbc60e9bd..7355dfd6ba912 100644
--- a/sklearn/metrics/tests/test_pairwise_distances_reduction.py
+++ b/sklearn/metrics/tests/test_pairwise_distances_reduction.py
@@ -13,10 +13,10 @@
 from sklearn.metrics._pairwise_distances_reduction import (
     BaseDistancesReductionDispatcher,
     ArgKmin,
+    ArgKminClassMode,
     RadiusNeighbors,
     sqeuclidean_row_norms,
 )
-
 from sklearn.metrics import euclidean_distances
 from sklearn.utils.fixes import sp_version, parse_version
 from sklearn.utils._testing import (
@@ -656,6 +656,124 @@ def test_argkmin_factory_method_wrong_usages():
         ArgKmin.compute(X=X, Y=Y, k=k, metric=metric, metric_kwargs=metric_kwargs)
 
 
+def test_argkmin_classmode_factory_method_wrong_usages():
+    rng = np.random.RandomState(1)
+    X = rng.rand(100, 10)
+    Y = rng.rand(100, 10)
+    k = 5
+    metric = "manhattan"
+
+    weights = "uniform"
+    labels = rng.randint(low=0, high=10, size=100)
+    unique_labels = np.unique(labels)
+
+    msg = (
+        "Only float64 or float32 datasets pairs are supported at this time, "
+        "got: X.dtype=float32 and Y.dtype=float64"
+    )
+    with pytest.raises(ValueError, match=msg):
+        ArgKminClassMode.compute(
+            X=X.astype(np.float32),
+            Y=Y,
+            k=k,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    msg = (
+        "Only float64 or float32 datasets pairs are supported at this time, "
+        "got: X.dtype=float64 and Y.dtype=int32"
+    )
+    with pytest.raises(ValueError, match=msg):
+        ArgKminClassMode.compute(
+            X=X,
+            Y=Y.astype(np.int32),
+            k=k,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    with pytest.raises(ValueError, match="k == -1, must be >= 1."):
+        ArgKminClassMode.compute(
+            X=X,
+            Y=Y,
+            k=-1,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    with pytest.raises(ValueError, match="k == 0, must be >= 1."):
+        ArgKminClassMode.compute(
+            X=X,
+            Y=Y,
+            k=0,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    with pytest.raises(ValueError, match="Unrecognized metric"):
+        ArgKminClassMode.compute(
+            X=X,
+            Y=Y,
+            k=k,
+            metric="wrong metric",
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    with pytest.raises(
+        ValueError, match=r"Buffer has wrong number of dimensions \(expected 2, got 1\)"
+    ):
+        ArgKminClassMode.compute(
+            X=np.array([1.0, 2.0]),
+            Y=Y,
+            k=k,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    with pytest.raises(ValueError, match="ndarray is not C-contiguous"):
+        ArgKminClassMode.compute(
+            X=np.asfortranarray(X),
+            Y=Y,
+            k=k,
+            metric=metric,
+            weights=weights,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    non_existent_weights_strategy = "non_existent_weights_strategy"
+    message = (
+        "Only the 'uniform' or 'distance' weights options are supported at this time. "
+        f"Got: weights='{non_existent_weights_strategy}'."
+    )
+    with pytest.raises(ValueError, match=message):
+        ArgKminClassMode.compute(
+            X=X,
+            Y=Y,
+            k=k,
+            metric=metric,
+            weights=non_existent_weights_strategy,
+            labels=labels,
+            unique_labels=unique_labels,
+        )
+
+    # TODO: introduce assertions on UserWarnings once the Euclidean specialisation
+    # of ArgKminClassMode is supported.
+
+
 def test_radius_neighbors_factory_method_wrong_usages():
     rng = np.random.RandomState(1)
     X = rng.rand(100, 10)
@@ -1222,3 +1340,36 @@ def test_sqeuclidean_row_norms(
     with pytest.raises(ValueError):
         X = np.asfortranarray(X)
         sqeuclidean_row_norms(X, num_threads=num_threads)
+
+
+def test_argkmin_classmode_strategy_consistent():
+    rng = np.random.RandomState(1)
+    X = rng.rand(100, 10)
+    Y = rng.rand(100, 10)
+    k = 5
+    metric = "manhattan"
+
+    weights = "uniform"
+    labels = rng.randint(low=0, high=10, size=100)
+    unique_labels = np.unique(labels)
+    results_X = ArgKminClassMode.compute(
+        X=X,
+        Y=Y,
+        k=k,
+        metric=metric,
+        weights=weights,
+        labels=labels,
+        unique_labels=unique_labels,
+        strategy="parallel_on_X",
+    )
+    results_Y = ArgKminClassMode.compute(
+        X=X,
+        Y=Y,
+        k=k,
+        metric=metric,
+        weights=weights,
+        labels=labels,
+        unique_labels=unique_labels,
+        strategy="parallel_on_Y",
+    )
+    assert_array_equal(results_X, results_Y)
diff --git a/sklearn/neighbors/tests/test_neighbors.py b/sklearn/neighbors/tests/test_neighbors.py
index b82d369f2c12c..4ef2831990f0a 100644
--- a/sklearn/neighbors/tests/test_neighbors.py
+++ b/sklearn/neighbors/tests/test_neighbors.py
@@ -675,15 +675,18 @@ def test_kneighbors_classifier_predict_proba(global_dtype):
     cls = neighbors.KNeighborsClassifier(n_neighbors=3, p=1)  # cityblock dist
     cls.fit(X, y)
     y_prob = cls.predict_proba(X)
-    real_prob = np.array(
-        [
-            [0, 2.0 / 3, 1.0 / 3],
-            [1.0 / 3, 2.0 / 3, 0],
-            [1.0 / 3, 0, 2.0 / 3],
-            [0, 1.0 / 3, 2.0 / 3],
-            [2.0 / 3, 1.0 / 3, 0],
-            [2.0 / 3, 1.0 / 3, 0],
-        ]
+    real_prob = (
+        np.array(
+            [
+                [0, 2, 1],
+                [1, 2, 0],
+                [1, 0, 2],
+                [0, 1, 2],
+                [2, 1, 0],
+                [2, 1, 0],
+            ]
+        )
+        / 3.0
     )
     assert_array_equal(real_prob, y_prob)
     # Check that it also works with non integer labels