compresses raw forcing netcdf files, cuts file size in about half

[quinnylee]

see issue #107 and discussion there for more details

Compression level was set to 1 because according to tests, higher compression levels provided negligible benefits to file size

[JoshCu]

What do you think about only enabling this if the files are larger than say 5 or 10gb? This new Nan interpolation step goes from taking 2s on a 217mb file to taking ~30s once it's compressed to 51mb.

For smaller files I'd rather things be faster but for larger datasets things the compression is far more valuable. The nan interpolation step rewrites the netcdf without overwriting it so it needs 2x the files size available to work.

[quinnylee]

Interesting, I hadn't noticed such a dramatic difference in performance. Maybe I got lucky and only tested compression on time periods and catchments where there were no nans...

I'll write something to enable compression if files are larger than 5 gb.

[JoshCu]

the machine I'm testing on is unusual, lots of very slow cores, so if something is serial it takes an eternity for me. Which is exactly what was happening. The problem was with the interpolation step, it opens the database after its been downloaded with no dask chunking so all the operations were running serially. I'm updating that which should alleviate this issue

[JoshCu]

I've messed with this a bunch more and even without interpolation I don't think we can compress the data to netcdf without slowing things down way too much. With the aorc data the save time for each chunk goes from 3ms to 300-400ms and the download speed drops to ~13mbs from ~500.

I've been trying out saving to zarr and that does seem to work and is actually faster than uncompressed netcdf saving too (the download part, I haven't tested out the forcing processing yet)

modifying the client calls in dask_utils.py to this speeds things up client = Client(threads_per_worker=1)

after the download things still break, but this is what I've got so far

def save_dataset(
    ds_to_save: xr.Dataset,
    target_path: Path,
    engine: Literal["netcdf4", "scipy", "h5netcdf"] = "h5netcdf",
):
    """
    Helper function to compute and save an xarray.Dataset (specifically, the raw
    forcing data) to a NetCDF file.
    Uses a temporary file and rename for atomicity.
    """
    x_len = ds_to_save.sizes["x"]
    y_len = ds_to_save.sizes["y"]
    time_len = ds_to_save.sizes["time"]

    encoding = {}
    compressor = Blosc(cname="zstd", clevel=3, shuffle=2)
    for var in list(ds_to_save.keys()):
        encoding[var] = {
            "compressor": compressor,  # higher compression levels provide negligible benefits to file size
        }
    for name, var in ds_to_save.data_vars.items():
        chunks = {"time": 144, "x": x_len, "y": y_len}
        ds_to_save[name] = var.chunk(chunks)

    if not target_path.parent.exists():
        target_path.parent.mkdir(parents=True, exist_ok=True)

    temp_file_path = target_path.with_name(target_path.name + ".saving.nc")
    if temp_file_path.exists():
        try:
            temp_file_path.unlink()
        except IsADirectoryError:
            shutil.rmtree(temp_file_path)

    client = Client.current()
    future: Future = client.compute(
        ds_to_save.to_zarr(
            store=temp_file_path,
            mode="w",
            compute=False,
            write_empty_chunks=False,
            encoding=encoding,
        )
    )  # type: ignore
    logger.debug(
        f"NetCDF write task submitted to Dask. Waiting for completion to {temp_file_path}..."
    )
    logger.info("For more detailed progress, see the Dask dashboard http://localhost:8787/status")
    progress(future)
    future.result()
    os.rename(str(temp_file_path), str(target_path))
    logger.info(f"Successfully saved data to: {target_path}")

[JoshCu]

I think I'm going to close this because proceeding with zarr seems to improve compression as well as the performance #173