From 75a4c91362eff2a8ebb0eed5b371c902cb535f60 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Fri, 13 Sep 2024 08:27:48 -0500 Subject: [PATCH] Add multi-variable rechunking (#1546) --- tests/geospatial/test_rechunking.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/geospatial/test_rechunking.py b/tests/geospatial/test_rechunking.py index fec62d7289..13712e0851 100644 --- a/tests/geospatial/test_rechunking.py +++ b/tests/geospatial/test_rechunking.py @@ -13,10 +13,17 @@ def test_era5_rechunking(client, gcs_url, scale): if scale == "small": # 101.83 GiB (small) time_range = slice("2020-01-01", "2023-01-01") + variables = ["sea_surface_temperature"] + elif scale == "medium": + # 2.12 TiB (medium) + time_range = slice(None) + variables = ["sea_surface_temperature"] else: - # 2.12 TiB (large) + # 4.24 TiB (large) + # This currently doesn't complete successfully. time_range = slice(None) - subset = ds.sea_surface_temperature.sel(time=time_range) + variables = ["sea_surface_temperature", "snow_depth"] + subset = ds[variables].sel(time=time_range) # Rechunk result = subset.chunk({"time": -1, "longitude": "auto", "latitude": "auto"})