WIP: lazy loading for netCDF backend

Author: maarten-ic

imas/backends/imas_core/al_context.py

@@ -10,6 +10,7 @@
 
 import numpy
 
+import imas
 from imas.backends.imas_core.imas_interface import ll_interface
 from imas.exception import LowlevelError
 from imas.ids_defs import (
@@ -280,6 +281,9 @@ def __init__(
         self.context = None
         """Potential weak reference to opened context."""
 
+    def get_child(self, child):
+        imas.backends.imas_core.db_entry_helpers._get_child(child, self)
+
     def get_context(self) -> ALContext:
         """Create and yield the actual ALContext."""
         if self.dbentry._db_ctx is not self.dbentry_ctx:

imas/backends/imas_core/db_entry_helpers.py

@@ -22,7 +22,7 @@ def get_children(
     structure: IDSStructure,
     ctx: ALContext,
     time_mode: int,
-    nbc_map: Optional[NBCPathMap],
+    nbc_map: Optional["NBCPathMap"],
 ) -> None:
     """Recursively get all children of an IDSStructure."""
     # NOTE: changes in this method must be propagated to _get_child and vice versa
@@ -77,15 +77,11 @@ def get_children(
                 getattr(structure, name)._IDSPrimitive__value = data
 
 
-def _get_child(child: IDSBase, ctx: Optional[LazyALContext]):
+def _get_child(child: IDSBase, ctx: LazyALContext):
     """Get a single child when required (lazy loading)."""
     # NOTE: changes in this method must be propagated to _get_children and vice versa
     #   Performance: this method is specialized for the lazy get
 
-    # ctx can be None when the parent structure does not exist in the on-disk DD version
-    if ctx is None:
-        return  # There is no data to be loaded
-
     time_mode = ctx.time_mode
     if time_mode == IDS_TIME_MODE_INDEPENDENT and child.metadata.type.is_dynamic:
         return  # skip dynamic (time-dependent) nodes
@@ -148,7 +144,7 @@ def put_children(
     ctx: ALContext,
     time_mode: int,
     is_slice: bool,
-    nbc_map: Optional[NBCPathMap],
+    nbc_map: Optional["NBCPathMap"],
     verify_maxoccur: bool,
 ) -> None:
     """Recursively put all children of an IDSStructure"""

imas/backends/netcdf/db_entry_nc.py

@@ -108,10 +108,6 @@ def get(
             else:
                 func = "get_sample"
             raise NotImplementedError(f"`{func}` is not available for netCDF files.")
-        if lazy:
-            raise NotImplementedError(
-                "Lazy loading is not implemented for netCDF files."
-            )
 
         # Check if the IDS/occurrence exists, and obtain the group it is stored in
         try:
@@ -123,7 +119,7 @@ def get(
 
         # Load data into the destination IDS
         if self._ds_factory.dd_version == destination._dd_version:
-            NC2IDS(group, destination, destination.metadata, None).run()
+            NC2IDS(group, destination, destination.metadata, None).run(lazy)
         else:
             # Construct relevant NBCPathMap, the one we get from DBEntry has the reverse
             # mapping from what we need. The imas_core logic does the mapping from
@@ -135,7 +131,7 @@ def get(
             nbc_map = ddmap.old_to_new if source_is_older else ddmap.new_to_old
             NC2IDS(
                 group, destination, self._ds_factory.new(ids_name).metadata, nbc_map
-            ).run()
+            ).run(lazy)
 
         return destination
 

imas/backends/netcdf/nc2ids.py

@@ -3,6 +3,7 @@
 from typing import Iterator, List, Optional, Tuple
 
 import netCDF4
+import numpy as np
 
 from imas.backends.netcdf import ids2nc
 from imas.backends.netcdf.nc_metadata import NCMetadata
@@ -100,6 +101,8 @@ def __init__(
         """NetCDF related metadata."""
         self.variables = list(group.variables)
         """List of variable names stored in the netCDF group."""
+
+        self._lazy_map = {}
         # Don't use masked arrays: they're slow and we'll handle most of the unset
         # values through the `:shape` arrays
         self.group.set_auto_mask(False)
@@ -113,18 +116,20 @@ def __init__(
                 "Mandatory variable `ids_properties.homogeneous_time` does not exist."
             )
         var = group["ids_properties.homogeneous_time"]
-        self._validate_variable(var, ids.ids_properties.homogeneous_time.metadata)
+        self._validate_variable(var, ids.metadata["ids_properties/homogeneous_time"])
         if var[()] not in [0, 1, 2]:
             raise InvalidNetCDFEntry(
                 f"Invalid value for ids_properties.homogeneous_time: {var[()]}. "
                 "Was expecting: 0, 1 or 2."
             )
         self.homogeneous_time = var[()] == IDS_TIME_MODE_HOMOGENEOUS
 
-    def run(self) -> None:
+    def run(self, lazy: bool) -> None:
         """Load the data from the netCDF group into the IDS."""
         self.variables.sort()
         self.validate_variables()
+        if lazy:
+            self.ids._set_lazy_context(LazyContext(self))
         for var_name in self.variables:
             if var_name.endswith(":shape"):
                 continue
@@ -157,6 +162,10 @@ def run(self) -> None:
                 target_metadata = metadata  # no conversion required
 
             var = self.group[var_name]
+            if lazy:
+                self._lazy_map[target_metadata.path_string] = var
+                continue
+
             if metadata.data_type is IDSDataType.STRUCT_ARRAY:
                 if "sparse" in var.ncattrs():
                     shapes = self.group[var_name + ":shape"][()]
@@ -342,3 +351,74 @@ def _validate_sparsity(
             raise variable_error(
                 shape_var, "dtype", shape_var.dtype, "any integer type"
             )
+
+
+class LazyContext:
+    def __init__(self, nc2ids, index=()):
+        self.nc2ids = nc2ids
+        self.index = index
+
+    def get_child(self, child):
+        metadata = child.metadata
+        path = metadata.path_string
+        data_type = metadata.data_type
+
+        var = self.nc2ids._lazy_map.get(path)
+        if data_type is IDSDataType.STRUCT_ARRAY:
+            # Determine size of the aos
+            if var is None:
+                size = 0
+            elif "sparse" in var.ncattrs():
+                size = self.group[var.name + ":shape"][self.index][0]
+            else:
+                # FIXME: extract dimension name from nc file?
+                dim = self.ncmeta.get_dimensions(
+                    metadata.path_string, self.homogeneous_time
+                )[-1]
+                size = self.group.dimensions[dim].size
+
+            child._set_lazy_context(
+                LazyArrayStructContext(self.nc2ids, self.index, size)
+            )
+
+        elif data_type is IDSDataType.STRUCTURE:
+            child._set_lazy_context(self)
+
+        else:  # Data elements
+            var = self.nc2ids._lazy_map.get(path)
+            if var is None:
+                return  # nothing to load
+
+            value = None
+            if "sparse" in var.ncattrs():
+                if metadata.ndim:
+                    shape_var = self.nc2ids.group[var.name + ":shape"]
+                    shape = shape_var[self.index]
+                    if shape.all():
+                        value = var[self.index + tuple(map(slice, shape))]
+                else:
+                    value = var[self.index]
+                    if value == getattr(var, "_FillValue", None):
+                        value = None  # Skip setting
+            else:
+                value = var[self.index]
+
+            if value is not None:
+                if isinstance(value, np.ndarray):
+                    # Convert the numpy array to a read-only view
+                    value = value.view()
+                    value.flags.writeable = False
+                # NOTE: bypassing IDSPrimitive.value.setter logic
+                child._IDSPrimitive__value = value
+
+
+class LazyArrayStructContext(LazyContext):
+    def __init__(self, nc2ids, index, size):
+        super().__init__(nc2ids, index)
+        self.size = size
+
+    def get_context(self):
+        return self  # IDSStructArray expects to get something with a size attribute
+
+    def iterate_to_index(self, index: int) -> LazyContext:
+        return LazyContext(self.nc2ids, self.index + (index,))

imas/ids_structure.py

@@ -6,18 +6,20 @@
 import logging
 from copy import deepcopy
 from types import MappingProxyType
-from typing import Generator, List, Optional, Union
+from typing import TYPE_CHECKING, Generator, List, Optional, Union
 
 from xxhash import xxh3_64
 
-from imas.backends.imas_core.al_context import LazyALContext
 from imas.ids_base import IDSBase, IDSDoc
 from imas.ids_identifiers import IDSIdentifier
 from imas.ids_metadata import IDSDataType, IDSMetadata
 from imas.ids_path import IDSPath
 from imas.ids_primitive import IDSPrimitive
 from imas.ids_struct_array import IDSStructArray
 
+if TYPE_CHECKING:
+    from imas.backends.imas_core.al_context import LazyALContext
+
 logger = logging.getLogger(__name__)
 
 
@@ -32,7 +34,7 @@ class IDSStructure(IDSBase):
 
     __doc__ = IDSDoc(__doc__)
     _children: "MappingProxyType[str, IDSMetadata]"
-    _lazy_context: Optional[LazyALContext]
+    _lazy_context: Optional["LazyALContext"]
 
     def __init__(self, parent: IDSBase, metadata: IDSMetadata):
         """Initialize IDSStructure from metadata specification
@@ -62,10 +64,8 @@ def __getattr__(self, name):
         child_meta = self._children[name]
         child = child_meta._node_type(self, child_meta)
         self.__dict__[name] = child  # bypass setattr logic below: avoid recursion
-        if self._lazy:  # lazy load the child
-            from imas.backends.imas_core.db_entry_helpers import _get_child
-
-            _get_child(child, self._lazy_context)
+        if self._lazy and self._lazy_context is not None:  # lazy load the child
+            self._lazy_context.get_child(child)
         return child
 
     def _assign_identifier(self, value: Union[IDSIdentifier, str, int]) -> None:
@@ -168,7 +168,7 @@ def __eq__(self, other) -> bool:
             return False  # Not equal if there is any difference
         return True  # Equal when there are no differences
 
-    def _set_lazy_context(self, ctx: LazyALContext) -> None:
+    def _set_lazy_context(self, ctx: "LazyALContext") -> None:
         """Called by DBEntry during a lazy get/get_slice.
 
         Set the context that we can use for retrieving our children.