Repository: jina-ai/serve
Pull request: feat: match with batching capability
This PR  adds batching to `d1.match(d2)` to load `d2` in batches.

Problem statement:
Add batching to match function
It may make sense only for `DocumentArrayMemmap` but it may be applicable to also `DocumentArray`.

The point is to have `batching` for the `Y` (The right hand side of the `match` operation) so that the Y can be loaded in memory in batches

Hints from issue discussion:
The Draft PR is https://github.com/jina-ai/jina/pull/3113

Commit message: fix: add types

Patch to review:
diff --git a/jina/math/helper.py b/jina/math/helper.py
index 9381f23a16df0..1c1c5c0343a44 100644
--- a/jina/math/helper.py
+++ b/jina/math/helper.py
@@ -49,3 +49,29 @@ def top_k(
         values = -values
 
     return values, idx
+
+
+def top_k_from_pair(
+    x_mat: 'np.ndarray',
+    x_inds: 'np.ndarray',
+    y_mat: 'np.ndarray',
+    y_inds: 'np.ndarray',
+    k: int,
+) -> Tuple['np.ndarray', 'np.ndarray']:
+    """
+    Finds values and indices of the k largest entries found in `x_ma` union `y_mat`
+    and the indices from `x_inds` union `y_inds`.
+
+    :param x_mat: numpy array of the first matrix
+    :param x_inds: numpy array of the indices of the first matrix
+    :param y_mat: numpy array of the second matrix
+    :param y_inds: numpy array of the indices of the second matrix
+    :param k: number of values to retrieve
+    :return: indices and distances of the best items in the pair
+    """
+    all_dists = np.hstack((x_mat, y_mat))
+    all_inds = np.hstack((x_inds, y_inds))
+    best_inds = np.argpartition(all_dists, kth=k, axis=1)
+    d_mat = np.take_along_axis(all_dists, best_inds, axis=1)
+    i_mat = np.take_along_axis(all_inds, best_inds, axis=1)
+    return d_mat[:, :k], i_mat[:, :k]
diff --git a/jina/types/arrays/neural_ops.py b/jina/types/arrays/neural_ops.py
index f67df82355f05..c9f226e9caff7 100644
--- a/jina/types/arrays/neural_ops.py
+++ b/jina/types/arrays/neural_ops.py
@@ -5,7 +5,7 @@
 
 from ... import Document
 from ...importer import ImportExtensions
-from ...math.helper import top_k, minmax_normalize
+from ...math.helper import top_k, minmax_normalize, top_k_from_pair
 
 if False:
     from .document import DocumentArray
@@ -25,6 +25,7 @@ def match(
         normalization: Optional[Tuple[int, int]] = None,
         use_scipy: bool = False,
         metric_name: Optional[str] = None,
+        batch_size: Union[None, int] = None,
     ) -> None:
         """Compute embedding based nearest neighbour in `another` for each Document in `self`,
         and store results in `matches`.
@@ -52,29 +53,117 @@ def match(
                                 all values will be rescaled into range `[a, b]`.
         :param use_scipy: use Scipy as the computation backend
         :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: if provided, then `darray` is loaded in chunks of, at most, batch_size elements. This option
+                           will be slower but more memory efficient. Specialy indicated if `darray` is a big
+                           DocumentArrayMemmap.
         """
 
-        X = np.stack(self.get_attributes('embedding'))
-        Y = np.stack(darray.get_attributes('embedding'))
-
-        if isinstance(metric, str):
+        if callable(metric):
+            cdist = metric
+        elif isinstance(metric, str):
             if use_scipy:
-                from scipy.spatial.distance import cdist
+                from scipy.spatial.distance import cdist as cdist
             else:
-                from ...math.distance import cdist
-            dists = cdist(X, Y, metric)
-        elif callable(metric):
-            dists = metric(X, Y)
+                from ...math.distance import cdist as cdist
         else:
             raise TypeError(
                 f'metric must be either string or a 2-arity function, received: {metric!r}'
             )
+        # brekpoint()
+        metric_name = metric_name or (metric.__name__ if callable(metric) else metric)
+
+        if batch_size:
+            self._match_online(
+                darray, cdist, limit, normalization, metric_name, batch_size
+            )
+        else:
+            self._match(darray, cdist, limit, normalization, metric_name)
 
+    def _match(self, darray, cdist, limit, normalization, metric_name):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        y_mat = np.stack(darray.get_attributes('embedding'))
+
+        dists = cdist(x_mat, y_mat, metric_name)
         dist, idx = top_k(dists, min(limit, len(darray)), descending=False)
         if normalization is not None:
             if isinstance(normalization, (tuple, list)):
                 dist = minmax_normalize(dist, normalization)
 
+        for _q, _ids, _dists in zip(self, idx, dist):
+            _q.matches.clear()
+            for _id, _dist in zip(_ids, _dists):
+                # Note, when match self with other, or both of them share the same Document
+                # we might have recursive matches .
+                # checkout https://github.com/jina-ai/jina/issues/3034
+                d = darray[int(_id)]
+                if d.id in self:
+                    d = Document(d, copy=True)
+                    d.pop('matches')
+                _q.matches.append(d, scores={metric_name: _dist}, copy=False)
+
+    def _match_online(
+        self, darray, cdist, limit, normalization, metric_name, batch_size
+    ):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: length of the chunks loaded into memory from darray.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        n_x = x_mat.shape[0]
+
+        def batch_generator(y_darray: 'DocumentArrayMemmap', n_batch: int):
+            for i in range(0, len(y_darray), n_batch):
+                y_mat = np.vstack(
+                    [y_darray[j].embedding for j in range(i, i + n_batch)]
+                )
+
+                yield y_mat, i
+
+        y_batch_generator = batch_generator(darray, batch_size)
+        top_dists = np.inf * np.ones((n_x, limit))
+        top_inds = np.zeros((n_x, limit), dtype=int)
+
+        for ybatch, ybatch_start_pos in y_batch_generator:
+            distances = cdist(x_mat, ybatch, metric_name)
+            dists, inds = top_k(distances, limit, descending=False)
+
+            if normalization is not None:
+                if isinstance(normalization, (tuple, list)):
+                    dist = minmax_normalize(dist, normalization)
+
+            inds = ybatch_start_pos + inds
+            top_dists, top_inds = top_k_from_pair(
+                dists, inds, top_dists, top_inds, limit
+            )
+
+        # sort final the final `top_dists` and `top_inds` per row
+        permutation = np.argsort(top_dists, axis=1)
+        dist = np.take_along_axis(top_dists, permutation, axis=1)
+        idx = np.take_along_axis(top_inds, permutation, axis=1)
+
         m_name = metric_name or (metric.__name__ if callable(metric) else metric)
         for _q, _ids, _dists in zip(self, idx, dist):
             _q.matches.clear()
diff --git a/tests/unit/types/arrays/test_neural_ops.py b/tests/unit/types/arrays/test_neural_ops.py
index 344cc3a25c77e..7f121ab1a0bb4 100644
--- a/tests/unit/types/arrays/test_neural_ops.py
+++ b/tests/unit/types/arrays/test_neural_ops.py
@@ -249,7 +249,7 @@ def test_scipy_dist(
 
 
 def test_2arity_function(docarrays_for_embedding_distance_computation):
-    def dotp(x, y):
+    def dotp(x, y, *args):
         return np.dot(x, np.transpose(y))
 
     D1, D2 = docarrays_for_embedding_distance_computation


Code review comments:
diff_hunk: @@ -52,29 +53,117 @@ def match(
                                 all values will be rescaled into range `[a, b]`.
         :param use_scipy: use Scipy as the computation backend
         :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: if provided, then `darray` is loaded in chunks of, at most, batch_size elements. This option
+                           will be slower but more memory efficient. Specialy indicated if `darray` is a big
+                           DocumentArrayMemmap.
         """
 
-        X = np.stack(self.get_attributes('embedding'))
-        Y = np.stack(darray.get_attributes('embedding'))
-
-        if isinstance(metric, str):
+        if callable(metric):
+            cdist = metric
+        elif isinstance(metric, str):
             if use_scipy:
-                from scipy.spatial.distance import cdist
+                from scipy.spatial.distance import cdist as cdist
             else:
-                from ...math.distance import cdist
-            dists = cdist(X, Y, metric)
-        elif callable(metric):
-            dists = metric(X, Y)
+                from ...math.distance import cdist as cdist
         else:
             raise TypeError(
                 f'metric must be either string or a 2-arity function, received: {metric!r}'
             )
+        # brekpoint()
path: jina/types/arrays/neural_ops.py
text: @user1:
remove

@author:
done

diff_hunk: @@ -52,29 +53,117 @@ def match(
                                 all values will be rescaled into range `[a, b]`.
         :param use_scipy: use Scipy as the computation backend
         :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: if provided, then `darray` is loaded in chunks of, at most, batch_size elements. This option
+                           will be slower but more memory efficient. Specialy indicated if `darray` is a big
+                           DocumentArrayMemmap.
         """
 
-        X = np.stack(self.get_attributes('embedding'))
-        Y = np.stack(darray.get_attributes('embedding'))
-
-        if isinstance(metric, str):
+        if callable(metric):
+            cdist = metric
+        elif isinstance(metric, str):
             if use_scipy:
-                from scipy.spatial.distance import cdist
+                from scipy.spatial.distance import cdist as cdist
             else:
-                from ...math.distance import cdist
-            dists = cdist(X, Y, metric)
-        elif callable(metric):
-            dists = metric(X, Y)
+                from ...math.distance import cdist as cdist
         else:
             raise TypeError(
                 f'metric must be either string or a 2-arity function, received: {metric!r}'
             )
+        # brekpoint()
+        metric_name = metric_name or (metric.__name__ if callable(metric) else metric)
+
+        if batch_size:
+            self._match_online(
+                darray, cdist, limit, normalization, metric_name, batch_size
+            )
+        else:
+            self._match(darray, cdist, limit, normalization, metric_name)
 
+    def _match(self, darray, cdist, limit, normalization, metric_name):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        y_mat = np.stack(darray.get_attributes('embedding'))
+
+        dists = cdist(x_mat, y_mat, metric_name)
         dist, idx = top_k(dists, min(limit, len(darray)), descending=False)
         if normalization is not None:
             if isinstance(normalization, (tuple, list)):
                 dist = minmax_normalize(dist, normalization)
 
+        for _q, _ids, _dists in zip(self, idx, dist):
+            _q.matches.clear()
+            for _id, _dist in zip(_ids, _dists):
+                # Note, when match self with other, or both of them share the same Document
+                # we might have recursive matches .
+                # checkout https://github.com/jina-ai/jina/issues/3034
+                d = darray[int(_id)]
+                if d.id in self:
+                    d = Document(d, copy=True)
+                    d.pop('matches')
+                _q.matches.append(d, scores={metric_name: _dist}, copy=False)
+
+    def _match_online(
+        self, darray, cdist, limit, normalization, metric_name, batch_size
+    ):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: length of the chunks loaded into memory from darray.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        n_x = x_mat.shape[0]
+
+        def batch_generator(y_darray: 'DocumentArrayMemmap', n_batch: int):
+            for i in range(0, len(y_darray), n_batch):
+                y_mat = np.vstack(
+                    [y_darray[j].embedding for j in range(i, i + n_batch)]
+                )
+
+                yield y_mat, i
+
+        y_batch_generator = batch_generator(darray, batch_size)
+        top_dists = np.inf * np.ones((n_x, limit))
+        top_inds = np.zeros((n_x, limit), dtype=int)
+
+        for ybatch, ybatch_start_pos in y_batch_generator:
+            distances = cdist(x_mat, ybatch, metric_name)
+            dists, inds = top_k(distances, limit, descending=False)
+
+            if normalization is not None:
path: jina/types/arrays/neural_ops.py
text: @user1:
put in one line
```suggestion
            if normalization is not None and isinstance(normalization, (tuple, list)):
```

@user1:
Or even `isinstance` already contains the first condition

@author:
done

diff_hunk: @@ -52,29 +53,117 @@ def match(
                                 all values will be rescaled into range `[a, b]`.
         :param use_scipy: use Scipy as the computation backend
         :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: if provided, then `darray` is loaded in chunks of, at most, batch_size elements. This option
+                           will be slower but more memory efficient. Specialy indicated if `darray` is a big
+                           DocumentArrayMemmap.
         """
 
-        X = np.stack(self.get_attributes('embedding'))
-        Y = np.stack(darray.get_attributes('embedding'))
-
-        if isinstance(metric, str):
+        if callable(metric):
+            cdist = metric
+        elif isinstance(metric, str):
             if use_scipy:
-                from scipy.spatial.distance import cdist
+                from scipy.spatial.distance import cdist as cdist
             else:
-                from ...math.distance import cdist
-            dists = cdist(X, Y, metric)
-        elif callable(metric):
-            dists = metric(X, Y)
+                from ...math.distance import cdist as cdist
         else:
             raise TypeError(
                 f'metric must be either string or a 2-arity function, received: {metric!r}'
             )
+        # brekpoint()
+        metric_name = metric_name or (metric.__name__ if callable(metric) else metric)
+
+        if batch_size:
+            self._match_online(
+                darray, cdist, limit, normalization, metric_name, batch_size
+            )
+        else:
+            self._match(darray, cdist, limit, normalization, metric_name)
 
+    def _match(self, darray, cdist, limit, normalization, metric_name):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        y_mat = np.stack(darray.get_attributes('embedding'))
+
+        dists = cdist(x_mat, y_mat, metric_name)
         dist, idx = top_k(dists, min(limit, len(darray)), descending=False)
         if normalization is not None:
             if isinstance(normalization, (tuple, list)):
                 dist = minmax_normalize(dist, normalization)
 
+        for _q, _ids, _dists in zip(self, idx, dist):
+            _q.matches.clear()
+            for _id, _dist in zip(_ids, _dists):
+                # Note, when match self with other, or both of them share the same Document
+                # we might have recursive matches .
+                # checkout https://github.com/jina-ai/jina/issues/3034
+                d = darray[int(_id)]
+                if d.id in self:
+                    d = Document(d, copy=True)
+                    d.pop('matches')
+                _q.matches.append(d, scores={metric_name: _dist}, copy=False)
+
+    def _match_online(
+        self, darray, cdist, limit, normalization, metric_name, batch_size
+    ):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: length of the chunks loaded into memory from darray.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        n_x = x_mat.shape[0]
+
+        def batch_generator(y_darray: 'DocumentArrayMemmap', n_batch: int):
+            for i in range(0, len(y_darray), n_batch):
+                y_mat = np.vstack(
+                    [y_darray[j].embedding for j in range(i, i + n_batch)]
+                )
+
+                yield y_mat, i
+
+        y_batch_generator = batch_generator(darray, batch_size)
+        top_dists = np.inf * np.ones((n_x, limit))
+        top_inds = np.zeros((n_x, limit), dtype=int)
+
+        for ybatch, ybatch_start_pos in y_batch_generator:
path: jina/types/arrays/neural_ops.py
text: @user1:
```suggestion
        for y_batch, y_batch_start_pos in y_batch_generator:
```

@user1:
Coherent naming

diff_hunk: @@ -52,29 +53,117 @@ def match(
                                 all values will be rescaled into range `[a, b]`.
         :param use_scipy: use Scipy as the computation backend
         :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: if provided, then `darray` is loaded in chunks of, at most, batch_size elements. This option
+                           will be slower but more memory efficient. Specialy indicated if `darray` is a big
+                           DocumentArrayMemmap.
         """
 
-        X = np.stack(self.get_attributes('embedding'))
-        Y = np.stack(darray.get_attributes('embedding'))
-
-        if isinstance(metric, str):
+        if callable(metric):
+            cdist = metric
+        elif isinstance(metric, str):
             if use_scipy:
-                from scipy.spatial.distance import cdist
+                from scipy.spatial.distance import cdist as cdist
             else:
-                from ...math.distance import cdist
-            dists = cdist(X, Y, metric)
-        elif callable(metric):
-            dists = metric(X, Y)
+                from ...math.distance import cdist as cdist
         else:
             raise TypeError(
                 f'metric must be either string or a 2-arity function, received: {metric!r}'
             )
+        # brekpoint()
+        metric_name = metric_name or (metric.__name__ if callable(metric) else metric)
+
+        if batch_size:
+            self._match_online(
+                darray, cdist, limit, normalization, metric_name, batch_size
+            )
+        else:
+            self._match(darray, cdist, limit, normalization, metric_name)
 
+    def _match(self, darray, cdist, limit, normalization, metric_name):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        y_mat = np.stack(darray.get_attributes('embedding'))
+
+        dists = cdist(x_mat, y_mat, metric_name)
         dist, idx = top_k(dists, min(limit, len(darray)), descending=False)
         if normalization is not None:
             if isinstance(normalization, (tuple, list)):
                 dist = minmax_normalize(dist, normalization)
 
+        for _q, _ids, _dists in zip(self, idx, dist):
+            _q.matches.clear()
+            for _id, _dist in zip(_ids, _dists):
+                # Note, when match self with other, or both of them share the same Document
+                # we might have recursive matches .
+                # checkout https://github.com/jina-ai/jina/issues/3034
+                d = darray[int(_id)]
+                if d.id in self:
+                    d = Document(d, copy=True)
+                    d.pop('matches')
+                _q.matches.append(d, scores={metric_name: _dist}, copy=False)
+
+    def _match_online(
+        self, darray, cdist, limit, normalization, metric_name, batch_size
+    ):
+        """
+        Computes the matches between self and `darray` loading `darray` into main memory in chunks of size `batch_size`.
+
+        :param darray: the other DocumentArray or DocumentArrayMemmap to match against
+        :param cdist: the distance metric
+        :param limit: the maximum number of matches, when not given
+                      all Documents in `another` are considered as matches
+        :param normalization: a tuple [a, b] to be used with min-max normalization,
+                                the min distance will be rescaled to `a`, the max distance will be rescaled to `b`
+                                all values will be rescaled into range `[a, b]`.
+        :param metric_name: if provided, then match result will be marked with this string.
+        :param batch_size: length of the chunks loaded into memory from darray.
+        """
+
+        x_mat = np.stack(self.get_attributes('embedding'))
+        n_x = x_mat.shape[0]
+
+        def batch_generator(y_darray: 'DocumentArrayMemmap', n_batch: int):
+            for i in range(0, len(y_darray), n_batch):
+                y_mat = np.vstack(
+                    [y_darray[j].embedding for j in range(i, i + n_batch)]
+                )
+
+                yield y_mat, i
+
+        y_batch_generator = batch_generator(darray, batch_size)
+        top_dists = np.inf * np.ones((n_x, limit))
+        top_inds = np.zeros((n_x, limit), dtype=int)
+
+        for ybatch, ybatch_start_pos in y_batch_generator:
+            distances = cdist(x_mat, ybatch, metric_name)
+            dists, inds = top_k(distances, limit, descending=False)
+
+            if normalization is not None:
+                if isinstance(normalization, (tuple, list)):
+                    dist = minmax_normalize(dist, normalization)
+
+            inds = ybatch_start_pos + inds
+            top_dists, top_inds = top_k_from_pair(
+                dists, inds, top_dists, top_inds, limit
+            )
+
+        # sort final the final `top_dists` and `top_inds` per row
+        permutation = np.argsort(top_dists, axis=1)
+        dist = np.take_along_axis(top_dists, permutation, axis=1)
+        idx = np.take_along_axis(top_inds, permutation, axis=1)
+
         m_name = metric_name or (metric.__name__ if callable(metric) else metric)
path: jina/types/arrays/neural_ops.py
text: @user1:
this part is common to `_match` and `_match_online` so u can redefine the interface of the private methods, and maybe also make these private methods private to the `match` method

@author:
makes sense, updated reusing the code