[ArrayManager] Implement concat with axis=1 (merge/join)

.github/workflows/ci.yml

@@ -153,6 +153,7 @@ jobs:
       run: |
         source activate pandas-dev
         pytest pandas/tests/frame/methods --array-manager
+        pytest pandas/tests/reshape/merge --array-manager
 
         # indexing iset related (temporary since other tests don't pass yet)
         pytest pandas/tests/frame/indexing/test_indexing.py::TestDataFrameIndexing::test_setitem_multi_index --array-manager

pandas/core/internals/array_manager.py

@@ -754,10 +754,12 @@ def reindex_indexer(
         # ignored keywords
         consolidate: bool = True,
         only_slice: bool = False,
+        # ArrayManager specific keywords
+        do_integrity_check: bool = True,
     ) -> T:
         axis = self._normalize_axis(axis)
         return self._reindex_indexer(
-            new_axis, indexer, axis, fill_value, allow_dups, copy
+            new_axis, indexer, axis, fill_value, allow_dups, copy, do_integrity_check
         )
 
     def _reindex_indexer(
@@ -768,6 +770,7 @@ def _reindex_indexer(
         fill_value=None,
         allow_dups: bool = False,
         copy: bool = True,
+        do_integrity_check: bool = True,
     ) -> T:
         """
         Parameters
@@ -822,7 +825,7 @@ def _reindex_indexer(
         new_axes = list(self._axes)
         new_axes[axis] = new_axis
 
-        return type(self)(new_arrays, new_axes)
+        return type(self)(new_arrays, new_axes, do_integrity_check=do_integrity_check)
 
     def take(self, indexer, axis: int = 1, verify: bool = True, convert: bool = True):
         """

pandas/core/internals/concat.py

@@ -31,6 +31,47 @@
     from pandas import Index
 
 
+def concatenate_array_managers(
+    mgrs_indexers, axes: List[Index], concat_axis: int, copy: bool
+) -> Manager:
+    """
+    Concatenate array managers into one.
+
+    Parameters
+    ----------
+    mgrs_indexers : list of (ArrayManager, {axis: indexer,...}) tuples
+    axes : list of Index
+    concat_axis : int
+    copy : bool
+
+    Returns
+    -------
+    ArrayManager
+    """
+    # reindex all arrays
+    mgrs = []
+    for mgr, indexers in mgrs_indexers:
+        for ax, indexer in indexers.items():
+            mgr = mgr.reindex_indexer(
+                axes[ax], indexer, axis=ax, allow_dups=True, do_integrity_check=False
+            )
+        mgrs.append(mgr)
+
+    # concatting along the rows -> concat the reindexed arrays
+    # TODO(ArrayManager) doesn't yet preserve the correct dtype
+    if concat_axis == 1:
+        arrays = [
+            concat_compat([mgrs[i].arrays[j] for i in range(len(mgrs))])
+            for j in range(len(mgrs[0].arrays))
+        ]
+        return ArrayManager(arrays, [axes[1], axes[0]], do_integrity_check=False)
+    # concatting along the columns -> combine reindexed arrays in a single manager
+    else:
+        assert concat_axis == 0
+        arrays = list(itertools.chain.from_iterable([mgr.arrays for mgr in mgrs]))
+        return ArrayManager(arrays, [axes[1], axes[0]], do_integrity_check=False)
+
+
 def concatenate_block_managers(
     mgrs_indexers, axes: List[Index], concat_axis: int, copy: bool
 ) -> Manager:
@@ -49,19 +90,7 @@ def concatenate_block_managers(
     BlockManager
     """
     if isinstance(mgrs_indexers[0][0], ArrayManager):
-
-        if concat_axis == 1:
-            # TODO for now only fastpath without indexers
-            mgrs = [t[0] for t in mgrs_indexers]
-            arrays = [
-                concat_compat([mgrs[i].arrays[j] for i in range(len(mgrs))], axis=0)
-                for j in range(len(mgrs[0].arrays))
-            ]
-            return ArrayManager(arrays, [axes[1], axes[0]])
-        elif concat_axis == 0:
-            mgrs = [t[0] for t in mgrs_indexers]
-            arrays = list(itertools.chain.from_iterable([mgr.arrays for mgr in mgrs]))
-            return ArrayManager(arrays, [axes[1], axes[0]])
+        return concatenate_array_managers(mgrs_indexers, axes, concat_axis, copy)
 
     concat_plans = [
         _get_mgr_concatenation_plan(mgr, indexers) for mgr, indexers in mgrs_indexers

pandas/tests/frame/methods/test_drop.py

@@ -155,7 +155,7 @@ def test_drop(self):
         assert return_value is None
         tm.assert_frame_equal(df, expected)
 
-    @td.skip_array_manager_not_yet_implemented
+    @td.skip_array_manager_not_yet_implemented  # TODO(ArrayManager) groupby
     def test_drop_multiindex_not_lexsorted(self):
         # GH#11640
 

pandas/tests/frame/methods/test_explode.py

@@ -1,14 +1,9 @@
 import numpy as np
 import pytest
 
-import pandas.util._test_decorators as td
-
 import pandas as pd
 import pandas._testing as tm
 
-# TODO(ArrayManager) concat with reindexing
-pytestmark = td.skip_array_manager_not_yet_implemented
-
 
 def test_error():
     df = pd.DataFrame(

pandas/tests/frame/methods/test_join.py

@@ -3,15 +3,10 @@
 import numpy as np
 import pytest
 
-import pandas.util._test_decorators as td
-
 import pandas as pd
 from pandas import DataFrame, Index, MultiIndex, date_range, period_range
 import pandas._testing as tm
 
-# TODO(ArrayManager) concat with reindexing
-pytestmark = td.skip_array_manager_not_yet_implemented
-
 
 @pytest.fixture
 def frame_with_period_index():
@@ -234,8 +229,9 @@ def test_join(self, multiindex_dataframe_random_data):
         b = frame.loc[frame.index[2:], ["B", "C"]]
 
         joined = a.join(b, how="outer").reindex(frame.index)
-        expected = frame.copy()
-        expected.values[np.isnan(joined.values)] = np.nan
+        expected = frame.copy().values
+        expected[np.isnan(joined.values)] = np.nan
+        expected = DataFrame(expected, index=frame.index, columns=frame.columns)
 
         assert not np.isnan(joined.values).all()
 

pandas/tests/io/formats/test_printing.py

@@ -121,7 +121,7 @@ def test_ambiguous_width(self):
         assert adjoined == expected
 
 
-@td.skip_array_manager_not_yet_implemented
+@td.skip_array_manager_not_yet_implemented  # TODO(ArrayManager) JSON
 class TestTableSchemaRepr:
     @classmethod
     def setup_class(cls):

pandas/tests/io/test_fsspec.py

@@ -247,7 +247,7 @@ def test_pickle_options(fsspectest):
     tm.assert_frame_equal(df, out)
 
 
-@td.skip_array_manager_not_yet_implemented
+@td.skip_array_manager_not_yet_implemented  # TODO(ArrayManager) JSON
 def test_json_options(fsspectest):
     df = DataFrame({"a": [0]})
     df.to_json("testmem://afile", storage_options={"test": "json_write"})

pandas/tests/reshape/merge/test_join.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+import pandas.util._test_decorators as td
+
 import pandas as pd
 from pandas import (
     Categorical,
@@ -547,6 +549,7 @@ def test_join_non_unique_period_index(self):
         )
         tm.assert_frame_equal(result, expected)
 
+    @td.skip_array_manager_not_yet_implemented  # TODO(ArrayManager) groupby
     def test_mixed_type_join_with_suffix(self):
         # GH #916
         df = DataFrame(np.random.randn(20, 6), columns=["a", "b", "c", "d", "e", "f"])

pandas/tests/reshape/merge/test_merge.py

@@ -277,17 +277,27 @@ def test_merge_copy(self):
         merged["d"] = "peekaboo"
         assert (right["d"] == "bar").all()
 
-    def test_merge_nocopy(self):
+    def test_merge_nocopy(self, using_array_manager):
         left = DataFrame({"a": 0, "b": 1}, index=range(10))
         right = DataFrame({"c": "foo", "d": "bar"}, index=range(10))
 
         merged = merge(left, right, left_index=True, right_index=True, copy=False)
 
-        merged["a"] = 6
-        assert (left["a"] == 6).all()
+        if using_array_manager:
+            # With ArrayManager, setting a column doesn't change the values inplace
+            # and thus does not propagate the changes to the original left/right
+            # dataframes -> need to check that no copy was made in a different way
+            # TODO(ArrayManager) we should be able to simplify this with a .loc
+            #  setitem test: merged.loc[0, "a"] = 10; assert left.loc[0, "a"] == 10
+            #  but this currently replaces the array (_setitem_with_indexer_split_path)
+            assert merged._mgr.arrays[0] is left._mgr.arrays[0]
+            assert merged._mgr.arrays[2] is right._mgr.arrays[0]
+        else:
+            merged["a"] = 6
+            assert (left["a"] == 6).all()
 
-        merged["d"] = "peekaboo"
-        assert (right["d"] == "peekaboo").all()
+            merged["d"] = "peekaboo"
+            assert (right["d"] == "peekaboo").all()
 
     def test_intelligently_handle_join_key(self):
         # #733, be a bit more 1337 about not returning unconsolidated DataFrame
@@ -1362,7 +1372,7 @@ def test_merge_take_missing_values_from_index_of_other_dtype(self):
         expected = expected.reindex(columns=["a", "key", "b"])
         tm.assert_frame_equal(result, expected)
 
-    def test_merge_readonly(self):
+    def test_merge_readonly(self, using_array_manager):
         # https://github.com/pandas-dev/pandas/issues/27943
         data1 = DataFrame(
             np.arange(20).reshape((4, 5)) + 1, columns=["a", "b", "c", "d", "e"]
@@ -1371,7 +1381,12 @@ def test_merge_readonly(self):
             np.arange(20).reshape((5, 4)) + 1, columns=["a", "b", "x", "y"]
         )
 
-        data1._mgr.blocks[0].values.flags.writeable = False
+        if using_array_manager:
+            for arr in data1._mgr.arrays:
+                arr.flags.writeable = False
+        else:
+            data1._mgr.blocks[0].values.flags.writeable = False
+
         data1.merge(data2)  # no error