From 33141a898de0b294f7be26392224863313d6b4dd Mon Sep 17 00:00:00 2001 From: Matthew Roeschke Date: Wed, 29 Dec 2021 10:47:38 -0800 Subject: [PATCH] TST: Mock pivot and pivot_table test for PerformanceWarning --- pandas/tests/frame/test_stack_unstack.py | 30 ++++++++++++++-------- pandas/tests/reshape/test_pivot.py | 32 ++++++++++++++---------- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/pandas/tests/frame/test_stack_unstack.py b/pandas/tests/frame/test_stack_unstack.py index 7efe4b08895bc..eecae31bec914 100644 --- a/pandas/tests/frame/test_stack_unstack.py +++ b/pandas/tests/frame/test_stack_unstack.py @@ -18,6 +18,7 @@ date_range, ) import pandas._testing as tm +from pandas.core.reshape import reshape as reshape_lib class TestDataFrameReshape: @@ -1826,19 +1827,26 @@ def test_unstack_unobserved_keys(self): tm.assert_frame_equal(recons, df) @pytest.mark.slow - def test_unstack_number_of_levels_larger_than_int32(self): + def test_unstack_number_of_levels_larger_than_int32(self, monkeypatch): # GH#20601 # GH 26314: Change ValueError to PerformanceWarning - df = DataFrame( - np.random.randn(2 ** 16, 2), index=[np.arange(2 ** 16), np.arange(2 ** 16)] - ) - msg = "The following operation may generate" - with tm.assert_produces_warning(PerformanceWarning, match=msg): - try: - df.unstack() - except MemoryError: - # Just checking the warning - return + + class MockUnstacker(reshape_lib._Unstacker): + def __init__(self, *args, **kwargs): + # __init__ will raise the warning + super().__init__(*args, **kwargs) + raise Exception("Don't compute final result.") + + with monkeypatch.context() as m: + m.setattr(reshape_lib, "_Unstacker", MockUnstacker) + df = DataFrame( + np.random.randn(2 ** 16, 2), + index=[np.arange(2 ** 16), np.arange(2 ** 16)], + ) + msg = "The following operation may generate" + with tm.assert_produces_warning(PerformanceWarning, match=msg): + with pytest.raises(Exception, match="Don't compute final result."): + df.unstack() def test_stack_order_with_unsorted_levels(self): # GH#16323 diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py index aa7af4548a770..a023adfb509a0 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -23,6 +23,7 @@ ) import pandas._testing as tm from pandas.api.types import CategoricalDtype as CDT +from pandas.core.reshape import reshape as reshape_lib from pandas.core.reshape.pivot import pivot_table @@ -1991,22 +1992,27 @@ def test_pivot_string_func_vs_func(self, f, f_numpy): tm.assert_frame_equal(result, expected) @pytest.mark.slow - def test_pivot_number_of_levels_larger_than_int32(self): + def test_pivot_number_of_levels_larger_than_int32(self, monkeypatch): # GH 20601 # GH 26314: Change ValueError to PerformanceWarning - df = DataFrame( - {"ind1": np.arange(2 ** 16), "ind2": np.arange(2 ** 16), "count": 0} - ) + class MockUnstacker(reshape_lib._Unstacker): + def __init__(self, *args, **kwargs): + # __init__ will raise the warning + super().__init__(*args, **kwargs) + raise Exception("Don't compute final result.") + + with monkeypatch.context() as m: + m.setattr(reshape_lib, "_Unstacker", MockUnstacker) + df = DataFrame( + {"ind1": np.arange(2 ** 16), "ind2": np.arange(2 ** 16), "count": 0} + ) - msg = "The following operation may generate" - with tm.assert_produces_warning(PerformanceWarning, match=msg): - try: - df.pivot_table( - index="ind1", columns="ind2", values="count", aggfunc="count" - ) - except MemoryError: - # Just checking the warning - return + msg = "The following operation may generate" + with tm.assert_produces_warning(PerformanceWarning, match=msg): + with pytest.raises(Exception, match="Don't compute final result."): + df.pivot_table( + index="ind1", columns="ind2", values="count", aggfunc="count" + ) def test_pivot_table_aggfunc_dropna(self, dropna): # GH 22159