66 DataFrame ,
77 MultiIndex ,
88 Series ,
9+ StringDtype ,
910 date_range ,
1011)
1112import pandas ._testing as tm
13+ from pandas .util .version import Version
1214
13- pytest .importorskip ("xarray" )
15+ xarray = pytest .importorskip ("xarray" )
1416
1517
1618class TestDataFrameToXArray :
@@ -29,13 +31,17 @@ def df(self):
2931 }
3032 )
3133
32- def test_to_xarray_index_types (self , index_flat , df , using_infer_string ):
34+ def test_to_xarray_index_types (self , index_flat , df , request ):
3335 index = index_flat
3436 # MultiIndex is tested in test_to_xarray_with_multiindex
3537 if len (index ) == 0 :
3638 pytest .skip ("Test doesn't make sense for empty index" )
37-
38- from xarray import Dataset
39+ elif Version (xarray .__version__ ) <= Version ("2024.9.0" ):
40+ request .applymarker (
41+ pytest .mark .xfail (
42+ reason = "Categorical column not preserved." ,
43+ )
44+ )
3945
4046 df .index = index [:4 ]
4147 df .index .name = "foo"
@@ -45,29 +51,22 @@ def test_to_xarray_index_types(self, index_flat, df, using_infer_string):
4551 assert len (result .coords ) == 1
4652 assert len (result .data_vars ) == 8
4753 tm .assert_almost_equal (list (result .coords .keys ()), ["foo" ])
48- assert isinstance (result , Dataset )
54+ assert isinstance (result , xarray . Dataset )
4955
5056 # idempotency
5157 # datetimes w/tz are preserved
5258 # column names are lost
5359 expected = df .copy ()
54- expected ["f" ] = expected ["f" ].astype (
55- object if not using_infer_string else "str"
56- )
5760 expected .columns .name = None
5861 tm .assert_frame_equal (result .to_dataframe (), expected )
5962
6063 def test_to_xarray_empty (self , df ):
61- from xarray import Dataset
62-
6364 df .index .name = "foo"
6465 result = df [0 :0 ].to_xarray ()
6566 assert result .sizes ["foo" ] == 0
66- assert isinstance (result , Dataset )
67+ assert isinstance (result , xarray . Dataset )
6768
6869 def test_to_xarray_with_multiindex (self , df , using_infer_string ):
69- from xarray import Dataset
70-
7170 # MultiIndex
7271 df .index = MultiIndex .from_product ([["a" ], range (4 )], names = ["one" , "two" ])
7372 result = df .to_xarray ()
@@ -76,7 +75,7 @@ def test_to_xarray_with_multiindex(self, df, using_infer_string):
7675 assert len (result .coords ) == 2
7776 assert len (result .data_vars ) == 8
7877 tm .assert_almost_equal (list (result .coords .keys ()), ["one" , "two" ])
79- assert isinstance (result , Dataset )
78+ assert isinstance (result , xarray . Dataset )
8079
8180 result = result .to_dataframe ()
8281 expected = df .copy ()
@@ -88,43 +87,48 @@ def test_to_xarray_with_multiindex(self, df, using_infer_string):
8887
8988
9089class TestSeriesToXArray :
91- def test_to_xarray_index_types (self , index_flat ):
90+ def test_to_xarray_index_types (self , index_flat , request ):
9291 index = index_flat
92+ if (
93+ isinstance (index .dtype , StringDtype )
94+ and index .dtype .storage == "pyarrow"
95+ and Version (xarray .__version__ ) > Version ("2024.9.0" )
96+ ):
97+ request .applymarker (
98+ pytest .mark .xfail (
99+ reason = "xarray calling reshape of ArrowExtensionArray" ,
100+ raises = NotImplementedError ,
101+ )
102+ )
93103 # MultiIndex is tested in test_to_xarray_with_multiindex
94104
95- from xarray import DataArray
96-
97105 ser = Series (range (len (index )), index = index , dtype = "int64" )
98106 ser .index .name = "foo"
99107 result = ser .to_xarray ()
100108 repr (result )
101109 assert len (result ) == len (index )
102110 assert len (result .coords ) == 1
103111 tm .assert_almost_equal (list (result .coords .keys ()), ["foo" ])
104- assert isinstance (result , DataArray )
112+ assert isinstance (result , xarray . DataArray )
105113
106114 # idempotency
107115 tm .assert_series_equal (result .to_series (), ser )
108116
109117 def test_to_xarray_empty (self ):
110- from xarray import DataArray
111-
112118 ser = Series ([], dtype = object )
113119 ser .index .name = "foo"
114120 result = ser .to_xarray ()
115121 assert len (result ) == 0
116122 assert len (result .coords ) == 1
117123 tm .assert_almost_equal (list (result .coords .keys ()), ["foo" ])
118- assert isinstance (result , DataArray )
124+ assert isinstance (result , xarray . DataArray )
119125
120126 def test_to_xarray_with_multiindex (self ):
121- from xarray import DataArray
122-
123127 mi = MultiIndex .from_product ([["a" , "b" ], range (3 )], names = ["one" , "two" ])
124128 ser = Series (range (6 ), dtype = "int64" , index = mi )
125129 result = ser .to_xarray ()
126130 assert len (result ) == 2
127131 tm .assert_almost_equal (list (result .coords .keys ()), ["one" , "two" ])
128- assert isinstance (result , DataArray )
132+ assert isinstance (result , xarray . DataArray )
129133 res = result .to_series ()
130134 tm .assert_series_equal (res , ser )
0 commit comments