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

asv_bench/benchmarks/stat_ops.py

@@ -11,6 +11,9 @@ class FrameOps:
     param_names = ["op", "dtype", "axis"]
 
     def setup(self, op, dtype, axis):
+        if dtype == "Int64":
+            # XXX only dealing with numpy arrays in ArrayManager right now
+            raise NotImplementedError
         if op == "mad" and dtype == "Int64":
             # GH-33036, GH#33600
             raise NotImplementedError

pandas/core/config_init.py

@@ -482,6 +482,12 @@ def use_inf_as_na_cb(key):
     cf.register_option(
         "use_inf_as_null", False, use_inf_as_null_doc, cb=use_inf_as_na_cb
     )
+    cf.register_option(
+        "data_manager",
+        "block",
+        "internal manager type",
+        validator=is_one_of_factory(["block", "array"]),
+    )
 
 cf.deprecate_option(
     "mode.use_inf_as_null", msg=use_inf_as_null_doc, rkey="mode.use_inf_as_na"

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,9 @@ def __init__(
         columns: Optional[Axes] = None,
         dtype: Optional[Dtype] = None,
         copy: bool = False,
+        # TODO setting default to "array" for testing purposes (the actual default
+        # needs to stay "block" initially of course for backwards compatibility)
+        manager: Optional[str] = None,
     ):
         if data is None:
             data = {}
@@ -455,7 +458,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 +567,16 @@ def __init__(
                     values, index, columns, dtype=values.dtype, copy=False
                 )
 
+        if manager is None:
+            manager = get_option("mode.data_manager")
+
+        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]])
+        # TODO check for case of manager="block" but mgr is ArrayManager
+
         NDFrame.__init__(self, mgr)
 
     # ----------------------------------------------------------------------
@@ -638,6 +651,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:
@@ -649,6 +664,8 @@ def _can_fast_transpose(self) -> bool:
         """
         Can we transpose this DataFrame without creating any new array objects.
         """
+        if isinstance(self._data, ArrayManager):
+            return False
         if self._data.any_extension_types:
             # TODO(EA2D) special case would be unnecessary with 2D EAs
             return False

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:
@@ -236,8 +238,9 @@ def _init_mgr(cls, mgr, axes, dtype=None, copy: bool = False) -> BlockManager:
             mgr = mgr.copy()
         if dtype is not None:
             # avoid further copies if we can
-            if len(mgr.blocks) > 1 or mgr.blocks[0].values.dtype != dtype:
-                mgr = mgr.astype(dtype=dtype)
+            # TODO
+            # if len(mgr.blocks) > 1 or mgr.blocks[0].values.dtype != dtype:
+            mgr = mgr.astype(dtype=dtype)
         return mgr
 
     # ----------------------------------------------------------------------
@@ -5372,6 +5375,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,21 @@ 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 = [
+                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
     ]