From bb12f5b702d060221fb6010f8606eef1a2f31826 Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Fri, 31 Jul 2026 15:26:06 +0200 Subject: [PATCH] Compress datasets exported to netcdf Datasets exported to netcdf were written without any compression. Enable the deflate (gzip) compression built into netcdf4/HDF5 when exporting. The encoding is tuned rather than using the defaults, since the defaults leave a large part of the achievable compression unused: * HDF5 auto chunking picks chunks of about 20 kB. Each chunk is compressed independently so this compresses noticeably worse than using chunks of a few MB. Chunks are now sized to about 4 MB. * The shuffle filter is only enabled for non complex data. Complex numbers are stored as an HDF5 compound type of two floats and shuffling the bytes of such a compound type roughly halves the compression ratio, while shuffling is what gives most of the gain for real valued data. * Dask backed datasets are rechunked to match the HDF5 chunks before writing. Otherwise every dask block triggers a read, decompress, update and recompress cycle of the same HDF5 chunk in the delayed writer. For typical QCoDeS data this reduces the file size by a factor of 1.9 for gridded data, 3.0 for complex data and 3.3 for non gridded data, without slowing down the export. The compression is lossless and the files remain regular netcdf files readable by any netcdf client. The level is configurable via the new config option `dataset.export_netcdf_compression_level` which defaults to 4 and can be set to 0 to write uncompressed files. Higher levels are slower with almost no additional size reduction. The temporary per row files written by the chunked export of large datasets are written uncompressed since they are immediately recombined into the final file. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8fa733cd-3f15-4e38-ab87-43a0fc14cbf8 --- docs/changes/newsfragments/8356.improved | 10 ++ ...Exporting-data-to-other-file-formats.ipynb | 4 +- src/qcodes/configuration/qcodesrc.json | 1 + src/qcodes/configuration/qcodesrc_schema.json | 7 + src/qcodes/dataset/data_set.py | 3 + .../dataset/exporters/export_to_xarray.py | 133 +++++++++++++++++- tests/dataset/test_dataset_export.py | 133 +++++++++++++++++- 7 files changed, 287 insertions(+), 4 deletions(-) create mode 100644 docs/changes/newsfragments/8356.improved diff --git a/docs/changes/newsfragments/8356.improved b/docs/changes/newsfragments/8356.improved new file mode 100644 index 000000000000..fc6e0cac9f9f --- /dev/null +++ b/docs/changes/newsfragments/8356.improved @@ -0,0 +1,10 @@ +Datasets exported to NetCDF are now compressed using the deflate (gzip) compression +built into NetCDF4/HDF5. Depending on the data this reduces the size of the exported +files by a factor of roughly 1.3 to 3.3 with no measurable slowdown of the export, and +no loss of precision since the compression is lossless. Compressed files remain regular +NetCDF files that can be read by any NetCDF client without any special handling. + +The compression level can be controlled with the new +``qcodes.config.dataset.export_netcdf_compression_level`` config option. It defaults to +``4`` and can be set to ``0`` to write uncompressed files as before. Higher levels are +significantly slower while giving almost no additional size reduction. diff --git a/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb b/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb index f9175d0da1dc..4c330ef7881b 100644 --- a/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb +++ b/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb @@ -178,7 +178,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The dataset can be exported using the `export` method. Currently exporting to netcdf and csv is supported." + "The dataset can be exported using the `export` method. Currently exporting to netcdf and csv is supported.\n", + "\n", + "NetCDF files are written with the deflate (gzip) compression that is built into NetCDF4/HDF5 enabled. This is lossless and typically reduces the file size by a factor of 1.3 to 3 depending on the data. The compression level can be changed with the `qcodes.config.dataset.export_netcdf_compression_level` config option and set to `0` to write uncompressed files.\n" ] }, { diff --git a/src/qcodes/configuration/qcodesrc.json b/src/qcodes/configuration/qcodesrc.json index 583f36bf8326..c501cdcee004 100644 --- a/src/qcodes/configuration/qcodesrc.json +++ b/src/qcodes/configuration/qcodesrc.json @@ -78,6 +78,7 @@ "export_name_elements": ["captured_run_id", "guid"], "export_chunked_export_of_large_files_enabled": false, "export_chunked_threshold": 1000, + "export_netcdf_compression_level": 4, "in_memory_cache": true, "load_from_exported_file": false }, diff --git a/src/qcodes/configuration/qcodesrc_schema.json b/src/qcodes/configuration/qcodesrc_schema.json index ddd4929be22c..bcb41807fa66 100644 --- a/src/qcodes/configuration/qcodesrc_schema.json +++ b/src/qcodes/configuration/qcodesrc_schema.json @@ -373,6 +373,13 @@ "default": 1000, "description": "Estimated size in MB above which the dataset will be exported in chuncks and recombined." }, + "export_netcdf_compression_level": { + "type": "integer", + "minimum": 0, + "maximum": 9, + "default": 4, + "description": "Deflate (gzip) compression level used when exporting datasets to netcdf. Set to 0 to write uncompressed files. Higher levels are slower with almost no additional size reduction, so the default of 4 is recommended." + }, "load_from_exported_file": { "description": "Flag to load metadata and raw data from exported file of type specified in export_type. If set to true, qcodes will try to import from file first, if it exists.", "type": "boolean", diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index f36ad008772f..5b8fd43eb0e2 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -1522,6 +1522,9 @@ def _export_as_netcdf(self, path: Path, file_name: str) -> Path: xarray_to_h5netcdf_with_complex_numbers( self.to_xarray_dataset(start=i + 1, end=i + 1), temp_path / file_name_template.format(i), + # these files are temporary and immediately recombined + # into the final file so compressing them only costs time + compression_level=0, ) files = tuple(temp_path.glob("*.nc")) data = xr.open_mfdataset(files) diff --git a/src/qcodes/dataset/exporters/export_to_xarray.py b/src/qcodes/dataset/exporters/export_to_xarray.py index 1734831dd820..ffa5604957e3 100644 --- a/src/qcodes/dataset/exporters/export_to_xarray.py +++ b/src/qcodes/dataset/exporters/export_to_xarray.py @@ -1,14 +1,16 @@ from __future__ import annotations +import itertools import logging import warnings from importlib.metadata import version from math import prod -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal import numpy as np from packaging import version as p_version +import qcodes from qcodes.dataset.linked_datasets.links import links_to_str from ..descriptions.versioning import serialization as serial @@ -29,6 +31,107 @@ _LOG = logging.getLogger(__name__) +# Target size in bytes of a single HDF5 chunk when compression is enabled. +# The chunk size matters a lot for how well the data compresses since each chunk +# is compressed independently. The auto chunking performed by h5py picks chunks +# that are significantly smaller than this which results in noticeably worse +# compression ratios for typical QCoDeS data. +_TARGET_CHUNK_SIZE_BYTES = 4 * 1024**2 + + +def _variable_chunksizes( + variable: xr.DataArray, target_bytes: int = _TARGET_CHUNK_SIZE_BYTES +) -> tuple[int, ...] | None: + """ + Calculate the HDF5 chunk shape to use for a given variable. + + The trailing dimensions are kept whole and the leading dimension is sized + such that a chunk is approximately ``target_bytes`` large. Returns None for + scalars and for variables with a zero sized dimension since those cannot be + chunked. + """ + shape = variable.shape + if len(shape) == 0 or 0 in shape: + return None + + bytes_per_leading_element = max(prod(shape[1:]), 1) * variable.dtype.itemsize + leading_chunk = target_bytes // max(bytes_per_leading_element, 1) + leading_chunk = max(1, min(int(leading_chunk), shape[0])) + return (leading_chunk, *shape[1:]) + + +def _netcdf_compression_encoding( + dataset: xr.Dataset, complevel: int +) -> dict[str, dict[str, Any]]: + """ + Build a per variable netcdf encoding dict enabling deflate compression. + + The shuffle filter is only enabled for non complex data. QCoDeS stores + complex numbers as an HDF5 compound type of two floats and byte shuffling + such a compound type destroys the byte patterns that deflate relies on, + roughly halving the compression ratio. + + Args: + dataset: The dataset that is about to be written. + complevel: Deflate compression level between 1 and 9. + + Returns: + A mapping from variable name to netcdf encoding options. Variables that + cannot be compressed, such as variable length strings, are omitted. + + """ + encoding: dict[str, dict[str, Any]] = {} + + for name in itertools.chain(dataset.data_vars, dataset.coords): + variable = dataset[name] + # Compression filters can only be applied to numeric and boolean data. + # Variable length strings and object arrays are written without filters. + if variable.dtype.kind not in "buifc": + continue + + variable_encoding: dict[str, Any] = { + "zlib": True, + "complevel": complevel, + "shuffle": variable.dtype.kind != "c", + } + chunksizes = _variable_chunksizes(variable) + if chunksizes is not None: + variable_encoding["chunksizes"] = chunksizes + encoding[str(name)] = variable_encoding + + return encoding + + +def _rechunk_to_match_encoding( + dataset: xr.Dataset, encoding: Mapping[str, Mapping[str, Any]] +) -> xr.Dataset: + """ + Align the dask chunks of a dataset with the HDF5 chunks it will be written to. + + Writing a dask backed dataset is only performed one dask block at a time. If + a block covers less than a full HDF5 chunk, that chunk has to be read, + decompressed, updated and recompressed again for every block, which is very + slow. Datasets that are not dask backed are returned unmodified. + """ + is_dask_backed = any( + getattr(dataset[name].data, "chunks", None) is not None + for name in dataset.variables + ) + if not is_dask_backed: + return dataset + + dim_chunks: dict[Hashable, int] = {} + for name, variable_encoding in encoding.items(): + chunksizes = variable_encoding.get("chunksizes") + if chunksizes is None: + continue + for dim, chunksize in zip(dataset[name].dims, chunksizes, strict=True): + dim_chunks[dim] = min(dim_chunks.get(dim, chunksize), chunksize) + + if not dim_chunks: + return dataset + return dataset.chunk(dim_chunks) + def _calculate_index_shape(idx: pd.Index | pd.MultiIndex) -> dict[Hashable, int]: # heavily inspired by xarray.core.dataset.from_dataframe @@ -415,11 +518,30 @@ def _paramspec_dict_with_extras( def xarray_to_h5netcdf_with_complex_numbers( - xarray_dataset: xr.Dataset, file_path: str | Path, compute: bool = True + xarray_dataset: xr.Dataset, + file_path: str | Path, + compute: bool = True, + compression_level: int | None = None, ) -> None: + """ + Write an xarray dataset to a netcdf file using the h5netcdf engine. + + Args: + xarray_dataset: The dataset to write. + file_path: Path of the netcdf file to write. + compute: If False the write is returned as a Dask delayed job which is + computed with a progress bar rather than written eagerly. + compression_level: Deflate compression level between 0 and 9 where 0 + means no compression. If None the level is read from + ``qcodes.config.dataset.export_netcdf_compression_level``. + + """ import cf_xarray as cf_xr from pandas import MultiIndex + if compression_level is None: + compression_level = int(qcodes.config.dataset.export_netcdf_compression_level) + has_multi_index = any( isinstance(xarray_dataset.indexes[index_name], MultiIndex) for index_name in xarray_dataset.indexes @@ -453,6 +575,12 @@ def xarray_to_h5netcdf_with_complex_numbers( xarray_too_old or h5netcdf_too_old ) + if compression_level > 0: + encoding = _netcdf_compression_encoding(internal_ds, compression_level) + internal_ds = _rechunk_to_match_encoding(internal_ds, encoding) + else: + encoding = {} + with warnings.catch_warnings(): # see http://xarray.pydata.org/en/stable/howdoi.html # for how to export complex numbers @@ -468,6 +596,7 @@ def xarray_to_h5netcdf_with_complex_numbers( engine="h5netcdf", invalid_netcdf=allow_invalid_netcdf, compute=compute, + encoding=encoding, ) if not compute and maybe_write_job is not None: # Dask and therefor tqdm.dask is slow to diff --git a/tests/dataset/test_dataset_export.py b/tests/dataset/test_dataset_export.py index bf1036e348f3..c3c328686da2 100644 --- a/tests/dataset/test_dataset_export.py +++ b/tests/dataset/test_dataset_export.py @@ -34,7 +34,11 @@ from qcodes.dataset.descriptions.versioning import serialization as serial from qcodes.dataset.export_config import DataExportType from qcodes.dataset.exporters.export_to_pandas import _generate_pandas_index -from qcodes.dataset.exporters.export_to_xarray import _calculate_index_shape +from qcodes.dataset.exporters.export_to_xarray import ( + _TARGET_CHUNK_SIZE_BYTES, + _calculate_index_shape, + _netcdf_compression_encoding, +) from qcodes.dataset.linked_datasets.links import links_to_str from qcodes.parameters import ManualParameter, Parameter, ParamSpecBase @@ -2326,3 +2330,130 @@ def test_incomplete_measurement_with_shared_setpoint( assert "signal_1d" in xr_ds.data_vars assert "signal_2d" in xr_ds.data_vars assert "x" in xr_ds.coords + + +def test_netcdf_compression_encoding_helper() -> None: + """Compression is applied per variable with shuffle disabled for complex data.""" + ds = xr.Dataset( + { + "real": (("x", "y"), np.zeros((100, 7))), + "complex": (("x", "y"), np.zeros((100, 7), dtype=np.complex128)), + "text": ("x", np.array(["a"] * 100, dtype=object)), + "scalar": ((), 1.0), + }, + coords={"x": np.arange(100.0), "y": np.arange(7.0)}, + ) + encoding = _netcdf_compression_encoding(ds, complevel=4) + + # object/string variables cannot be compressed and are left out + assert "text" not in encoding + assert set(encoding) == {"real", "complex", "scalar", "x", "y"} + + assert encoding["real"] == { + "zlib": True, + "complevel": 4, + "shuffle": True, + "chunksizes": (100, 7), + } + # shuffle destroys the byte layout of the complex compound type + assert encoding["complex"]["shuffle"] is False + assert encoding["complex"]["chunksizes"] == (100, 7) + # scalars cannot be chunked + assert "chunksizes" not in encoding["scalar"] + + +def test_netcdf_compression_encoding_chunks_are_capped_to_target_size() -> None: + """The leading dimension is sized to give chunks of about the target size.""" + n_columns = 128 + ds = xr.Dataset({"z": (("x", "y"), np.zeros((10000, n_columns)))}) + encoding = _netcdf_compression_encoding(ds, complevel=4) + + chunksizes = encoding["z"]["chunksizes"] + assert chunksizes[1] == n_columns + assert chunksizes[0] * n_columns * 8 <= _TARGET_CHUNK_SIZE_BYTES + assert (chunksizes[0] + 1) * n_columns * 8 > _TARGET_CHUNK_SIZE_BYTES + + +@pytest.mark.parametrize("complevel", [0, 4]) +def test_export_netcdf_compression_config( + tmp_path_factory: TempPathFactory, mock_dataset_grid: DataSet, complevel: int +) -> None: + """The compression level from the config is applied to the exported file.""" + h5py = pytest.importorskip("h5py") + tmp_path = tmp_path_factory.mktemp("export_netcdf_compression") + qcodes.config.dataset.export_netcdf_compression_level = complevel + + mock_dataset_grid.export(export_type="netcdf", path=tmp_path, prefix="qcodes_") + file_path = mock_dataset_grid.export_info.export_paths["nc"] + + with h5py.File(file_path, "r") as file: + z = file["z"] + if complevel == 0: + assert z.compression is None + else: + assert z.compression == "gzip" + assert z.compression_opts == complevel + assert z.shuffle is True + assert z.chunks == (10, 5) + + loaded_ds = xr.load_dataset(file_path, engine="h5netcdf") + assert_allclose(loaded_ds.z, mock_dataset_grid.to_xarray_dataset().z) + + +def test_export_netcdf_compression_of_complex_data( + tmp_path_factory: TempPathFactory, mock_dataset_numpy_complex: DataSet +) -> None: + """Complex data is compressed but not shuffled and round trips exactly.""" + h5py = pytest.importorskip("h5py") + tmp_path = tmp_path_factory.mktemp("export_netcdf_compression_complex") + qcodes.config.dataset.export_netcdf_compression_level = 4 + + mock_dataset_numpy_complex.export( + export_type="netcdf", path=tmp_path, prefix="qcodes_" + ) + file_path = mock_dataset_numpy_complex.export_info.export_paths["nc"] + + with h5py.File(file_path, "r") as file: + z = file["z"] + assert z.compression == "gzip" + assert z.shuffle is False + + original = mock_dataset_numpy_complex.to_xarray_dataset() + loaded_ds = xr.load_dataset(file_path, engine="h5netcdf") + assert_array_equal(loaded_ds.z, original.z) + + +def test_export_netcdf_compression_reduces_file_size( + tmp_path_factory: TempPathFactory, experiment: Experiment +) -> None: + """Compression gives a significantly smaller file for realistic data.""" + dataset = new_data_set("compressible") + xparam = ParamSpecBase("x", "numeric") + yparam = ParamSpecBase("y", "numeric") + zparam = ParamSpecBase("z", "numeric") + dataset.set_interdependencies( + InterDependencies_(dependencies={zparam: (xparam, yparam)}) + ) + dataset.mark_started() + y_values = np.linspace(-1, 1, 200) + for i in range(100): + # quantized as a real instrument would return it + z_values = np.round(np.sin(i / 10) * np.exp(-(y_values**2)) * 32767) / 32767 + dataset.add_results( + [{"x": i, "y": y, "z": z} for y, z in zip(y_values, z_values, strict=True)] + ) + dataset.mark_completed() + + sizes = {} + for complevel in (0, 4): + qcodes.config.dataset.export_netcdf_compression_level = complevel + path = tmp_path_factory.mktemp(f"compression_size_{complevel}") + dataset.export(export_type="netcdf", path=path, prefix="qcodes_") + sizes[complevel] = Path(dataset.export_info.export_paths["nc"]).stat().st_size + + assert sizes[4] < sizes[0] / 1.3 + + loaded_ds = xr.load_dataset( + dataset.export_info.export_paths["nc"], engine="h5netcdf" + ) + assert_array_equal(loaded_ds.z, dataset.to_xarray_dataset().z)