Faster fastpath

[hmaarrfk]

I noticed that slicing into datetime64[ns] arrays was really starting to limit our overall performance.

What is strangest is that it is faster to do it on a lazy array than an in-memory array.

I think it comes down to the code on the main branch triggers:

xarray/xarray/core/variable.py

Line 220 in aaa778c

def _possibly_convert_objects(values):

which does a numpy->pandas -> numpy conversion for safety.

But that seems a little over the top.

My benchmark of interest is:

Indexing into a 1D array with a scalar to get a single entry from it out.

This often occurs when we slice into our larger datasets.

on main

python software_timestamp_benchmark.py
software_timestamp.nc
h5netcdf  lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 8664.24it/s]
h5netcdf  computed  : 100%|███████████| 100000/100000 [00:16<00:00, 5976.98it/s]
netcdf4   lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 8453.77it/s]
netcdf4   computed  : 100%|███████████| 100000/100000 [00:16<00:00, 5924.96it/s]
software_timestamp_float64.nc
h5netcdf  lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 9042.25it/s]
h5netcdf  computed  : 100%|██████████| 100000/100000 [00:07<00:00, 14157.30it/s]
netcdf4   lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 9005.22it/s]
netcdf4   computed  : 100%|██████████| 100000/100000 [00:07<00:00, 13935.46it/s]

on this branch:

$ python software_timestamp_benchmark.py
software_timestamp.nc
h5netcdf  lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 8951.20it/s]
h5netcdf  computed  : 100%|██████████| 100000/100000 [00:06<00:00, 14793.26it/s]
netcdf4   lazy      : 100%|███████████| 100000/100000 [00:11<00:00, 8774.30it/s]
netcdf4   computed  : 100%|██████████| 100000/100000 [00:06<00:00, 14671.77it/s]
software_timestamp_float64.nc
h5netcdf  lazy      : 100%|███████████| 100000/100000 [00:10<00:00, 9364.32it/s]
h5netcdf  computed  : 100%|██████████| 100000/100000 [00:06<00:00, 14955.89it/s]
netcdf4   lazy      : 100%|███████████| 100000/100000 [00:10<00:00, 9324.61it/s]
netcdf4   computed  : 100%|██████████| 100000/100000 [00:06<00:00, 14697.09it/s]

The benchmark

import xarray
# from tqdm import tqdm
tqdm = lambda x, *args, **kwargs: x

# dataset = xarray.open_dataset('.nc')
for filename in [
    # 'software_timestamp_float64.nc',
    'software_timestamp.nc',
]:
    print(filename)
    for engine in [
        'h5netcdf', 
        # 'netcdf4',
    ]:
        for compute in [
            # False, 
            True,
        ]:
            dataset = xarray.open_dataset(filename, engine='h5netcdf')
            N_frames = dataset.sizes['frame_number']
            N_iters = 100_000
    
            software_timestamp = dataset['software_timestamp']
            if compute:
                software_timestamp = software_timestamp.compute()
                compute_str = "computed"
            else:
                compute_str = "lazy"
            desc = f"{engine:10s}{compute_str:10s}"
            for i in tqdm(range(N_iters), desc=desc):
                software_timestamp.isel(frame_number=i % N_frames)
            dataset.close()

I can try to expand the benchmark, though it is difficult to "see" these slowdowns with toy examples sometimes.

software_timestamp.zip

• Closes #xxxx
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

xref: #2799
xref: #7045

[max-sixty]

Small point relative to avoiding _maybe_wrap_data:

Wouldn't we only want to get ndim if fastpath is True? So we could only change the comparison to 0 to a comparison to None on L272?

[hmaarrfk]

Thank you for the thorough and quick review. you are correct.

[dcherian]

Looks like both pd.Index and ExtensionArray define ndim so should be OK. It'd be nice to add this as a comment.

[max-sixty]

Merging pending tests. I'm not 100% on the implications of skipping _maybe_wrap_data in some cases, so if anyone knows whether this was required, we can reassess. Thanks @hmaarrfk