Performance Issue - DataArray.drop_vars increase runtime significantly.

[KenTSUI-dev]

In one of my project trying to process and visualise some 4D modelling data from Delft3D FM in ArcGIS Online, I found that DaskArray.drop_vars doubles the runtime of my code. From my understanding, dropping variables should not be a costly task. Though I am not able to copy my project here for your testing, I tried to produce a minimal working example below to demonstrate the situation. I would like to know if it is expected (i.e. dropping variables is costly) and how to avoid it?

import xarray as xr
import dask
import dask.array as da
from dask.diagnostics import ProgressBar
import numpy as np



def drop_vars_performance1():
    """
    without drop_vars
    """
    dim_time, dim_layer, dim_lat, dim_lon  = 100, 20, 800, 800
    lat = np.linspace(23.5, 24.5, dim_lat)
    lon = np.linspace(113.5, 114.5, dim_lon)
    time = np.arange(dim_time)
    layer = np.arange(dim_layer)
    data = da.random.random((dim_time, dim_layer, dim_lat, dim_lon), chunks=(1, 1, -1, -1))


    xarr = xr.DataArray(
        data=data,
        dims=['time','layer', 'lat', 'lon'],
        coords={
            'lat': lat,
            'lon': lon,
            'time': time,
            'layer': layer,
            'dummy1': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
            'dummy2': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
            'dummy3': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
        }
    )

    job = xarr.to_netcdf(r'D:/temp/drop_vars_performance_1.nc', compute=False)

    with ProgressBar():
        job.compute()


def drop_vars_performance2():
    """
    with drop_vars
    """
    dim_time, dim_layer, dim_lat, dim_lon  = 100, 20, 800, 800
    lat = np.linspace(23.5, 24.5, dim_lat)
    lon = np.linspace(113.5, 114.5, dim_lon)
    time = np.arange(dim_time)
    layer = np.arange(dim_layer)
    data = da.random.random((dim_time, dim_layer, dim_lat, dim_lon), chunks=(1, 1, -1, -1))

    xarr = xr.DataArray(
        data=data,
        dims=['time','layer', 'lat', 'lon'],
        coords={
            'lat': lat,
            'lon': lon,
            'time': time,
            'layer': layer,
            'dummy1': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
            'dummy2': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
            'dummy3': (('lat', 'lon'), np.random.random((dim_lat, dim_lon))),
        }
    )

    xarr = xarr.drop_vars({'dummy1', 'dummy2', 'dummy3',})

    job = xarr.to_netcdf(r'D:/temp/drop_vars_performance_2.nc', compute=False)

    with ProgressBar():
        job.compute()


if __name__ == "__main__":
    print(f"xarray version: {xr.__version__}")
    print(f"dask version: {dask.__version__}")
    print("drop_vars_performance1:")
    drop_vars_performance1()
    print("drop_vars_performance2:")
    drop_vars_performance2()

output:

Python 3.9.13 (main, Aug 25 2022, 23:51:50) [MSC v.1916 64 bit (AMD64)]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.8.0 -- An enhanced Interactive Python. Type '?' for help.
PyDev console: using IPython 8.8.0
Python 3.9.13 (main, Aug 25 2022, 23:51:50) [MSC v.1916 64 bit (AMD64)] on win32

xarray version: 2022.9.0
dask version: 2022.9.2
drop_vars_performance1:
[########################################] | 100% Completed | 124.24 s
drop_vars_performance2:
[########################################] | 100% Completed | 179.83 s

[kmuehlbauer]

@KenTSUI-dev
The coordinates aren't dropped in your current code. You would need to assign a return Dataset.

xarr = xarr.drop_vars({'dummy1', 'dummy2', 'dummy3',})

I'm in no way an expert in dask performance diagnosing, but the timings vary from run to run. Even when run separately timings differ from run to run. I'm interested how proper diagnosing is done in such cases, too.

[KenTSUI-dev]

Thanks. I have updated and re-run the code. The performance issue is still observed.

[keewis]

I consistently observe that the second computation takes longer, and it doesn't matter which version is called last. I'm no expert, but I believe that this is a hardware issue (a full disk cache that has to be cleared before writing more data, maybe?).

Can you try calling drop_vars_performance2 before drop_vars_performance1 to verify?

[kmuehlbauer]

I consistently observe that the second computation takes longer, and it doesn't matter which version is called last. I'm no expert, but I believe that this is a hardware issue (a full disk cache that has to be cleared before writing more data, maybe?).

I have observed the same. But even running each function separately gives me timings which considerably change from run to run for both of them. This would need a dask-expert which lectures the correct performance profiling.