[DataFrame] Implemented __setitem__, select_dtypes, and astype

python/ray/dataframe/dataframe.py

@@ -15,7 +15,8 @@
     is_bool_dtype,
     is_list_like,
     is_numeric_dtype,
-    is_timedelta64_dtype)
+    is_timedelta64_dtype,
+    _get_dtype_from_object)
 from pandas.core.indexing import check_bool_indexer
 
 import warnings
@@ -977,9 +978,34 @@ def assign(self, **kwargs):
             "github.com/ray-project/ray.")
 
     def astype(self, dtype, copy=True, errors='raise', **kwargs):
-        raise NotImplementedError(
-            "To contribute to Pandas on Ray, please visit "
-            "github.com/ray-project/ray.")
+        if errors == 'raise':
+            try:
+                pd.DataFrame().astype(dtype)
+            except (ValueError, TypeError):
+                return self
+        if isinstance(dtype, dict):
+            new_rows = _map_partitions(lambda df: df.astype(dtype=dtype,
+                                                            copy=True,
+                                                            errors='ignore',
+                                                            **kwargs),
+                                       self._row_partitions)
+            if copy:
+                return DataFrame(row_partitions=new_rows,
+                                 columns=self.columns,
+                                 index=self.index)
+            self._row_partitions = new_rows
+        else:
+            new_blocks = [_map_partitions(lambda d: d.astype(dtype=dtype,
+                                                             copy=True,
+                                                             errors='ignore',
+                                                             **kwargs),
+                                          block)
+                          for block in self._block_partitions]
+            if copy:
+                return DataFrame(block_partitions=new_blocks,
+                                 columns=self.columns,
+                                 index=self.index)
+            self._block_partitions = new_blocks
 
     def at_time(self, time, asof=False):
         raise NotImplementedError(
@@ -2418,9 +2444,31 @@ def rdiv(self, other, axis='columns', level=None, fill_value=None):
     def reindex(self, labels=None, index=None, columns=None, axis=None,
                 method=None, copy=True, level=None, fill_value=np.nan,
                 limit=None, tolerance=None):
-        raise NotImplementedError(
-            "To contribute to Pandas on Ray, please visit "
-            "github.com/ray-project/ray.")
+        if not columns:
+            return DataFrame()
+        col_idx = [self.columns.get_loc(columns[i])
+                   for i in range(len(columns))
+                   if columns[i] in self.columns]
+
+        if not copy:
+            raise NotImplementedError(
+                "To contribute to Pandas on Ray, please visit "
+                "github.com/ray-project/ray.")
+
+        axis = pd.DataFrame()._get_axis_number(axis) if (axis) else 0
+        if axis == 1 or columns:
+            def row_helper(df, col_idx):
+                df = df.reindex(columns=col_idx, copy=True)
+                return df
+            new_rows = _map_partitions(row_helper,
+                                       self._row_partitions, col_idx)
+            return DataFrame(row_partitions=new_rows,
+                             columns=columns,
+                             index=self.index)
+        else:
+            raise NotImplementedError(
+                "To contribute to Pandas on Ray, please visit "
+                "github.com/ray-project/ray.")
 
     def reindex_axis(self, labels, axis=0, method=None, level=None, copy=True,
                      limit=None, fill_value=np.nan):
@@ -2688,9 +2736,33 @@ def select(self, crit, axis=0):
             "github.com/ray-project/ray.")
 
     def select_dtypes(self, include=None, exclude=None):
-        raise NotImplementedError(
-            "To contribute to Pandas on Ray, please visit "
-            "github.com/ray-project/ray.")
+        # Validates arguments for whether both include and exclude are None or
+        # if they are disjoint. Also invalidates string dtypes.
+        pd.DataFrame().select_dtypes(include, exclude)
+
+        if include and not is_list_like(include):
+            include = [include]
+        elif not include:
+            include = []
+
+        if exclude and not is_list_like(exclude):
+            exclude = [exclude]
+        elif not exclude:
+            exclude = []
+
+        sel = tuple(map(set, (include, exclude)))
+
+        include, exclude = map(
+            lambda x: set(map(_get_dtype_from_object, x)), sel)
+
+        dtypes = self.dtypes
+        indicate = [i for i in range(len(dtypes))
+                    if ((len(exclude) != 0 and any(map(
+                        lambda x: issubclass(dtypes[i].type, x), exclude)))
+                    or (len(include) != 0 and not any(map(
+                        lambda x: issubclass(dtypes[i].type, x), include))))]
+
+        return self.drop(columns=self.columns[indicate], inplace=False)
 
     def sem(self, axis=None, skipna=None, level=None, ddof=1,
             numeric_only=None, **kwargs):
@@ -3260,7 +3332,7 @@ def _getitem_array(self, key):
             new_parts = _map_partitions(lambda df: df[key],
                                         self._col_partitions)
             columns = self.columns
-            index = self.index[key]
+            index = self._col_metadata[key].index
 
             return DataFrame(col_partitions=new_parts,
                              columns=columns,
@@ -3317,9 +3389,16 @@ def __getattr__(self, key):
             raise e
 
     def __setitem__(self, key, value):
-        raise NotImplementedError(
-            "To contribute to Pandas on Ray, please visit "
-            "github.com/ray-project/ray.")
+        if not isinstance(key, str):
+            raise NotImplementedError(
+                "To contribute to Pandas on Ray, please visit "
+                "github.com/ray-project/ray.")
+        if key not in self.columns:
+            self.insert(loc=len(self.columns), column=key, value=value)
+        else:
+            loc = self.columns.get_loc(key)
+            self.__delitem__(key)
+            self.insert(loc=loc, column=key, value=value)
 
     def __len__(self):
         """Gets the length of the dataframe.

python/ray/dataframe/test/test_dataframe.py

@@ -905,10 +905,28 @@ def test_assign():
 
 
 def test_astype():
-    ray_df = create_test_dataframe()
+    td = TestData()
+    ray_df_frame = from_pandas(td.frame, 2)
+    our_df_casted = ray_df_frame.astype(np.int32)
+    expected_df_casted = pd.DataFrame(td.frame.values.astype(np.int32),
+                                      index=td.frame.index,
+                                      columns=td.frame.columns)
 
-    with pytest.raises(NotImplementedError):
-        ray_df.astype(None)
+    assert(ray_df_equals_pandas(our_df_casted, expected_df_casted))
+
+    our_df_casted = ray_df_frame.astype(np.float64)
+    expected_df_casted = pd.DataFrame(td.frame.values.astype(np.float64),
+                                      index=td.frame.index,
+                                      columns=td.frame.columns)
+
+    assert(ray_df_equals_pandas(our_df_casted, expected_df_casted))
+
+    our_df_casted = ray_df_frame.astype(str)
+    expected_df_casted = pd.DataFrame(td.frame.values.astype(str),
+                                      index=td.frame.index,
+                                      columns=td.frame.columns)
+
+    assert(ray_df_equals_pandas(our_df_casted, expected_df_casted))
 
 
 def test_at_time():
@@ -2147,10 +2165,13 @@ def test_rdiv():
 
 
 def test_reindex():
-    ray_df = create_test_dataframe()
-
-    with pytest.raises(NotImplementedError):
-        ray_df.reindex()
+    trd = rdf.DataFrame({'a': list('abc'),
+                         'b': list(range(1, 4)),
+                         'c': np.arange(3, 6).astype('u1'),
+                         'd': np.arange(4.0, 7.0, dtype='float64'),
+                         'e': [True, False, True],
+                         'f': pd.date_range('now', periods=3).values})
+    assert(trd.reindex(columns=['a', 'b'])['b'].equals(trd['b']))
 
 
 def test_reindex_axis():
@@ -2524,10 +2545,25 @@ def test_select():
 
 
 def test_select_dtypes():
-    ray_df = create_test_dataframe()
-
-    with pytest.raises(NotImplementedError):
-        ray_df.select_dtypes()
+    df = pd.DataFrame({'test1': list('abc'),
+                       'test2': np.arange(3, 6).astype('u1'),
+                       'test3': np.arange(8.0, 11.0, dtype='float64'),
+                       'test4': [True, False, True],
+                       'test5': pd.date_range('now', periods=3).values,
+                       'test6': list(range(5, 8))})
+    include = np.float, 'integer'
+    exclude = np.bool_,
+    rd = from_pandas(df, 2)
+    r = rd.select_dtypes(include=include, exclude=exclude)
+
+    e = df[["test2", "test3", "test6"]]
+    assert(ray_df_equals_pandas(r, e))
+
+    try:
+        rdf.DataFrame().select_dtypes()
+        assert(False)
+    except ValueError:
+        assert(True)
 
 
 def test_sem():