diff --git a/qcodes/dataset/sqlite/database.py b/qcodes/dataset/sqlite/database.py index cd9431cb37a..2b487ad4913 100644 --- a/qcodes/dataset/sqlite/database.py +++ b/qcodes/dataset/sqlite/database.py @@ -26,7 +26,13 @@ perform_db_upgrade, ) from qcodes.dataset.sqlite.initial_schema import init_db -from qcodes.utils.types import complex_types, numpy_floats, numpy_ints +from qcodes.utils.types import ( + complex_type_union, + complex_types, + numpy_complex_map_size2type, + numpy_floats, + numpy_ints, +) JournalMode = Literal["DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"] @@ -54,10 +60,32 @@ def _convert_array(text: bytes) -> np.ndarray: return np.lib.format.read_array(io.BytesIO(text), allow_pickle=False) -def _convert_complex(text: bytes) -> np.complexfloating: - out = io.BytesIO(text) - out.seek(0) - return np.load(out)[0] +def _convert_complex(text: bytes) -> complex_type_union: + """ + See this: + https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-3-0 + + np.load and np.lib.format.read_array parse the npy header + Skipping the header is 70 times faster + + For backward compatibility with qcodes <=0.34: + npy format has an header len evenly divisible by 64, and QCoDeS saved complexes as + single value arrays + HEADER HEADER binary value + |--- n x 64 --||- len % 64-| + + For qcodes >0.34: + QCoDeS saves complexes directly as binary value: there is no header + + NOTE: Will break with complex512 whose lenght is 64 bytes + """ + try: + value_size = len(text) % 64 + return np.frombuffer( + text[-value_size:], dtype=numpy_complex_map_size2type[value_size] + ).item() + except KeyError as exc: + raise ValueError(f"Cannot parse {str(text)}") from exc this_session_default_encoding = sys.getdefaultencoding() @@ -108,11 +136,15 @@ def _adapt_float(fl: float) -> float | str: return float(fl) -def _adapt_complex(value: complex | np.complexfloating) -> sqlite3.Binary: - out = io.BytesIO() - np.save(out, np.array([value])) - out.seek(0) - return sqlite3.Binary(out.read()) +def _adapt_complex(value: complex_type_union) -> sqlite3.Binary: + # np.save and np.lib.format.write_array add an header that is useless for a single + # value + # Avoiding the header is 15 times faster. + return sqlite3.Binary( + ( + value if isinstance(value, np.complexfloating) else np.complex_(value) + ).tobytes() + ) def connect(name: str | Path, debug: bool = False, version: int = -1) -> ConnectionPlus: diff --git a/qcodes/tests/dataset/test_dataset_basic.py b/qcodes/tests/dataset/test_dataset_basic.py index 9b27dd573ac..8e7bf9e7b44 100644 --- a/qcodes/tests/dataset/test_dataset_basic.py +++ b/qcodes/tests/dataset/test_dataset_basic.py @@ -24,7 +24,11 @@ from qcodes.dataset.descriptions.rundescriber import RunDescriber from qcodes.dataset.guids import parse_guid from qcodes.dataset.sqlite.connection import atomic, path_to_dbfile -from qcodes.dataset.sqlite.database import _convert_array, get_DB_location +from qcodes.dataset.sqlite.database import ( + _convert_array, + _convert_complex, + get_DB_location, +) from qcodes.dataset.sqlite.queries import _rewrite_timestamps, _unicode_categories from qcodes.tests.common import error_caused_by from qcodes.tests.dataset.helper_functions import verify_data_dict @@ -615,6 +619,15 @@ def test_backward_compat__adapt_array_v0_33(): assert arr == _convert_array(out.read()) +def test_backward_compat__adapt_complex_v0_33(): + for dtype in complex_types: + value = dtype(1j) + out = io.BytesIO() + np.save(out, np.array([value])) + out.seek(0) + assert value == _convert_complex(out.read()) + + def test_missing_keys(dataset): """ Test that we can now have partial results with keys missing. This is for diff --git a/qcodes/utils/types.py b/qcodes/utils/types.py index bf39784039f..3ba66cef1cc 100644 --- a/qcodes/utils/types.py +++ b/qcodes/utils/types.py @@ -54,6 +54,7 @@ """ numpy_concrete_complex = (np.complex64, np.complex128) +numpy_complex_map_size2type = {np.dtype(t).itemsize: t for t in numpy_concrete_complex} """ Complex types with fixed sizes. """