| 183 | tm.assert_almost_equal(result, expected) |
| 184 | |
| 185 | def test_2d_datetime64(self): |
| 186 | # 2005/01/01 - 2006/01/01 |
| 187 | arr = ( |
| 188 | np.random.default_rng(2).integers(11_045_376, 11_360_736, (5, 3)) |
| 189 | * 100_000_000_000 |
| 190 | ) |
| 191 | arr = arr.view(dtype="datetime64[ns]") |
| 192 | indexer = [0, 2, -1, 1, -1] |
| 193 | |
| 194 | # axis=0 |
| 195 | result = algos.take_nd(arr, indexer, axis=0) |
| 196 | expected = arr.take(indexer, axis=0) |
| 197 | expected.view(np.int64)[[2, 4], :] = iNaT |
| 198 | tm.assert_almost_equal(result, expected) |
| 199 | |
| 200 | result = algos.take_nd(arr, indexer, axis=0, fill_value=datetime(2007, 1, 1)) |
| 201 | expected = arr.take(indexer, axis=0) |
| 202 | expected[[2, 4], :] = datetime(2007, 1, 1) |
| 203 | tm.assert_almost_equal(result, expected) |
| 204 | |
| 205 | # axis=1 |
| 206 | result = algos.take_nd(arr, indexer, axis=1) |
| 207 | expected = arr.take(indexer, axis=1) |
| 208 | expected.view(np.int64)[:, [2, 4]] = iNaT |
| 209 | tm.assert_almost_equal(result, expected) |
| 210 | |
| 211 | result = algos.take_nd(arr, indexer, axis=1, fill_value=datetime(2007, 1, 1)) |
| 212 | expected = arr.take(indexer, axis=1) |
| 213 | expected[:, [2, 4]] = datetime(2007, 1, 1) |
| 214 | tm.assert_almost_equal(result, expected) |
| 215 | |
| 216 | def test_take_axis_0(self): |
| 217 | arr = np.arange(12).reshape(4, 3) |