ENH: set name attribute during DataFrame.apply computation, close #983

pandas-dev · Apr 10, 2012 · cb730a7 · cb730a7
1 parent e1583c1
commit cb730a7
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 6 deletions.
diff --git a/pandas/core/frame.py b/pandas/core/frame.py
@@ -3082,7 +3082,8 @@ def _apply_standard(self, func, axis, ignore_failures=False):
             values = self.values
             dummy = Series(np.nan, index=self._get_axis(axis),
                            dtype=values.dtype)
-            result = lib.reduce(values, func, axis=axis, dummy=dummy)
+            result = lib.reduce(values, func, axis=axis, dummy=dummy,
+                                labels=self._get_agg_axis(axis))
             return Series(result, index=self._get_agg_axis(axis))
         except Exception:
             pass

diff --git a/pandas/src/reduce.pyx b/pandas/src/reduce.pyx
@@ -8,9 +8,10 @@ cdef class Reducer:
     '''
     cdef:
         Py_ssize_t increment, chunksize, nresults
-        object arr, dummy, f
+        object arr, dummy, f, labels
 
-    def __init__(self, object arr, object f, axis=1, dummy=None):
+    def __init__(self, object arr, object f, axis=1, dummy=None,
+                 labels=None):
         n, k = arr.shape
 
         if axis == 0:
@@ -31,6 +32,7 @@ cdef class Reducer:
         self.f = f
         self.arr = arr
         self.dummy = self._check_dummy(dummy)
+        self.labels = labels
 
     def _check_dummy(self, dummy=None):
         if dummy is None:
@@ -51,14 +53,28 @@ cdef class Reducer:
             Py_ssize_t i
             flatiter it
             object res
+            bint set_label = 0
+            ndarray[object] labels
 
         arr = self.arr
         chunk = self.dummy
 
         dummy_buf = chunk.data
         chunk.data = arr.data
+
+        set_label = self.labels is not None
+
+        if set_label:
+            if not np.issubdtype(self.labels.dtype, object):
+                labels = self.labels.astype('O')
+            else:
+                labels = self.labels
+
         try:
             for i in range(self.nresults):
+                if set_label:
+                    chunk.name = labels[i]
+
                 res = self.f(chunk)
                 if i == 0:
                     result = self._get_result_array(res)
@@ -228,6 +244,6 @@ cdef class Slider:
         self.buf.shape[0] = self.orig_len
         self.buf.data = self.orig_data
 
-def reduce(arr, f, axis=0, dummy=None):
-    reducer = Reducer(arr, f, axis=axis, dummy=dummy)
+def reduce(arr, f, axis=0, dummy=None, labels=None):
+    reducer = Reducer(arr, f, axis=axis, dummy=dummy, labels=labels)
     return reducer.get_result()
diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py
@@ -3882,6 +3882,15 @@ def test_apply_convert_objects(self):
         result = data.apply(lambda x: x, axis=1)
         assert_frame_equal(result, data)
 
+    def test_apply_attach_name(self):
+        result = self.frame.apply(lambda x: x.name)
+        expected = Series(self.frame.columns, index=self.frame.columns)
+        assert_series_equal(result, expected)
+
+        result = self.frame.apply(lambda x: x.name, axis=1)
+        expected = Series(self.frame.index, index=self.frame.index)
+        assert_series_equal(result, expected)
+
     def test_applymap(self):
         applied = self.frame.applymap(lambda x: x * 2)
         assert_frame_equal(applied, self.frame * 2)
@@ -5040,7 +5049,6 @@ def test_stale_cached_series_bug_473(self):
         exp = Y['g'].sum()
         self.assert_(isnull(Y['g']['c']))
 
-
 if __name__ == '__main__':
     # unittest.main()
     import nose