Significant performance improvements for complex scalars

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:

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

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.
 """