from typing import Type
from . import BaseNdArray
from .dense import BaseDenseNdArray
from .dense.numpy import DenseNdArray
from .sparse import BaseSparseNdArray
from .sparse.scipy import SparseNdArray
from ...proto import jina_pb2
__all__ = ['NdArray']
[docs]class NdArray(BaseNdArray):
"""
:class:`NdArray` is one of the **primitive data type** in Jina.
It offers a Pythonic interface to allow users access and manipulate
:class:`jina.jina_pb2.NdArrayProto` object without working with Protobuf itself.
A generic view of the Protobuf NdArray, unifying the view of DenseNdArray and SparseNdArray
This class should be used in nearly all the Jina context.
Simple usage:
.. highlight:: python
.. code-block:: python
# start from empty proto
a = NdArray()
# start from an existig proto
a = NdArray(doc.embedding)
# set value
a.value = np.random.random([10, 5])
# get value
print(a.value)
# set value to a TF sparse tensor
a.is_sparse = True
a.value = SparseTensor(...)
print(a.value)
Advanced usage:
:class:`NdArray` also takes a dense NdArray and a sparse NdArray constructor
as arguments. You can consider them as the backend for dense and sparse NdArray. The combination
is your choice, it could be:
.. highlight:: python
.. code-block:: python
# numpy (dense) + scipy (sparse)
from .dense.numpy import DenseNdArray
from .sparse.scipy import SparseNdArray
NdArray(dense_cls=DenseNdArray, sparse_cls=SparseNdArray)
# numpy (dense) + pytorch (sparse)
from .dense.numpy import DenseNdArray
from .sparse.pytorch import SparseNdArray
NdArray(dense_cls=DenseNdArray, sparse_cls=SparseNdArray)
# numpy (dense) + tensorflow (sparse)
from .dense.numpy import DenseNdArray
from .sparse.tensorflow import SparseNdArray
NdArray(dense_cls=DenseNdArray, sparse_cls=SparseNdArray)
Once you set `sparse_cls`, it will only accept the data type in that particular type.
That is, you can not use a :class:`NdArray` equipped with Tensorflow sparse to
set/get Pytorch or Scipy sparse matrices.
:param proto: the protobuf message, when not given then create a new one via :meth:`get_null_proto`
:param is_sparse: if the ndarray is sparse, can be changed later
:param dense_cls: the to-be-used class for DenseNdArray when `is_sparse=False`
:param sparse_cls: the to-be-used class for SparseNdArray when `is_sparse=True`
:param args: additional positional arguments stored as member and used for the parent initialization
:param kwargs: additional key value arguments stored as member and used for the parent initialization
"""
def __init__(
self,
proto: 'jina_pb2.NdArrayProto' = None,
is_sparse: bool = False,
dense_cls: Type['BaseDenseNdArray'] = DenseNdArray,
sparse_cls: Type['BaseSparseNdArray'] = SparseNdArray,
*args,
**kwargs
):
self.is_sparse = is_sparse
self.dense_cls = dense_cls
self.sparse_cls = sparse_cls
super().__init__(proto, *args, **kwargs)
self._args = args
self._kwargs = kwargs
[docs] def null_proto(self):
"""
Get the new protobuf representation.
:return: ndarray proto instance
"""
return jina_pb2.NdArrayProto()
@property
def value(self):
"""
Get the value of protobuf and return in corresponding type.
:return: value
"""
stype = self._pb_body.WhichOneof('content')
if stype == 'dense':
return self.dense_cls(self._pb_body.dense).value
elif stype == 'sparse':
return self.sparse_cls(self._pb_body.sparse, **self._kwargs).value
@value.setter
def value(self, value):
"""
Set the value of protobuf and with :param:`value`.
:param value: value to set
"""
if self.is_sparse:
self.sparse_cls(self._pb_body.sparse).value = value
else:
self.dense_cls(self._pb_body.dense).value = value
def _build_content_dict(self):
return {'value': self.value, 'is_sparse': self.is_sparse}