ENH: concat with keys for MultiIndex works. A bit more testing / docs required

Author: wesm

pandas/core/index.py

@@ -54,6 +54,7 @@ class Index(np.ndarray):
     _backfill = lib.backfill_object
 
     name = None
+
     def __new__(cls, data, dtype=None, copy=False, name=None):
         if isinstance(data, np.ndarray):
             if dtype is None and issubclass(data.dtype.type, np.integer):

pandas/tools/merge.py

@@ -6,7 +6,8 @@
 
 from pandas.core.frame import DataFrame, _merge_doc
 from pandas.core.groupby import get_group_index
-from pandas.core.index import Index, MultiIndex, _get_combined_index
+from pandas.core.index import (Index, MultiIndex, _get_combined_index,
+                               _ensure_index)
 from pandas.core.internals import (IntBlock, BoolBlock, BlockManager,
                                    make_block, _consolidate)
 from pandas.util.decorators import cache_readonly, Appender, Substitution
@@ -588,25 +589,31 @@ def _get_all_block_kinds(blockmaps):
 # Concatenate DataFrame objects
 
 def concat(objs, axis=0, join='outer', join_axes=None, ignore_index=False,
-           keys=None, names=None, levels=None, verify_integrity=False):
+           keys=None, levels=None, names=None, verify_integrity=False):
     """
-    Concatenate DataFrame objects row or column wise
+    Concatenate pandas objects along a particular axis with optional set logic
+    along the other axes. Can also add a layer of hierarchical indexing on the
+    concatenation axis, which may be useful if the labels are the same (or
+    overlapping) on the passed axis number
 
     Parameters
     ----------
-    objs : list of DataFrame objects
-    axis : {0, 1}, default 0
+    objs : list of DataFrame (or other pandas) objects
+    axis : {0, 1, ...}, default 0
         The axis to concatenate along
     join : {'inner', 'outer'}, default 'outer'
         How to handle indexes on other axis(es)
     join_index : index-like
     verify_integrity : boolean, default False
         Check whether the new concatenated axis contains duplicates. This can
         be very expensive relative to the actual data concatenation
+    keys : sequence-like or list of sequences
+    levels :
+    names :
 
     Returns
     -------
-    concatenated : DataFrame
+    concatenated : type of objects
     """
     op = _Concatenator(objs, axis=axis, join_axes=join_axes,
                        ignore_index=ignore_index, join=join,
@@ -722,8 +729,16 @@ def _concat_blocks(self, blocks):
                                 'DataFrames')
             return make_block(concat_values, blocks[0].items, self.new_axes[0])
         else:
-            concat_items = _concat_indexes([b.items for b in blocks])
-            # TODO: maybe want to "take" from the new columns?
+            all_items = [b.items for b in blocks]
+            if self.axis == 0 and self.keys is not None:
+                offsets = np.r_[0, [len(x._data.axes[self.axis]) for
+                                    x in self.objs]]
+                indexer = np.concatenate([offsets[i] + b.ref_locs
+                                          for i, b in enumerate(blocks)])
+                concat_items = self.new_axes[0].take(indexer)
+            else:
+                concat_items = _concat_indexes(all_items)
+
             return make_block(concat_values, concat_items, self.new_axes[0])
 
     def _concat_single_item(self, item):
@@ -837,18 +852,23 @@ def _concat_frames_hierarchical(frames, keys, names, levels, axis=0):
 def _concat_indexes(indexes):
     return indexes[0].append(indexes[1:])
 
-def _make_concat_multiindex(indexes, keys, levels, names):
-    single_level = len(levels) == 1
+def _make_concat_multiindex(indexes, keys, levels=None, names=None):
+    single_level = levels is None or len(levels) == 1
+
+    if single_level:
+        zipped = [keys]
+        if names is None:
+            names = [None]
+    else:
+        zipped = zip(*keys)
+        if names is None:
+            names = [None] * len(keys)
 
     if not _all_indexes_same(indexes):
         label_list = []
 
         # things are potentially different sizes, so compute the exact labels
         # for each level and pass those to MultiIndex.from_arrays
-        if single_level:
-            zipped = [keys]
-        else:
-            zipped = zip(*keys)
 
         for hlevel in zipped:
             to_concat = []
@@ -874,20 +894,21 @@ def _make_concat_multiindex(indexes, keys, levels, names):
 
     names.append(indexes[0].name)
 
-    new_levels = list(levels)
+    if levels is None:
+        if single_level:
+            new_levels = [_ensure_index(keys)]
+        else:
+            new_levels = [_ensure_index(k) for k in keys]
+    else:
+        new_levels = list(levels)
 
     # do something a bit more speedy
     new_levels.append(new_index)
 
     # construct labels
     labels = []
 
-    if single_level:
-        zipped = [keys]
-    else:
-        zipped = zip(*keys)
-
-    for hlevel, level in zip(zipped, levels):
+    for hlevel, level in zip(zipped, new_levels[:-1]):
         mapped = level.get_indexer(hlevel)
         labels.append(np.repeat(mapped, n))
 

pandas/tools/tests/test_merge.py

@@ -742,6 +742,42 @@ def test_append_missing_column_proper_upcast(self):
         self.assert_(appended['B'].dtype == 'O')
 
     def test_concat_with_group_keys(self):
+        df = DataFrame(np.random.randn(4, 3))
+        df2 = DataFrame(np.random.randn(4, 4))
+
+        # axis=0
+        df = DataFrame(np.random.randn(3, 4))
+        df2 = DataFrame(np.random.randn(4, 4))
+
+        result = concat([df, df2], keys=[0, 1])
+        exp_index = MultiIndex.from_arrays([[0, 0, 0, 1, 1, 1, 1],
+                                            [0, 1, 2, 0, 1, 2, 3]])
+        expected = DataFrame(np.r_[df.values, df2.values],
+                             index=exp_index)
+        tm.assert_frame_equal(result, expected)
+
+        result = concat([df, df], keys=[0, 1])
+        exp_index2 = MultiIndex.from_arrays([[0, 0, 0, 1, 1, 1],
+                                            [0, 1, 2, 0, 1, 2]])
+        expected = DataFrame(np.r_[df.values, df.values],
+                             index=exp_index2)
+        tm.assert_frame_equal(result, expected)
+
+        # axis=1
+        df = DataFrame(np.random.randn(4, 3))
+        df2 = DataFrame(np.random.randn(4, 4))
+
+        result = concat([df, df2], keys=[0, 1], axis=1)
+        expected = DataFrame(np.c_[df.values, df2.values],
+                             columns=exp_index)
+        tm.assert_frame_equal(result, expected)
+
+        result = concat([df, df], keys=[0, 1], axis=1)
+        expected = DataFrame(np.c_[df.values, df.values],
+                             columns=exp_index2)
+        tm.assert_frame_equal(result, expected)
+
+    def test_concat_keys_and_levels(self):
         pass
 
     def test_crossed_dtypes_weird_corner(self):