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)

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

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.~

@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!

Performance module:metrics module:neighbors cython


diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7aea0d58..35082befe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 **Note**: Numbers like (\#123) point to closed Pull Requests on the fractal-tasks-core repository.

+# 1.4.1
+
+* Tasks:
+    * Remove overlap checking for output ROIs in Cellpose task to address performance issues (\#889).

 # 1.4.0

diff --git a/fractal_tasks_core/tasks/cellpose_segmentation.py b/fractal_tasks_core/tasks/cellpose_segmentation.py
index 9ee23b8e8..499ecc0b7 100644
--- a/fractal_tasks_core/tasks/cellpose_segmentation.py
+++ b/fractal_tasks_core/tasks/cellpose_segmentation.py
@@ -44,7 +44,6 @@
 from fractal_tasks_core.roi import (
     find_overlaps_in_ROI_indices,
 )
-from fractal_tasks_core.roi import get_overlapping_pairs_3D
 from fractal_tasks_core.roi import is_ROI_table_valid
 from fractal_tasks_core.roi import load_region
 from fractal_tasks_core.tables import write_table
@@ -561,15 +560,6 @@ def cellpose_segmentation(

             bbox_dataframe_list.append(bbox_df)

-            overlap_list = get_overlapping_pairs_3D(
-                bbox_df, full_res_pxl_sizes_zyx
-            )
-            if len(overlap_list) > 0:
-                logger.warning(
-                    f"ROI {indices} has "
-                    f"{len(overlap_list)} bounding-box pairs overlap"
-                )
-
         # Compute and store 0-th level to disk
         da.array(new_label_img).to_zarr(
             url=mask_zarr,

diff --git a/tests/tasks/test_workflows_cellpose_segmentation.py b/tests/tasks/test_workflows_cellpose_segmentation.py
index 2c3c04216..1ce511d30 100644
--- a/tests/tasks/test_workflows_cellpose_segmentation.py
+++ b/tests/tasks/test_workflows_cellpose_segmentation.py
@@ -602,49 +602,6 @@ def test_workflow_bounding_box(
     assert "encoding-version" in table_group.attrs.asdict().keys()

-def test_workflow_bounding_box_with_overlap(
-    tmp_path: Path,
-    zenodo_zarr: list[str],
-    caplog: pytest.LogCaptureFixture,
-    monkeypatch: MonkeyPatch,
-):
-    monkeypatch.setattr(
-        "fractal_tasks_core.tasks.cellpose_segmentation.cellpose.core.use_gpu",
-        patched_cellpose_core_use_gpu,
-    )
-
-    monkeypatch.setattr(
-        "fractal_tasks_core.tasks.cellpose_segmentation.segment_ROI",
-        patched_segment_ROI_overlapping_organoids,
-    )
-
-    # Setup caplog fixture, see
-    # https://docs.pytest.org/en/stable/how-to/logging.html#caplog-fixture
-    caplog.set_level(logging.WARNING)
-
-    # Use pre-made 3D zarr
-    zarr_dir = tmp_path / "tmp_out/"
-    zarr_urls = prepare_3D_zarr(str(zarr_dir), zenodo_zarr)
-    debug(zarr_dir)
-    debug(zarr_urls)
-
-    # Per-FOV labeling
-    channel = CellposeChannel1InputModel(
-        wavelength_id="A01_C01", normalize=CellposeCustomNormalizer()
-    )
-    for zarr_url in zarr_urls:
-        cellpose_segmentation(
-            zarr_url=zarr_url,
-            channel=channel,
-            level=3,
-            relabeling=True,
-            diameter_level0=80.0,
-            output_ROI_table="bbox_table",
-        )
-        debug(caplog.text)
-        assert "bounding-box pairs overlap" in caplog.text
-
-
 def test_cellpose_within_masked_bb_with_overlap(
     tmp_path: Path,
     zenodo_zarr: list[str],
diff --git a/tests/test_unit_ROIs.py b/tests/test_unit_ROIs.py
index eb8cc3bf4..3806bcdb2 100644
--- a/tests/test_unit_ROIs.py
+++ b/tests/test_unit_ROIs.py
@@ -30,6 +30,7 @@
     find_overlaps_in_ROI_indices,
 )
 from fractal_tasks_core.roi import get_image_grid_ROIs
+from fractal_tasks_core.roi import get_overlapping_pairs_3D
 from fractal_tasks_core.roi import get_single_image_ROI
 from fractal_tasks_core.roi import is_ROI_table_valid
 from fractal_tasks_core.roi import (
@@ -755,3 +756,32 @@ def test_create_roi_table_from_df_list_with_label_repeats():
         ]
     )
     np.testing.assert_allclose(output_array, roi_table.X)
+
+
+def test_get_overlapping_pairs_3D():
+    common_columns = {
+        "y_micrometer": [0.0, 0.0],
+        "z_micrometer": [0.0, 0.0],
+        "len_x_micrometer": [1.0, 1.0],
+        "len_y_micrometer": [1.0, 1.0],
+        "len_z_micrometer": [1.0, 1.0],
+        "label": [1, 2],
+    }
+    df_overlapping = pd.DataFrame(
+        {"x_micrometer": [0.0, 0.5], **common_columns}
+    )
+    full_res_pxl_sizes_zyx = [1.0, 1.0, 1.0]
+    df_non_overlapping = pd.DataFrame(
+        {"x_micrometer": [0.0, 2.0], **common_columns}
+    )
+    res_overlapping = get_overlapping_pairs_3D(
+        df_overlapping,
+        full_res_pxl_sizes_zyx,
+    )
+    assert len(res_overlapping) == 1
+
+    res_non_overlapping = get_overlapping_pairs_3D(
+        df_non_overlapping,
+        full_res_pxl_sizes_zyx,
+    )
+    assert len(res_non_overlapping) == 0

Review/optimize use of `get_overlapping_pairs_3D` in Cellpose task
Branching from #764 (fixed in principle with #778)

After a fix, it'd be useful to have a rough benchmark of `get_overlapping_pairs_3D for` one of those many-labels cases (say 1000 labels for a given `i_ROI`). Since this function does nothing else than printing a warning, I would be very strict in how long a runtime we can accept.

If it turns out that it is slow, we can easily improve it in a trivial way (simply stop after the first overlap is found) or also in more systematic ways (the current approach is close to being the definition of a slow Python function: a quadratic-scaling for loop, which calls a pure-Python function and then even appends to a list).

If the function is not slow, then there's no special need for a refactor.

Another thing to consider here:
This (as well as the original implementation) will only check for potential bounding box overlaps within a given ROI that's being processed, right?

A future improvement may be to check for overlaps across ROIs as well, as in theory, they could also overlap. Given that this is just for warnings, not sure how relevant this is though. Our masking should still work for this after all.
> A future improvement may be to check for overlaps across ROIs as well, as in theory, they could also overlap. Given that this is just for warnings, not sure how relevant this is though. Our masking should still work for this after all.

For the record, making this change is quite straightforward, since we should only move the check from within the ROI loop to its end (EDIT: this diff is not working, but the idea remains valid):
```diff
diff --git a/fractal_tasks_core/tasks/cellpose_segmentation.py b/fractal_tasks_core/tasks/cellpose_segmentation.py
index ea78661a..403630a1 100644
--- a/fractal_tasks_core/tasks/cellpose_segmentation.py
+++ b/fractal_tasks_core/tasks/cellpose_segmentation.py
@@ -541,15 +541,6 @@ def cellpose_segmentation(

             bbox_dataframe_list.append(bbox_df)

-            overlap_list = get_overlapping_pairs_3D(
-                bbox_df, full_res_pxl_sizes_zyx
-            )
-            if len(overlap_list) > 0:
-                logger.warning(
-                    f"ROI {indices} has "
-                    f"{len(overlap_list)} bounding-box pairs overlap"
-                )
-
         # Compute and store 0-th level to disk
         da.array(new_label_img).to_zarr(
             url=mask_zarr,
@@ -582,6 +573,15 @@ def cellpose_segmentation(
             bbox_dataframe_list = [empty_bounding_box_table()]
         # Concatenate all ROI dataframes
         df_well = pd.concat(bbox_dataframe_list, axis=0, ignore_index=True)
+
+        overlap_list = get_overlapping_pairs_3D(
+            bbox_dataframe_list, full_res_pxl_sizes_zyx
+        )
+        if len(overlap_list) > 0:
+            logger.warning(
+                f"{len(overlap_list)} bounding-box pairs overlap"
+            )
+
         df_well.index = df_well.index.astype(str)
         # Extract labels and drop them from df_well
         labels = pd.DataFrame(df_well["label"].astype(str))
```

---

However we should not do this without actual benchmark, or we risk introducing a very bad scaling. If there are `N` organoids with `M` labels each (under the fake assumption that each organoid has the same number of internal objects), we have:
1. Current `main`: `(N*(N-1)/2) * (M*(M-1)/2)` overlap checks
2. After #778: N*M*(M-1)/2 overlap checks
3. After the patch described in this comment: (N*M)*(N*M-1)/2 overlap checks

E.g. for N=138 organoids (ref #764) and 100 labels (random guess), we'd have:
1. 46*10^6 checks
2. 683k checks
3. 95*10^6 checks

Note that:
> If it turns out that it is slow, we can easily improve it in a trivial way (simply stop after the first overlap is found) or also in more systematic ways (the current approach is close to being the definition of a slow Python function: a quadratic-scaling for loop, which calls a pure-Python function and then even appends to a list).
Agreed on not doing this now. More something to test if this gets benchmarked further :)

I'd say we go with the fix now as you implemented it here. And consider most of this benchmarking in the context of a future Cellpose task refactor to use new OME-Zarr reader/writer strategy
I just debugged this with some users who (mistakenly) had output ROI tables set for single cell segmentation where they have tens of thousands of segmented objects. In that case, this performance became super prohibitive, even in the slightly optimized version.

I'd say we should remove this in the next Cellpose refactor. The warning is not very useful, but sometimes highly performance critcial.

For the extreme example, for 1 user, their whole segmentation task ran in under a minute, but then it took almost 9 hours (!!) to do these overlap checks for 28 thousand label objects:
```
2024-10-16 10:28:52,329; INFO; START cellpose_segmentation task
[... Detail logs]
2024-10-16 10:28:53,449; INFO; [segment_ROI] START | x: <class 'numpy.ndarray'>, (1, 1, 5400, 5120) | do_3D=False | model.diam_mean=17.0 | diameter=10.0 | advanced_cellpose_model_params.flow_threshold=0.4 | normalize.type='default'
2024-10-16 10:29:27,081; INFO; [segment_ROI] END   | Elapsed: 33.627 s | mask.shape=(1, 5400, 5120), mask.dtype=dtype('uint16') (then <class 'numpy.uint32'>), np.max(mask)=28360 | model.diam_mean=17.0 | diameter=10.0 | advanced_cellpose_model_params.flow_threshold=0.4
2024-10-16 10:29:27,134; INFO; ROI had num_labels_roi=28360, num_labels_tot={'num_labels_tot': 28360}
[This is were the cellpose segmentation has ended]
2024-10-16 19:20:10,278; WARNING; ROI [0, 1, 0, 5400, 0, 5120] has 18940 bounding-box pairs overlap
```


diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42432fd8..404fb9ac 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,8 @@ and the versioning aims to respect [Semantic Versioning](http://semver.org/spec/
 ### Changed
 - Repair Header image in Readme
   [#587](https://github.com/OpenEnergyPlatform/open-MaStR/pull/587)
+- Increase XML parsing speed
+  [#598](https://github.com/OpenEnergyPlatform/open-MaStR/pull/598)
 - Update bug issue templates
   [#593](https://github.com/OpenEnergyPlatform/open-MaStR/pull/593)
 ### Removed
diff --git a/CITATION.cff b/CITATION.cff
index 591a84ea..25c61bab 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -30,6 +30,10 @@ authors:
     alias: "@nesnoj"
     affiliation: "Reiner Lemoine Institut"
     orcid: " https://orcid.org/0000-0002-8563-5261"
+  - family-names: "Imbrisca"
+    given-names: "Alexandra-Andreea"
+    alias: "@AlexandraImbrisca"
+    affiliation: "Technical University of Munich"
 title: "open-MaStR"
 type: software
 license: AGPL-3.0
diff --git a/open_mastr/xml_download/utils_write_to_database.py b/open_mastr/xml_download/utils_write_to_database.py
index 4917e9d9..d0e9b524 100644
--- a/open_mastr/xml_download/utils_write_to_database.py
+++ b/open_mastr/xml_download/utils_write_to_database.py
@@ -1,3 +1,4 @@
+from io import StringIO
 from shutil import Error
 from zipfile import ZipFile

@@ -5,14 +6,14 @@
 import numpy as np
 import pandas as pd
 import sqlalchemy
-from sqlalchemy import select
+from sqlalchemy import inspect, select
 from sqlalchemy.sql import text
+from sqlalchemy.sql.sqltypes import Date, DateTime

 from open_mastr.utils.config import setup_logger
 from open_mastr.utils.helpers import data_to_include_tables
 from open_mastr.utils.orm import tablename_mapping
 from open_mastr.xml_download.utils_cleansing_bulk import cleanse_bulk_data
-from io import StringIO

 def write_mastr_xml_to_database(
@@ -26,74 +27,79 @@ def write_mastr_xml_to_database(
     include_tables = data_to_include_tables(data, mapping="write_xml")

     with ZipFile(zipped_xml_file_path, "r") as f:
-        files_list = f.namelist()
-        files_list = correct_ordering_of_filelist(files_list)
+        files_list = correct_ordering_of_filelist(f.namelist())
+
         for file_name in files_list:
-            # xml_tablename is the beginning of the filename without the number in lowercase
-            xml_tablename = file_name.split("_")[0].split(".")[0].lower()
-
-            if is_table_relevant(
-                xml_tablename=xml_tablename, include_tables=include_tables
-            ):
-                sql_tablename = tablename_mapping[xml_tablename]["__name__"]
-
-                if is_first_file(file_name):
-                    create_database_table(engine=engine, xml_tablename=xml_tablename)
-                    print(
-                        f"Table '{sql_tablename}' is filled with data '{xml_tablename}' "
-                        "from the bulk download."
-                    )
-                print(f"File '{file_name}' is parsed.")
+            xml_table_name = extract_xml_table_name(file_name)

-                df = preprocess_table_for_writing_to_database(
-                    f=f,
-                    file_name=file_name,
-                    xml_tablename=xml_tablename,
-                    bulk_download_date=bulk_download_date,
-                )
+            if not is_table_relevant(xml_table_name, include_tables):
+                continue

-                # Convert date and datetime columns into the datatype datetime
-                df = cast_date_columns_to_datetime(xml_tablename, df)
+            sql_table_name = extract_sql_table_name(xml_table_name)

-                if bulk_cleansing:
-                    df = cleanse_bulk_data(df, zipped_xml_file_path)
+            if is_first_file(file_name):
+                create_database_table(engine, xml_table_name)
+                print(
+                    f"Table '{sql_table_name}' is filled with data '{xml_table_name}' "
+                    "from the bulk download."
+                )
+            print(f"File '{file_name}' is parsed.")
+
+            df = read_xml_file(f, file_name)
+            df = process_table_before_insertion(
+                df,
+                xml_table_name,
+                zipped_xml_file_path,
+                bulk_download_date,
+                bulk_cleansing,
+            )

-                add_table_to_database(
-                    df=df,
-                    xml_tablename=xml_tablename,
-                    sql_tablename=sql_tablename,
-                    if_exists="append",
-                    engine=engine,
+            if engine.dialect.name == "sqlite":
+                add_table_to_sqlite_database(df, xml_table_name, sql_table_name, engine)
+            else:
+                add_table_to_non_sqlite_database(
+                    df, xml_table_name, sql_table_name, engine
                 )
+
     print("Bulk download and data cleansing were successful.")

-def is_table_relevant(xml_tablename: str, include_tables: list) -> bool:
+def extract_xml_table_name(file_name: str) -> str:
+    """Extract the XML table name from the file name."""
+    return file_name.split("_")[0].split(".")[0].lower()
+
+
+def extract_sql_table_name(xml_table_name: str) -> str:
+    """Extract the SQL table name from the xml table name."""
+    return tablename_mapping[xml_table_name]["__name__"]
+
+
+def is_table_relevant(xml_table_name: str, include_tables: list) -> bool:
     """Checks if the table contains relevant data and if the user wants to
     have it in the database."""
     # few tables are only needed for data cleansing of the xml files and contain no
     # information of relevance
     try:
         boolean_write_table_to_sql_database = (
-            tablename_mapping[xml_tablename]["__class__"] is not None
+            tablename_mapping[xml_table_name]["__class__"] is not None
         )
     except KeyError:
         print(
-            f"Table '{xml_tablename}' is not supported by your open-mastr version and "
+            f"Table '{xml_table_name}' is not supported by your open-mastr version and "
             f"will be skipped."
         )
         return False
     # check if the table should be written to sql database (depends on user input)
-    include_count = include_tables.count(xml_tablename)
+    include_count = include_tables.count(xml_table_name)

     return include_count == 1 and boolean_write_table_to_sql_database

-def create_database_table(engine: sqlalchemy.engine.Engine, xml_tablename: str) -> None:
-    orm_class = tablename_mapping[xml_tablename]["__class__"]
-    # drop the content from table
+def create_database_table(
+    engine: sqlalchemy.engine.Engine, xml_table_name: str
+) -> None:
+    orm_class = tablename_mapping[xml_table_name]["__class__"]
     orm_class.__table__.drop(engine, checkfirst=True)
-    # create table schema
     orm_class.__table__.create(engine)

@@ -105,8 +111,10 @@ def is_first_file(file_name: str) -> bool:
     )

-def cast_date_columns_to_datetime(xml_tablename: str, df: pd.DataFrame) -> pd.DataFrame:
-    sqlalchemy_columnlist = tablename_mapping[xml_tablename][
+def cast_date_columns_to_datetime(
+    xml_table_name: str, df: pd.DataFrame
+) -> pd.DataFrame:
+    sqlalchemy_columnlist = tablename_mapping[xml_table_name][
         "__class__"
     ].__table__.columns.items()
     for column in sqlalchemy_columnlist:
@@ -117,15 +125,31 @@ def cast_date_columns_to_datetime(xml_tablename: str, df: pd.DataFrame) -> pd.Da
     return df

+def cast_date_columns_to_string(xml_table_name: str, df: pd.DataFrame) -> pd.DataFrame:
+    column_list = tablename_mapping[xml_table_name][
+        "__class__"
+    ].__table__.columns.items()
+    for column in column_list:
+        column_name = column[0]
+
+        if not (column[0] in df.columns and is_date_column(column, df)):
+            continue
+
+        df[column_name] = pd.to_datetime(df[column_name], errors="coerce")
+
+        if type(column[1].type) is Date:
+            df[column_name] = (
+                df[column_name].dt.strftime("%Y-%m-%d").replace("NaT", None)
+            )
+        elif type(column[1].type) is DateTime:
+            df[column_name] = (
+                df[column_name].dt.strftime("%Y-%m-%d %H:%M:%S.%f").replace("NaT", None)
+            )
+    return df
+
+
 def is_date_column(column, df: pd.DataFrame) -> bool:
-    return (
-        type(column[1].type)
-        in [
-            sqlalchemy.sql.sqltypes.Date,
-            sqlalchemy.sql.sqltypes.DateTime,
-        ]
-        and column[0] in df.columns
-    )
+    return type(column[1].type) in [Date, DateTime] and column[0] in df.columns

 def correct_ordering_of_filelist(files_list: list) -> list:
@@ -154,49 +178,35 @@ def correct_ordering_of_filelist(files_list: list) -> list:
     return files_list

-def preprocess_table_for_writing_to_database(
-    f: ZipFile,
-    file_name: str,
-    xml_tablename: str,
-    bulk_download_date: str,
-) -> pd.DataFrame:
-    data = f.read(file_name)
-    try:
-        df = pd.read_xml(data, encoding="UTF-16", compression="zip")
-    except lxml.etree.XMLSyntaxError as err:
-        df = handle_xml_syntax_error(data.decode("utf-16"), err)
-
-    df = add_zero_as_first_character_for_too_short_string(df)
-    df = change_column_names_to_orm_format(df, xml_tablename)
-
-    # Add Column that refers to the source of the data
-    df["DatenQuelle"] = "bulk"
-    df["DatumDownload"] = bulk_download_date
-    return df
+def read_xml_file(f: ZipFile, file_name: str) -> pd.DataFrame:
+    """Read the xml file from the zip file and return it as a DataFrame."""
+    with f.open(file_name) as xml_file:
+        try:
+            return pd.read_xml(xml_file, encoding="UTF-16", parser="etree")
+        except lxml.etree.XMLSyntaxError as error:
+            return handle_xml_syntax_error(xml_file.read().decode("utf-16"), error)

 def change_column_names_to_orm_format(
-    df: pd.DataFrame, xml_tablename: str
+    df: pd.DataFrame, xml_table_name: str
 ) -> pd.DataFrame:
-    if tablename_mapping[xml_tablename]["replace_column_names"]:
+    if tablename_mapping[xml_table_name]["replace_column_names"]:
         df.rename(
-            columns=tablename_mapping[xml_tablename]["replace_column_names"],
+            columns=tablename_mapping[xml_table_name]["replace_column_names"],
             inplace=True,
         )
     return df

-def add_table_to_database(
+def add_table_to_non_sqlite_database(
     df: pd.DataFrame,
-    xml_tablename: str,
-    sql_tablename: str,
-    if_exists: str,
+    xml_table_name: str,
+    sql_table_name: str,
     engine: sqlalchemy.engine.Engine,
 ) -> None:
     # get a dictionary for the data types
-
     table_columns_list = list(
-        tablename_mapping[xml_tablename]["__class__"].__table__.columns
+        tablename_mapping[xml_table_name]["__class__"].__table__.columns
     )
     dtypes_for_writing_sql = {
         column.name: column.type
@@ -204,16 +214,22 @@ def add_table_to_database(
         if column.name in df.columns
     }

-    add_missing_columns_to_table(engine, xml_tablename, column_list=df.columns.tolist())
+    # Convert date and datetime columns into the datatype datetime.
+    df = cast_date_columns_to_datetime(xml_table_name, df)
+
+    add_missing_columns_to_table(
+        engine, xml_table_name, column_list=df.columns.tolist()
+    )
+
     for _ in range(10000):
         try:
             with engine.connect() as con:
                 with con.begin():
                     df.to_sql(
-                        sql_tablename,
+                        sql_table_name,
                         con=con,
                         index=False,
-                        if_exists=if_exists,
+                        if_exists="append",
                         dtype=dtypes_for_writing_sql,
                     )
                     break
@@ -224,7 +240,7 @@ def add_table_to_database(
         except sqlalchemy.exc.IntegrityError:
             # error resulting from Unique constraint failed
             df = write_single_entries_until_not_unique_comes_up(
-                df=df, xml_tablename=xml_tablename, engine=engine
+                df, xml_table_name, engine
             )

@@ -261,14 +277,14 @@ def add_zero_as_first_character_for_too_short_string(df: pd.DataFrame) -> pd.Dat

 def write_single_entries_until_not_unique_comes_up(
-    df: pd.DataFrame, xml_tablename: str, engine: sqlalchemy.engine.Engine
+    df: pd.DataFrame, xml_table_name: str, engine: sqlalchemy.engine.Engine
 ) -> pd.DataFrame:
     """
     Remove from dataframe these rows, which are already existing in the database table
     Parameters
     ----------
     df
-    xml_tablename
+    xml_table_name
     engine

     Returns
@@ -276,7 +292,7 @@ def write_single_entries_until_not_unique_comes_up(
     Filtered dataframe
     """

-    table = tablename_mapping[xml_tablename]["__class__"].__table__
+    table = tablename_mapping[xml_table_name]["__class__"].__table__
     primary_key = next(c for c in table.columns if c.primary_key)

     with engine.connect() as con:
@@ -295,14 +311,14 @@ def write_single_entries_until_not_unique_comes_up(
         labels=key_list, errors="ignore"
     )  # drop primary keys that already exist in the table
     df = df.reset_index()
-    print(f"{len_df_before-len(df)} entries already existed in the database.")
+    print(f"{len_df_before - len(df)} entries already existed in the database.")

     return df

 def add_missing_columns_to_table(
     engine: sqlalchemy.engine.Engine,
-    xml_tablename: str,
+    xml_table_name: str,
     column_list: list,
 ) -> None:
     """
@@ -312,8 +328,8 @@ def add_missing_columns_to_table(
     Parameters
     ----------
     engine
-    xml_tablename
-    df
+    xml_table_name
+    column_list

     Returns
     -------
@@ -323,7 +339,7 @@ def add_missing_columns_to_table(

     # get the columns name from the existing database
     inspector = sqlalchemy.inspect(engine)
-    table_name = tablename_mapping[xml_tablename]["__class__"].__table__.name
+    table_name = tablename_mapping[xml_table_name]["__class__"].__table__.name
     columns = inspector.get_columns(table_name)
     column_names_from_database = [column["name"] for column in columns]

@@ -392,3 +408,59 @@ def find_nearest_brackets(xml_string: str, position: int) -> tuple[int, int]:
             continue

     raise Error("An error occured when parsing the xml file. Maybe it is corrupted?")
+
+
+def process_table_before_insertion(
+    df: pd.DataFrame,
+    xml_table_name: str,
+    zipped_xml_file_path: str,
+    bulk_download_date: str,
+    bulk_cleansing: bool,
+) -> pd.DataFrame:
+    df = add_zero_as_first_character_for_too_short_string(df)
+    df = change_column_names_to_orm_format(df, xml_table_name)
+
+    # Add Column that refers to the source of the data
+    df["DatenQuelle"] = "bulk"
+    df["DatumDownload"] = bulk_download_date
+
+    if bulk_cleansing:
+        df = cleanse_bulk_data(df, zipped_xml_file_path)
+    return df
+
+
+def add_table_to_sqlite_database(
+    df: pd.DataFrame,
+    xml_table_name: str,
+    sql_table_name: str,
+    engine: sqlalchemy.engine.Engine,
+) -> None:
+    column_list = df.columns.tolist()
+    add_missing_columns_to_table(engine, xml_table_name, column_list)
+
+    # Convert NaNs to None.
+    df = df.where(pd.notnull(df), None)
+
+    # Convert date columns to strings. Dates are not supported directly by SQLite.
+    df = cast_date_columns_to_string(xml_table_name, df)
+
+    # Create SQL statement for bulk insert. ON CONFLICT DO NOTHING prevents duplicates.
+    insert_stmt = f"INSERT INTO {sql_table_name} ({','.join(column_list)}) VALUES ({','.join(['?' for _ in column_list])}) ON CONFLICT DO NOTHING"
+
+    for _ in range(10000):
+        try:
+            with engine.connect() as con:
+                with con.begin():
+                    con.connection.executemany(insert_stmt, df.to_numpy())
+                    break
+        except sqlalchemy.exc.DataError as err:
+            delete_wrong_xml_entry(err, df)
+        except sqlalchemy.exc.IntegrityError:
+            # error resulting from Unique constraint failed
+            df = write_single_entries_until_not_unique_comes_up(
+                df, xml_table_name, engine
+            )
+        except:
+            # If any unexpected error occurs, we'll switch back to the non-SQLite method.
+            add_table_to_non_sqlite_database(df, xml_table_name, sql_table_name, engine)
+            break

diff --git a/tests/xml_download/test_utils_write_to_database.py b/tests/xml_download/test_utils_write_to_database.py
index dc5aaacd..e72e5501 100644
--- a/tests/xml_download/test_utils_write_to_database.py
+++ b/tests/xml_download/test_utils_write_to_database.py
@@ -1,25 +1,35 @@
+import os
+import sqlite3
 import sys
+from datetime import datetime
+from os.path import expanduser
 from zipfile import ZipFile

+import numpy as np
+import pandas as pd
+import pytest
+from sqlalchemy import create_engine, inspect
+from sqlalchemy.sql import text
+
 from open_mastr.utils import orm
-from open_mastr.xml_download.utils_cleansing_bulk import (
-    replace_mastr_katalogeintraege,
-)
+from open_mastr.utils.orm import RetrofitUnits, NuclearExtended, tablename_mapping
 from open_mastr.xml_download.utils_write_to_database import (
-    cast_date_columns_to_datetime,
-    preprocess_table_for_writing_to_database,
-    add_table_to_database,
+    add_missing_columns_to_table,
     add_zero_as_first_character_for_too_short_string,
+    cast_date_columns_to_string,
+    change_column_names_to_orm_format,
     correct_ordering_of_filelist,
+    create_database_table,
+    extract_sql_table_name,
+    extract_xml_table_name,
+    is_date_column,
+    is_first_file,
+    is_table_relevant,
+    process_table_before_insertion,
+    read_xml_file,
+    add_table_to_non_sqlite_database,
+    add_table_to_sqlite_database,
 )
-import os
-from os.path import expanduser
-import sqlite3
-from sqlalchemy import create_engine
-import pandas as pd
-import pytest
-import numpy as np
-from datetime import datetime

 # Check if xml file exists
 _xml_file_exists = False
@@ -70,103 +80,93 @@ def engine_testdb():
     yield create_engine(testdb_url)

-@pytest.mark.skipif(
-    not _xml_file_exists, reason="The zipped xml file could not be found."
-)
-def test_preprocess_table_for_writing_to_database(zipped_xml_file_path):
-    # Prepare
-    file_name = "EinheitenKernkraft.xml"
-    xml_tablename = "einheitenkernkraft"
-    bulk_download_date = zipped_xml_file_path.split("_")[-1].replace(".zip", "")
-    # Check if bulk_download_date is derived correctly like 20220323
-    assert len(bulk_download_date) == 8
+def test_extract_xml_table_name():
+    file_name = "Netzanschlusspunkte_31.xml"
+    assert extract_xml_table_name(file_name) == "netzanschlusspunkte"

-    with ZipFile(zipped_xml_file_path, "r") as f:
-        # Act
-        df = preprocess_table_for_writing_to_database(
-            f=f,
-            file_name=file_name,
-            xml_tablename=xml_tablename,
-            bulk_download_date=bulk_download_date,
-        )

-    # Assert
-    assert df["DatenQuelle"].unique().tolist() == ["bulk"]
-    assert df["DatumDownload"].unique().tolist() == [bulk_download_date]
-    assert df["Gemeindeschluessel"].apply(len).unique().tolist() == [8]
-    assert df["Postleitzahl"].apply(len).unique().tolist() == [5]
+def text_extract_sql_table_name():
+    xml_table_name = "netzanschlusspunkte"
+    assert extract_sql_table_name(xml_table_name) == "network_connection_points"

-@pytest.mark.skipif(
-    not _xml_file_exists, reason="The zipped xml file could not be found."
-)
-def test_add_table_to_database(zipped_xml_file_path, engine_testdb):
-    # Prepare
+def test_is_table_relevant():
+    include_tables = ["anlagengasspeicher", "marktakteure"]
+    assert is_table_relevant("anlagengasspeicher", include_tables) is True
+    assert is_table_relevant("netzanschlusspunkte", include_tables) is False
+
+
+def test_create_database_table(engine_testdb):
     orm.Base.metadata.create_all(engine_testdb)
-    file_name = "EinheitenKernkraft.xml"
-    xml_tablename = "einheitenkernkraft"
-    sql_tablename = "nuclear_extended"
-    bulk_download_date = zipped_xml_file_path.split("_")[-1].replace(".zip", "")
-    # Check if bulk_download_date is derived correctly like 20220323
-    assert len(bulk_download_date) == 8
-    with ZipFile(zipped_xml_file_path, "r") as f:
-        df_write = preprocess_table_for_writing_to_database(
-            f=f,
-            file_name=file_name,
-            xml_tablename=xml_tablename,
-            bulk_download_date=bulk_download_date,
-        )
+    xml_table_name = "einheitenkernkraft"
+    sql_table_name = "nuclear_extended"

-    # Convert date and datetime columns into the datatype datetime
-    df_write = cast_date_columns_to_datetime(xml_tablename, df_write)
+    create_database_table(engine_testdb, xml_table_name)

-    # Katalogeintraege: int -> string value
-    df_write = replace_mastr_katalogeintraege(
-        zipped_xml_file_path=zipped_xml_file_path, df=df_write
-    )
+    assert inspect(engine_testdb).has_table(sql_table_name) is True

-    # Act
-    add_table_to_database(
-        df=df_write,
-        xml_tablename=xml_tablename,
-        sql_tablename=sql_tablename,
-        if_exists="replace",
-        engine=engine_testdb,
+
+def test_is_first_file():
+    assert is_first_file("EinheitenKernkraft.xml") is True
+    assert is_first_file("EinheitenKernkraft_1.xml") is True
+    assert is_first_file("EinheitenKernkraft_2.xml") is False
+
+
+def test_cast_date_columns_to_string():
+    initial_df = pd.DataFrame(
+        {
+            "EegMastrNummer": [1, 2, 3],
+            "Registrierungsdatum": [
+                datetime(2024, 3, 11).date(),
+                datetime(1999, 2, 1).date(),
+                np.datetime64("nat"),
+            ],
+            "DatumLetzteAktualisierung": [
+                datetime(2022, 3, 22),
+                datetime(2020, 1, 2, 10, 12, 46),
+                np.datetime64("nat"),
+            ],
+        }
     )
-    with engine_testdb.connect() as con:
-        with con.begin():
-            df_read = pd.read_sql_table(table_name=sql_tablename, con=con)
-    # Drop the empty columns which come from orm and are not filled by bulk download
-    df_read.dropna(how="all", axis=1, inplace=True)
-    # Rename LokationMaStRNummer -> LokationMastrNummer unresolved error
-    df_write.rename(
-        columns={"LokationMaStRNummer": "LokationMastrNummer"}, inplace=True
-    )
-    # Reorder columns of the df_write since the order of columns doesn't play a role
-    df_read = df_read[df_write.columns.tolist()]
-    # Cast the dtypes of df_write like df_read dtypes
-    # Datatypes are changed during inserting via sqlalchemy and orm
-    df_read = df_read.astype(dict(zip(df_write.columns.tolist(), df_write.dtypes)))
-    # Assert
+    expected_df = pd.DataFrame(
+        {
+            "EegMastrNummer": [1, 2, 3],
+            "Registrierungsdatum": ["2024-03-11", "1999-02-01", np.nan],
+            "DatumLetzteAktualisierung": [
+                "2022-03-22 00:00:00.000000",
+                "2020-01-02 10:12:46.000000",
+                np.nan,
+            ],
+        }
+    )
+
     pd.testing.assert_frame_equal(
-        df_write.select_dtypes(exclude="datetime"),
-        df_read.select_dtypes(exclude="datetime"),
+        expected_df, cast_date_columns_to_string("anlageneegwasser", initial_df)
     )

-def test_add_zero_as_first_character_for_too_short_string():
-    # Prepare
-    df_raw = pd.DataFrame(
-        {"ID": [0, 1, 2], "Gemeindeschluessel": [9162000, np.nan, 19123456]}
-    )
-    df_correct = pd.DataFrame(
-        {"ID": [0, 1, 2], "Gemeindeschluessel": ["09162000", np.nan, "19123456"]}
+def test_is_date_column():
+    columns = RetrofitUnits.__table__.columns.items()
+    df = pd.DataFrame(
+        {
+            "Id": [1],
+            "DatumLetzteAktualisierung": [datetime(2022, 3, 22)],
+            "WiederinbetriebnahmeDatum": [datetime(2024, 3, 11).date()],
+        }
     )

-    # Act
-    df_edited = add_zero_as_first_character_for_too_short_string(df_raw)
-    # Assert
-    pd.testing.assert_frame_equal(df_edited, df_correct)
+    date_column = list(filter(lambda col: col[0] == "Id", columns))[0]
+    assert is_date_column(date_column, df) is False
+
+    datetime_column = list(
+        filter(lambda col: col[0] == "DatumLetzteAktualisierung", columns)
+    )[0]
+    assert is_date_column(datetime_column, df) is True
+
+    date_column = list(
+        filter(lambda col: col[0] == "WiederinbetriebnahmeDatum", columns)
+    )[0]
+    assert is_date_column(date_column, df) is True

 def test_correct_ordering_of_filelist():
@@ -215,24 +215,181 @@ def test_correct_ordering_of_filelist():
     ]

-def test_cast_date_columns_to_datetime():
+@pytest.mark.skipif(
+    not _xml_file_exists, reason="The zipped xml file could not be found."
+)
+def test_read_xml_file(zipped_xml_file_path):
+    with ZipFile(zipped_xml_file_path, "r") as f:
+        df = read_xml_file(f, "EinheitenKernkraft.xml")
+
+    assert df.shape[0] > 0
+
+    # Since the file is from the latest download, its content can vary over time. To make sure that the table is
+    # correctly created, we check that all of its columns are associated are included in our mapping.
+    for column in df.columns:
+        if column in tablename_mapping["einheitenkernkraft"]["replace_column_names"]:
+            column = tablename_mapping["einheitenkernkraft"]["replace_column_names"][
+                column
+            ]
+        assert column in NuclearExtended.__table__.columns.keys()
+
+
+def test_add_zero_as_first_character_for_too_short_string():
+    # Prepare
     df_raw = pd.DataFrame(
+        {"ID": [0, 1, 2], "Gemeindeschluessel": [9162000, np.nan, 19123456]}
+    )
+    df_correct = pd.DataFrame(
+        {"ID": [0, 1, 2], "Gemeindeschluessel": ["09162000", np.nan, "19123456"]}
+    )
+
+    # Act
+    df_edited = add_zero_as_first_character_for_too_short_string(df_raw)
+    # Assert
+    pd.testing.assert_frame_equal(df_edited, df_correct)
+
+
+def test_change_column_names_to_orm_format():
+    initial_df = pd.DataFrame(
         {
-            "ID": [0, 1, 2],
-            "Registrierungsdatum": ["2022-03-22", "2020-01-02", "2022-03-35"],
+            "VerknuepfteEinheitenMaStRNummern": ["test1", "test2"],
+            "NetzanschlusspunkteMaStRNummern": [1, 2],
         }
     )
-    df_replaced = pd.DataFrame(
+    expected_df = pd.DataFrame(
         {
-            "ID": [0, 1, 2],
-            "Registrierungsdatum": [
-                datetime(2022, 3, 22),
-                datetime(2020, 1, 2),
-                np.datetime64("nat"),
-            ],
+            "VerknuepfteEinheiten": ["test1", "test2"],
+            "Netzanschlusspunkte": [1, 2],
+        }
+    )
+
+    pd.testing.assert_frame_equal(
+        expected_df, change_column_names_to_orm_format(initial_df, "lokationen")
+    )
+
+
+def test_process_table_before_insertion(zipped_xml_file_path):
+    bulk_download_date = datetime.now().date().strftime("%Y%m%d")
+    initial_df = pd.DataFrame(
+        {
+            "Gemeindeschluessel": [9162000, 19123456],
+            "Postleitzahl": [1234, 54321],
+            "NameKraftwerk": ["test1", "test2"],
+            "LokationMaStRNummer": ["test3", "test4"],
+        }
+    )
+    expected_df = pd.DataFrame(
+        {
+            "Gemeindeschluessel": ["09162000", "19123456"],
+            "Postleitzahl": ["01234", "54321"],
+            "NameKraftwerk": ["test1", "test2"],
+            "LokationMastrNummer": ["test3", "test4"],
+            "DatenQuelle": ["bulk", "bulk"],
+            "DatumDownload": [bulk_download_date, bulk_download_date],
         }
     )

     pd.testing.assert_frame_equal(
-        df_replaced, cast_date_columns_to_datetime("anlageneegwasser", df_raw)
+        expected_df,
+        process_table_before_insertion(
+            initial_df,
+            "einheitenkernkraft",
+            zipped_xml_file_path,
+            bulk_download_date,
+            bulk_cleansing=False,
+        ),
+    )
+
+
+def test_add_missing_columns_to_table(engine_testdb):
+    with engine_testdb.connect() as con:
+        with con.begin():
+            # We must recreate the table to be sure that the new colum is not present.
+            con.execute(text("DROP TABLE IF EXISTS gas_consumer"))
+            create_database_table(engine_testdb, "einheitengasverbraucher")
+
+            initial_data_in_db = pd.DataFrame(
+                {
+                    "EinheitMastrNummer": ["id1"],
+                    "DatumLetzteAktualisierung": [datetime(2022, 2, 2)],
+                }
+            )
+            initial_data_in_db.to_sql(
+                "gas_consumer", con=con, if_exists="append", index=False
+            )
+
+    add_missing_columns_to_table(
+        engine_testdb, "einheitengasverbraucher", ["NewColumn"]
     )
+
+    expected_df = pd.DataFrame(
+        {
+            "EinheitMastrNummer": ["id1"],
+            "DatumLetzteAktualisierung": [datetime(2022, 2, 2)],
+            "NewColumn": [None],
+        }
+    )
+    with engine_testdb.connect() as con:
+        with con.begin():
+            actual_df = pd.read_sql_table("gas_consumer", con=con)
+            # The actual_df will contain more columns than the expected_df, so we can't use assert_frame_equal.
+            assert expected_df.index.isin(actual_df.index).all()
+
+
+@pytest.mark.parametrize(
+    "add_table_to_database_function",
+    [add_table_to_sqlite_database, add_table_to_non_sqlite_database],
+)
+def test_add_table_to_sqlite_database(engine_testdb, add_table_to_database_function):
+    with engine_testdb.connect() as con:
+        with con.begin():
+            # We must recreate the table to be sure that no other data is present.
+            con.execute(text("DROP TABLE IF EXISTS gsgk_eeg"))
+            create_database_table(
+                engine_testdb, "anlageneeggeothermiegrubengasdruckentspannung"
+            )
+
+    df = pd.DataFrame(
+        {
+            "Registrierungsdatum": ["2022-02-02", "2024-03-20"],
+            "EegMastrNummer": ["id1", "id2"],
+            "DatumLetzteAktualisierung": [
+                "2022-12-02 10:10:10.000300",
+                "2024-10-10 00:00:00.000000",
+            ],
+            "AusschreibungZuschlag": [True, False],
+            "Netzbetreiberzuordnungen": ["test1", "test2"],
+            "InstallierteLeistung": [1.0, 100.4],
+        }
+    )
+    expected_df = pd.DataFrame(
+        {
+            "InstallierteLeistung": [1.0, 100.4],
+            "AnlageBetriebsstatus": [None, None],
+            "Registrierungsdatum": [datetime(2022, 2, 2), datetime(2024, 3, 20)],
+            "EegMastrNummer": ["id1", "id2"],
+            "Meldedatum": [np.datetime64("NaT"), np.datetime64("NaT")],
+            "DatumLetzteAktualisierung": [
+                datetime(2022, 12, 2, 10, 10, 10, 300),
+                datetime(2024, 10, 10),
+            ],
+            "EegInbetriebnahmedatum": [np.datetime64("NaT"), np.datetime64("NaT")],
+            "VerknuepfteEinheit": [None, None],
+            "AnlagenschluesselEeg": [None, None],
+            "AusschreibungZuschlag": [True, False],
+            "AnlagenkennzifferAnlagenregister": [None, None],
+            "AnlagenkennzifferAnlagenregister_nv": [None, None],
+            "Netzbetreiberzuordnungen": ["test1", "test2"],
+            "DatenQuelle": [None, None],
+            "DatumDownload": [np.datetime64("NaT"), np.datetime64("NaT")],
+        }
+    )
+
+    add_table_to_database_function(
+        df, "anlageneeggeothermiegrubengasdruckentspannung", "gsgk_eeg", engine_testdb
+    )
+    with engine_testdb.connect() as con:
+        with con.begin():
+            pd.testing.assert_frame_equal(
+                expected_df, pd.read_sql_table("gsgk_eeg", con=con)
+            )

Increase parsing speed
This task contains several steps:

1. Search different ways that might increase parsing speed. Parsing is done right now by the `pandas.read_xml` method [here](https://github.com/OpenEnergyPlatform/open-MaStR/blob/d17a92c17912f7f96814d309859dc495f942c07a/open_mastr/xml_download/utils_write_to_database.py#L158). Several alternatives are:
  * `polars`, `duckdb`, `pyspark` might have xml parsers and might be faster
  * use plain xml parsing from python (without pandas)
  * ...

2. Writing to sqlite database right now is done by `pandas.to_sql` [here](https://github.com/OpenEnergyPlatform/open-MaStR/blob/d17a92c17912f7f96814d309859dc495f942c07a/open_mastr/xml_download/utils_write_to_database.py#L182C5-L182C26). There might be other faster methods depending on step 1.
3. Construct a benchmark in an own repository. Use a benchmark xml file from the Marktstammdatenregister and test different implementations for parsing them.
4. Decide for a best method and implement it in `open-mastr`

Hi! I started working on this task and decided to change the steps slightly:

1. Construct the benchmark
    - Use the Marktstammdatenregister to construct a few datasets of various size - ✅  ([link](https://github.com/AlexandraImbrisca/open-MaStR/tree/develop/benchmark))
    - Create a script to automate the calculation and comparison of the parsing speed between various optimisations - ✅  ([link](https://github.com/AlexandraImbrisca/open-MaStR/pull/6))
2. Explore faster methods of parsing the XML
     - Research the options and implement the changes
     - Run the benchmark and analyse the results
3. Explore faster methods of writing to the sqlite database
     - Research the options and implement the changes
     - Run the benchmark and analyse the results
4. Decide on the best method and add it to this repository
I was at DACH Energy Informatics Conference and took two points from there:
- Many researchers use open-mastr
- The feature request I heard most often was the question, if we can decrease the time it needs to download and parse the data
I think people will be really happy if this issue is successful 😃
> I was at DACH Energy Informatics Conference and took two points from there:
>
> * Many researchers use open-mastr
> * The feature request I heard most often was the question, if we can decrease the time it needs to download and parse the data
>   I think people will be really happy if this issue is successful 😃

Sounds great. I think we cannot do much about the dl speed but I'm really looking forward to the parsing enhancement @AlexandraImbrisca :smile:


diff --git a/RELEASE.md b/RELEASE.md
index 7af26f6ca4..c19418f0ca 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -2,6 +2,7 @@

 ## Major features and improvements
 * Implemented `KedroDataCatalog.to_config()` method that converts the catalog instance into a configuration format suitable for serialization.
+* Improve OmegaConfigLoader performance

 ## Bug fixes and other changes
 * Added validation to ensure dataset versions consistency across catalog.
diff --git a/kedro/config/omegaconf_config.py b/kedro/config/omegaconf_config.py
index 8d82ebf360..3f40624a4b 100644
--- a/kedro/config/omegaconf_config.py
+++ b/kedro/config/omegaconf_config.py
@@ -126,7 +126,8 @@ def __init__(  # noqa: PLR0913
         self.base_env = base_env or ""
         self.default_run_env = default_run_env or ""
         self.merge_strategy = merge_strategy or {}
-
+        self._globals_oc: DictConfig | None = None
+        self._runtime_params_oc: DictConfig | None = None
         self.config_patterns = {
             "catalog": ["catalog*", "catalog*/**", "**/catalog*"],
             "parameters": ["parameters*", "parameters*/**", "**/parameters*"],
@@ -346,12 +347,11 @@ def load_and_merge_dir_config(
                 OmegaConf.merge(*aggregate_config, self.runtime_params), resolve=True
             )

+        merged_config_container = OmegaConf.to_container(
+            OmegaConf.merge(*aggregate_config), resolve=True
+        )
         return {
-            k: v
-            for k, v in OmegaConf.to_container(
-                OmegaConf.merge(*aggregate_config), resolve=True
-            ).items()
-            if not k.startswith("_")
+            k: v for k, v in merged_config_container.items() if not k.startswith("_")
         }

     @staticmethod
@@ -436,9 +436,12 @@ def _get_globals_value(self, variable: str, default_value: Any = _NO_VALUE) -> A
             raise InterpolationResolutionError(
                 "Keys starting with '_' are not supported for globals."
             )
-        globals_oc = OmegaConf.create(self._globals)
+
+        if not self._globals_oc:
+            self._globals_oc = OmegaConf.create(self._globals)
+
         interpolated_value = OmegaConf.select(
-            globals_oc, variable, default=default_value
+            self._globals_oc, variable, default=default_value
         )
         if interpolated_value != _NO_VALUE:
             return interpolated_value
@@ -449,9 +452,11 @@ def _get_globals_value(self, variable: str, default_value: Any = _NO_VALUE) -> A

     def _get_runtime_value(self, variable: str, default_value: Any = _NO_VALUE) -> Any:
         """Return the runtime params values to the resolver"""
-        runtime_oc = OmegaConf.create(self.runtime_params)
+        if not self._runtime_params_oc:
+            self._runtime_params_oc = OmegaConf.create(self.runtime_params)
+
         interpolated_value = OmegaConf.select(
-            runtime_oc, variable, default=default_value
+            self._runtime_params_oc, variable, default=default_value
         )
         if interpolated_value != _NO_VALUE:
             return interpolated_value

Improve OmegaConfigLoader performance when global/variable interpolations are involved
## Description

Extending on the investigation in #3893 , OmegaConfigLoader lags in resolving catalog configurations when global/variable interpolations are involved.

## Context

**Some previous observations:**
https://github.com/kedro-org/kedro/issues/3893#issuecomment-2460629752
https://github.com/kedro-org/kedro/issues/3893#issuecomment-2465507178

## Steps to Reproduce

Run stress test which creates a single catalog with variable interpolations - https://github.com/kedro-org/kedro/blob/test/ocl-bottleneck/kedro_benchmarks/temp_investigate_ocl/ocl_plot_variables.py

## Expected Result

Reduce the time spent on below methods (when interpolations are involved, the bottleneck seems to be `OmegaConf.to_container`)

- [OmegaConf.load](https://omegaconf.readthedocs.io/en/2.3_branch/_modules/omegaconf/omegaconf.html#OmegaConf)
- [OmegaConf.to_container](https://omegaconf.readthedocs.io/en/2.3_branch/api_reference.html#omegaconf.OmegaConf.to_container)
- [OmegaConf.merge](https://omegaconf.readthedocs.io/en/2.3_branch/api_reference.html#omegaconf.OmegaConf.merge)

## Actual Result

https://github.com/kedro-org/kedro/issues/3893#issuecomment-2465507178

## Your Environment

<img width="245" alt="image" src="https://github.com/user-attachments/assets/07e3d00f-ad2b-41c6-863a-29e7d20d37c1">

* Kedro version used (`pip show kedro` or `kedro -V`): 0.19.9
* Python version used (`python -V`): 3.11
* Operating system and version: macOS and Sonoma 14.7.1

During the investigation, I found that there was a slow bit in `_set_globals_value`. I didn't spent enough time to fix it, but with a quick fix it improves roughly from 1.5s -> 0.9s, but there are probably more.
Particularly, there is an obvious slow path `global_oc` get created and destroyed for every reference of `$globals`

https://github.com/kedro-org/kedro/blob/01c095b2537c1cfbafe6443a021e60ec94bfa1af/kedro/config/omegaconf_config.py#L433-L439


diff --git a/src/tlo/methods/contraception.py b/src/tlo/methods/contraception.py
index ab6c633f4c..09cb394804 100644
--- a/src/tlo/methods/contraception.py
+++ b/src/tlo/methods/contraception.py
@@ -1146,12 +1146,12 @@ def __init__(self, module, person_id, new_contraceptive):

         self.TREATMENT_ID = "Contraception_Routine"
         self.ACCEPTED_FACILITY_LEVEL = _facility_level
-
-    @property
-    def EXPECTED_APPT_FOOTPRINT(self):
-        """Return the expected appt footprint based on contraception method and whether the HSI has been rescheduled."""
-        person_id = self.target
         current_method = self.sim.population.props.loc[person_id].co_contraception
+        self.EXPECTED_APPT_FOOTPRINT = self._get_appt_footprint(current_method)
+
+
+    def _get_appt_footprint(self, current_method: str):
+        """Return the appointment footprint based on contraception method and whether the HSI has been rescheduled."""
         if self._number_of_times_run > 0:  # if it is to re-schedule due to unavailable consumables
             return self.make_appt_footprint({})
         # if to switch to a method
@@ -1165,9 +1165,11 @@ def EXPECTED_APPT_FOOTPRINT(self):
         elif self.new_contraceptive in ['male_condom', 'other_modern', 'pill']:
             return self.make_appt_footprint({'PharmDispensing': 1})
         else:
+            warnings.warn(
+                "Assumed empty footprint for Contraception Routine appt because couldn't find actual case.",
+                stacklevel=3,
+            )
             return self.make_appt_footprint({})
-            warnings.warn(UserWarning("Assumed empty footprint for Contraception Routine appt because couldn't find"
-                                      "actual case."))

     def apply(self, person_id, squeeze_factor):
         """If the relevant consumable is available, do change in contraception and log it"""
@@ -1266,6 +1268,8 @@ def apply(self, person_id, squeeze_factor):
         ):
             self.reschedule()

+        return self._get_appt_footprint(current_method)
+
     def post_apply_hook(self):
         self._number_of_times_run += 1

Population dataframe accesses in `HSI_Contraception_FamilyPlanningAppt.EXPECTED_APPT_FOOTPRINT` are performance bottleneck
Look at the most recent successful profiling run results

https://github-pages.ucl.ac.uk/TLOmodel-profiling/_static/profiling_html/schedule_1226_dbf33d69edadbacdc2dad1fc32c93eb771b03f9b.html

calls to the `HSI_Contraception_FamilyPlanningAppt.EXPECTED_APPT_FOOTPRINT` property are consuming a lot of run time, with a total of 2570s spent in evaluating this property out of a total run time of 27 400s (9.3%). There are three functions calling `HSI_Contraception_FamilyPlanningAppt.EXPECTED_APPT_FOOTPRINT` with a significant amount of time recorded in each call path

- `HealthSystem.check_hsi_event_is_valid` (1110s)
- `HSI_Contraception_FamilyPlanningAppt.initialise` (545s)
- `HSI_Contraception_FamilyPlanningAppt.run` (911s)

and in all cases the vast majority of the time in the `HSI_Contraception_FamilyPlanningAppt.EXPECTED_APPT_FOOTPRINT` call is spent in `.loc` indexing operations on the population dataframe.

Looking at the implementation of the property

https://github.com/UCL/TLOmodel/blob/5b4f328e5fc0df469eec40bcf1110a555326aa80/src/tlo/methods/contraception.py#L1150-L1170

we can see this must be arising from the access to the current contraception method property for the individual the HSI event is targetting in the line

```Python
current_method = self.sim.population.props.loc[person_id].co_contraception
```

<details>
<summary>
This pattern of the `EXPECTED_APPT_FOOTPRINT` being a property rather than an attribute seems to also be used by a few other HSI events
</summary>

`HSI_Hiv_StartOrContinueTreatment`

https://github.com/UCL/TLOmodel/blob/5b4f328e5fc0df469eec40bcf1110a555326aa80/src/tlo/methods/hiv.py#L2968-L2986

`HSI_Tb_StartTreatment`

https://github.com/UCL/TLOmodel/blob/5b4f328e5fc0df469eec40bcf1110a555326aa80/src/tlo/methods/tb.py#L2175-L2183

`HsiBaseVaccine`

https://github.com/UCL/TLOmodel/blob/5b4f328e5fc0df469eec40bcf1110a555326aa80/src/tlo/methods/epi.py#L407-L431
</details>

I'm assuming this has been implemented like this to allow the expected appointment footprint to be based on the target's properties at the point at which it is run rather than when it is scheduled? Tagging @tbhallett, @BinglingICL and @EvaJanouskova as I think this was originally implemented in #826.

One thing I'm not sure about is how this pattern interacts with this logic in `HSI_Event.initialise`

https://github.com/UCL/TLOmodel/blob/5b4f328e5fc0df469eec40bcf1110a555326aa80/src/tlo/methods/hsi_event.py#L333-L338

which reassigns `self.EXPECTED_APPT_FOOTPRINT` in certain cases, in which case it would end up as a plain (static) attribute rather than a property. As `HSI_Event.initialise` appears to be called in `HealthSystem.schedule_hsi_event`, that is at the time the HSI event is scheduled, this means in some cases even when `EXPECTED_APPT_FOOTPRINT` is defined as a property it will be 'locked-in' as a static attribute with values taken from population dataframe at point at which it is scheduled.

As accessing the population dataframe dynamically when `EXPECTED_APPT_FOOTPRINT` is a property seems to be a performance bottleneck, and given it does seem to be guaranteed this will actually be done at the point the HSI event is run rather than scheduled, it seems that enforcing that `EXPECTED_APPT_FOOTPRINT` is an attribute which is contructed when the HSI event is constructed (when scheduling) might be a good idea to both keep the behaviour consistent and allow us to reduce the population dataframes by just doing once on construction rather than each access of `EXPECTED_APPT_FOOTPRINT`.

@BinglingICL, I think it is okay to set the EXPECTED_APPT_FOOTPRINT with the same logic when HSI is scheduled instead of when it is run, isn't it?
> [@BinglingICL](https://github.com/BinglingICL), I think it is okay to set the EXPECTED_APPT_FOOTPRINT with the same logic when HSI is scheduled instead of when it is run, isn't it?

Hi Eva, it seems so to me, if you meant the way to assign appt footprint within the property.