Pulling subset of larger dataset from s3

[sgavrilovicarup]

Hello everyone and thanks for your help!

I am having issues pulling a subset of a larger dataset. The entire dataset (240MB) is downloaded to extract a small subset (90KB) of data. The dataset is publicly available NASA GDDP climate data stored on s3. I am hoping there is a way around this, like a mask you can pass to only download what you need (this is what I do with large rasters on s3). A simple reproducible example of my code is below.

import xarray as xr
import fsspec

lat_indices=[389, 390, 391, 392, 393, 394, 395]
lon_indices=[947, 948, 949, 950, 951, 952, 953, 954, 955]

# link to view bucket https://nex-gddp-cmip6.s3.us-west-2.amazonaws.com/index.html#NEX-GDDP-CMIP6/ACCESS-CM2/ssp245/r1i1p1f1/tasmax/
net_cdf_url='s3://nex-gddp-cmip6/NEX-GDDP-CMIP6/ACCESS-CM2/ssp245/r1i1p1f1/tasmax/tasmax_day_ACCESS-CM2_ssp245_r1i1p1f1_gn_2049.nc'

fs = fsspec.filesystem('s3', anon=True)

ds = xr.open_dataset(fs.open(net_cdf_url), engine='h5netcdf', chunks={'time': 'auto', 'lat': 'auto', 'lon': 'auto'})

print('Entire ds')
print(ds)

subset = ds['tasmax'].isel(lat=slice(lat_indices[0], lat_indices[-1]), lon=slice(lon_indices[0], lon_indices[-1]), time=slice(None))

subset.persist()

print(f"ds kilobytes {ds.nbytes/1e3}")
print(f"subset kilobytes {subset.nbytes/1e3}")

print('subset ds')
print(subset)

[dcherian]

This is selecting the whole dataset...

 .isel(lat=slice(lat_indices[0], lat_indices[-1]), lon=slice(lon_indices[0], lon_indices[-1]), time=slice(None))

[sgavrilovicarup]

Hi Deepak, I provided a small list of lat & lon indices above. I don't believe this selects the whole dataset. Unless I am missing something? My goal is download the full time series data at only the indices given.

[TomNicholas]

When I run this example it prints

ds kilobytes 1261459.24
subset kilobytes 70.08

as I would expect. Can you double-check this? And if you get something different can you please share your environment details?

[sgavrilovicarup]

Hi Tom, happy holidays and thanks for responding.

No that is correct and as expected. However, my issue is when I review my network logs, I am seeing 240MB of data downloaded, i.e., the entire dataset from s3. With lazy loading, I expected it to only download the 70KB that I selected (the subset only), and not the whole dataset. My question is, is this the intended behavior? And if so, is there any way to extract a subset of a dataset from s3 without downloading the entire dataset?

For example, in another workflow, I have 20GB rasters on s3. When I need a small area, I only download 250KB from the 20GB raster (using masking in rasterio). In other words, I do not need to download the 20GB raster to extract a tiny area of interest. Makes sense?

[TomNicholas]

Yes that does make sense. What is the chunking of the netCDF file on disk? To me it looks like it's 1 chunk per time point, which I'm finding out by using chunks={}.

ds = xr.open_dataset(fs.open(net_cdf_url), engine='h5netcdf', chunks={})

In [10]: ds
Out[10]: 
<xarray.Dataset>
Dimensions:  (time: 365, lat: 600, lon: 1440)
Coordinates:
  * time     (time) datetime64[ns] 2049-01-01T12:00:00 ... 2049-12-31T12:00:00
  * lat      (lat) float64 -59.88 -59.62 -59.38 -59.12 ... 89.38 89.62 89.88
  * lon      (lon) float64 0.125 0.375 0.625 0.875 ... 359.1 359.4 359.6 359.9
Data variables:
    tasmax   (time, lat, lon) float32 dask.array<chunksize=(1, 600, 1440), meta=np.ndarray>

Opening the data like this creates a ds['tasmax'] object which wraps a dask array with one chunk per time step. (Note that this is a different chunking pattern to if you open using chunks='auto'.)

Your isel is selecting a few spatial points, but across all timesteps. My understanding is that the dask array backing the the ds['tasmax'] object needs to open, uncompress (possibly), load, then subset all of those chunks in order to persist all the data points you've requested. In other words you have to load a whole chunk to get at any values in a chunk. Try opening using chunks={} then subsetting using .isel(..., time=0) instead and see if the network only downloads 1/365th of that data.

[sgavrilovicarup]

Thanks for the suggestion! When I open with chunks={} and call ds['tasmax'].chunks I get

((1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1), (600,), (1440,))

When I try pulling only one time step as you suggested, it still downloads the entire dataset.