PERF: column_arrays

pandas/core/internals/managers.py

@@ -981,28 +981,35 @@ def iget_values(self, i: int) -> ArrayLike:
         return values
 
     @property
-    def column_arrays(self) -> list[np.ndarray]:
+    def column_arrays(self) -> list[ArrayLike]:
         """
         Used in the JSON C code to access column arrays.
         This optimizes compared to using `iget_values` by converting each
-        block.values to a np.ndarray only once up front
         """
-        # special casing datetimetz to avoid conversion through object dtype
-        arrays = [
-            blk.values._ndarray
-            if isinstance(blk, DatetimeTZBlock)
-            else np.asarray(blk.values)
-            for blk in self.blocks
-        ]
-        result = []
-        for i in range(len(self.items)):
-            arr = arrays[self.blknos[i]]
-            if arr.ndim == 2:
-                values = arr[self.blklocs[i]]
+        # This is an optimized equivalent to
+        #  result = [self.iget_values(i) for i in range(len(self.items))]
+        result: list[ArrayLike | None] = [None] * len(self.items)
+        for blk in self.blocks:
+            is_dtz = isinstance(blk, DatetimeTZBlock)
+            mgr_locs = blk._mgr_locs
+            values = blk.values
+            if values.ndim == 1:
+                # TODO(EA2D): special casing not needed with 2D EAs
+                result[mgr_locs.indexer] = values
+
             else:
-                values = arr
-            result.append(values)
-        return result
+                for i, loc in enumerate(mgr_locs):
+                    result[loc] = values[i]
+                    if is_dtz:
+                        # special casing datetimetz to avoid conversion through
+                        #  object dtype
+                        # error: Item "ndarray[Any, Any]" of
+                        # "Optional[Union[ExtensionArray, ndarray[Any, Any]]]"
+                        # has no attribute "_ndarray"
+                        result[loc] = result[loc]._ndarray  # type: ignore[union-attr]
+        # error: Incompatible return value type (got "List[None]",
+        # expected "List[ndarray[Any, Any]]")
+        return result  # type: ignore[return-value]
 
     def iset(self, loc: int | slice | np.ndarray, value: ArrayLike):
         """