@@ -1169,18 +1169,6 @@ def assert_dtype_equal(a, b):
11691169 )
11701170
11711171
1172- def assert_pandas_batch_types (
1173- batch : pandas .DataFrame , expected_types : list [type ]
1174- ) -> None :
1175- assert batch .dtypes is not None
1176-
1177- pandas_dtypes = batch .dtypes
1178- # pd.string is represented as an np.object
1179- # np.dtype string is not the same as pd.string (python)
1180- for pandas_dtype , expected_type in zip (pandas_dtypes , expected_types ):
1181- assert_dtype_equal (pandas_dtype .type , numpy .dtype (expected_type ).type )
1182-
1183-
11841172def test_pandas_dtypes (conn_cnx ):
11851173 with conn_cnx (
11861174 session_parameters = {
@@ -1191,12 +1179,18 @@ def test_pandas_dtypes(conn_cnx):
11911179 cur .execute (
11921180 "select 1::integer, 2.3::double, 'foo'::string, current_timestamp()::timestamp where 1=0"
11931181 )
1194- expected_types = [numpy .int64 , float , object , numpy .datetime64 ]
1195- assert_pandas_batch_types (cur .fetch_pandas_all (), expected_types )
1196-
11971182 batches = cur .get_result_batches ()
1183+ batch = batches [0 ].to_pandas ()
1184+
1185+ assert batch .dtypes is not None
11981186 assert batches [0 ].to_arrow () is not True
1199- assert_pandas_batch_types (batches [0 ].to_pandas (), expected_types )
1187+
1188+ pandas_dtypes = batch .dtypes
1189+ expected_types = [numpy .int64 , float , object , numpy .datetime64 ]
1190+ # pd.string is represented as an np.object
1191+ # np.dtype string is not the same as pd.string (python)
1192+ for i , typ in enumerate (expected_types ):
1193+ assert_dtype_equal (pandas_dtypes [i ].type , numpy .dtype (typ ).type )
12001194
12011195
12021196def test_timestamp_tz (conn_cnx ):
0 commit comments