@@ -2360,61 +2360,73 @@ def test_get_engine_auto_error_message(self):
2360
2360
2361
2361
@pytest .mark .parametrize ("option" , [True , False ])
2362
2362
@pytest .mark .parametrize ("func" , ["read_sql" , "read_sql_query" ])
2363
- def test_read_sql_nullable_dtypes (self , string_storage , func , option ):
2363
+ @pytest .mark .parametrize ("dtype_backend" , ["pandas" , "pyarrow" ])
2364
+ def test_read_sql_nullable_dtypes (
2365
+ self , string_storage , func , option , dtype_backend
2366
+ ):
2364
2367
# GH#50048
2365
2368
table = "test"
2366
2369
df = self .nullable_data ()
2367
2370
df .to_sql (table , self .conn , index = False , if_exists = "replace" )
2368
2371
2369
2372
with pd .option_context ("mode.string_storage" , string_storage ):
2370
- if option :
2371
- with pd .option_context ("mode.nullable_dtypes" , True ):
2372
- result = getattr (pd , func )(f"Select * from { table } " , self .conn )
2373
- else :
2374
- result = getattr (pd , func )(
2375
- f"Select * from { table } " , self .conn , use_nullable_dtypes = True
2376
- )
2377
- expected = self .nullable_expected (string_storage )
2373
+ with pd .option_context ("mode.dtype_backend" , dtype_backend ):
2374
+ if option :
2375
+ with pd .option_context ("mode.nullable_dtypes" , True ):
2376
+ result = getattr (pd , func )(f"Select * from { table } " , self .conn )
2377
+ else :
2378
+ result = getattr (pd , func )(
2379
+ f"Select * from { table } " , self .conn , use_nullable_dtypes = True
2380
+ )
2381
+ expected = self .nullable_expected (string_storage , dtype_backend )
2378
2382
tm .assert_frame_equal (result , expected )
2379
2383
2380
2384
with pd .option_context ("mode.string_storage" , string_storage ):
2381
- iterator = getattr (pd , func )(
2382
- f"Select * from { table } " ,
2383
- self .conn ,
2384
- use_nullable_dtypes = True ,
2385
- chunksize = 3 ,
2386
- )
2387
- expected = self .nullable_expected (string_storage )
2388
- for result in iterator :
2389
- tm .assert_frame_equal (result , expected )
2385
+ with pd .option_context ("mode.dtype_backend" , dtype_backend ):
2386
+ iterator = getattr (pd , func )(
2387
+ f"Select * from { table } " ,
2388
+ self .conn ,
2389
+ use_nullable_dtypes = True ,
2390
+ chunksize = 3 ,
2391
+ )
2392
+ expected = self .nullable_expected (string_storage , dtype_backend )
2393
+ for result in iterator :
2394
+ tm .assert_frame_equal (result , expected )
2390
2395
2391
2396
@pytest .mark .parametrize ("option" , [True , False ])
2392
2397
@pytest .mark .parametrize ("func" , ["read_sql" , "read_sql_table" ])
2393
- def test_read_sql_nullable_dtypes_table (self , string_storage , func , option ):
2398
+ @pytest .mark .parametrize ("dtype_backend" , ["pandas" , "pyarrow" ])
2399
+ def test_read_sql_nullable_dtypes_table (
2400
+ self , string_storage , func , option , dtype_backend
2401
+ ):
2394
2402
# GH#50048
2395
2403
table = "test"
2396
2404
df = self .nullable_data ()
2397
2405
df .to_sql (table , self .conn , index = False , if_exists = "replace" )
2398
2406
2399
2407
with pd .option_context ("mode.string_storage" , string_storage ):
2400
- if option :
2401
- with pd .option_context ("mode.nullable_dtypes" , True ):
2402
- result = getattr (pd , func )(table , self .conn )
2403
- else :
2404
- result = getattr (pd , func )(table , self .conn , use_nullable_dtypes = True )
2405
- expected = self .nullable_expected (string_storage )
2408
+ with pd .option_context ("mode.dtype_backend" , dtype_backend ):
2409
+ if option :
2410
+ with pd .option_context ("mode.nullable_dtypes" , True ):
2411
+ result = getattr (pd , func )(table , self .conn )
2412
+ else :
2413
+ result = getattr (pd , func )(
2414
+ table , self .conn , use_nullable_dtypes = True
2415
+ )
2416
+ expected = self .nullable_expected (string_storage , dtype_backend )
2406
2417
tm .assert_frame_equal (result , expected )
2407
2418
2408
2419
with pd .option_context ("mode.string_storage" , string_storage ):
2409
- iterator = getattr (pd , func )(
2410
- table ,
2411
- self .conn ,
2412
- use_nullable_dtypes = True ,
2413
- chunksize = 3 ,
2414
- )
2415
- expected = self .nullable_expected (string_storage )
2416
- for result in iterator :
2417
- tm .assert_frame_equal (result , expected )
2420
+ with pd .option_context ("mode.dtype_backend" , dtype_backend ):
2421
+ iterator = getattr (pd , func )(
2422
+ table ,
2423
+ self .conn ,
2424
+ use_nullable_dtypes = True ,
2425
+ chunksize = 3 ,
2426
+ )
2427
+ expected = self .nullable_expected (string_storage , dtype_backend )
2428
+ for result in iterator :
2429
+ tm .assert_frame_equal (result , expected )
2418
2430
2419
2431
def nullable_data (self ) -> DataFrame :
2420
2432
return DataFrame (
@@ -2430,7 +2442,7 @@ def nullable_data(self) -> DataFrame:
2430
2442
}
2431
2443
)
2432
2444
2433
- def nullable_expected (self , storage ) -> DataFrame :
2445
+ def nullable_expected (self , storage , dtype_backend ) -> DataFrame :
2434
2446
2435
2447
string_array : StringArray | ArrowStringArray
2436
2448
string_array_na : StringArray | ArrowStringArray
@@ -2443,7 +2455,7 @@ def nullable_expected(self, storage) -> DataFrame:
2443
2455
string_array = ArrowStringArray (pa .array (["a" , "b" , "c" ]))
2444
2456
string_array_na = ArrowStringArray (pa .array (["a" , "b" , None ]))
2445
2457
2446
- return DataFrame (
2458
+ df = DataFrame (
2447
2459
{
2448
2460
"a" : Series ([1 , np .nan , 3 ], dtype = "Int64" ),
2449
2461
"b" : Series ([1 , 2 , 3 ], dtype = "Int64" ),
@@ -2455,6 +2467,18 @@ def nullable_expected(self, storage) -> DataFrame:
2455
2467
"h" : string_array_na ,
2456
2468
}
2457
2469
)
2470
+ if dtype_backend == "pyarrow" :
2471
+ pa = pytest .importorskip ("pyarrow" )
2472
+
2473
+ from pandas .arrays import ArrowExtensionArray
2474
+
2475
+ df = DataFrame (
2476
+ {
2477
+ col : ArrowExtensionArray (pa .array (df [col ], from_pandas = True ))
2478
+ for col in df .columns
2479
+ }
2480
+ )
2481
+ return df
2458
2482
2459
2483
def test_chunksize_empty_dtypes (self ):
2460
2484
# GH#50245
@@ -2578,8 +2602,14 @@ class Test(BaseModel):
2578
2602
2579
2603
assert list (df .columns ) == ["id" , "string_column" ]
2580
2604
2581
- def nullable_expected (self , storage ) -> DataFrame :
2582
- return super ().nullable_expected (storage ).astype ({"e" : "Int64" , "f" : "Int64" })
2605
+ def nullable_expected (self , storage , dtype_backend ) -> DataFrame :
2606
+ df = super ().nullable_expected (storage , dtype_backend )
2607
+ if dtype_backend == "pandas" :
2608
+ df = df .astype ({"e" : "Int64" , "f" : "Int64" })
2609
+ else :
2610
+ df = df .astype ({"e" : "int64[pyarrow]" , "f" : "int64[pyarrow]" })
2611
+
2612
+ return df
2583
2613
2584
2614
@pytest .mark .parametrize ("func" , ["read_sql" , "read_sql_table" ])
2585
2615
def test_read_sql_nullable_dtypes_table (self , string_storage , func ):
@@ -2613,8 +2643,14 @@ def setup_driver(cls):
2613
2643
def test_default_type_conversion (self ):
2614
2644
pass
2615
2645
2616
- def nullable_expected (self , storage ) -> DataFrame :
2617
- return super ().nullable_expected (storage ).astype ({"e" : "Int64" , "f" : "Int64" })
2646
+ def nullable_expected (self , storage , dtype_backend ) -> DataFrame :
2647
+ df = super ().nullable_expected (storage , dtype_backend )
2648
+ if dtype_backend == "pandas" :
2649
+ df = df .astype ({"e" : "Int64" , "f" : "Int64" })
2650
+ else :
2651
+ df = df .astype ({"e" : "int64[pyarrow]" , "f" : "int64[pyarrow]" })
2652
+
2653
+ return df
2618
2654
2619
2655
2620
2656
@pytest .mark .db
0 commit comments