File size: 899 Bytes
7885a28 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
from pandas import (
Categorical,
Index,
Series,
)
from pandas.core.arrays import BaseMaskedArray
def get_array(obj, col=None):
"""
Helper method to get array for a DataFrame column or a Series.
Equivalent of df[col].values, but without going through normal getitem,
which triggers tracking references / CoW (and we might be testing that
this is done by some other operation).
"""
if isinstance(obj, Index):
arr = obj._values
elif isinstance(obj, Series) and (col is None or obj.name == col):
arr = obj._values
else:
assert col is not None
icol = obj.columns.get_loc(col)
assert isinstance(icol, int)
arr = obj._get_column_array(icol)
if isinstance(arr, BaseMaskedArray):
return arr._data
elif isinstance(arr, Categorical):
return arr
return getattr(arr, "_ndarray", arr)
|