-
Notifications
You must be signed in to change notification settings - Fork 237
Expand file tree
/
Copy pathdataframe.py
More file actions
43 lines (32 loc) · 1.27 KB
/
dataframe.py
File metadata and controls
43 lines (32 loc) · 1.27 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
from typing import TYPE_CHECKING, Type
if TYPE_CHECKING: # pragma: no cover
from pandas import DataFrame
from docarray.typing import T
class DataframeIOMixin:
"""Save/load from :class:`pandas.dataframe`
.. note::
These functions require you to install `pandas`
"""
def to_dataframe(self, **kwargs) -> 'DataFrame':
"""Export itself to a :class:`pandas.DataFrame` object.
:param kwargs: the extra kwargs will be passed to :meth:`pandas.DataFrame.from_dict`.
:return: a :class:`pandas.DataFrame` object
"""
from pandas import DataFrame
return DataFrame.from_dict(self.to_list(), **kwargs)
@classmethod
def from_dataframe(cls: Type['T'], df: 'DataFrame', *args, **kwargs) -> 'T':
"""Import a :class:`DocumentArray` from a :class:`pandas.DataFrame` object.
:param df: a :class:`pandas.DataFrame` object.
:return: a :class:`DocumentArray` object
"""
da = cls(**kwargs)
from docarray import Document
for m in df.to_dict(orient='records'):
# drop nan
da.append(
Document(
{k: v for k, v in m.items() if (not isinstance(v, float) or v == v)}
)
)
return da