POC: ArrayManager -- array-based data manager for columnar store

Author: jorisvandenbossche

pandas/core/frame.py

@@ -128,7 +128,7 @@
 from pandas.core.indexes.multi import MultiIndex, maybe_droplevels
 from pandas.core.indexes.period import PeriodIndex
 from pandas.core.indexing import check_bool_indexer, convert_to_index_sliceable
-from pandas.core.internals import BlockManager
+from pandas.core.internals import ArrayManager, BlockManager
 from pandas.core.internals.construction import (
     arrays_to_mgr,
     dataclasses_to_dicts,
@@ -446,6 +446,7 @@ def __init__(
         columns: Optional[Axes] = None,
         dtype: Optional[Dtype] = None,
         copy: bool = False,
+        manager: str = "array",
     ):
         if data is None:
             data = {}
@@ -455,7 +456,7 @@ def __init__(
         if isinstance(data, DataFrame):
             data = data._mgr
 
-        if isinstance(data, BlockManager):
+        if isinstance(data, (BlockManager, ArrayManager)):
             if index is None and columns is None and dtype is None and copy is False:
                 # GH#33357 fastpath
                 NDFrame.__init__(
@@ -564,6 +565,11 @@ def __init__(
                     values, index, columns, dtype=values.dtype, copy=False
                 )
 
+        if manager == "array" and not isinstance(mgr, ArrayManager):
+            # TODO proper initialization
+            df = DataFrame(mgr, manager="block")
+            arrays = [arr.copy() for arr in df._iter_column_arrays()]
+            mgr = ArrayManager(arrays, [mgr.axes[1], mgr.axes[0]])
         NDFrame.__init__(self, mgr)
 
     # ----------------------------------------------------------------------
@@ -638,6 +644,8 @@ def _is_homogeneous_type(self) -> bool:
         ...    "B": np.array([1, 2], dtype=np.int64)})._is_homogeneous_type
         False
         """
+        if isinstance(self._mgr, ArrayManager):
+            return False
         if self._mgr.any_extension_types:
             return len({block.dtype for block in self._mgr.blocks}) == 1
         else:

pandas/core/generic.py

@@ -100,7 +100,7 @@
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.core.indexes.period import Period, PeriodIndex
 import pandas.core.indexing as indexing
-from pandas.core.internals import BlockManager
+from pandas.core.internals import ArrayManager, BlockManager
 from pandas.core.missing import find_valid_index
 from pandas.core.ops import _align_method_FRAME
 from pandas.core.shared_docs import _shared_docs
@@ -197,7 +197,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     _deprecations: FrozenSet[str] = frozenset(["get_values", "tshift"])
     _metadata: List[str] = []
     _is_copy = None
-    _mgr: BlockManager
+    _mgr: Union[BlockManager, ArrayManager]
     _attrs: Dict[Optional[Hashable], Any]
     _typ: str
 
@@ -206,7 +206,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
 
     def __init__(
         self,
-        data: BlockManager,
+        data: Union[BlockManager, ArrayManager],
         copy: bool = False,
         attrs: Optional[Mapping[Optional[Hashable], Any]] = None,
     ):
@@ -223,7 +223,9 @@ def __init__(
         object.__setattr__(self, "_flags", Flags(self, allows_duplicate_labels=True))
 
     @classmethod
-    def _init_mgr(cls, mgr, axes, dtype=None, copy: bool = False) -> BlockManager:
+    def _init_mgr(
+        cls, mgr, axes, dtype=None, copy: bool = False
+    ) -> Union[BlockManager, ArrayManager]:
         """ passed a manager and a axes dict """
         for a, axe in axes.items():
             if axe is not None:
@@ -5372,6 +5374,8 @@ def _protect_consolidate(self, f):
         Consolidate _mgr -- if the blocks have changed, then clear the
         cache
         """
+        if isinstance(self._mgr, ArrayManager):
+            return f()
         blocks_before = len(self._mgr.blocks)
         result = f()
         if len(self._mgr.blocks) != blocks_before:

pandas/core/internals/__init__.py

@@ -16,6 +16,7 @@
 from pandas.core.internals.concat import concatenate_block_managers
 from pandas.core.internals.managers import (
     BlockManager,
+    ArrayManager,
     SingleBlockManager,
     create_block_manager_from_arrays,
     create_block_manager_from_blocks,

pandas/core/internals/concat.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 import copy
+import itertools
 from typing import Dict, List
 
 import numpy as np
@@ -26,7 +27,7 @@
 import pandas.core.algorithms as algos
 from pandas.core.arrays import ExtensionArray
 from pandas.core.internals.blocks import make_block
-from pandas.core.internals.managers import BlockManager
+from pandas.core.internals.managers import ArrayManager, BlockManager
 
 
 def concatenate_block_managers(
@@ -46,6 +47,23 @@ def concatenate_block_managers(
     -------
     BlockManager
     """
+    # breakpoint()
+
+    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 = [
+                np.concatenate([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]])
+        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]])
+
     concat_plans = [
         _get_mgr_concatenation_plan(mgr, indexers) for mgr, indexers in mgrs_indexers
     ]