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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
|
"""
Property-based tests for roundtripping between xarray and pandas objects.
"""
from functools import partial
import numpy as np
import pandas as pd
import pytest
import xarray as xr
pytest.importorskip("hypothesis")
import hypothesis.extra.numpy as npst # isort:skip
import hypothesis.extra.pandas as pdst # isort:skip
import hypothesis.strategies as st # isort:skip
from hypothesis import given # isort:skip
from xarray.tests import has_pyarrow
numeric_dtypes = st.one_of(
npst.unsigned_integer_dtypes(endianness="="),
npst.integer_dtypes(endianness="="),
npst.floating_dtypes(endianness="="),
)
numeric_series = numeric_dtypes.flatmap(lambda dt: pdst.series(dtype=dt))
@st.composite
def dataframe_strategy(draw):
tz = draw(st.timezones())
dtype = pd.DatetimeTZDtype(unit="ns", tz=tz)
datetimes = st.datetimes(
min_value=pd.Timestamp("1677-09-21T00:12:43.145224193"),
max_value=pd.Timestamp("2262-04-11T23:47:16.854775807"),
timezones=st.just(tz),
)
df = pdst.data_frames(
[
pdst.column("datetime_col", elements=datetimes),
pdst.column("other_col", elements=st.integers()),
],
index=pdst.range_indexes(min_size=1, max_size=10),
)
return draw(df).astype({"datetime_col": dtype})
an_array = npst.arrays(
dtype=numeric_dtypes,
shape=npst.array_shapes(max_dims=2), # can only convert 1D/2D to pandas
)
@st.composite
def datasets_1d_vars(draw) -> xr.Dataset:
"""Generate datasets with only 1D variables
Suitable for converting to pandas dataframes.
"""
# Generate an index for the dataset
idx = draw(pdst.indexes(dtype="u8", min_size=0, max_size=100))
# Generate 1-3 variables, 1D with the same length as the index
vars_strategy = st.dictionaries(
keys=st.text(),
values=npst.arrays(dtype=numeric_dtypes, shape=len(idx)).map(
partial(xr.Variable, ("rows",))
),
min_size=1,
max_size=3,
)
return xr.Dataset(draw(vars_strategy), coords={"rows": idx})
@given(st.data(), an_array)
def test_roundtrip_dataarray(data, arr) -> None:
names = data.draw(
st.lists(st.text(), min_size=arr.ndim, max_size=arr.ndim, unique=True).map(
tuple
)
)
coords = {name: np.arange(n) for (name, n) in zip(names, arr.shape, strict=True)}
original = xr.DataArray(arr, dims=names, coords=coords)
roundtripped = xr.DataArray(original.to_pandas())
xr.testing.assert_identical(original, roundtripped)
@given(datasets_1d_vars())
def test_roundtrip_dataset(dataset) -> None:
df = dataset.to_dataframe()
assert isinstance(df, pd.DataFrame)
roundtripped = xr.Dataset(df)
xr.testing.assert_identical(dataset, roundtripped)
@given(numeric_series, st.text())
def test_roundtrip_pandas_series(ser, ix_name) -> None:
# Need to name the index, otherwise Xarray calls it 'dim_0'.
ser.index.name = ix_name
arr = xr.DataArray(ser)
roundtripped = arr.to_pandas()
pd.testing.assert_series_equal(ser, roundtripped)
xr.testing.assert_identical(arr, roundtripped.to_xarray())
# Dataframes with columns of all the same dtype - for roundtrip to DataArray
numeric_homogeneous_dataframe = numeric_dtypes.flatmap(
lambda dt: pdst.data_frames(columns=pdst.columns(["a", "b", "c"], dtype=dt))
)
@pytest.mark.xfail
@given(numeric_homogeneous_dataframe)
def test_roundtrip_pandas_dataframe(df) -> None:
# Need to name the indexes, otherwise Xarray names them 'dim_0', 'dim_1'.
df.index.name = "rows"
df.columns.name = "cols"
arr = xr.DataArray(df)
roundtripped = arr.to_pandas()
pd.testing.assert_frame_equal(df, roundtripped)
xr.testing.assert_identical(arr, roundtripped.to_xarray())
@given(df=dataframe_strategy())
def test_roundtrip_pandas_dataframe_datetime(df) -> None:
# Need to name the indexes, otherwise Xarray names them 'dim_0', 'dim_1'.
df.index.name = "rows"
df.columns.name = "cols"
dataset = xr.Dataset.from_dataframe(df)
roundtripped = dataset.to_dataframe()
roundtripped.columns.name = "cols" # why?
pd.testing.assert_frame_equal(df, roundtripped)
xr.testing.assert_identical(dataset, roundtripped.to_xarray())
@pytest.mark.parametrize(
"extension_array",
[
pd.Categorical(["a", "b", "c"]),
pd.array(["a", "b", "c"], dtype="string"),
pd.arrays.IntervalArray(
[pd.Interval(0, 1), pd.Interval(1, 5), pd.Interval(2, 6)]
),
pd.arrays.TimedeltaArray._from_sequence(pd.TimedeltaIndex(["1h", "2h", "3h"])),
pd.arrays.DatetimeArray._from_sequence(
pd.DatetimeIndex(["2023-01-01", "2023-01-02", "2023-01-03"], freq="D")
),
np.array([1, 2, 3], dtype="int64"),
]
+ ([pd.array([1, 2, 3], dtype="int64[pyarrow]")] if has_pyarrow else []),
ids=["cat", "string", "interval", "timedelta", "datetime", "numpy"]
+ (["pyarrow"] if has_pyarrow else []),
)
@pytest.mark.parametrize("is_index", [True, False])
def test_roundtrip_1d_pandas_extension_array(extension_array, is_index) -> None:
df = pd.DataFrame({"arr": extension_array})
if is_index:
df = df.set_index("arr")
arr = xr.Dataset.from_dataframe(df)["arr"]
roundtripped = arr.to_pandas()
df_arr_to_test = df.index if is_index else df["arr"]
assert (df_arr_to_test == roundtripped).all()
# `NumpyExtensionArray` types are not roundtripped, including `StringArray` which subtypes.
if isinstance(extension_array, pd.arrays.NumpyExtensionArray): # type: ignore[attr-defined]
assert isinstance(arr.data, np.ndarray)
else:
assert (
df_arr_to_test.dtype
== (roundtripped.index if is_index else roundtripped).dtype
)
xr.testing.assert_identical(arr, roundtripped.to_xarray())
|
