From c31216a336b39caa80a26e2e0799cfb612069328 Mon Sep 17 00:00:00 2001 From: Mikhail Astafev Date: Mon, 27 Jul 2026 16:29:57 +0200 Subject: [PATCH] Refactor split raw data storage into DataSetInSeparateSqliteDbFile subclass Move the per-dataset SQLite storage behaviour out of the config-flag conditionals in DataSet and into a dedicated DataSetInSeparateSqliteDbFile subclass, selected at the factory boundary. - Base DataSet gains generic results-backend hooks (no-ops by default): _setup_results_backend_on_load/_on_new_run/_on_start, plus the class attribute _creates_results_table_in_main_db. The generic 'results may live on a separate connection' routing (_data_conn, get_parameter_data, add_results) stays in the base and activates only when _raw_data_conn is populated. - DataSetInSeparateSqliteDbFile(DataSet) overrides only those hooks to create/ connect the per-dataset file and record raw_data_db_path. - Class selection moves to the factory boundary: new_data_set() and Measurement instantiate the subclass when the feature is enabled; a new _load_dataset_from_run_id() helper picks the subclass for existing runs that record a raw_data_db_path. Routed through experiment_container.data_set, extract_runs, and _get_datasetprotocol_from_guid. - Note: direct DataSet(run_id=...) construction on a split run no longer auto-detects the raw backend; use _load_dataset_from_run_id / load_by_id. Updated make_shadow_dataset accordingly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/qcodes/dataset/data_set.py | 203 ++++++++++++++------ src/qcodes/dataset/database_extract_runs.py | 8 +- src/qcodes/dataset/experiment_container.py | 9 +- src/qcodes/dataset/measurements.py | 16 +- tests/dataset/test_dataset_basic.py | 6 +- tests/dataset/test_raw_data_storage.py | 29 ++- 6 files changed, 203 insertions(+), 68 deletions(-) diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index d9dea16293c9..68a8ef27ec31 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -234,6 +234,33 @@ class DataSet(BaseDataSet): ) background_sleep_time = 1e-3 + #: Whether this dataset stores its results table in the main database. + #: Subclasses that keep raw data in a separate backend (e.g. + #: :class:`DataSetInSeparateSqliteDbFile`) set this to ``False`` so that no + #: results table is created in the main database. + _creates_results_table_in_main_db: bool = True + + def _setup_results_backend_on_load(self, *, read_only: bool) -> None: + """Hook: set up the results backend when loading an existing run. + + No-op for a plain :class:`DataSet` (results live in the main database). + Subclasses use this to e.g. connect to a per-dataset raw data file. + """ + + def _setup_results_backend_on_new_run(self) -> None: + """Hook: record results-backend bookkeeping for a newly created run. + + No-op for a plain :class:`DataSet`. Subclasses use this to e.g. record + the location of the per-dataset raw data file in the ``runs`` table. + """ + + def _setup_results_backend_on_start(self) -> None: + """Hook: create/open the results backend when the run is started. + + No-op for a plain :class:`DataSet`. Subclasses use this to e.g. create + the per-dataset raw data file and its results table. + """ + def __init__( self, path_to_db: str | None = None, @@ -324,24 +351,10 @@ def __init__( self._export_info = ExportInfo.from_str( self.metadata.get("export_info", "") ) - # If this dataset was saved with raw data in a separate db, - # re-open that connection for reads. The path is stored in a - # dedicated runs-table column, not in the user-facing metadata. - raw_db_path = get_raw_data_db_path_for_run(self.conn, self.run_id) - self._raw_data_db_path = raw_db_path - if raw_db_path is not None: - if Path(raw_db_path).is_file(): - self._raw_data_conn = connect_to_raw_data_db( - raw_db_path, read_only=read_only - ) - elif self._started: - raise FileNotFoundError( - f"Raw data file for dataset {self.guid} not found at " - f"'{raw_db_path}'. The per-dataset SQLite file may " - f"have been moved or deleted." - ) - # else: the dataset was never started, so the raw data file has - # not been created yet - there is simply no data to connect to. + # Delegate any results-backend setup (e.g. connecting to a + # per-dataset raw data file) to subclasses. For a plain DataSet + # this is a no-op since results live in the main database. + self._setup_results_backend_on_load(read_only=read_only) else: # Actually perform all the side effects needed for the creation # of a new dataset. Note that a dataset is created (in the DB) @@ -350,9 +363,10 @@ def __init__( if exp_id is None: exp_id = get_default_experiment_id(self.conn) name = name or "dataset" - # When raw data is stored in a separate backend (e.g. a per-dataset - # SQLite file), no results table is created in the main database - - # only the run metadata is kept there. This mirrors how + # Subclasses that store results outside the main database (e.g. in + # a per-dataset SQLite file) set ``_creates_results_table_in_main_db`` + # to False, so no results table is created here - only the run + # metadata is kept in the main database. This mirrors how # ``DataSetInMem`` records runs without a results table. _, run_id, __ = create_run( self.conn, @@ -362,7 +376,7 @@ def __init__( parameters=None, values=values, metadata=metadata, - create_run_table=not is_raw_data_storage_enabled(), + create_run_table=self._creates_results_table_in_main_db, ) # this is really the UUID (an ever increasing count in the db) self._run_id = run_id @@ -382,17 +396,10 @@ def __init__( self._parent_dataset_links = [] self._export_info = ExportInfo({}) - if is_raw_data_storage_enabled(): - # Record the raw-data backend location up front. This marks the - # run as a split-storage dataset (so it can be told apart from a - # ``DataSetInMem`` run, which also has no results table) even - # before it is started and before the raw data file is created. - # The path is stored in a dedicated column, not in the - # user-facing metadata. - raw_path_str = str(get_raw_data_db_path(self.guid)) - self._raw_data_db_path = raw_path_str - with atomic(self.conn) as aconn: - set_raw_data_db_path_for_run(aconn, self.run_id, raw_path_str) + # Let subclasses record any results-backend bookkeeping for the new + # run (e.g. the location of the per-dataset raw data file). No-op + # for a plain DataSet. + self._setup_results_backend_on_new_run() assert self.path_to_db is not None if _WRITERS.get(self.path_to_db) is None: queue: Queue[Any] = Queue() @@ -775,37 +782,21 @@ def _perform_start_actions(self, start_bg_writer: bool) -> None: Perform the actions that must take place once the run has been started """ paramspecs = new_to_old(self._rundescriber.interdeps).paramspecs - raw_data_enabled = is_raw_data_storage_enabled() for spec in paramspecs: add_parameter( spec, conn=self.conn, run_id=self.run_id, - # The results table only lives in the main database when raw - # data storage is disabled; with it enabled the parameter - # columns are created in the per-dataset raw data file below. - insert_into_results_table=not raw_data_enabled, + # The results table only lives in the main database for a plain + # DataSet; subclasses that use a separate backend create the + # parameter columns there in _setup_results_backend_on_start. + insert_into_results_table=self._creates_results_table_in_main_db, ) - # When raw data split is enabled, create a per-dataset SQLite file - # for results data with the full results table. - if raw_data_enabled: - # The raw-data path was already recorded at dataset creation time; - # reuse it so both locations stay in sync. - raw_path_str = self._raw_data_db_path or str( - get_raw_data_db_path(self.guid) - ) - raw_db_path = Path(raw_path_str) - self._raw_data_conn = create_raw_data_db( - raw_db_path, - self.table_name, - self._rundescriber.interdeps.paramspecs, - ) - if self._raw_data_db_path != raw_path_str: - self._raw_data_db_path = raw_path_str - with atomic(self.conn) as aconn: - set_raw_data_db_path_for_run(aconn, self.run_id, raw_path_str) + # Let subclasses create/open their results backend (e.g. a per-dataset + # SQLite file with the full results table). No-op for a plain DataSet. + self._setup_results_backend_on_start() desc_str = serial.to_json_for_storage(self.description) @@ -1779,6 +1770,97 @@ def _estimate_ds_size(self) -> float: return row_size * len(self) / 1024 / 1024 +class DataSetInSeparateSqliteDbFile(DataSet): + """A :class:`DataSet` whose raw measurement data is stored in an + individual, per-dataset SQLite file while all metadata remains in the main + database. + + This is the concrete backend behind the ``dataset.raw_data_to_separate_db`` + config option. The per-dataset file is named ``.db`` and lives in the + folder given by ``dataset.raw_data_path``. Only the results-backend setup + differs from a plain :class:`DataSet`: the generic "results live on a + separate connection" routing (``_data_conn``, ``get_parameter_data``, + ``add_results``) is provided by the base class and activated once + ``_raw_data_conn`` is populated here. + + The path to the per-dataset file is recorded in a dedicated + ``raw_data_db_path`` column of the ``runs`` table (not in the user-facing + metadata), which is also how such runs are recognised when loading. + """ + + _creates_results_table_in_main_db = False + + def _setup_results_backend_on_load(self, *, read_only: bool) -> None: + # Re-open the per-dataset raw data file for reads. The path is stored + # in a dedicated runs-table column, not in the user-facing metadata. + raw_db_path = get_raw_data_db_path_for_run(self.conn, self.run_id) + self._raw_data_db_path = raw_db_path + if raw_db_path is not None: + if Path(raw_db_path).is_file(): + self._raw_data_conn = connect_to_raw_data_db( + raw_db_path, read_only=read_only + ) + elif self._started: + raise FileNotFoundError( + f"Raw data file for dataset {self.guid} not found at " + f"'{raw_db_path}'. The per-dataset SQLite file may " + f"have been moved or deleted." + ) + # else: the dataset was never started, so the raw data file has + # not been created yet - there is simply no data to connect to. + + def _setup_results_backend_on_new_run(self) -> None: + # Record the raw-data backend location up front. This marks the run as + # a split-storage dataset (so it can be told apart from a DataSetInMem + # run, which also has no results table) even before it is started and + # before the raw data file is created. + raw_path_str = str(get_raw_data_db_path(self.guid)) + self._raw_data_db_path = raw_path_str + with atomic(self.conn) as aconn: + set_raw_data_db_path_for_run(aconn, self.run_id, raw_path_str) + + def _setup_results_backend_on_start(self) -> None: + # Create the per-dataset SQLite file with the full results table. The + # path was already recorded at creation time; reuse it so both stay in + # sync. + raw_path_str = self._raw_data_db_path or str(get_raw_data_db_path(self.guid)) + raw_db_path = Path(raw_path_str) + self._raw_data_conn = create_raw_data_db( + raw_db_path, + self.table_name, + self._rundescriber.interdeps.paramspecs, + ) + if self._raw_data_db_path != raw_path_str: + self._raw_data_db_path = raw_path_str + with atomic(self.conn) as aconn: + set_raw_data_db_path_for_run(aconn, self.run_id, raw_path_str) + + +def _load_dataset_from_run_id( + conn: AtomicConnection | None = None, + run_id: int | None = None, + *, + path_to_db: str | None = None, + read_only: bool = False, +) -> DataSet: + """Instantiate the appropriate DataSet class for an existing run. + + Returns a :class:`DataSetInSeparateSqliteDbFile` when the run records a + ``raw_data_db_path`` (raw data stored in a separate SQLite file), otherwise + a plain :class:`DataSet`. This is the recommended way to load an existing + run into a :class:`DataSet` object when not going through the public + ``load_by_*`` functions, since it selects the correct subclass. + """ + if run_id is None: + raise ValueError("run_id must be provided") + conn = conn_from_dbpath_or_conn(conn, path_to_db, read_only=read_only) + if get_raw_data_db_path_for_run(conn, run_id) is not None: + return DataSetInSeparateSqliteDbFile( + run_id=run_id, conn=conn, read_only=read_only + ) + return DataSet(run_id=run_id, conn=conn, read_only=read_only) + + # public api def load_by_run_spec( *, @@ -2094,7 +2176,7 @@ def _get_datasetprotocol_from_guid( # raw_data_db_path; anything else without a results table is an # in-memory (netcdf-backed) dataset. elif get_raw_data_db_path_for_run(conn, run_id) is not None: - d = DataSet(conn=conn, run_id=run_id) + d = DataSetInSeparateSqliteDbFile(conn=conn, run_id=run_id) else: d = DataSetInMem._load_from_db(conn=conn, guid=guid) @@ -2139,7 +2221,12 @@ def new_data_set( """ # note that passing `conn` is a secret feature that is unfortunately used # in `Runner` to pass a connection from an existing `Experiment`. - d = DataSet( + # Choose the concrete class based on whether split raw data storage is + # enabled in the config. + dataset_class = ( + DataSetInSeparateSqliteDbFile if is_raw_data_storage_enabled() else DataSet + ) + d = dataset_class( path_to_db=None, run_id=None, conn=conn, diff --git a/src/qcodes/dataset/database_extract_runs.py b/src/qcodes/dataset/database_extract_runs.py index deb74982a5bc..050ecdbf6f7c 100644 --- a/src/qcodes/dataset/database_extract_runs.py +++ b/src/qcodes/dataset/database_extract_runs.py @@ -11,7 +11,7 @@ from opentelemetry import trace from tqdm.auto import tqdm -from qcodes.dataset.data_set import DataSet, load_by_id +from qcodes.dataset.data_set import DataSet, _load_dataset_from_run_id, load_by_id from qcodes.dataset.data_set_in_memory import load_from_netcdf from qcodes.dataset.dataset_helpers import _add_run_to_runs_table from qcodes.dataset.experiment_container import _create_exp_if_needed @@ -132,7 +132,9 @@ def extract_runs_into_db( # Finally insert the runs for run_id in run_ids: _extract_single_dataset_into_db( - DataSet(run_id=run_id, conn=source_conn), target_conn, target_exp_id + _load_dataset_from_run_id(source_conn, run_id), + target_conn, + target_exp_id, ) finally: source_conn.close() @@ -378,7 +380,7 @@ def _copy_dataset_as_is( target_exp_id: int, ) -> Literal["copied_as_is", "failed"]: try: - dataset_obj = DataSet(run_id=dataset.run_id, conn=source_conn) + dataset_obj = _load_dataset_from_run_id(source_conn, dataset.run_id) with atomic(target_conn) as target_conn_atomic: _extract_single_dataset_into_db( dataset_obj, target_conn_atomic, target_exp_id diff --git a/src/qcodes/dataset/experiment_container.py b/src/qcodes/dataset/experiment_container.py index ad1f658e4a3c..5fd714750498 100644 --- a/src/qcodes/dataset/experiment_container.py +++ b/src/qcodes/dataset/experiment_container.py @@ -5,7 +5,12 @@ from typing import TYPE_CHECKING, Any from warnings import warn -from qcodes.dataset.data_set import DataSet, load_by_id, new_data_set +from qcodes.dataset.data_set import ( + DataSet, + _load_dataset_from_run_id, + load_by_id, + new_data_set, +) from qcodes.dataset.experiment_settings import _set_default_experiment_id from qcodes.dataset.sqlite.connection import AtomicConnection, path_to_dbfile from qcodes.dataset.sqlite.database import ( @@ -168,7 +173,7 @@ def data_set(self, counter: int) -> DataSet: """ run_id = get_runid_from_expid_and_counter(self.conn, self.exp_id, counter) - return DataSet(run_id=run_id, conn=self.conn) + return _load_dataset_from_run_id(self.conn, run_id) def data_sets(self) -> list[DataSetProtocol]: """Get all the datasets of this experiment""" diff --git a/src/qcodes/dataset/measurements.py b/src/qcodes/dataset/measurements.py index 3437bbb8a7c6..88e7099a8cf1 100644 --- a/src/qcodes/dataset/measurements.py +++ b/src/qcodes/dataset/measurements.py @@ -26,7 +26,12 @@ import qcodes as qc import qcodes.validators as vals -from qcodes.dataset.data_set import DataSet, load_by_guid +from qcodes.dataset._raw_data_storage import is_raw_data_storage_enabled +from qcodes.dataset.data_set import ( + DataSet, + DataSetInSeparateSqliteDbFile, + load_by_guid, +) from qcodes.dataset.data_set_in_memory import DataSetInMem from qcodes.dataset.data_set_protocol import ( DataSetProtocol, @@ -636,7 +641,14 @@ def __enter__(self) -> DataSaver: conn = None if self._dataset_class is DataSetType.DataSet: - self.ds = DataSet( + # When split raw data storage is enabled, use the subclass that + # writes results to a per-dataset SQLite file. + dataset_cls = ( + DataSetInSeparateSqliteDbFile + if is_raw_data_storage_enabled() + else DataSet + ) + self.ds = dataset_cls( name=self.name, exp_id=exp_id, conn=conn, diff --git a/tests/dataset/test_dataset_basic.py b/tests/dataset/test_dataset_basic.py index e0e45bb3b73d..33530da505e6 100644 --- a/tests/dataset/test_dataset_basic.py +++ b/tests/dataset/test_dataset_basic.py @@ -20,7 +20,7 @@ new_data_set, new_experiment, ) -from qcodes.dataset.data_set import DataSet +from qcodes.dataset.data_set import DataSet, _load_dataset_from_run_id from qcodes.dataset.data_set_protocol import CompletedError from qcodes.dataset.descriptions.dependencies import InterDependencies_ from qcodes.dataset.descriptions.rundescriber import RunDescriber @@ -54,7 +54,9 @@ def make_shadow_dataset(dataset: DataSet): database file. """ - return DataSet(path_to_db=dataset.path_to_db, run_id=dataset.run_id) + return _load_dataset_from_run_id( + path_to_db=dataset.path_to_db, run_id=dataset.run_id + ) @pytest.mark.usefixtures("experiment") diff --git a/tests/dataset/test_raw_data_storage.py b/tests/dataset/test_raw_data_storage.py index f7d4c3396580..84d87c59ef8b 100644 --- a/tests/dataset/test_raw_data_storage.py +++ b/tests/dataset/test_raw_data_storage.py @@ -30,7 +30,11 @@ purge_orphaned_datasets, update_raw_data_paths, ) -from qcodes.dataset.data_set import DataSet, load_by_id +from qcodes.dataset.data_set import ( + DataSet, + DataSetInSeparateSqliteDbFile, + load_by_id, +) from qcodes.dataset.database_extract_runs import ( export_datasets_and_create_metadata_db, extract_runs_into_db, @@ -192,6 +196,22 @@ def test_raw_data_conn_is_set(self) -> None: assert ds._raw_data_conn is not None self._close_ds(ds) + def test_new_data_set_returns_separate_file_subclass(self) -> None: + """When split is enabled, new_data_set returns the dedicated subclass.""" + ds = new_data_set("test-split") + assert isinstance(ds, DataSetInSeparateSqliteDbFile) + self._close_ds(ds) + + def test_loaded_dataset_is_separate_file_subclass(self) -> None: + """Loading a split dataset returns the dedicated subclass.""" + ds, _ = self._make_dataset_with_data(n_rows=3) + run_id = ds.run_id + self._close_ds(ds) + + loaded = load_by_id(run_id) + assert isinstance(loaded, DataSetInSeparateSqliteDbFile) + self._close_ds(loaded) + def test_raw_data_file_created(self, tmp_path: Path) -> None: """A per-dataset SQLite file should be created.""" ds, _ = self._make_dataset_with_data() @@ -388,6 +408,13 @@ def test_missing_raw_data_file_raises(self, tmp_path: Path) -> None: @pytest.mark.usefixtures("experiment") class TestDataSetWithoutSplitRawData: + def test_new_data_set_returns_plain_dataset(self) -> None: + """When split is disabled, new_data_set returns a plain DataSet.""" + ds = new_data_set("test-no-split") + assert isinstance(ds, DataSet) + assert not isinstance(ds, DataSetInSeparateSqliteDbFile) + ds.conn.close() + def test_raw_data_conn_is_none(self) -> None: """When split is disabled, _raw_data_conn should be None.""" ds = new_data_set("test-no-split")