@@ -47,6 +47,9 @@ def transform(self, X, y):
4747 def _gets_props (self , array ):
4848 props = {}
4949 props ["type" ] = array .__class__ .__name__
50+ if props ["type" ].lower () in ("series" , "dataframe" ):
51+ suffix = "dask-" if is_dask_container (array ) else "pandas-"
52+ props ["type" ] = suffix + props ["type" ]
5053 props ["columns" ] = getattr (array , "columns" , None )
5154 props ["name" ] = getattr (array , "name" , None )
5255 props ["dtypes" ] = getattr (array , "dtypes" , None )
@@ -56,13 +59,34 @@ def _transfrom_one(self, array, props):
5659 type_ = props ["type" ].lower ()
5760 if type_ == "list" :
5861 ret = array .tolist ()
59- elif type_ == "dataframe" :
62+ elif type_ == "pandas- dataframe" :
6063 import pandas as pd
64+
6165 ret = pd .DataFrame (array , columns = props ["columns" ])
6266 ret = ret .astype (props ["dtypes" ])
63- elif type_ == "series" :
67+ elif type_ == "pandas- series" :
6468 import pandas as pd
69+
6570 ret = pd .Series (array , dtype = props ["dtypes" ], name = props ["name" ])
71+ elif type_ == "dask-dataframe" :
72+ from dask import dataframe
73+
74+ if is_dask_container (array ):
75+ ret = dataframe .from_dask_array (
76+ array , columns = props ["columns" ]
77+ )
78+ else :
79+ ret = dataframe .from_array (array , columns = props ["columns" ])
80+ ret = ret .astype (props ["dtypes" ])
81+ elif type_ == "dask-series" :
82+ from dask import dataframe
83+
84+ if is_dask_container (array ):
85+ ret = dataframe .from_dask_array (array )
86+ else :
87+ ret = dataframe .from_array (array )
88+ ret = ret .astype (props ["dtypes" ])
89+ ret = ret .rename (props ["name" ])
6690 else :
6791 ret = array
6892 return ret
0 commit comments