|
5 | 5 |
|
6 | 6 | import pandas as pd
|
7 | 7 | from pandas import DataFrame
|
| 8 | +import pandas._testing as tm |
8 | 9 | from pandas.tests.copy_view.util import get_array
|
9 | 10 |
|
10 | 11 |
|
@@ -93,3 +94,49 @@ def test_switch_options():
|
93 | 94 | subset.iloc[0, 0] = 0
|
94 | 95 | # df updated with CoW disabled
|
95 | 96 | assert df.iloc[0, 0] == 0
|
| 97 | + |
| 98 | + |
| 99 | +@td.skip_array_manager_invalid_test |
| 100 | +@pytest.mark.parametrize("dtype", [np.intp, np.int8]) |
| 101 | +@pytest.mark.parametrize( |
| 102 | + "locs, arr", |
| 103 | + [ |
| 104 | + ([0], np.array([-1, -2, -3])), |
| 105 | + ([1], np.array([-1, -2, -3])), |
| 106 | + ([5], np.array([-1, -2, -3])), |
| 107 | + ([0, 1], np.array([[-1, -2, -3], [-4, -5, -6]]).T), |
| 108 | + ([0, 2], np.array([[-1, -2, -3], [-4, -5, -6]]).T), |
| 109 | + ([0, 1, 2], np.array([[-1, -2, -3], [-4, -5, -6], [-4, -5, -6]]).T), |
| 110 | + ([1, 2], np.array([[-1, -2, -3], [-4, -5, -6]]).T), |
| 111 | + ([1, 3], np.array([[-1, -2, -3], [-4, -5, -6]]).T), |
| 112 | + ([1, 3], np.array([[-1, -2, -3], [-4, -5, -6]]).T), |
| 113 | + ], |
| 114 | +) |
| 115 | +def test_iset_splits_blocks_inplace(using_copy_on_write, locs, arr, dtype): |
| 116 | + # Nothing currently calls iset with |
| 117 | + # more than 1 loc with inplace=True (only happens with inplace=False) |
| 118 | + # but ensure that it works |
| 119 | + df = DataFrame( |
| 120 | + { |
| 121 | + "a": [1, 2, 3], |
| 122 | + "b": [4, 5, 6], |
| 123 | + "c": [7, 8, 9], |
| 124 | + "d": [10, 11, 12], |
| 125 | + "e": [13, 14, 15], |
| 126 | + "f": ["a", "b", "c"], |
| 127 | + }, |
| 128 | + ) |
| 129 | + arr = arr.astype(dtype) |
| 130 | + df_orig = df.copy() |
| 131 | + df2 = df.copy(deep=None) # Trigger a CoW (if enabled, otherwise makes copy) |
| 132 | + df2._mgr.iset(locs, arr, inplace=True) |
| 133 | + |
| 134 | + tm.assert_frame_equal(df, df_orig) |
| 135 | + |
| 136 | + if using_copy_on_write: |
| 137 | + for i, col in enumerate(df.columns): |
| 138 | + if i not in locs: |
| 139 | + assert np.shares_memory(get_array(df, col), get_array(df2, col)) |
| 140 | + else: |
| 141 | + for col in df.columns: |
| 142 | + assert not np.shares_memory(get_array(df, col), get_array(df2, col)) |
0 commit comments