diff --git a/plottr/apps/inspectr.py b/plottr/apps/inspectr.py index 10d81c18..39d192a3 100644 --- a/plottr/apps/inspectr.py +++ b/plottr/apps/inspectr.py @@ -452,7 +452,8 @@ def loadDB(self) -> None: overview: Optional[Dict[int, Any]] = None if self.use_fast_sql: try: - overview = get_db_overview(self.path) + overview = get_db_overview(self.path, + extra_columns=['inspectr_tag']) except Exception as e: LOGGER.warning(f"Fast SQL overview failed, falling back to " f"qcodes API: {e}") diff --git a/plottr/data/_qcodes_db_overview.py b/plottr/data/_qcodes_db_overview.py new file mode 100644 index 00000000..68433b34 --- /dev/null +++ b/plottr/data/_qcodes_db_overview.py @@ -0,0 +1,267 @@ +""" +Local (vendored) copy of the QCoDeS ``get_db_overview`` implementation. + +This module is a private fallback used by :mod:`plottr.data.qcodes_db_overview` +for QCoDeS versions that do not yet expose ``qcodes.dataset.get_db_overview``. +Its content is kept as an exact copy of the upstream QCoDeS implementation, +except that the (private) QCoDeS SQLite imports are done lazily inside +:func:`get_db_overview` rather than at module import time. + +:func:`get_db_overview` issues a single ``JOIN`` query against the ``runs`` and +``experiments`` tables to collect run metadata (experiment/sample names, time +stamps, record counts, guids, ...) without instantiating a ``DataSet`` object +per run. This avoids the expensive ``experiments()`` + ``data_sets()`` +enumeration and makes it possible to list the contents of databases with many +thousands of runs almost instantly. +""" + +from __future__ import annotations + +import datetime +import json +import logging +from contextlib import closing, nullcontext +from typing import TYPE_CHECKING + +from typing_extensions import TypedDict + +if TYPE_CHECKING: + from collections.abc import Sequence + from pathlib import Path + + from qcodes.dataset.sqlite.connection import AtomicConnection + +log = logging.getLogger(__name__) + + +class RunOverviewDict(TypedDict): + """ + Lightweight overview of a single run. + + Contains only cheap-to-query metadata: no snapshot, no data and no full + ``DataSet`` object. Extra ad-hoc metadata columns requested via the + ``extra_columns`` argument of :func:`get_db_overview` are added to the + dictionary under their column name in addition to the keys documented here. + """ + + #: ``run_id`` of the run. + run_id: int + #: Name of the experiment the run belongs to. + experiment: str + #: Sample name of the experiment the run belongs to. + sample: str + #: Name of the run. + name: str + #: Local date the run was started, formatted as ``YYYY-MM-DD`` (empty + #: string if unknown). + started_date: str + #: Local time the run was started, formatted as ``HH:MM:SS`` (empty string + #: if unknown). + started_time: str + #: Local date the run was completed, formatted as ``YYYY-MM-DD`` (empty + #: string if the run has not completed). + completed_date: str + #: Local time the run was completed, formatted as ``HH:MM:SS`` (empty + #: string if the run has not completed). + completed_time: str + #: Best-effort number of data points in the run, see :func:`get_db_overview`. + records: int + #: guid of the run. + guid: str + + +def _format_timestamp(ts: float | None) -> tuple[str, str]: + """ + Convert a unix timestamp into ``(date, time)`` strings in local time. + + Returns a pair of empty strings if the timestamp is missing or invalid. + """ + if ts is None or ts == 0: + return "", "" + try: + dt = datetime.datetime.fromtimestamp(ts) + except (OSError, ValueError, OverflowError): + return "", "" + return dt.strftime("%Y-%m-%d"), dt.strftime("%H:%M:%S") + + +def _records_from_run_description(run_description_json: str | None) -> int: + """ + Extract a data-point count from the ``shapes`` field of a run description. + + A QCoDeS run description may contain a ``shapes`` mapping from dependent + parameter names to their shape tuples. The count returned here is the sum + over all dependent parameters of the product of their shape dimensions. + Returns ``0`` if the run description is missing, cannot be parsed or does + not contain shape information. + """ + if not run_description_json: + return 0 + try: + desc = json.loads(run_description_json) + except (json.JSONDecodeError, TypeError): + return 0 + shapes = desc.get("shapes") if isinstance(desc, dict) else None + if not shapes: + return 0 + total = 0 + for shape in shapes.values(): + if isinstance(shape, (list, tuple)) and len(shape) > 0: + n = 1 + for dim in shape: + n *= dim + total += n + return total + + +def get_db_overview( + path_to_db: str | Path | None = None, + *, + conn: AtomicConnection | None = None, + start_run_id: int = 0, + extra_columns: Sequence[str] | None = None, +) -> dict[int, RunOverviewDict]: + """ + Get a lightweight overview of the runs in a QCoDeS database. + + This uses a single SQL ``JOIN`` query on the ``runs`` and ``experiments`` + tables to fetch run metadata, avoiding the much more expensive + ``experiments()`` + ``data_sets()`` enumeration that instantiates a + ``DataSet`` object per run. It is therefore well suited for listing the + contents of databases with many runs. The (potentially large) snapshot of + each run is deliberately not read, as it would slow down building the + overview significantly. + + The reported number of ``records`` is a best-effort estimate of the number + of data points in a run and may be less precise than + ``DataSet.number_of_results``: + + * For completed runs the shape information stored in the run description is + preferred (it is the authoritative final count), falling back to the + number of rows in the results table. + * For runs that are still in progress the number of rows in the results + table is preferred (it grows as data is added), falling back to the run + description shapes. + * If neither is available the count is reported as ``0`` (unknown). + + Only one of ``path_to_db`` and ``conn`` should be supplied. If a + ``path_to_db`` is given the database is opened in read-only mode and the + connection is closed again before returning. + + Args: + path_to_db: Path to the database file. Opened read-only if given. + conn: An existing connection to use instead of ``path_to_db``. It is + left open by this function. + start_run_id: Only return runs whose ``run_id`` is strictly greater + than this value. Use ``0`` (the default) to get all runs, or pass + the last known ``run_id`` to fetch only newly added runs. + extra_columns: Names of additional ``runs``-table columns to include in + each :class:`RunOverviewDict`. Columns that do not exist in the + ``runs`` table of the given database are silently skipped. This is + useful for reading ad-hoc metadata columns added via + ``DataSet.add_metadata``. + + Returns: + A dictionary mapping ``run_id`` to a :class:`RunOverviewDict`. + + """ + # These are (technically private) QCoDeS SQLite helpers; import them lazily + # so that importing this module does not depend on QCoDeS internals. + import sqlite3 + + from qcodes.dataset.sqlite.database import conn_from_dbpath_or_conn + from qcodes.dataset.sqlite.query_helpers import is_column_in_table + + overview: dict[int, RunOverviewDict] = {} + + created_conn = conn is None + connection = conn_from_dbpath_or_conn( + conn=conn, path_to_db=path_to_db, read_only=True + ) + manager = closing(connection) if created_conn else nullcontext(connection) + + with manager as c: + valid_extra_columns = [ + col for col in (extra_columns or []) if is_column_in_table(c, "runs", col) + ] + extra_select = "".join(f", r.{col}" for col in valid_extra_columns) + + # ``run_description`` is queried to derive the record count for + # completed runs; the (potentially large) snapshot is deliberately + # excluded. + query = f""" + SELECT r.run_id, e.name, e.sample_name, r.name, + r.run_timestamp, r.completed_timestamp, + r.guid, r.result_table_name, + r.run_description{extra_select} + FROM runs r + JOIN experiments e ON r.exp_id = e.exp_id + WHERE r.run_id > ? + ORDER BY r.run_id + """ + + try: + rows = c.execute(query, (start_run_id,)).fetchall() + except sqlite3.Error as e: + log.warning("Could not query database overview: %s", e) + return overview + + # ``result_counter`` in the runs table is the run's ordinal within its + # experiment, not a data-point count, so it is not usable here. For the + # ``array`` paramtype a single INSERT can also contain many data points. + # The real number of data points is therefore the row count of the + # results table, queried separately. + result_tables = {row[7] for row in rows if row[7]} + row_counts: dict[str, int] = {} + for table in result_tables: + try: + (count,) = c.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone() + except sqlite3.Error: + continue # results table may not exist (yet) + row_counts[table] = count + + n_fixed = 9 # number of columns selected before ``extra_columns`` + for row in rows: + run_id = row[0] + started_date, started_time = _format_timestamp(row[4]) + completed_date, completed_time = _format_timestamp(row[5]) + result_table = row[7] or "" + is_completed = row[5] is not None and row[5] != 0 + + # The record count is a best-effort data-point count. For completed + # runs the run-description shapes are the authoritative final count; + # for in-progress runs the live results-table row count is preferred + # as it grows while data is added. ``0`` means "unknown". + if is_completed: + records = _records_from_run_description(row[8]) + if records == 0: + records = row_counts.get(result_table, 0) + else: + records = row_counts.get(result_table, 0) + if records == 0: + records = _records_from_run_description(row[8]) + + entry: RunOverviewDict = { + "run_id": run_id, + "experiment": row[1] or "", + "sample": row[2] or "", + "name": row[3] or "", + "started_date": started_date, + "started_time": started_time, + "completed_date": completed_date, + "completed_time": completed_time, + "records": records, + "guid": row[6] or "", + } + if valid_extra_columns: + extra = { + col: row[n_fixed + i] for i, col in enumerate(valid_extra_columns) + } + # The keys of ``extra`` are only known at runtime (they are the + # user-supplied ``extra_columns``), so they cannot be part of + # the closed ``RunOverviewDict`` definition. + entry.update(extra) # type: ignore[typeddict-item] + + overview[run_id] = entry + + return overview diff --git a/plottr/data/qcodes_db_overview.py b/plottr/data/qcodes_db_overview.py index 6e421a25..95422026 100644 --- a/plottr/data/qcodes_db_overview.py +++ b/plottr/data/qcodes_db_overview.py @@ -1,200 +1,19 @@ """ plottr.data.qcodes_db_overview — Fast database overview queries. -This module provides optimized functions for listing QCoDeS dataset metadata -without loading full DataSet objects. It uses direct SQLite queries on the -QCoDeS database schema, avoiding the expensive experiments()/data_sets() -enumeration. - -**Intended for eventual contribution to QCoDeS.** The queries here rely on the -stable QCoDeS database schema (runs + experiments tables) which has not changed -across many QCoDeS versions. +This module exposes :func:`get_db_overview` (and :class:`RunOverviewDict`), +which provide a fast, lightweight listing of the runs in a QCoDeS database +without loading full ``DataSet`` objects. + +The implementation has been contributed to QCoDeS. When a QCoDeS version that +exposes ``qcodes.dataset.get_db_overview`` is installed, that implementation is +used; otherwise the vendored copy in :mod:`plottr.data._qcodes_db_overview` +(an exact copy of the upstream implementation) is used as a fallback. """ -import json -import sys -import time -import logging -from contextlib import closing -from typing import Dict, Optional, Tuple - -from typing_extensions import TypedDict - -from qcodes.dataset.sqlite.database import conn_from_dbpath_or_conn - -logger = logging.getLogger(__name__) - - -def _records_from_run_description(run_description_json: Optional[str]) -> int: - """Extract record count from run_description shapes field. - - QCoDeS run_description may contain a ``shapes`` dict mapping dependent - parameter names to their shape tuples. The total data-point count is the - product of shape dimensions summed across all parameter trees — matching - the semantics of ``DataSet.number_of_results``. - """ - if not run_description_json: - return 0 - try: - desc = json.loads(run_description_json) - shapes = desc.get('shapes') - if not shapes: - return 0 - total = 0 - for shape in shapes.values(): - if isinstance(shape, (list, tuple)) and len(shape) > 0: - n = 1 - for dim in shape: - n *= dim - # Each parameter tree contributes n_values * n_params_in_tree - # But shapes only has dependent params, and number_of_results - # counts all values including axes. For display purposes, - # the product of the shape is the most useful number. - total += n - return total - except (json.JSONDecodeError, TypeError, KeyError): - return 0 - - -class RunOverviewDict(TypedDict): - """Lightweight run overview — no snapshot, no data, no full DataSet.""" - run_id: int - experiment: str - sample: str - name: str - started_date: str - started_time: str - completed_date: str - completed_time: str - records: int - guid: str - inspectr_tag: str - - -def _format_timestamp(ts: Optional[float]) -> Tuple[str, str]: - """Convert a unix timestamp float to (date, time) strings.""" - if ts is None or ts == 0: - return '', '' - try: - t = time.localtime(ts) - return time.strftime('%Y-%m-%d', t), time.strftime('%H:%M:%S', t) - except (OSError, ValueError, OverflowError): - return '', '' - - -def get_db_overview(db_path: str, - start_run_id: int = 0, - ) -> Dict[int, RunOverviewDict]: - """Get a lightweight overview of all runs in a QCoDeS database. - - Uses a single SQL JOIN query to fetch run metadata from the ``runs`` and - ``experiments`` tables, avoiding the expensive ``experiments()`` + - ``data_sets()`` enumeration that QCoDeS uses internally. - - For a database with 1500 runs, this completes in ~10ms vs 15+ minutes - with the standard QCoDeS API. - - :param db_path: path to the .db file. - :param start_run_id: only return runs with run_id > start_run_id. - Use 0 to get all runs. Pass the last known run_id for incremental - refresh. - :returns: dict mapping run_id to RunOverviewDict. - """ - overview: Dict[int, RunOverviewDict] = {} - - if sys.version_info >= (3, 11): - conn = conn_from_dbpath_or_conn(conn=None, path_to_db=db_path, read_only=True) - else: - conn = conn_from_dbpath_or_conn(conn=None, path_to_db=db_path) - - with closing(conn) as c: - # Check which ad-hoc metadata columns exist in the runs table. - # QCoDeS stores metadata added via ds.add_metadata() as extra columns. - try: - col_info = c.execute('PRAGMA table_info(runs)').fetchall() - col_names = {col[1] for col in col_info} - except Exception: - col_names = set() - - has_inspectr_tag = 'inspectr_tag' in col_names - - # Build query: include inspectr_tag column if it exists. - # Includes run_description to extract shape info for record count. - # Deliberately excludes snapshot (large blob). - tag_col = ", r.inspectr_tag" if has_inspectr_tag else "" - query = f""" - SELECT r.run_id, e.name, e.sample_name, r.name, - r.run_timestamp, r.completed_timestamp, - r.result_counter, r.guid, r.result_table_name, - r.run_description{tag_col} - FROM runs r - JOIN experiments e ON r.exp_id = e.exp_id - WHERE r.run_id > ? - ORDER BY r.run_id - """ - - try: - rows = c.execute(query, (start_run_id,)).fetchall() - except Exception as e: - logger.warning(f"Could not query database overview: {e}") - return overview - - # Build a map of actual row counts from each results table. - # result_counter in the runs table counts INSERT calls, not data points. - # For array paramtype one INSERT can contain thousands of data points, - # so result_counter can be much smaller than the real data point count. - results_tables: set[str] = set() - for row in rows: - tbl = row[8] # result_table_name - if tbl: - results_tables.add(tbl) - row_counts: dict[str, int] = {} - for tbl in results_tables: - try: - cnt = c.execute( - f'SELECT COUNT(*) FROM "{tbl}"' - ).fetchone() - row_counts[tbl] = cnt[0] if cnt else 0 - except Exception: - pass # table may not exist - - tag_col_idx = 10 if has_inspectr_tag else -1 - for row in rows: - run_id = row[0] - started_date, started_time = _format_timestamp(row[4]) - completed_date, completed_time = _format_timestamp(row[5]) - tag = row[tag_col_idx] if tag_col_idx > 0 and len(row) > tag_col_idx and row[tag_col_idx] else '' - result_table = row[8] or '' - is_completed = row[5] is not None and row[5] != 0 - - # Determine record count. - # For completed datasets: prefer shape metadata (authoritative - # final count) over results table rows. - # For active (incomplete) datasets: prefer results table rows - # (live count that grows as data is added). - # Fall back to result_counter if nothing else is available. - if is_completed: - records = _records_from_run_description(row[9]) - if records == 0: - records = row_counts.get(result_table, 0) - else: - records = row_counts.get(result_table, 0) - if records == 0: - records = _records_from_run_description(row[9]) - if records == 0: - records = row[6] or 0 - overview[run_id] = RunOverviewDict( - run_id=run_id, - experiment=row[1] or '', - sample=row[2] or '', - name=row[3] or '', - started_date=started_date, - started_time=started_time, - completed_date=completed_date, - completed_time=completed_time, - records=records, - guid=row[7] or '', - inspectr_tag=tag, - ) +try: + from qcodes.dataset import RunOverviewDict, get_db_overview +except ImportError: + from ._qcodes_db_overview import RunOverviewDict, get_db_overview - return overview +__all__ = ["RunOverviewDict", "get_db_overview"] diff --git a/test/pytest/test_qcodes_data.py b/test/pytest/test_qcodes_data.py index df0de7d2..a188c0b6 100644 --- a/test/pytest/test_qcodes_data.py +++ b/test/pytest/test_qcodes_data.py @@ -372,7 +372,7 @@ def test_counts_result_rows(self, tmp_path): def test_records_from_shapes(self): """Shape info in run_description should produce correct count.""" import json - from plottr.data.qcodes_db_overview import _records_from_run_description + from plottr.data._qcodes_db_overview import _records_from_run_description desc = json.dumps({"version": 3, "shapes": {"dep1": [100, 50]}}) assert _records_from_run_description(desc) == 5000