def test_arrow_extension_type():
from pandas.core.arrays._arrow_utils import ArrowPeriodType
p1 = ArrowPeriodType("D")
p2 = ArrowPeriodType("D")
p3 = ArrowPeriodType("M")
assert p1.freq == "D"
assert p1 == p2
assert not p1 == p3
assert hash(p1) == hash(p2)
assert not hash(p1) == hash(p3)
def __arrow_array__(self, type=None):
"""
Convert myself into a pyarrow Array.
"""
import pyarrow
from pandas.core.arrays._arrow_utils import ArrowPeriodType
if type is not None:
if pyarrow.types.is_integer(type):
return pyarrow.array(self._ndarray,
mask=self.isna(),
type=type)
elif isinstance(type, ArrowPeriodType):
# ensure we have the same freq
if self.freqstr != type.freq:
raise TypeError(
"Not supported to convert PeriodArray to array with different "
f"'freq' ({self.freqstr} vs {type.freq})")
else:
raise TypeError(
f"Not supported to convert PeriodArray to '{type}' type")
period_type = ArrowPeriodType(self.freqstr)
storage_array = pyarrow.array(self._ndarray,
mask=self.isna(),
type="int64")
return pyarrow.ExtensionArray.from_storage(period_type, storage_array)
def test_arrow_array(data, freq):
import pyarrow as pa
from pandas.core.arrays._arrow_utils import ArrowPeriodType
periods = period_array(data, freq=freq)
result = pa.array(periods)
assert isinstance(result.type, ArrowPeriodType)
assert result.type.freq == freq
expected = pa.array(periods.asi8, type="int64")
assert result.storage.equals(expected)
# convert to its storage type
result = pa.array(periods, type=pa.int64())
assert result.equals(expected)
# unsupported conversions
with pytest.raises(TypeError):
pa.array(periods, type="float64")
with pytest.raises(TypeError, match="different 'freq'"):
pa.array(periods, type=ArrowPeriodType("T"))