From 3c817ad908a3bdfa209c3aa09a8835036941b7f0 Mon Sep 17 00:00:00 2001 From: Devin Petersohn Date: Mon, 16 Apr 2018 08:50:56 -0700 Subject: [PATCH] Add slice functionality (#1832) --- python/ray/dataframe/dataframe.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/python/ray/dataframe/dataframe.py b/python/ray/dataframe/dataframe.py index b68d050db..8a51b5c6e 100644 --- a/python/ray/dataframe/dataframe.py +++ b/python/ray/dataframe/dataframe.py @@ -2945,9 +2945,7 @@ class DataFrame(object): # see if we can slice the rows indexer = self._row_metadata.convert_to_index_sliceable(key) if indexer is not None: - raise NotImplementedError("To contribute to Pandas on Ray, please" - "visit github.com/ray-project/ray.") - # return self._getitem_slice(indexer) + return self._getitem_slice(indexer) if isinstance(key, (pd.Series, np.ndarray, pd.Index, list)): return self._getitem_array(key) @@ -3015,6 +3013,15 @@ class DataFrame(object): lambda df: df.__getitem__(index), self._col_partitions[part]) + def _getitem_slice(self, key): + new_cols = _map_partitions(lambda df: df[key], + self._col_partitions) + + index = self.index[key] + return DataFrame(col_partitions=new_cols, + index=index, + columns=self.columns) + def __getattr__(self, key): """After regular attribute access, looks up the name in the columns