Add a callback to get_parameter_data to follow data loading

docs/changes/newsfragments/4688.new

@@ -0,0 +1,6 @@
+Add a callback to 1get_parameter1.
+This is usefull to track the progress of the data download.
+Since sqlite3 does not allow to keep track of the data loading progress,
+we compute how many sqlite requests correspond to a certain percentage of
+progress which is dictated by a config parameter "callback_percent".
+Then we perform x sql request instead of one, running the callback everytime.

qcodes/configuration/qcodesrc.json

@@ -69,6 +69,7 @@
         "use_threads": false,
         "dond_plot": false,
         "dond_show_progress": false,
+        "callback_percent" : 5.0,
         "export_automatic": false,
         "export_type": null,
         "export_prefix": "qcodes_",

qcodes/configuration/qcodesrc_schema.json

@@ -322,6 +322,11 @@
                     "default": false,
                     "description": "Should dond functions show a progress bar during the measurement"
                 },
+                "callback_percent": {
+                    "type": "number",
+                    "default": 5.0,
+                    "description": "If user wants to callback a function while loading data, the callback is done every callback_percent"
+                },
                 "export_automatic": {
                     "type": "boolean",
                     "default": false,

qcodes/dataset/data_set.py

@@ -762,6 +762,7 @@ def get_parameter_data(
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        callback: Callable[[float], None] | None = None,
     ) -> ParameterData:
         """
         Returns the values stored in the :class:`.DataSet` for the specified parameters
@@ -801,6 +802,8 @@ def get_parameter_data(
                 if None
             end: end value of selection range (by results count); ignored if
                 None
+            callback: Function called during the data loading every
+                config.dataset.callback_percent.
 
         Returns:
             Dictionary from requested parameters to Dict of parameter names
@@ -813,7 +816,7 @@ def get_parameter_data(
         else:
             valid_param_names = self._validate_parameters(*params)
         return get_parameter_data(self.conn, self.table_name,
-                                  valid_param_names, start, end)
+                                  valid_param_names, start, end, callback)
 
     def to_pandas_dataframe_dict(
         self,

qcodes/dataset/data_set_in_memory.py

@@ -8,7 +8,7 @@
 import warnings
 from collections.abc import Mapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 import numpy as np
 
@@ -844,6 +844,7 @@ def get_parameter_data(
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        callback: Callable[[float], None] | None = None,
     ) -> ParameterData:
         self._warn_if_set(*params, start=start, end=end)
         return self.cache.data()

qcodes/dataset/data_set_protocol.py

@@ -7,6 +7,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     List,
     Protocol,
@@ -231,6 +232,7 @@ def get_parameter_data(
         *params: str | ParamSpec | ParameterBase,
         start: int | None = None,
         end: int | None = None,
+        callback: Callable[[float], None] | None = None,
     ) -> ParameterData:
         pass
 

qcodes/dataset/sqlite/queries.py

@@ -17,6 +17,7 @@
 from typing_extensions import TypedDict
 
 import qcodes as qc
+from qcodes import config
 from qcodes.dataset.descriptions.dependencies import InterDependencies_
 from qcodes.dataset.descriptions.param_spec import ParamSpec, ParamSpecBase
 from qcodes.dataset.descriptions.rundescriber import RunDescriber
@@ -174,6 +175,7 @@ def get_parameter_data(
     columns: Sequence[str] = (),
     start: int | None = None,
     end: int | None = None,
+    callback: Callable[[float], None] | None = None,
 ) -> dict[str, dict[str, np.ndarray]]:
     """
     Get data for one or more parameters and its dependencies. The data
@@ -199,6 +201,8 @@ def get_parameter_data(
             are returned.
         start: start of range; if None, then starts from the top of the table
         end: end of range; if None, then ends at the bottom of the table
+        callback: Function called during the data loading every
+            config.dataset.callback_percent.
     """
     rundescriber = get_rundescriber_from_result_table_name(conn, table_name)
 
@@ -214,7 +218,8 @@ def get_parameter_data(
             rundescriber,
             output_param,
             start,
-            end)
+            end,
+            callback)
     return output
 
 
@@ -225,6 +230,7 @@ def get_shaped_parameter_data_for_one_paramtree(
     output_param: str,
     start: int | None,
     end: int | None,
+    callback: Callable[[float], None] | None = None,
 ) -> dict[str, np.ndarray]:
     """
     Get the data for a parameter tree and reshape it according to the
@@ -241,7 +247,8 @@ def get_shaped_parameter_data_for_one_paramtree(
         rundescriber,
         output_param,
         start,
-        end
+        end,
+        callback
     )
     if rundescriber.shapes is not None:
         shape = rundescriber.shapes.get(output_param)
@@ -288,10 +295,11 @@ def get_parameter_data_for_one_paramtree(
     output_param: str,
     start: int | None,
     end: int | None,
+    callback: Callable[[float], None] | None = None,
 ) -> tuple[dict[str, np.ndarray], int]:
     interdeps = rundescriber.interdeps
     data, paramspecs, n_rows = _get_data_for_one_param_tree(
-        conn, table_name, interdeps, output_param, start, end
+        conn, table_name, interdeps, output_param, start, end, callback
     )
     if not paramspecs[0].name == output_param:
         raise ValueError("output_param should always be the first "
@@ -389,6 +397,7 @@ def _get_data_for_one_param_tree(
     output_param: str,
     start: int | None,
     end: int | None,
+    callback: Callable[[float], None] | None = None,
 ) -> tuple[list[tuple[Any, ...]], list[ParamSpecBase], int]:
     output_param_spec = interdeps._id_to_paramspec[output_param]
     # find all the dependencies of this param
@@ -401,7 +410,8 @@ def _get_data_for_one_param_tree(
                                     output_param,
                                     *dependency_names,
                                     start=start,
-                                    end=end)
+                                    end=end,
+                                    callback=callback)
     n_rows = len(res)
     return res, paramspecs, n_rows
 
@@ -441,6 +451,7 @@ def get_parameter_tree_values(
     *other_param_names: str,
     start: int | None = None,
     end: int | None = None,
+    callback: Callable[[float], None] | None = None,
 ) -> list[tuple[Any, ...]]:
     """
     Get the values of one or more columns from a data table. The rows
@@ -461,27 +472,97 @@ def get_parameter_tree_values(
         end: The (1-indexed) result to include as the last result to be
             returned. None is equivalent to "all the rest". If start > end,
             nothing is returned.
+        callback: Function called during the data loading every
+            config.dataset.callback_percent.
 
     Returns:
         A list of list. The outer list index is row number, the inner list
         index is parameter value (first toplevel_param, then other_param_names)
     """
 
+    cursor = conn.cursor()
+
+    offset: int | np.ndarray
+
     offset = max((start - 1), 0) if start is not None else 0
     limit = max((end - offset), 0) if end is not None else -1
 
     if start is not None and end is not None and start > end:
         limit = 0
 
+    # start and end currently not working with callback
+    if start is None and end is None and callback is not None:
+
+        # Since sqlite3 does not allow to keep track of the data loading
+        # progress, we compute how many sqlite request correspond to
+        # a progress of config.dataset.callback_percent
+
+        # First, we get the number of dependent parameters
+        rd = get_rundescriber_from_result_table_name(conn, result_table_name)._to_dict()
+        # New qcodes
+        if 'interdependencies_' in rd.keys():
+            nb_param_dependent = len(rd['interdependencies_']['dependencies'])
+        # Old qcodes
+        else:
+            nb_param_dependent = len([i for i in rd['interdependencies']['paramspecs'] if i['name']==toplevel_param_name][0]['depends_on'])
+
+        # Second, we get the number of points
+        sql_callback = f"""
+                        SELECT MAX(id)
+                        FROM "{result_table_name}"
+                        """
+        cursor.execute(sql_callback, ())
+        rows = cursor.fetchall()
+        max_id = int(rows[0][0])
+        nb_point = int(max_id/nb_param_dependent)
+
+        # Third, we get the number of rows corresponding to a download of
+        # config.dataset.callback_percent
+        if nb_point>=100:
+            limit = int(max_id/100*config.dataset.callback_percent/2)
+
+            offset = np.arange(0, nb_point, limit)
+            # Ensure that the last call gets all the points
+            if offset[-1]!=nb_point:
+                offset = np.append(offset, nb_point)
+
+            iteration = 100/(len(offset)-1)
+
+    # Create the base sql query
     columns = [toplevel_param_name] + list(other_param_names)
     sql = f"""
-    SELECT {','.join(columns)} FROM "{result_table_name}"
-    WHERE {toplevel_param_name} IS NOT NULL
-    LIMIT ? OFFSET ?
-    """
-    cursor = conn.cursor()
-    cursor.execute(sql, (limit, offset))
-    return many_many(cursor, *columns)
+           SELECT {','.join(columns)} FROM "{result_table_name}"
+           WHERE {toplevel_param_name} IS NOT NULL
+           LIMIT ? OFFSET ?
+           """
+
+    # Request if no callback
+    if isinstance(offset, int) and callback is None:
+
+        cursor.execute(sql, (limit, offset))
+        res = many_many(cursor, *columns)
+
+    # Request if callback
+    elif isinstance(offset, np.ndarray) and callback is not None:
+
+        # 0
+        progress = 0.
+        callback(progress)
+
+        # 1
+        cursor.execute(sql, (limit, offset[0]))
+        res  = many_many(cursor, *columns)
+        progress += iteration
+        callback(progress)
+
+        # others
+        for i in range(1, len(offset)-1):
+            cursor.execute(sql, (limit, offset[i]))
+            res.extend(many_many(cursor, *columns))
+            progress += iteration
+            callback(progress)
+
+    return res
 
 
 @deprecate(alternative="get_parameter_data")
@@ -2068,7 +2149,8 @@ def load_new_data_for_rundescriber(
             rundescriber=rundescriber,
             output_param=meas_parameter,
             start=start,
-            end=None
+            end=None,
+            callback=None,
         )
         new_data_dict[meas_parameter] = new_data
         updated_read_status[meas_parameter] = start + n_rows_read - 1

qcodes/dataset/sqlite/query_helpers.py

@@ -6,7 +6,7 @@
 
 import itertools
 import sqlite3
-from typing import Any, List, Mapping, Sequence, Union
+from typing import Any, List, Mapping, Sequence, Tuple, Union
 
 import numpy as np
 from numpy import ndarray