File tree 1 file changed +13
-9
lines changed
1 file changed +13
-9
lines changed Original file line number Diff line number Diff line change @@ -614,16 +614,20 @@ def test_s3_roundtrip_for_dir(
614
614
# read_table uses the new Arrow Datasets API since pyarrow 1.0.0
615
615
# Previous behaviour was pyarrow partitioned columns become 'category' dtypes
616
616
# These are added to back of dataframe on read. In new API category dtype is
617
- # only used if partition field is string.
618
- legacy_read_table = LooseVersion (pyarrow .__version__ ) < LooseVersion ("1.0.0" )
619
- if partition_col and legacy_read_table :
620
- partition_col_type = "category"
621
- else :
622
- partition_col_type = "int32"
623
-
624
- expected_df [partition_col ] = expected_df [partition_col ].astype (
625
- partition_col_type
617
+ # only used if partition field is string, but this changed again to use
618
+ # category dtype for all types (not only strings) in pyarrow 2.0.0
619
+ pa10 = (LooseVersion (pyarrow .__version__ ) >= LooseVersion ("1.0.0" )) and (
620
+ LooseVersion (pyarrow .__version__ ) < LooseVersion ("2.0.0" )
626
621
)
622
+ if partition_col :
623
+ if pa10 :
624
+ partition_col_type = "int32"
625
+ else :
626
+ partition_col_type = "category"
627
+
628
+ expected_df [partition_col ] = expected_df [partition_col ].astype (
629
+ partition_col_type
630
+ )
627
631
628
632
check_round_trip (
629
633
df_compat ,
You can’t perform that action at this time.
0 commit comments