Repository: QCoDeS/Qcodes
Pull request #5320 (2023-08-17 03:23:53)

load_by_id optionally loads data from exported file
<!--

Thanks for submitting a pull request against QCoDeS.

To help us effectively merge your pr please consider the following check list.

- [ ] Make sure that the pull request contains a short description of the changes made.
- [ ] If you are submitting a new feature please document it. This can be in the form of inline
      docstrings, an example notebook or restructured text files.
- [ ] Please include automatic tests for the changes made when possible.

Unless your change is a small or trivial fix please add a small changelog entry:

- [ ] Create a file in the docs\changes\newsfragments folder with a short description of the change.

This file should be in the format number.categoryofcontribution. Here the number should either be the number
of the pull request. To get the number of the pull request one must
first the pull request and then subsequently update the number. The category of contribution should be
one of ``breaking``, ``new``, ``improved``, ``new_driver`` ``improved_driver``, ``underthehood``.

If this fixes a known bug reported against QCoDeS:

- [ ] Please include a string in the following form ``closes #xxx`` where ``xxx``` is the number of the bug fixed.

Please have a look at [the contributing guide](https://qcodes.github.io/Qcodes/community/contributing.html)
for more information.

If you are in doubt about any of this please ask and we will be happy to help.

-->
We've been having issues using Plottr in tandem with measurements because they access the sqlite database at the same time. We think that simultaneous read/write operations from both processes have been causing either one of them to crash. A potential solution I came up with is to have load_by_id load the dataset from an exported netcdf file, if it exists.
Please take a look and let me know if you think this is a reasonable solution, I'm currently testing it to see if it actually solves the problem.


Patch:
diff --git a/docs/changes/newsfragments/5320.improved b/docs/changes/newsfragments/5320.improved
new file mode 100644
index 000000000..fe8dd2c94
--- /dev/null
+++ b/docs/changes/newsfragments/5320.improved
@@ -0,0 +1,3 @@
+`load_by_run_spec` and `load_by_id` may now optionally load the dataset as a DataSetInMem from an exported netcdf file. This avoids reading
+from the database potentially resulting in better performance. This option is off by default at the moment but can be turned on by setting
+`qcodes.config.dataset.load_from_exported_file` to True.
diff --git a/requirements.txt b/requirements.txt
index e3c67603d..28b927d92 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,14 +9,14 @@ argon2-cffi~=23.1.0
 argon2-cffi-bindings~=21.2.0
 asttokens~=2.4.0
 attrs~=23.2.0
-autodocsumm==0.2.12
+autodocsumm==0.2.11
 azure-core~=1.29.0
 azure-identity~=1.15.0
 Babel~=2.14.0
 backcall==0.2.0
 beautifulsoup4~=4.12.0
 bleach~=6.1.0
-broadbean==0.12.0
+broadbean==0.11.0
 cachetools~=5.3.0
 certifi~=2023.11.17
 cffi~=1.16.0
diff --git a/src/qcodes/configuration/qcodesrc.json b/src/qcodes/configuration/qcodesrc.json
index f796ed145..583f36bf8 100644
--- a/src/qcodes/configuration/qcodesrc.json
+++ b/src/qcodes/configuration/qcodesrc.json
@@ -78,7 +78,8 @@
         "export_name_elements": ["captured_run_id", "guid"],
         "export_chunked_export_of_large_files_enabled": false,
         "export_chunked_threshold": 1000,
-        "in_memory_cache": true
+        "in_memory_cache": true,
+        "load_from_exported_file": false
     },
     "telemetry":
     {
diff --git a/src/qcodes/configuration/qcodesrc_schema.json b/src/qcodes/configuration/qcodesrc_schema.json
index 448b4639c..ddd4929be 100644
--- a/src/qcodes/configuration/qcodesrc_schema.json
+++ b/src/qcodes/configuration/qcodesrc_schema.json
@@ -373,6 +373,11 @@
                     "default": 1000,
                     "description": "Estimated size in MB above which the dataset will be exported in chuncks and recombined."
                 },
+                "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",
+                    "default": false
+                },
                 "in_memory_cache": {
                     "type": "boolean",
                     "default": true,
diff --git a/src/qcodes/dataset/__init__.py b/src/qcodes/dataset/__init__.py
index 79a44aba6..3d724195a 100644
--- a/src/qcodes/dataset/__init__.py
+++ b/src/qcodes/dataset/__init__.py
@@ -10,7 +10,7 @@ from .data_set import (
     load_by_run_spec,
     new_data_set,
 )
-from .data_set_in_memory import load_from_netcdf
+from .data_set_in_memory import load_from_file, load_from_netcdf
 from .data_set_protocol import DataSetProtocol, DataSetType
 from .database_extract_runs import extract_runs_into_db
 from .descriptions.dependencies import InterDependencies_, ParamSpecTree
@@ -103,6 +103,7 @@ __all__ = [
     "load_by_run_spec",
     "load_experiment",
     "load_experiment_by_name",
+    "load_from_file",
     "load_from_netcdf",
     "load_last_experiment",
     "load_or_create_experiment",
diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py
index 573c10003..c9b9b420e 100644
--- a/src/qcodes/dataset/data_set.py
+++ b/src/qcodes/dataset/data_set.py
@@ -33,6 +33,9 @@ from qcodes.dataset.descriptions.versioning.converters import new_to_old, old_to
 from qcodes.dataset.descriptions.versioning.rundescribertypes import Shapes
 from qcodes.dataset.descriptions.versioning.v0 import InterDependencies
 from qcodes.dataset.experiment_settings import get_default_experiment_id
+from qcodes.dataset.export_config import (
+    DataExportType,
+)
 from qcodes.dataset.guids import filter_guids_by_parts, generate_guid, parse_guid
 from qcodes.dataset.linked_datasets.links import Link, links_to_str, str_to_links
 from qcodes.dataset.sqlite.connection import ConnectionPlus, atomic, atomic_transaction
@@ -84,7 +87,7 @@ from qcodes.utils import (
 )
 
 from .data_set_cache import DataSetCacheWithDBBackend
-from .data_set_in_memory import DataSetInMem
+from .data_set_in_memory import DataSetInMem, load_from_file
 from .descriptions.versioning import serialization as serial
 from .exporters.export_info import ExportInfo
 from .exporters.export_to_csv import dataframe_to_csv
@@ -1589,8 +1592,9 @@ def load_by_run_spec(
     specs of the runs found will be printed.
 
     If the raw data is in the database this will be loaded as a
-    :class:`qcodes.dataset.data_set.DataSet`
-    otherwise it will be loaded as a :class:`.DataSetInMemory`
+    :class:`DataSet`. Otherwise it will be loaded as a :class:`.DataSetInMemory`
+    If the raw data is exported to netcdf and ``qcodes.config.dataset.load_from_exported_file`` is
+    set to True. this will be loaded from file as a :class:`DataSetInMemory`. regardless.
 
     Args:
         captured_run_id: The ``run_id`` that was originally assigned to this
@@ -1709,9 +1713,11 @@ def load_by_id(run_id: int, conn: ConnectionPlus | None = None) -> DataSetProtoc
     data to another db file. We recommend using :func:`.load_by_run_spec` which
     does not have this issue and is significantly more flexible.
 
+
     If the raw data is in the database this will be loaded as a
-    :class:`qcodes.dataset.data_set.DataSet` otherwise it will be
-    loaded as a :class:`.DataSetInMemory`
+    :class:`DataSet`. Otherwise it will be loaded as a :class:`.DataSetInMemory`
+    If the raw data is exported to netcdf and ``qcodes.config.dataset.load_from_exported_file`` is
+    set to True. this will be loaded from file as a :class:`DataSetInMemory`. regardless.
 
     Args:
         run_id: run id of the dataset
@@ -1747,8 +1753,9 @@ def load_by_guid(guid: str, conn: ConnectionPlus | None = None) -> DataSetProtoc
     is specified in the config.
 
     If the raw data is in the database this will be loaded as a
-    :class:`qcodes.dataset.data_set.DataSet`
-    otherwise it will be loaded as a :class:`.DataSetInMemory`
+    :class:`DataSet`. Otherwise it will be loaded as a :class:`.DataSetInMemory`
+    If the raw data is exported to netcdf and ``qcodes.config.dataset.load_from_exported_file`` is
+    set to True. this will be loaded from file as a :class:`DataSetInMemory`. regardless.
 
     Args:
         guid: guid of the dataset
@@ -1788,9 +1795,11 @@ def load_by_counter(
     data to another db file. We recommend using :func:`.load_by_run_spec` which
     does not have this issue and is significantly more flexible.
 
+
     If the raw data is in the database this will be loaded as a
-    :class:`qcodes.dataset.data_set.DataSet`
-    otherwise it will be loaded as a :class:`.DataSetInMemory`
+    :class:`DataSet`. Otherwise it will be loaded as a :class:`.DataSetInMemory`
+    If the raw data is exported to netcdf and ``qcodes.config.dataset.load_from_exported_file`` is
+    set to True. this will be loaded from file as a :class:`DataSetInMemory`. regardless.
 
     Args:
         counter: counter of the dataset within the given experiment
@@ -1799,7 +1808,7 @@ def load_by_counter(
           connection to the DB file specified in the config is made
 
     Returns:
-        :class:`qcodes.dataset.data_set.DataSet` or
+        :class:`DataSet` or
         :class:`.DataSetInMemory` of the given counter in
         the given experiment
     """
@@ -1821,15 +1830,41 @@ def load_by_counter(
 def _get_datasetprotocol_from_guid(guid: str, conn: ConnectionPlus) -> DataSetProtocol:
     run_id = get_runid_from_guid(conn, guid)
     if run_id is None:
-        raise NameError(f"No run with GUID: {guid} found in database.")
+        raise NameError("No run with GUID: %s found in database.", guid)
+
+    if qcodes.config.dataset.load_from_exported_file:
+        export_info = _get_datasetprotocol_export_info(run_id=run_id, conn=conn)
+
+        export_file_path = export_info.export_paths.get(
+            DataExportType.NETCDF.value, None
+        )
+
+        if export_file_path is not None:
+            try:
+                d: DataSetProtocol = load_from_file(export_file_path)
+
+            except (ValueError, FileNotFoundError) as e:
+                log.warning("Cannot load data from file: %s", e)
+
+            else:
+                return d
+
     result_table_name = _get_result_table_name_by_guid(conn, guid)
     if _check_if_table_found(conn, result_table_name):
-        d: DataSetProtocol = DataSet(conn=conn, run_id=run_id)
+        d = DataSet(conn=conn, run_id=run_id)
     else:
         d = DataSetInMem._load_from_db(conn=conn, guid=guid)
+
     return d
 
 
+def _get_datasetprotocol_export_info(run_id: int, conn: ConnectionPlus) -> ExportInfo:
+    metadata = get_metadata_from_run_id(conn=conn, run_id=run_id)
+    export_info_str = metadata.get("export_info", "")
+    export_info = ExportInfo.from_str(export_info_str)
+    return export_info
+
+
 def new_data_set(
     name: str,
     exp_id: int | None = None,
diff --git a/src/qcodes/dataset/data_set_in_memory.py b/src/qcodes/dataset/data_set_in_memory.py
index 463214290..6d6fd7319 100644
--- a/src/qcodes/dataset/data_set_in_memory.py
+++ b/src/qcodes/dataset/data_set_in_memory.py
@@ -23,6 +23,7 @@ from qcodes.dataset.descriptions.param_spec import ParamSpec, ParamSpecBase
 from qcodes.dataset.descriptions.rundescriber import RunDescriber
 from qcodes.dataset.descriptions.versioning.converters import new_to_old
 from qcodes.dataset.descriptions.versioning.rundescribertypes import Shapes
+from qcodes.dataset.export_config import DataExportType
 from qcodes.dataset.guids import generate_guid
 from qcodes.dataset.linked_datasets.links import Link, links_to_str
 from qcodes.dataset.sqlite.connection import ConnectionPlus, atomic
@@ -913,3 +914,36 @@ def load_from_netcdf(
         The loaded dataset.
     """
     return DataSetInMem._load_from_netcdf(path=path, path_to_db=path_to_db)
+
+
+def load_from_file(
+    path: Path | str, path_to_db: Path | str | None = None
+) -> DataSetInMem:
+    """
+    Create an in-memory dataset from a file.
+    The file is expected to contain a QCoDeS dataset that
+    has been exported using the QCoDeS export functions.
+    Currently, this only supports loading from netcdf files.
+
+    Args:
+        path: Path to the file to import.
+        path_to_db: Optional path to a database where this dataset may be
+            exported to. If not supplied the path can be given at export time
+            or the dataset exported to the default db as set in the QCoDeS config.
+
+    Returns:
+        The loaded dataset.
+    """
+    path = Path(path)
+    if not path.is_file():
+        raise FileNotFoundError(f"File {path} not found.")
+
+    if DataExportType.NETCDF.value == path.suffix.replace(".", ""):
+        return DataSetInMem._load_from_netcdf(path=path, path_to_db=path_to_db)
+
+    else:
+        raise ValueError(
+            f"Loading file of type {path.suffix} is not supported."
+            "Please refer to https://github.com/QCoDeS/Qcodes/issues"
+            "to find existing or create a new feature request."
+        )


Test patch:
diff --git a/tests/dataset/test_dataset_in_memory.py b/tests/dataset/test_dataset_in_memory.py
index efed3c0a7..23849fbb8 100644
--- a/tests/dataset/test_dataset_in_memory.py
+++ b/tests/dataset/test_dataset_in_memory.py
@@ -11,8 +11,9 @@ import xarray as xr
 from hypothesis import HealthCheck, given, settings
 from numpy.testing import assert_almost_equal
 
+import qcodes
 from qcodes.dataset import load_by_id, load_by_run_spec
-from qcodes.dataset.data_set_in_memory import DataSetInMem
+from qcodes.dataset.data_set_in_memory import DataSetInMem, load_from_file
 from qcodes.dataset.data_set_protocol import DataSetType
 from qcodes.dataset.sqlite.connection import ConnectionPlus, atomic_transaction
 from qcodes.station import Station
@@ -449,6 +450,69 @@ def test_load_from_db(meas_with_registered_param, DMM, DAC, tmp_path) -> None:
     compare_datasets(ds, loaded_ds)
 
 
+def test_load_from_file(meas_with_registered_param, DMM, DAC, tmp_path) -> None:
+    qcodes.config["dataset"]["export_prefix"] = "my_export_prefix"
+    qcodes.config["dataset"]["export_type"] = "netcdf"
+
+    Station(DAC, DMM)
+    with meas_with_registered_param.run(
+        dataset_class=DataSetType.DataSetInMem
+    ) as datasaver:
+        for set_v in np.linspace(0, 25, 10):
+            DAC.ch1.set(set_v)
+            get_v = DMM.v1()
+            datasaver.add_result((DAC.ch1, set_v), (DMM.v1, get_v))
+
+    ds: DataSetInMem = datasaver.dataset
+    ds.add_metadata("foo", "bar")
+    ds.export(path=tmp_path)
+    ds.add_metadata("metadata_added_after_export", "42")
+
+    export_file_path = ds.export_info.export_paths.get("nc")
+    assert export_file_path is not None
+    loaded_ds: DataSetInMem = load_from_file(export_file_path)
+
+    assert isinstance(loaded_ds, DataSetInMem)
+    assert loaded_ds.snapshot == ds.snapshot
+    assert loaded_ds.export_info == ds.export_info
+    assert loaded_ds.metadata == ds.metadata
+
+    assert "export_info" in loaded_ds.metadata.keys()
+    assert "metadata_added_after_export" in loaded_ds.metadata.keys()
+    assert loaded_ds.metadata["foo"] == "bar"
+    assert loaded_ds.metadata["metadata_added_after_export"] == "42"
+
+    compare_datasets(ds, loaded_ds)
+
+
+def test_load_from_file_by_id(meas_with_registered_param, DMM, DAC, tmp_path) -> None:
+    qcodes.config["dataset"]["export_prefix"] = "my_export_prefix"
+    qcodes.config["dataset"]["export_type"] = "netcdf"
+    qcodes.config["dataset"]["load_from_exported_file"] = True
+    assert qcodes.config.dataset.load_from_exported_file is True
+
+    Station(DAC, DMM)
+    with meas_with_registered_param.run() as datasaver:
+        for set_v in np.linspace(0, 25, 10):
+            DAC.ch1.set(set_v)
+            get_v = DMM.v1()
+            datasaver.add_result((DAC.ch1, set_v), (DMM.v1, get_v))
+
+    ds: DataSetInMem = datasaver.dataset
+    assert not isinstance(ds, DataSetInMem)
+    ds.export(path=tmp_path)
+
+    # Load from file
+    loaded_ds_from_file = load_by_id(ds.run_id)
+    assert isinstance(loaded_ds_from_file, DataSetInMem)
+
+    # Load from db
+    qcodes.config["dataset"]["load_from_exported_file"] = False
+    assert qcodes.config.dataset.load_from_exported_file is False
+    loaded_ds_from_db = load_by_id(ds.run_id)
+    assert not isinstance(loaded_ds_from_db, DataSetInMem)
+
+
 def test_load_from_netcdf_legacy_version(non_created_db) -> None:
     # Qcodes 0.26 exported netcdf files did not contain
     # the parent dataset links and used a different engine to write data