MCPcopy Index your code
hub / github.com/pydata/xarray / test_chunk

Method test_chunk

xarray/tests/test_dataset.py:1254–1313  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

1252
1253 @requires_dask
1254 def test_chunk(self) -> None:
1255 data = create_test_data()
1256 for v in data.variables.values():
1257 assert isinstance(v.data, np.ndarray)
1258 assert data.chunks == {}
1259
1260 reblocked = data.chunk()
1261 for k, v in reblocked.variables.items():
1262 if k in reblocked.dims:
1263 assert isinstance(v.data, np.ndarray)
1264 else:
1265 assert isinstance(v.data, da.Array)
1266
1267 expected_chunks: dict[Hashable, tuple[int, ...]] = {
1268 "dim1": (8,),
1269 "dim2": (9,),
1270 "dim3": (10,),
1271 }
1272 assert reblocked.chunks == expected_chunks
1273
1274 # test kwargs form of chunks
1275 assert data.chunk(expected_chunks).chunks == expected_chunks
1276
1277 def get_dask_names(ds):
1278 return {k: v.data.name for k, v in ds.items()}
1279
1280 orig_dask_names = get_dask_names(reblocked)
1281
1282 reblocked = data.chunk({"time": 5, "dim1": 5, "dim2": 5, "dim3": 5})
1283 # time is not a dim in any of the data_vars, so it
1284 # doesn't get chunked
1285 expected_chunks = {"dim1": (5, 3), "dim2": (5, 4), "dim3": (5, 5)}
1286 assert reblocked.chunks == expected_chunks
1287
1288 # make sure dask names change when rechunking by different amounts
1289 # regression test for GH3350
1290 new_dask_names = get_dask_names(reblocked)
1291 for k, v in new_dask_names.items():
1292 assert v != orig_dask_names[k]
1293
1294 reblocked = data.chunk(expected_chunks)
1295 assert reblocked.chunks == expected_chunks
1296
1297 # reblock on already blocked data
1298 orig_dask_names = get_dask_names(reblocked)
1299 reblocked = reblocked.chunk(expected_chunks)
1300 new_dask_names = get_dask_names(reblocked)
1301 assert reblocked.chunks == expected_chunks
1302 assert_identical(reblocked, data)
1303 # rechunking with same chunk sizes should not change names
1304 for k, v in new_dask_names.items():
1305 assert v == orig_dask_names[k]
1306
1307 with pytest.raises(
1308 ValueError,
1309 match=re.escape(
1310 "chunks keys ('foo',) not found in data dimensions ('dim2', 'dim3', 'time', 'dim1')"
1311 ),

Callers

nothing calls this directly

Calls 5

create_test_dataFunction · 0.90
assert_identicalFunction · 0.90
itemsMethod · 0.80
valuesMethod · 0.45
chunkMethod · 0.45

Tested by

no test coverage detected