@@ -1886,46 +1886,48 @@ def test_agg_lambda_pyarrow_to_same_data_type():
1886
1886
def test_agg_lambda_float64_pyarrow_dtype_conversion ():
1887
1887
# test numpy dtype conversion back to pyarrow dtype
1888
1888
# complexes, floats, ints, uints, object
1889
- df = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1889
+ df = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100.0 , 200 , 255.3873 ]})
1890
1890
df ["B" ] = df ["B" ].astype ("float64[pyarrow]" )
1891
1891
gb = df .groupby ("A" )
1892
1892
result = gb .agg (lambda x : x )
1893
1893
1894
- expected = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1895
- expected ["B" ] = expected ["B" ].astype ("float64 [pyarrow]" )
1894
+ expected = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100.0 , 200 , 255.3873 ]})
1895
+ expected ["B" ] = expected ["B" ].astype ("double [pyarrow]" )
1896
1896
expected .set_index ("A" , inplace = True )
1897
1897
1898
1898
tm .assert_frame_equal (result , expected )
1899
1899
assert result ["B" ].dtype == expected ["B" ].dtype
1900
1900
1901
1901
1902
- def test_agg_lambda_complex128_pyarrow_dtype_conversion ():
1902
+ def test_agg_lambda_int64_pyarrow_dtype_conversion ():
1903
+ # test numpy dtype conversion back to pyarrow dtype
1904
+ # complexes, floats, ints, uints, object
1903
1905
df = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1904
1906
df ["B" ] = df ["B" ].astype ("int64[pyarrow]" )
1905
1907
gb = df .groupby ("A" )
1906
- result = gb .agg (lambda x : complex ( x . sum (), x . count ()) )
1908
+ result = gb .agg (lambda x : x )
1907
1909
1908
- expected = DataFrame (
1909
- {
1910
- "A" : ["c1" , "c2" , "c3" ],
1911
- "B" : [complex (100 , 1 ), complex (200 , 1 ), complex (255 , 1 )],
1912
- }
1913
- )
1914
- expected ["B" ] = expected ["B" ].astype ("complex128" )
1910
+ expected = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1911
+ expected ["B" ] = expected ["B" ].astype ("int64[pyarrow]" )
1915
1912
expected .set_index ("A" , inplace = True )
1916
1913
1917
1914
tm .assert_frame_equal (result , expected )
1918
1915
assert result ["B" ].dtype == expected ["B" ].dtype
1919
1916
1920
1917
1921
- def test_agg_lambda_int64_pyarrow_dtype_conversion ():
1918
+ def test_agg_lambda_complex128_pyarrow_dtype_conversion ():
1922
1919
df = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1923
1920
df ["B" ] = df ["B" ].astype ("int64[pyarrow]" )
1924
1921
gb = df .groupby ("A" )
1925
- result = gb .agg (lambda x : x )
1922
+ result = gb .agg (lambda x : complex ( x . sum (), x . count ()) )
1926
1923
1927
- expected = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1928
- expected ["B" ] = expected ["B" ].astype ("int64[pyarrow]" )
1924
+ expected = DataFrame (
1925
+ {
1926
+ "A" : ["c1" , "c2" , "c3" ],
1927
+ "B" : [complex (100 , 1 ), complex (200 , 1 ), complex (255 , 1 )],
1928
+ }
1929
+ )
1930
+ expected ["B" ] = expected ["B" ].astype ("complex128" )
1929
1931
expected .set_index ("A" , inplace = True )
1930
1932
1931
1933
tm .assert_frame_equal (result , expected )
@@ -1975,8 +1977,6 @@ def test_agg_lambda_bool_pyarrow_dtype_conversion():
1975
1977
1976
1978
1977
1979
def test_agg_lambda_object_pyarrow_dtype_conversion ():
1978
- import pyarrow as pa
1979
-
1980
1980
df = DataFrame ({"A" : ["c1" , "c2" , "c3" ], "B" : [100 , 200 , 255 ]})
1981
1981
df ["B" ] = df ["B" ].astype ("int64[pyarrow]" )
1982
1982
gb = df .groupby ("A" )
@@ -1986,9 +1986,7 @@ def test_agg_lambda_object_pyarrow_dtype_conversion():
1986
1986
{"A" : ["c1" , "c2" , "c3" ], "B" : [{"number" : 1 }, {"number" : 1 }, {"number" : 1 }]}
1987
1987
)
1988
1988
1989
- pyarrow_type = pa .struct ({"number" : pa .int64 ()})
1990
- pandas_pyarrow_dtype = pd .ArrowDtype (pyarrow_type )
1991
- expected ["B" ] = expected ["B" ].astype (pandas_pyarrow_dtype )
1989
+ expected ["B" ] = expected ["B" ].astype ("object" )
1992
1990
expected .set_index ("A" , inplace = True )
1993
1991
1994
1992
tm .assert_frame_equal (result , expected )
0 commit comments