Repository: jina-ai/serve
Pull request: feat: support for sparse arrays in documents
This PR wants to complete sparse support, as discussed in #2182 .

Problem statement:
Allow sparse vectors to be supported by Jina Primitive Data Types
# Feature Request

This ticket is a sub compound of implementing sparse vector functionality to Jina. Currently, there is some low-level code implementation to support sparse vectors.  Over the course of this epic we will implement the remaining logic. Main goal is to be able to build a fully neural search system where embedding representations are `sparse vectors`. This will support certain usecase like TD-IDF, and provide the functionitilty for graph embedding in the future. 

# Acceptance Criteria

- We need to be able to store `sparse` vectors in `blob` field of `Document content`

- We need to be able to store `sparse` vectors in `embedding` field of `Document`


<p></p>

Hints from issue discussion:
Link to subfeature: https://jinaai.productboard.com/feature-board/planning/features/7331485

Commit message: fix: update format

Patch to review:
diff --git a/jina/types/document/__init__.py b/jina/types/document/__init__.py
index 12475263bb4a5..dcd3b928e2fdb 100644
--- a/jina/types/document/__init__.py
+++ b/jina/types/document/__init__.py
@@ -10,6 +10,7 @@
 from typing import Union, Dict, Optional, TypeVar, Any, Tuple, List
 
 import numpy as np
+import scipy
 from google.protobuf import json_format
 from google.protobuf.field_mask_pb2 import FieldMask
 
@@ -27,6 +28,7 @@
 from ...logging import default_logger
 from ...proto import jina_pb2
 
+
 __all__ = ['Document', 'DocumentContentType', 'DocumentSourceType']
 DIGEST_SIZE = 8
 
@@ -492,6 +494,15 @@ def _update_ndarray(self, k, v):
         elif isinstance(v, NdArray):
             NdArray(getattr(self._pb_body, k)).is_sparse = v.is_sparse
             NdArray(getattr(self._pb_body, k)).value = v.value
+        elif scipy.sparse.issparse(v) == True:
+            from ..ndarray.sparse.scipy import SparseNdArray
+
+            protbuff_updater = NdArray(
+                is_sparse=True,
+                sparse_cls=SparseNdArray,
+                proto=getattr(self._pb_body, k),
+            )
+            protbuff_updater.value = v
         else:
             raise TypeError(f'{k} is in unsupported type {typename(v)}')
 
diff --git a/tests/unit/types/ndarray/test_ndarray.py b/tests/unit/types/ndarray/test_ndarray.py
index f3ce83e5db67b..8acc12e2cb0b8 100644
--- a/tests/unit/types/ndarray/test_ndarray.py
+++ b/tests/unit/types/ndarray/test_ndarray.py
@@ -158,6 +158,45 @@ def test_generic():
     np.testing.assert_equal(b.value, c)
 
 
+def test_sparse_scipy_formats():
+    from jina.types.ndarray.generic import NdArray
+    import scipy.sparse as sp
+    from scipy.sparse import coo_matrix, bsr_matrix, csr_matrix, csc_matrix
+    from jina import Document
+
+    row = np.array([0, 0, 1, 2, 2, 2])
+    col = np.array([0, 2, 2, 0, 1, 2])
+    data = np.array([1, 2, 3, 4, 5, 6])
+
+    X = sp.coo_matrix((data, (row, col)), shape=(4, 10))
+    d = Document()
+    d.embedding = X
+    # The following print statement is true because data is converted to coo_matrix
+    print('d.embedding == sp.coo_matrix: ', d.embedding == sp.coo_matrix)
+    np.testing.assert_array_equal(d.embedding.todense(), X.todense())
+
+    X = sp.bsr_matrix((data, (row, col)), shape=(3, 3))
+    d = Document()
+    d.embedding = X
+    # The following print statement is False because data is converted to coo_matrix
+    print('d.embedding == sp.bsr_matrix: ', d.embedding == sp.bsr_matrix)
+    np.testing.assert_array_equal(d.embedding.todense(), X.todense())
+
+    X = sp.csr_matrix((data, (row, col)), shape=(4, 10))
+    d = Document()
+    d.embedding = X
+    # The following print statement is False because data is converted to coo_matrix
+    print('d.embedding == sp.csr_matrix: ', d.embedding == sp.csr_matrix)
+    np.testing.assert_array_equal(d.embedding.todense(), X.todense())
+
+    X = sp.csc_matrix((data, (row, col)), shape=(4, 10))
+    d = Document()
+    d.embedding = X
+    # The following print statement is False because data is converted to coo_matrix
+    print('d.embedding == sp.csc_matrix: ', d.embedding == sp.csc_matrix)
+    np.testing.assert_array_equal(d.embedding.todense(), X.todense())
+
+
 @pytest.mark.parametrize('shape', [[10], [7, 8], [7, 8, 9]])
 def test_dummy_numpy_sparse(shape):
     a = np.random.random(shape)


Code review comments:
diff_hunk: @@ -492,6 +494,15 @@ def _update_ndarray(self, k, v):
         elif isinstance(v, NdArray):
             NdArray(getattr(self._pb_body, k)).is_sparse = v.is_sparse
             NdArray(getattr(self._pb_body, k)).value = v.value
+        elif scipy.sparse.issparse(v) == True:
path: jina/types/document/__init__.py
text: @user1:
style comment, `== True` is not needed