Skip to content

Commit c6a98c0

Browse files
author
Kei
committed
Update tests
1 parent 4ef96f7 commit c6a98c0

File tree

1 file changed

+19
-21
lines changed

1 file changed

+19
-21
lines changed

pandas/tests/groupby/aggregate/test_aggregate.py

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1886,46 +1886,48 @@ def test_agg_lambda_pyarrow_to_same_data_type():
18861886
def test_agg_lambda_float64_pyarrow_dtype_conversion():
18871887
# test numpy dtype conversion back to pyarrow dtype
18881888
# 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]})
18901890
df["B"] = df["B"].astype("float64[pyarrow]")
18911891
gb = df.groupby("A")
18921892
result = gb.agg(lambda x: x)
18931893

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]")
18961896
expected.set_index("A", inplace=True)
18971897

18981898
tm.assert_frame_equal(result, expected)
18991899
assert result["B"].dtype == expected["B"].dtype
19001900

19011901

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
19031905
df = DataFrame({"A": ["c1", "c2", "c3"], "B": [100, 200, 255]})
19041906
df["B"] = df["B"].astype("int64[pyarrow]")
19051907
gb = df.groupby("A")
1906-
result = gb.agg(lambda x: complex(x.sum(), x.count()))
1908+
result = gb.agg(lambda x: x)
19071909

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]")
19151912
expected.set_index("A", inplace=True)
19161913

19171914
tm.assert_frame_equal(result, expected)
19181915
assert result["B"].dtype == expected["B"].dtype
19191916

19201917

1921-
def test_agg_lambda_int64_pyarrow_dtype_conversion():
1918+
def test_agg_lambda_complex128_pyarrow_dtype_conversion():
19221919
df = DataFrame({"A": ["c1", "c2", "c3"], "B": [100, 200, 255]})
19231920
df["B"] = df["B"].astype("int64[pyarrow]")
19241921
gb = df.groupby("A")
1925-
result = gb.agg(lambda x: x)
1922+
result = gb.agg(lambda x: complex(x.sum(), x.count()))
19261923

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")
19291931
expected.set_index("A", inplace=True)
19301932

19311933
tm.assert_frame_equal(result, expected)
@@ -1975,8 +1977,6 @@ def test_agg_lambda_bool_pyarrow_dtype_conversion():
19751977

19761978

19771979
def test_agg_lambda_object_pyarrow_dtype_conversion():
1978-
import pyarrow as pa
1979-
19801980
df = DataFrame({"A": ["c1", "c2", "c3"], "B": [100, 200, 255]})
19811981
df["B"] = df["B"].astype("int64[pyarrow]")
19821982
gb = df.groupby("A")
@@ -1986,9 +1986,7 @@ def test_agg_lambda_object_pyarrow_dtype_conversion():
19861986
{"A": ["c1", "c2", "c3"], "B": [{"number": 1}, {"number": 1}, {"number": 1}]}
19871987
)
19881988

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")
19921990
expected.set_index("A", inplace=True)
19931991

19941992
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)