from typing import TYPE_CHECKING, Tuple, Sequence, Optional

import numpy as np

if TYPE_CHECKING:

from ..types import ArrayType

from .. import Document

def unravel(docs: Sequence['Document'], field: str) -> Optional['ArrayType']:

_first = getattr(docs[0], field)

if _first is None:

# failed to unravel, return as a list

r = [getattr(d, field) for d in docs]

if any(_rr is not None for _rr in r):

return r

else:

return None

framework, is_sparse = get_array_type(_first)

all_fields = [getattr(d, field) for d in docs]

cls_type = type(_first)

if framework == 'python':

return cls_type(all_fields)

elif framework == 'numpy':

return np.stack(all_fields)

elif framework == 'tensorflow':

import tensorflow as tf

return tf.stack(all_fields)

elif framework == 'torch':

import torch

return torch.stack(all_fields)

elif framework == 'paddle':

import paddle

return paddle.stack(all_fields)

elif framework == 'scipy':

import scipy.sparse

return cls_type(scipy.sparse.vstack(all_fields))

def ravel(value: 'ArrayType', docs: Sequence['Document'], field: str) -> None:

"""Ravel :attr:`value` into ``doc.field`` of each documents

:param docs: the docs to set

:param field: the field of the doc to set

:param value: the value to be set on ``doc.field``

"""

use_get_row = False

if hasattr(value, 'getformat'):

# for scipy only

sp_format = value.getformat()

if sp_format in {'bsr', 'coo'}:

# for BSR and COO, they dont implement [j, ...] in scipy

# but they offer get_row() API which implicitly translate the

# sparse row into CSR format, hence needs to convert back

# not very efficient, but this is the best we can do.

use_get_row = True

if use_get_row:

emb_shape0 = value.shape[0]

for d, j in zip(docs, range(emb_shape0)):

row = getattr(value.getrow(j), f'to{sp_format}')()

setattr(d, field, row)

elif isinstance(value, (list, tuple)):

for d, j in zip(docs, value):

setattr(d, field, j)

else:

emb_shape0 = value.shape[0]

for d, j in zip(docs, range(emb_shape0)):

setattr(d, field, value[j, ...])

def get_array_type(array: 'ArrayType') -> Tuple[str, bool]:

"""Get the type of ndarray without importing the framework

:param array: any array, scipy, numpy, tf, torch, etc.

:return: a tuple where the first element represents the framework, the second represents if it is sparse array

"""

module_tags = array.__class__.__module__.split('.')

class_name = array.__class__.__name__

if isinstance(array, (list, tuple)):

return 'python', False

if 'numpy' in module_tags:

return 'numpy', False

if 'jina' in module_tags:

if class_name == 'NdArray':

return 'jina', False # sparse or not is irrelevant

if 'docarray_pb2' in module_tags:

if class_name == 'NdArrayProto':

return 'jina_proto', False # sparse or not is irrelevant

if 'tensorflow' in module_tags:

if class_name == 'SparseTensor':

return 'tensorflow', True

if class_name == 'Tensor' or class_name == 'EagerTensor':

return 'tensorflow', False

if 'torch' in module_tags and class_name == 'Tensor':

return 'torch', array.is_sparse

if 'paddle' in module_tags and class_name == 'Tensor':

# Paddle does not support sparse tensor on 11/8/2021

# https://github.com/PaddlePaddle/Paddle/issues/36697

return 'paddle', False

if 'scipy' in module_tags and 'sparse' in module_tags:

return 'scipy', True

raise TypeError(f'can not determine the array type: {module_tags}.{class_name}')

def to_numpy_array(value) -> 'np.ndarray':

"""Return the value always in :class:`numpy.ndarray` regardless the framework type.

:return: the value in :class:`numpy.ndarray`.

"""

v = value

framework, is_sparse = get_array_type(value)

if is_sparse:

if hasattr(v, 'todense'):

v = v.todense()

elif hasattr(v, 'to_dense'):

v = v.to_dense()

elif framework == 'tensorflow':

import tensorflow as tf

if isinstance(v, tf.SparseTensor):

v = tf.sparse.to_dense(v)

if hasattr(v, 'numpy'):

v = v.numpy()

return v