Implements head/tail based on iloc, and fixes bug in getitem. (#1057)

Author: sighingnow

mars/dataframe/indexing/__init__.py

@@ -14,20 +14,22 @@
 
 
 def _install():
-    from .iloc import iloc
+    from .iloc import iloc, head, tail
     from .loc import loc
     from .set_index import set_index
     from .getitem import dataframe_getitem, series_getitem
     from ..operands import DATAFRAME_TYPE, SERIES_TYPE
 
-    for cls in DATAFRAME_TYPE:
+    for cls in DATAFRAME_TYPE + SERIES_TYPE:
         setattr(cls, 'iloc', property(iloc))
         setattr(cls, 'loc', property(loc))
+        setattr(cls, 'head', head)
+        setattr(cls, 'tail', tail)
+
+    for cls in DATAFRAME_TYPE:
         setattr(cls, 'set_index', set_index)
         setattr(cls, '__getitem__', dataframe_getitem)
     for cls in SERIES_TYPE:
-        setattr(cls, 'iloc', property(iloc))
-        setattr(cls, 'loc', property(loc))
         setattr(cls, '__getitem__', series_getitem)
 
 

mars/dataframe/indexing/getitem.py

@@ -276,11 +276,11 @@ def tile_with_mask(cls, op):
             for idx, df_chunk in zip(out_chunk_indexes, df_chunks):
                 mask_chunk = mask_chunks[df_chunk.index[0]]
                 index_value = parse_index(out_df.index_value.to_pandas(), df_chunk)
-                out_chunk = op.copy().reset_key().new_chunk([df_chunk, mask_chunk],
-                                                            shape=(np.nan, df_chunk.shape[1]), index=idx,
+                out_chunk = op.copy().reset_key().new_chunk([df_chunk, mask_chunk], index=idx,
+                                                            shape=(np.nan, df_chunk.shape[1]),
+                                                            dtypes=df_chunk.dtypes,
                                                             index_value=index_value,
-                                                            columns_value=df_chunk.columns_value,
-                                                            dtypes=df_chunk.dtypes)
+                                                            columns_value=df_chunk.columns_value)
                 out_chunks.append(out_chunk)
 
         else:
@@ -292,8 +292,10 @@ def tile_with_mask(cls, op):
                     chunk_op = op.copy().reset_key()
                     chunk_op._mask = op.mask.iloc[nsplits_acc[idx]:nsplits_acc[idx+1]]
                     out_chunk = chunk_op.new_chunk([in_chunk], index=in_chunk.index,
-                                                   shape=(np.nan, in_chunk.shape[1]), dtypes=in_chunk.dtypes,
-                                                   index_value=in_df.index_value, columns_value=in_chunk.columns_value)
+                                                   shape=(np.nan, in_chunk.shape[1]),
+                                                   dtypes=in_chunk.dtypes,
+                                                   index_value=in_df.index_value,
+                                                   columns_value=in_chunk.columns_value)
                     out_chunks.append(out_chunk)
 
             nsplits = ((np.nan,) * in_df.chunk_shape[0], in_df.nsplits[1])
@@ -363,7 +365,7 @@ def execute(cls, ctx, op):
                 mask = ctx[op.inputs[1].key]
             else:
                 mask = op.mask
-            ctx[op.outputs[0].key] = df[mask]
+            ctx[op.outputs[0].key] = df[mask.reindex_like(ctx[op.inputs[0].key]).fillna(False)]
 
 
 _list_like_types = (list, np.ndarray, SERIES_TYPE, pd.Series, TENSOR_TYPE)
@@ -382,7 +384,9 @@ def dataframe_getitem(df, item):
             if col_name not in columns:
                 raise KeyError('%s not in columns' % col_name)
         op = DataFrameIndex(col_names=item, object_type=ObjectType.dataframe)
-    elif isinstance(item, _list_like_types) and astensor(item).dtype == np.bool:
+    elif isinstance(item, _list_like_types):
+        # NB: don't enforce the dtype of `item` to be `bool` since it may be unknown
+        print('use mask: item = ', item)
         op = DataFrameIndex(mask=item, object_type=ObjectType.dataframe)
     else:
         if item not in columns:

mars/dataframe/indexing/iloc.py

@@ -432,3 +432,11 @@ def execute(cls, ctx, op):
 
 def iloc(a):
     return DataFrameIloc(a)
+
+
+def head(a, n=5):
+    return DataFrameIloc(a)[0:n]
+
+
+def tail(a, n=5):
+    return DataFrameIloc(a)[-n:]

mars/dataframe/indexing/tests/test_indexing_execution.py

@@ -456,3 +456,41 @@ def testSeriesGetitem(self):
         series5 = series[selected]
         pd.testing.assert_series_equal(
             self.executor.execute_dataframe(series5, concat=True)[0], data[selected])
+
+    def testHead(self):
+        data = pd.DataFrame(np.random.rand(10, 5), columns=['c1', 'c2', 'c3', 'c4', 'c5'])
+        df = md.DataFrame(data, chunk_size=2)
+
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(), concat=True)[0], data.head())
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(3), concat=True)[0], data.head(3))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(-3), concat=True)[0], data.head(-3))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(8), concat=True)[0], data.head(8))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(-8), concat=True)[0], data.head(-8))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(13), concat=True)[0], data.head(13))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.head(-13), concat=True)[0], data.head(-13))
+
+    def testTail(self):
+        data = pd.DataFrame(np.random.rand(10, 5), columns=['c1', 'c2', 'c3', 'c4', 'c5'])
+        df = md.DataFrame(data, chunk_size=2)
+
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(), concat=True)[0], data.tail())
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(3), concat=True)[0], data.tail(3))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(-3), concat=True)[0], data.tail(-3))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(8), concat=True)[0], data.tail(8))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(-8), concat=True)[0], data.tail(-8))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(13), concat=True)[0], data.tail(13))
+        pd.testing.assert_frame_equal(
+            self.executor.execute_dataframe(df.tail(-13), concat=True)[0], data.tail(-13))