@@ -541,9 +541,10 @@ def test_arrow_roundtrip(dtype, string_storage, using_infer_string):
541
541
else :
542
542
assert isinstance (result ["a" ].dtype , pd .StringDtype )
543
543
expected = df .astype (pd .StringDtype (string_storage , na_value = dtype .na_value ))
544
- expected .columns = expected .columns .astype (
545
- pd .StringDtype (string_storage , na_value = np .nan )
546
- )
544
+ if using_infer_string :
545
+ expected .columns = expected .columns .astype (
546
+ pd .StringDtype (string_storage , na_value = np .nan )
547
+ )
547
548
tm .assert_frame_equal (result , expected )
548
549
# ensure the missing value is represented by NA and not np.nan or None
549
550
assert result .loc [2 , "a" ] is result ["a" ].dtype .na_value
@@ -571,9 +572,10 @@ def test_arrow_load_from_zero_chunks(dtype, string_storage, using_infer_string):
571
572
else :
572
573
assert isinstance (result ["a" ].dtype , pd .StringDtype )
573
574
expected = df .astype (pd .StringDtype (string_storage , na_value = dtype .na_value ))
574
- expected .columns = expected .columns .astype (
575
- pd .StringDtype (string_storage , na_value = np .nan )
576
- )
575
+ if using_infer_string :
576
+ expected .columns = expected .columns .astype (
577
+ pd .StringDtype (string_storage , na_value = np .nan )
578
+ )
577
579
tm .assert_frame_equal (result , expected )
578
580
579
581
0 commit comments