Local and S3 windowed-reading Sentinel-2 rasters is 10x slower with rioxarray than xarray

[Bencpr]

Hello all !

I am currently doing some speed tests with reading Sentinel 2 rasters located on a private S3 bucket. I have both standard and COG versions of them.

After somes tests it turns out reading data from S3 using windowed reading takes me about 10x more time with rioxarray.open_rasterio() than with xarray.open_rasterio() .
I also included rasterio.open() in the tests, which is even longer than rioxarray.

import os
from time import time
import xarray as xr
import rioxarray as rioxr
from rasterio.windows import Window

os.environ['AWS_S3_ENDPOINT'] = XXXXX
os.environ['AWS_ACCESS_KEY_ID'] = XXXXX
os.environ['AWS_SECRET_ACCESS_KEY'] = XXXXX
os.environ['AWS_VIRTUAL_HOSTING'] =  "False" # need for /vsis3/

# path to S3 raster
path = "'/vsis3/path/to/file.tif"

# prepare different n_wins window offsets to be read
n_wins = 200
col_offs = [int(i) for i in np.linspace(0, 8000, n_wins)]
row_offs = [int(i) for i in np.linspace(400, 8400, n_wins)]
h, w = 256, 256

# repeat n_tries and average time elapsed
n_tries = 5

readers = ("xarray, "rioxarray", "rasterio")

for reader in readers:

  # time elapsed for n_tries
  time_tries = 0

  for n in range(n_tries):
      time_try = 0
      start = time()
  
      for col_off, row_off in zip(col_offs, row_offs):
  
          if reader == "xarray":
              win = xr.open_rasterio(path).isel(x=slice(row_off, h), y=slice(col_off, w))
  
          if reader == "rioxarray":
              win = rioxr.open_rasterio(path).isel(x=slice(row_off, h), y=slice(col_off, w))
  
          if reader == "rasterio":
              with rasterio.open(path) as src:
                  win = src.read(window=Window(col_off=col_off, row_off=row_off, height=h, width=w))
  
      end = time()
      time_try = end - start
      time_tries += time_try
  
  total_time = time_tries / n_tries
  print(f'reader {reader} - elapsed {mode}: {final_time} for {n_wins} windows with {n_tries} tries')

Using a COG tif I get the following outputs:

reader xarray - time : 0.142 for 200 windows with 5 tries
reader rioxarray - time : 14.408 for 200 windows with 5 tries
reader rasterio - time : 22.988 for 200 windows with 5 tries

So apparently xarray is 10x faster ? What am I missing here?
I may have included mistakes...or I didn't configure stuff well.

Thanks in advance for any advice :)

Best,
Benjamin

[snowman2]

rioxarray/xarray both use lazy loading to load data from the rasters. Calling win.load() will likely get you performance closer to rasterio.

rioxarray loads metadata from the files, so that is likely why it takes longer than xarray.

Also, can you provide the output of rioxarray.show_versions?

EDIT: To avoid the cost of re-opening the files, I recommend modifying the tests to:

for reader in readers:
    if reader == "xarray":
        xds = xr.open_rasterio(path)

    if reader == "rioxarray":
        xds = rioxr.open_rasterio(path)

  # time elapsed for n_tries
  time_tries = 0

  for n in range(n_tries):
      time_try = 0
      start = time()
  
      for col_off, row_off in zip(col_offs, row_offs):
  
         if reader in ("xarray", "rioxarray"):
             win = xds.isel(x=slice(row_off, h), y=slice(col_off, w)).load()
         if reader == "rasterio":
             with rasterio.open(path) as src:
                 win = src.read(window=Window(col_off=col_off, row_off=row_off, height=h, width=w))
  
      end = time()
      time_try = end - start
      time_tries += time_try
  
  total_time = time_tries / n_tries
  print(f'reader {reader} - elapsed {mode}: {final_time} for {n_wins} windows with {n_tries} tries')

[Bencpr]

Thanks for your answer.

I added a .load() on each reading as follows

if reader == "xarray":
    win = xr.open_rasterio(path).isel(x=slice(row_off, h), y=slice(col_off, w)).load()

if reader == "rioxarray":
    win = rioxr.open_rasterio(path).isel(x=slice(row_off, h), y=slice(col_off, w)).load()

if reader == "rasterio":
    with rasterio.open(path) as src:
        win = src.read(window=Window(col_off=col_off, row_off=row_off, height=h, width=w))

And get the following:

reader xarray - time : 0.523 for 200 windows with 5 tries
reader rioxarray - time : 15.173 for 200 windows with 5 tries
reader rasterio - time : 21.328 for 200 windows with 5 tries

xarray stills reads way faster apparently.

Here is the output of rioxarray.show_versions():

rioxarray (0.5.0) deps:
  rasterio: 1.1.3
    xarray: 0.17.0
      GDAL: 3.0.4

Other python deps:
     scipy: 1.6.1
    pyproj: 3.1.0

System:
    python: 3.8.2 (default, Jul 16 2020, 14:00:26)  [GCC 9.3.0]
executable: /home/name/.pyenv/versions/datalab/bin/python
   machine: Linux-5.4.72-microsoft-standard-WSL2-x86_64-with-glibc2.29

Any advice ? Again thank you !

[snowman2]

Can you update to the latest pyproj (3.3.0) and rioxarray (0.9.0) and try again?

[snowman2]

Also, see the updated answer above.

[Bencpr]

With the new versions of pyproj and rioxarray, it's much better :

reader xarray - time : 0.494 for 200 windows with 5 tries
reader rioxarray - time : 1.241 for 200 windows with 5 tries

But still better for xarray, maybe due to metadata reading ?

Thanks for the tips to avoid the cost of re-opening the file ! But actually re-openening is closer to my use case, as I will face potentially hundreds of files and slicing different windows in them. Btw do you recommend to lazy open them, even in case of many files?

[snowman2]

I am thinking this issue is related to: #234
rioxarray.set_options(export_grid_mapping=False) may also help improve performance.

do you recommend to lazy open them, even in case of many files?

I would recommend trying to avoid re-opening files if you can and it makes sense.

[Bencpr]

Thank you @snowman2 ! All good now.