From 82b274300d29756b5bc3cdd3f4294d4696e290e9 Mon Sep 17 00:00:00 2001
From: Florian Blanchet <florian.blanchet@aalto.fi>
Date: Mon, 20 Jun 2022 15:56:50 +0300
Subject: [PATCH] Do not use numpy.ndarray for sqlite3 converters and adapters
 on complex scalars

---
 qcodes/dataset/sqlite/database.py          | 52 +++++++++++++++++-----
 qcodes/tests/dataset/test_dataset_basic.py | 15 ++++++-
 qcodes/utils/types.py                      |  1 +
 3 files changed, 57 insertions(+), 11 deletions(-)

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