Repository: ecmwf/earthkit-data
Pull request #497 (2024-10-24 15:14:00)

Feature/to xarray from stream

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) Report
Attention: Patch coverage is `88.04348%` with `11 lines` in your changes missing coverage. Please review.
> Project coverage is 90.88%. Comparing base [(`3781045`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/37810453390ed76008510b9c42c86518f8c7fe4f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) to head [(`dd1611e`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/dd1611e1b3fc8548151fb3ef68f7d4cd6893d19f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

| [Files with missing lines](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) | Patch % | Lines |
|---|---|---|
| [tests/xr\_engine/xr\_engine\_fixtures.py](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?src=pr&el=tree&filepath=tests%2Fxr_engine%2Fxr_engine_fixtures.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-dGVzdHMveHJfZW5naW5lL3hyX2VuZ2luZV9maXh0dXJlcy5weQ==) | 56.52% | [7 Missing and 3 partials :warning: ](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) |
| [tests/xr\_engine/test\_xr\_memory.py](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?src=pr&el=tree&filepath=tests%2Fxr_engine%2Ftest_xr_memory.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf#diff-dGVzdHMveHJfZW5naW5lL3Rlc3RfeHJfbWVtb3J5LnB5) | 98.21% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #497      +/-   ##
===========================================
- Coverage    90.91%   90.88%   -0.04%     
===========================================
  Files          144      145       +1     
  Lines        10118    10202      +84     
  Branches       479      487       +8     
===========================================
+ Hits          9199     9272      +73     
- Misses         762      769       +7     
- Partials       157      161       +4     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/497?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).


Patch:
diff --git a/docs/examples/xarray_engine_field_dims.ipynb b/docs/examples/xarray_engine_field_dims.ipynb
index c83a8f2..79fedaa 100644
--- a/docs/examples/xarray_engine_field_dims.ipynb
+++ b/docs/examples/xarray_engine_field_dims.ipynb
@@ -35,7 +35,7 @@
    "id": "2f23c9c6-09b3-477f-8d2f-534312bc835f",
    "metadata": {
     "editable": true,
-    "raw_mimetype": "",
+    "raw_mimetype": "text/restructuredtext",
     "slideshow": {
      "slide_type": ""
     },
diff --git a/src/earthkit/data/core/fieldlist.py b/src/earthkit/data/core/fieldlist.py
index 2463cfc..6e650b7 100644
--- a/src/earthkit/data/core/fieldlist.py
+++ b/src/earthkit/data/core/fieldlist.py
@@ -829,8 +829,8 @@ class FieldList(Index):
 
         Parameters
         ----------
-        fields: list
-            List of :obj:`Field` objects.
+        fields: iterable
+            Iterable of :obj:`Field` objects.
 
         Returns
         -------
@@ -839,7 +839,7 @@ class FieldList(Index):
         """
         from earthkit.data.indexing.fieldlist import SimpleFieldList
 
-        return SimpleFieldList(fields)
+        return SimpleFieldList([f for f in fields])
 
     @staticmethod
     def from_numpy(array, metadata):
@@ -966,7 +966,6 @@ class FieldList(Index):
             r[0] = vals
             for i, f in enumerate(it, start=1):
                 r[i] = _vals(f)
-
             return r
 
     def to_numpy(self, **kwargs):
diff --git a/src/earthkit/data/indexing/tensor.py b/src/earthkit/data/indexing/tensor.py
index 894363e..84dc295 100644
--- a/src/earthkit/data/indexing/tensor.py
+++ b/src/earthkit/data/indexing/tensor.py
@@ -367,6 +367,17 @@ class FieldListTensor(TensorCore):
 
         return cls(source, user_coords, field_coords, field_dims, flatten_values)
 
+    def clear(self):
+        self.source = None
+        self._user_coords = None
+        self._user_shape = None
+        self._user_dims = None
+        self._field_coords = None
+        self._field_shape = None
+        self._field_dims = None
+        self._full_shape = None
+        self.flatten_values = None
+
     @flatten_arg
     def to_numpy(self, index=None, **kwargs):
         if index is not None:
diff --git a/src/earthkit/data/readers/grib/codes.py b/src/earthkit/data/readers/grib/codes.py
index f1b0d5a..c8fd72f 100644
--- a/src/earthkit/data/readers/grib/codes.py
+++ b/src/earthkit/data/readers/grib/codes.py
@@ -302,7 +302,7 @@ class GribField(Field):
         )
 
     def write(self, f, **kwargs):
-        r"""Writes the message to a file object.
+        r"""Write the message to a file object.
 
         Parameters
         ----------
@@ -317,7 +317,7 @@ class GribField(Field):
         write(f, self, **kwargs)
 
     def message(self):
-        r"""Returns a buffer containing the encoded message.
+        r"""Return a buffer containing the encoded message.
 
         Returns
         -------
diff --git a/src/earthkit/data/readers/grib/memory.py b/src/earthkit/data/readers/grib/memory.py
index 7702bba..ca6b9b8 100644
--- a/src/earthkit/data/readers/grib/memory.py
+++ b/src/earthkit/data/readers/grib/memory.py
@@ -155,6 +155,9 @@ class GribFieldInMemory(GribField):
             GribCodesHandle(handle, None, None), use_metadata_cache=get_use_grib_metadata_cache()
         )
 
+    def _release(self):
+        self._handle = None
+
     def copy(self, **kwargs):
         return NewGribFieldInMemory(self, **kwargs)
 
diff --git a/src/earthkit/data/readers/grib/xarray.py b/src/earthkit/data/readers/grib/xarray.py
index 5c941b6..2e2a683 100644
--- a/src/earthkit/data/readers/grib/xarray.py
+++ b/src/earthkit/data/readers/grib/xarray.py
@@ -192,7 +192,7 @@ class XarrayMixIn:
                 "level" and "level_type" roles.
                 - "level_and_type": Use a single dimension for combined level and type of level.
             * squeeze: bool, None
-                Remove dimensions which has only one valid values. Not applies to dimension in
+                Remove dimensions which have only one valid value. Not applies to dimensions in
                 ``ensure_dims``. Its default value (None) expands
                 to True unless the ``profile`` overwrites it.
             * add_valid_time_coord: bool, None
@@ -245,6 +245,22 @@ class XarrayMixIn:
                 A dictionary of attribute to rename. Default is None.
             * remapping: dict, None
                 Define new metadata keys for indexing. Default is None.
+            * lazy_load: bool, None
+                If True, the resulting DataSet will load data lazily from the
+                underlying data source. If False, a DataSet holding all the data in memory
+                and decoupled from the backend source will be created.
+                Using ``lazy_load=False`` with ``release_source=True`` can provide optimised
+                memory usage in certain cases. The default value of ``lazy_load`` (None)
+                expands to True unless the ``profile`` overwrites it.
+            * release_source: bool, None
+                Only used when ``lazy_load=False``. If True, memory held in the input fields are
+                released as soon as their values are copied into the resulting DataSet. This is
+                done per field to avoid memory spikes. The release operation is currently
+                only supported for GRIB fields stored entirely in memory, e.g. when read from a
+                :ref:`stream <streams>`. When a field does not support the release operation, this
+                option is ignored. Having run :obj:`to_xarray` the input data becomes unusable,
+                so use this option carefully. The default value of ``release_source`` (None) expands
+                to False unless the ``profile`` overwrites it.
             * strict: bool, None
                 If True, perform stricter checks on hypercube consistency. Its default value (None) expands
                 to False unless the ``profile`` overwrites it.
@@ -287,9 +303,11 @@ class XarrayMixIn:
         --------
         >>> import earthkit.data
         >>> fs = earthkit.data.from_source("file", "test6.grib")
+        >>> ds = fs.to_xarray(time_dim_mode="forecast")
+        >>> # also possible to use the xarray_open_dataset_kwargs
         >>> ds = fs.to_xarray(
         ...     xarray_open_dataset_kwargs={
-        ...         "backend_kwargs": {"ignore_keys": ["number"]}
+        ...         "backend_kwargs": {"time_dim_mode": "forecast"}
         ...     }
         ... )
 
@@ -349,7 +367,7 @@ class XarrayMixIn:
         # print(f"{kwargs=}")
         # print(f"{xarray_open_dataset_kwargs=}")
 
-        from earthkit.data.utils.xarray.engine import from_earthkit
+        from earthkit.data.utils.xarray.builder import from_earthkit
 
         return from_earthkit(self, **xarray_open_dataset_kwargs)
 
diff --git a/src/earthkit/data/sources/array_list.py b/src/earthkit/data/sources/array_list.py
index 5675383..01a0eb0 100644
--- a/src/earthkit/data/sources/array_list.py
+++ b/src/earthkit/data/sources/array_list.py
@@ -85,6 +85,10 @@ class ArrayField(Field):
     def handle(self):
         return self._metadata._handle
 
+    def _release(self):
+        self._array = None
+        self.__metadata = None
+
     def __getstate__(self) -> dict:
         ret = {}
         ret["_array"] = self._array
diff --git a/src/earthkit/data/sources/stream.py b/src/earthkit/data/sources/stream.py
index 501f36a..32b55f8 100644
--- a/src/earthkit/data/sources/stream.py
+++ b/src/earthkit/data/sources/stream.py
@@ -9,6 +9,7 @@
 
 import itertools
 import logging
+from functools import cached_property
 
 from earthkit.data.core.fieldlist import FieldList
 from earthkit.data.readers import stream_reader
@@ -31,15 +32,13 @@ class StreamMemorySource(MemoryBaseSource):
         if not isinstance(stream, Stream):
             raise ValueError(f"Invalid stream={stream}")
         self._stream = stream
-        self._reader_ = None
 
-    @property
+    @cached_property
     def _reader(self):
-        if self._reader_ is None:
-            self._reader_ = stream_reader(self, self._stream.stream, True, **self._kwargs)
-            if self._reader_ is None:
-                raise TypeError(f"could not create reader for stream={self._stream}")
-        return self._reader_
+        reader = stream_reader(self, self._stream.stream, True, **self._kwargs)
+        if reader is None:
+            raise TypeError(f"could not create reader for stream={self._stream}")
+        return reader
 
     def mutate(self):
         source = self._reader.mutate_source()
@@ -52,9 +51,10 @@ class StreamMemorySource(MemoryBaseSource):
 class StreamSource(Source):
     def __init__(self, stream, *, read_all=False, **kwargs):
         super().__init__()
-        self._reader_ = None
         self._stream = self._wrap_stream(stream)
         self.memory = read_all
+
+        # TODO: remove this check in a future release
         for k in ["group_by", "batch_size"]:
             if k in kwargs:
                 raise ValueError(f"Invalid argument '{k}' for StreamSource. Deprecated since 0.8.0.")
@@ -74,13 +74,12 @@ class StreamSource(Source):
                 return StreamFieldList(self._reader, **self._kwargs)
         return self
 
-    @property
+    @cached_property
     def _reader(self):
-        if self._reader_ is None:
-            self._reader_ = stream_reader(self, self._stream.stream, False, **self._kwargs)
-            if self._reader_ is None:
-                raise TypeError(f"could not create reader for stream={self._stream.stream}")
-        return self._reader_
+        reader = stream_reader(self, self._stream.stream, False, **self._kwargs)
+        if reader is None:
+            raise TypeError(f"could not create reader for stream={self._stream.stream}")
+        return reader
 
     def batched(self, n):
         """Iterate through the stream in batches of ``n``.
@@ -133,13 +132,6 @@ class StreamSource(Source):
 
         return stream
 
-    def _status(self):
-        """For testing purposes."""
-        return {
-            "reader": self._reader_ is not None,
-            "stream": self._stream._stream is not None,
-        }
-
 
 class MultiStreamSource(Source):
     def __init__(self, sources, read_all=False, **kwargs):
@@ -183,10 +175,6 @@ class MultiStreamSource(Source):
                 raise TypeError(f"Invalid source={s}")
         return r
 
-    def _status(self):
-        """For testing purposes."""
-        return [s._status() for s in self.sources]
-
 
 class StreamFieldList(FieldList, Source):
     def __init__(self, source, **kwargs):
@@ -208,6 +196,12 @@ class StreamFieldList(FieldList, Source):
     def __getstate__(self):
         raise NotImplementedError("StreamFieldList cannot be pickled")
 
+    def to_xarray(self, **kwargs):
+        from earthkit.data.core.fieldlist import FieldList
+
+        fields = [f for f in self]
+        return FieldList.from_fields(fields).to_xarray(**kwargs)
+
 
 class Stream:
     def __init__(self, stream=None, maker=None, **kwargs):
diff --git a/src/earthkit/data/utils/diag.py b/src/earthkit/data/utils/diag.py
index ee48238..567ed3c 100644
--- a/src/earthkit/data/utils/diag.py
+++ b/src/earthkit/data/utils/diag.py
@@ -11,11 +11,29 @@ import time
 from collections import defaultdict
 
 
-class TimeDiag:
+class DiagCore:
     def __init__(self, name=""):
         self.name = name
+
+    def _build_label(self, label):
+        if label:
+            label = f"[{label:10}] "
+
+        if self.name:
+            name = f"[{self.name:10}] "
+            if label:
+                label = f"{name}{label}"
+            else:
+                label = f"{name}"
+
+        return label
+
+
+class TimeDiag(DiagCore):
+    def __init__(self, name="", **kwargs):
         self.start = time.time()
         self.prev = self.start
+        super().__init__(name=name, **kwargs)
 
     def elapsed(self):
         return time.time() - self.start
@@ -24,14 +42,7 @@ class TimeDiag:
         curr = time.time()
         delta = curr - self.prev
         self.prev = curr
-        if label:
-            label = f"[{label:10}] "
-        if self.name:
-            if label:
-                label = f"[{self.name}] {label}"
-            else:
-                label = f"[{self.name}]"
-
+        label = self._build_label(label)
         s = f"{label}elapsed={self.elapsed():.3f}s delta={delta:.3f}s"
 
         if as_str:
@@ -40,14 +51,13 @@ class TimeDiag:
             print(s)
 
 
-class MemoryDiag:
-    def __init__(self, name="", peak=False):
+class MemoryDiag(DiagCore):
+    def __init__(self, name="", peak=False, **kwargs):
         import os
         import platform
 
         import psutil
 
-        self.name = name
         self.proc = psutil.Process()
         self.prev = 0
         self.scale = 1
@@ -59,6 +69,8 @@ class MemoryDiag:
         except Exception:
             pass
 
+        super().__init__(name=name, **kwargs)
+
     def scale_to_mbytes(self, v):
         return v * self.scale
 
@@ -85,14 +97,7 @@ class MemoryDiag:
         m = self.current()
         _delta = m - self.prev
         self.prev = m
-        if label:
-            label = f"[{label:10}] "
-
-        if self.name:
-            if label:
-                label = f"[{self.name}] {label}"
-            else:
-                label = f"[{self.name}]"
+        label = self._build_label(label)
 
         s = ""
         if delta:
@@ -109,15 +114,17 @@ class MemoryDiag:
             print(s)
 
 
-class Diag:
-    def __init__(self, name="", peak=False):
-        self.time = TimeDiag(name)
-        self.memory = MemoryDiag(name, peak=peak)
+class Diag(DiagCore):
+    def __init__(self, name="", peak=False, **kwargs):
+        self.time = TimeDiag("")
+        self.memory = MemoryDiag("", peak=peak)
+        super().__init__(name=name, **kwargs)
 
     def __call__(self, label=""):
-        if label:
-            label = f"[{label:10}] "
-        return f"{label}{self.time(as_str=True)} {self.memory(as_str=True)}"
+        label = str(label)
+        label = self._build_label(label)
+        s = f"{label}{self.time(as_str=True)} {self.memory(as_str=True)}"
+        print(s)
 
     def peak(self):
         return self.memory.peak()
@@ -127,12 +134,6 @@ def metadata_cache_diag(fieldlist):
     r = defaultdict(int)
     for f in fieldlist:
         collect_field_metadata_cache_diag(f, r)
-        # try:
-        #     md_cache = f._diag()
-        #     for k in ["metadata_cache_hits", "metadata_cache_misses", "metadata_cache_size"]:
-        #         r[k] += md_cache[k]
-        # except Exception:
-        #     pass
     return r
 
 
diff --git a/src/earthkit/data/utils/xarray/builder.py b/src/earthkit/data/utils/xarray/builder.py
index 76573a1..385bbda 100644
--- a/src/earthkit/data/utils/xarray/builder.py
+++ b/src/earthkit/data/utils/xarray/builder.py
@@ -8,7 +8,8 @@
 #
 
 import logging
-from functools import cached_property
+from abc import ABCMeta
+from abc import abstractmethod
 
 import numpy
 import xarray
@@ -16,13 +17,16 @@ import xarray.core.indexing as indexing
 
 from earthkit.data.utils import ensure_iterable
 
+from .profile import Profile
+
 LOG = logging.getLogger(__name__)
 
 
 class VariableBuilder:
-    def __init__(self, var_dims, data, local_attr_keys, tensor, remapping):
+    def __init__(self, name, var_dims, data_maker, local_attr_keys, tensor, remapping):
+        self.name = name
         self.var_dims = var_dims
-        self.data = data
+        self.data_maker = data_maker
         self._attrs = {}
         self.local_keys = ensure_iterable(local_attr_keys)
         self.tensor = tensor
@@ -33,7 +37,8 @@ class VariableBuilder:
             "message": self.tensor.source[0].metadata().override()._handle.get_buffer(),
         }
         self._attrs["_earthkit"] = attrs
-        return xarray.Variable(self.var_dims, self.data, attrs=self._attrs)
+        data = self.data_maker(self.tensor, self.var_dims, self.name)
+        return xarray.Variable(self.var_dims, data, attrs=self._attrs)
 
     def load_attrs(self, keys, strict=True):
         keys = keys + self.local_keys
@@ -155,7 +160,7 @@ class TensorBackendArray(xarray.backends.common.BackendArray):
         return result
 
 
-class TensorBackendBuilder:
+class BackendDataBuilder(metaclass=ABCMeta):
     def __init__(
         self,
         ds,
@@ -171,15 +176,14 @@ class TensorBackendBuilder:
         self.dtype = profile.dtype
         self.array_module = profile.array_module
 
-        # these coords inside the the tensor are called user_coords
-        # Note: in the tensor the corresponding dims are called user_dims
+        # Note: these coords inside the tensor are called user_coords and
+        # the corresponding dims are called user_dims
         self.tensor_coords = {}
 
         # coords describing the field dimensions
         self.field_coords = {}
 
         self.grid = self._ensure_grid(grid)
-
         if self.profile.add_geo_coords:
             self.field_coords = self._make_field_coords()
 
@@ -213,37 +217,46 @@ class TensorBackendBuilder:
 
             _dims, _vals = tensor.make_valid_datetime()
             if _dims is not None and _vals is not None:
-                # if _dims is None or _vals is None:
-                #     raise ValueError("valid_time coord could not be created")
                 self.tensor_coords["valid_time"] = Coord.make("valid_time", _vals, dims=_dims)
 
     def build(self):
-        t_vars = {}
         # we assume each variable forms a full cube
-        for variable in self.profile.variables:
-            t_vars[variable] = self.make_variable(self.ds, self.dims, self.profile.variable_key, variable)
+        var_builders = self.pre_build_variables()
 
         # build variable and global attributes
-        xr_attrs = self.profile.attrs.builder.build(self.ds, t_vars, rename=True)
-
+        xr_attrs = self.profile.attrs.builder.build(self.ds, var_builders, rename=True)
         xr_coords = self.coords()
-        xr_vars = {self.profile.rename_variable(k): v.build() for k, v in t_vars.items()}
 
+        # build variables
+        xr_vars = {self.profile.rename_variable(k): v.build() for k, v in var_builders.items()}
+
+        # build dataset
         dataset = xarray.Dataset(xr_vars, coords=xr_coords, attrs=xr_attrs)
+
         if self.profile.rename_dims_map():
             dataset = dataset.rename(self.profile.rename_dims_map())
 
+        if "source" not in dataset.encoding:
+            dataset.encoding["source"] = None
+
         return dataset
 
-    def make_variable(self, ds, dims, key, name):
-        ds_var = ds.sel(**{key: name})
+    @abstractmethod
+    def build_values(self, *args, **kwargs):
+        pass
 
+    @abstractmethod
+    def pre_build_variables(self):
+        pass
+
+    def pre_build_variable(self, ds_var, dims, name):
         tensor_dims, tensor_coords, tensor_coords_component, extra_tensor_attrs = self.prepare_tensor(
             ds_var, dims, name
         )
         tensor_dim_keys = [d.key for d in tensor_dims]
 
-        # print("tensor_dims", tensor_dims)
+        LOG.debug(f"{tensor_dims=} {tensor_coords=} {tensor_coords_component=} {extra_tensor_attrs=}")
+        LOG.debug(f"{tensor_dim_keys=}")
 
         tensor = ds_var.to_tensor(
             *tensor_dim_keys,
@@ -271,42 +284,12 @@ class TensorBackendBuilder:
         var_dims.extend(tensor.field_dims)
 
         self.collect_date_coords(tensor)
-
-        # print(f" full_dims={tensor.full_dims}")
-        # print(f" full_shape={tensor.full_shape}")
-
-        # self.collect_coords(tensor)
-        # var_dims = self.var_dims(tensor)
-        # var_dims = [d.key for d in tensor_dims]
-        # print(f"var_dims={var_dims}")
-
-        backend_array = TensorBackendArray(
-            tensor,
-            var_dims,
-            tensor.full_shape,
-            self.array_module,
-            self.dtype,
-            name,
-        )
-
-        data = indexing.LazilyIndexedArray(backend_array)
-
-        # Get metadata keys which are common for all fields, and not listed in dataset attrs
-        # kk = [k for k in self.profile.index_keys if k not in self.attributes]
-        # var_attrs = ds.common_attributes_other(ds_var, kk)
-
-        # var_attrs = {
-        #     k: ds_var.index(k)[0] for k in self.profile.var_attributes() if len(ds_var.index(k)) >= 1
-        # }
-        # var_attrs.update(tensor_attrs)
-
-        # var_attrs = self.profile.remap(var_attrs)
-
+        data_maker = self.build_values
         remapping = self.profile.remapping.build()
 
-        var = VariableBuilder(var_dims, data, extra_tensor_attrs, tensor, remapping)
-        # var = xarray.Variable(var_dims, data, attrs=var_attrs)
-        return var
+        var_builder = VariableBuilder(name, var_dims, data_maker, extra_tensor_attrs, tensor, remapping)
+
+        return var_builder
 
     def prepare_tensor(self, ds, dims, name):
         tensor_dims = []
@@ -314,21 +297,17 @@ class TensorBackendBuilder:
         tensor_coords_component = {}
         extra_tensor_attrs = []
 
-        # print(f"prepare_tensor: {name=} {dims=}")
-        # First check if the dims/coords are consistent with the tensors of the previous variables
-
-        # from .fieldlist import unique_values
+        LOG.debug(f"{name=} {dims=}")
 
-        # remapping = self.profile.remapping.build()
+        # First check if the dims/coords are consistent with the tensors of the previous variables
         vals, component_vals = ds.unique_values([d.key for d in dims], component=True)
 
-        # print("unique_values", vals)
-        # print("ensure_dims", self.profile.dims.ensure_dims)
+        LOG.debug(f"unique_values={vals}")
+        LOG.debug(f"ensure_dims={self.profile.dims.ensure_dims}")
 
         for d in dims:
-            # num = len(ds.index(d.key))
             num = len(vals[d.key])
-            # print(f"  {d.key=} {vals[d.key]}")
+            # LOG.debug(f"  {d.key=} {vals[d.key]}")
             if num == 0:
                 continue
                 # if d.name not in self.profile.dims.ensure_dims:
@@ -342,13 +321,12 @@ class TensorBackendBuilder:
                     extra_tensor_attrs.append(d.key)
                 elif num > 1 or not self.profile.dims.squeeze or d.name in self.profile.dims.ensure_dims:
                     tensor_dims.append(d)
-                    # tensor_coords[d.key] = ds.index(d.key)
                     tensor_coords[d.key] = vals[d.key]
                     if d.key in component_vals:
                         tensor_coords_component[d.key] = component_vals[d.key]
                     self.check_tensor_coords(name, d.key, tensor_coords)
 
-        # TODO:  check if fieldlist forms a full hypercube with respect to the the dims/coordinate
+        # TODO:  check if fieldlist forms a full hypercube with respect to the the dims/coordinates
         return tensor_dims, tensor_coords, tensor_coords_component, extra_tensor_attrs
 
     def check_tensor_coords(self, var_name, coord_name, tensor_coords):
@@ -358,6 +336,73 @@ class TensorBackendBuilder:
             check_coords(var_name, coord_name, tensor_coords, self.tensor_coords)
 
 
+class TensorBackendDataBuilder(BackendDataBuilder):
+    """Build a dataset using memory backend.
+
+    Each variable contain a TensorBackendArray object that loads the data on demand from
+    the GRIB using the associated TensorFieldList. This solution can work bot for GRIB data
+    stored in memory or on disk.
+    """
+
+    def pre_build_variables(self):
+        """Generate a builder for each variable"""
+        builders = {}
+
+        # we assume each variable forms a full cube
+        for name in self.profile.variables:
+            ds_var = self.ds.sel(**{self.profile.variable_key: name})
+            builders[name] = self.pre_build_variable(ds_var, self.dims, name)
+
+        return builders
+
+    def build_values(self, tensor, var_dims, name):
+        """Generate the data object stored in the xarray variable"""
+        backend_array = TensorBackendArray(
+            tensor,
+            var_dims,
+            tensor.full_shape,
+            self.array_module,
+            self.dtype,
+            name,
+        )
+
+        data = indexing.LazilyIndexedArray(backend_array)
+        return data
+
+
+class MemoryBackendDataBuilder(BackendDataBuilder):
+    """Build a dataset using memory backend.
+
+    Each variable contain data values in memory copied from the input fields. As soon as a
+    field's values are copied, the field is immediately released and its memory is freed if the
+    field supports this operation. Fields released in this way are not available for further access.
+    """
+
+    def pre_build_variables(self):
+        """Generate a builder for each variable"""
+        builders = {}
+        groups = self.ds.group(self.profile.variable_key, self.profile.variables)
+
+        # we assume each variable forms a full cube
+        for name in groups:
+            ds_var = groups[name]
+            builders[name] = self.pre_build_variable(ds_var, self.dims, name)
+
+        return builders
+
+    def build_values(self, tensor, var_dims, name):
+        """Generate the data object stored in the xarray variable"""
+
+        # At this point all the fields must be a ReleasableField.
+        # We mark the fields so that their data will be released on the next
+        # values access.
+        if self.profile.release_source:
+            for f in tensor.source:
+                f.keep = False
+
+        return tensor.to_numpy(dtype=self.dtype)
+
+
 class DatasetBuilder:
     def __init__(
         self,
@@ -365,74 +410,58 @@ class DatasetBuilder:
         profile="mars",
         **kwargs,
     ):
-        """
-        auto_split: bool
-            When it is True and the data does not form a complete hypercube automatically
-            tries to split it into multiple hypercubes and returns a list of datasets (one
-            dataset per hypercube). When it is False and the data does not form a complete hypercube
-            rases an error. Default is False.
-        split_dims: str, or iterable of str, None
-            Dimension or list of dimensions to use for splitting the data into multiple hypercubes.
-            Default is None.
-        """
+        self.profile = Profile.make(profile, **kwargs)
         self.ds = ds
+
+        if self.profile.lazy_load:
+            self.builder = TensorBackendDataBuilder
+        else:
+            self.builder = MemoryBackendDataBuilder
+
         self.kwargs = kwargs
-        self.profile_name = profile
         self.grids = {}
 
-    @cached_property
-    def profile(self):
-        from .profile import Profile
+    # @cached_property
+    # def profile(self):
+    #     from .profile import Profile
 
-        return Profile.make(self.profile_name, **self.kwargs)
+    #     return Profile.make(self.profile_name, **self.kwargs)
 
     def parse(self):
         assert not hasattr(self.ds, "_ek_builder")
-        from .fieldlist import WrappedFieldList
-
-        # from .profile import Profile
-        # profile = Profile.make(self.profile, **self.kwargs)
+        from .fieldlist import XArrayInputFieldList
 
         remapping = self.profile.remapping.build()
-        # print(f"{remapping=}")
-        # print(f"{profile.remapping=}")
-
-        # create a new fieldlist and ensure all the required metadata is kept in memory
-        ds = WrappedFieldList(self.ds, keys=self.profile.index_keys, remapping=remapping)
-        # print(f"{remapping=}")
-        # print(f"ds: {ds.indices()}")
-
-        # print(f"ds: {ds.index('level_and_type')}")
-        # print("components", ds.index("level_and_type", component=True))
 
-        # global attributes are keys which are the same for all the fields
-        # attributes = {k: v[0] for k, v in ds_ori.indices().items() if len(v) == 1}
-        # global_attrs = ds.common_indices()
+        LOG.debug(f"{remapping=}")
+        LOG.debug(f"{self.profile.remapping=}")
 
-        # LOG.info(f"{attributes=}")
+        # create a new fieldlist for optimised access to unique values
+        ds = XArrayInputFieldList(self.ds, keys=self.profile.index_keys, remapping=remapping)
 
-        # if hasattr(ds, "path"):
-        #     global_attrs["ekds_source"] = ds.path
+        # LOG.debug(f"{ds.db=}")
+        LOG.debug(f"before update: {self.profile.dim_keys=}")
 
-        # print("parse dims", profile.dim_keys)
-
-        # print("ds keys", ds.db[0].keys())
-        # print("ds indices", ds._md_indices)
         self.profile.update(ds)
 
-        # print("parse dims", profile.dim_keys)
+        LOG.debug(f"after update: {self.profile.dim_keys=}")
 
         # the data is only sorted once
-        ds_sorted = ds.order_by(self.profile.sort_keys)
+        ds = ds.order_by(self.profile.sort_keys)
+
+        if not self.profile.lazy_load and self.profile.release_source:
+            ds.make_releasable()
 
-        return ds, ds_sorted
+        return ds
 
     def grid(self, ds):
         grids = ds.index("md5GridSection")
+
         if len(grids) != 1:
             raise ValueError(f"Expected one grid, got {len(grids)}")
         grid = grids[0]
         key = (grid, self.profile.flatten_values)
+
         if key not in self.grids:
             from .grid import TensorGrid
 
@@ -442,52 +471,51 @@ class DatasetBuilder:
 
 class SingleDatasetBuilder(DatasetBuilder):
     def __init__(self, *args, **kwargs):
-        auto_split = kwargs.get("auto_split", False)
         split_dims = kwargs.get("split_dims", None)
-        if auto_split or split_dims:
+        if split_dims:
             raise ValueError("SingleDatasetMaker does not support splitting")
 
         super().__init__(*args, **kwargs)
 
     def build(self):
-        ds, ds_sorted = self.parse()
+        ds_sorted = self.parse()
         dims = self.profile.dims.to_list()
-        # print("SingleDatasetBuilder.build dims", dims)
-        builder = TensorBackendBuilder(
+        LOG.debug(f"{dims=}")
+        builder = self.builder(
             ds_sorted,
             self.profile,
             dims,
-            grid=self.grid(ds),
+            grid=self.grid(ds_sorted),
         )
-
-        return builder.build()
+        r = builder.build()
+        return r
 
 
 class SplitDatasetBuilder(DatasetBuilder):
     def __init__(self, *args, backend_kwargs=None, **kwargs):
-        self.auto_split = backend_kwargs.pop("auto_split", False)
+        """
+        split_dims: str, or iterable of str, None
+            Dimension or list of dimensions to use for splitting the data into multiple hypercubes.
+            Default is None.
+        """
         self.split_dims = backend_kwargs.pop("split_dims", None)
         self.backend_kwargs = dict(**backend_kwargs)
         self.xr_open_dataset_kwargs = dict(**kwargs)
 
-        if not self.auto_split and not self.split_dims:
-            raise ValueError("SplitDatasetMaker requires auto_split or split_dims")
-
-        # auto_split is experimental and not yet supported
-        if self.auto_split:
-            raise ValueError("auto_split is not yet supported")
+        if not self.split_dims:
+            raise ValueError("SplitDatasetMaker requires split_dims")
 
         super().__init__(*args, split_dims=self.split_dims, **backend_kwargs)
 
     def build(self):
         from .splitter import Splitter
 
-        _, ds_sorted = self.parse()
-        splitter = Splitter.make(self.auto_split, self.split_dims)
+        ds_sorted = self.parse()
+        splitter = Splitter.make(self.split_dims)
         datasets = []
         for s_dims, s_ds in splitter.split(ds_sorted, self.profile):
-            # print(f"split_dims: {s_dims}   {type(s_ds)}")
-            builder = TensorBackendBuilder(
+            LOG.debug(f"{s_dims=} type of s_ds={type(s_ds)}")
+            builder = self.builder(
                 s_ds,
                 self.profile,
                 s_dims,
@@ -501,3 +529,23 @@ class SplitDatasetBuilder(DatasetBuilder):
             s_ds._ek_builder = None
 
         return datasets[0] if len(datasets) == 1 else datasets
+
+
+def from_earthkit(ds, **kwargs):
+    """Create an xarray dataset from an earthkit fieldlist."""
+    backend_kwargs = kwargs.get("backend_kwargs", {})
+    split_dims = backend_kwargs.get("split_dims", None)
+
+    assert kwargs["engine"] == "earthkit"
+
+    if not split_dims:
+        backend_kwargs.pop("split_dims", None)
+
+        if len(kwargs) == 1:
+            assert kwargs["engine"] == "earthkit"
+            return SingleDatasetBuilder(ds, **backend_kwargs).build()
+        else:
+            return xarray.open_dataset(ds, **kwargs)
+    else:
+        backend_kwargs = kwargs.pop("backend_kwargs", {})
+        return SplitDatasetBuilder(ds, backend_kwargs=backend_kwargs, **kwargs).build()
diff --git a/src/earthkit/data/utils/xarray/defaults.yaml b/src/earthkit/data/utils/xarray/defaults.yaml
index b514f41..7039d0b 100644
--- a/src/earthkit/data/utils/xarray/defaults.yaml
+++ b/src/earthkit/data/utils/xarray/defaults.yaml
@@ -35,6 +35,8 @@ dtype: float64
 array_module: numpy
 
 # other
+lazy_load: true
+release_source: false
 strict: false
 errors: raise
 
diff --git a/src/earthkit/data/utils/xarray/dim.py b/src/earthkit/data/utils/xarray/dim.py
index c3a70a7..537695b 100644
--- a/src/earthkit/data/utils/xarray/dim.py
+++ b/src/earthkit/data/utils/xarray/dim.py
@@ -751,7 +751,7 @@ class Dims:
         return (
             f'Variable-related keys "{keys}"" cannot be dimensions. Such a key'
             " must be specified as the variable_key. The variable_key can only "
-            f'be set to a single key, its current value is "{self.profile.variable_key}"'
+            f'be set to a single key. Its current value is "{self.profile.variable_key}"'
         )
 
     def register_remapping(self, remapping, patch=None):
diff --git a/src/earthkit/data/utils/xarray/engine.py b/src/earthkit/data/utils/xarray/engine.py
index 46f1a53..6e2515f 100644
--- a/src/earthkit/data/utils/xarray/engine.py
+++ b/src/earthkit/data/utils/xarray/engine.py
@@ -15,24 +15,6 @@ from xarray.backends import BackendEntrypoint
 LOG = logging.getLogger(__name__)
 
 
-def from_earthkit(ds, **kwargs):
-    backend_kwargs = kwargs.get("backend_kwargs", {})
-    auto_split = backend_kwargs.get("auto_split", False)
-    split_dims = backend_kwargs.get("split_dims", None)
-
-    assert kwargs["engine"] == "earthkit"
-
-    if not auto_split and not split_dims:
-        backend_kwargs.pop("auto_split", None)
-        backend_kwargs.pop("split_dims", None)
-        return xarray.open_dataset(ds, **kwargs)
-    else:
-        from .builder import SplitDatasetBuilder
-
-        backend_kwargs = kwargs.pop("backend_kwargs", {})
-        return SplitDatasetBuilder(ds, backend_kwargs=backend_kwargs, **kwargs).build()
-
-
 class EarthkitBackendEntrypoint(BackendEntrypoint):
     def open_dataset(
         self,
@@ -64,6 +46,8 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
         rename_attrs=None,
         remapping=None,
         flatten_values=None,
+        lazy_load=None,
+        release_source=None,
         strict=None,
         dtype=None,
         array_module=None,
@@ -166,7 +150,7 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
               "level" and "level_type" roles.
             - "level_and_type": Use a single dimension for combined level and type of level.
         squeeze: bool, None
-            Remove dimensions which has only one valid values. Not applies to dimension in
+            Remove dimensions which have only one valid value. Not applies to dimensions in
             ``ensure_dims``. Its default value (None) expands
             to True unless the ``profile`` overwrites it.
         add_valid_time_coord: bool, None
@@ -219,6 +203,22 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
             A dictionary of attribute to rename. Default is None.
         remapping: dict, None
             Define new metadata keys for indexing. Default is None.
+        lazy_load: bool, None
+            If True, the resulting DataSet will load data lazily from the
+            underlying data source. If False, a DataSet holding all the data in memory
+            and decoupled from the backend source will be created.
+            Using ``lazy_load=False`` with ``release_source=True`` can provide optimised
+            memory usage in certain cases. The default value of ``lazy_load`` (None)
+            expands to True unless the ``profile`` overwrites it.
+        release_source: bool, None
+            Only used when ``lazy_load=False``. If True, memory held in the input fields are
+            released as soon as their values are copied into the resulting DataSet. This is
+            done per field to avoid memory spikes. The release operation is currently
+            only supported for GRIB fields stored entirely in memory, e.g. when read from a
+            :ref:`stream <streams>`. When a field does not support the release operation, this
+            option is ignored. Having run :obj:`to_xarray` the input data becomes unusable,
+            so use this option carefully. The default value of ``release_source`` (None) expands
+            to False unless the ``profile`` overwrites it.
         strict: bool, None
             If True, perform stricter checks on hypercube consistency. Its default value (None) expands
             to False unless the ``profile`` overwrites it.
@@ -254,6 +254,8 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
             remapping=remapping,
             decode_times=decode_times,
             decode_timedelta=decode_timedelta,
+            lazy_load=lazy_load,
+            release_source=release_source,
             strict=strict,
             dtype=dtype,
             array_module=array_module,
diff --git a/src/earthkit/data/utils/xarray/fieldlist.py b/src/earthkit/data/utils/xarray/fieldlist.py
index 8abb609..c7c00bf 100644
--- a/src/earthkit/data/utils/xarray/fieldlist.py
+++ b/src/earthkit/data/utils/xarray/fieldlist.py
@@ -101,7 +101,7 @@ class IndexDB:
         return f"IndexDB(_index={self._index}, component={self._component})"
 
 
-class WrappedFieldList(FieldList):
+class XArrayInputFieldList(FieldList):
     def __init__(self, fieldlist, keys=None, db=None, remapping=None):
         super().__init__()
         self.ds = fieldlist
@@ -134,6 +134,25 @@ class WrappedFieldList(FieldList):
     def __len__(self):
         return len(self.ds)
 
+    def make_releasable(self):
+        print("Making releasable")
+        self.ds = FieldList.from_fields([ReleasableField(f) for f in self.ds])
+
+    def group(self, key, values):
+        groups = defaultdict(list)
+        for f in self.ds:
+            v = str(f.metadata(key, default=None))
+            if v in values:
+                groups[v].append(f)
+
+        for k, v in groups.items():
+            if not v:
+                continue
+            db = None
+            groups[k] = XArrayInputFieldList(FieldList.from_fields(v), db=db, remapping=self.remapping)
+
+        return groups
+
     def sel(self, *args, **kwargs):
         assert "remapping" not in kwargs
         assert "patches" not in kwargs
@@ -141,14 +160,17 @@ class WrappedFieldList(FieldList):
         db = None
         # if not args and self.db and all(k in self.db for k in kwargs):
         #     db = self.db.filter(**kwargs)
-        return WrappedFieldList(ds, db=db, remapping=self.remapping)
+        return XArrayInputFieldList(ds, db=db, remapping=self.remapping)
 
     def order_by(self, *args, **kwargs):
         assert "remapping" not in kwargs
         assert "patches" not in kwargs
-        return WrappedFieldList(
-            self.ds.order_by(*args, remapping=self.remapping, **kwargs), db=self.db, remapping=self.remapping
+        ds = XArrayInputFieldList(
+            self.ds.order_by(*args, remapping=self.remapping, **kwargs),
+            db=self.db,
+            remapping=self.remapping,
         )
+        return ds
 
     def unique_values(self, names, component=False):
         if isinstance(names, str):
@@ -193,3 +215,42 @@ class WrappedFieldList(FieldList):
             return indices, components
         else:
             return indices
+
+
+class ReleasableField:
+    def __init__(self, field):
+        self.field = field
+        self.keep = True
+        self.released = False
+
+    def to_numpy(self, *args, **kwargs):
+        if self.released:
+            return None
+
+        r = self.field.to_numpy(*args, **kwargs)
+        if not self.keep:
+            self._release()
+        return r
+
+    def to_array(self, *args, **kwargs):
+        if self.released:
+            return None
+
+        r = self.field.to_array(*args, **kwargs)
+        if not self.keep:
+            self._release()
+        return r
+
+    def _release(self):
+        if not self.released:
+            if hasattr(self.field, "_release"):
+                self.field._release()
+            self.released = True
+
+    def metadata(self, *args, **kwargs):
+        if self.released:
+            return None
+        return self.field.metadata(*args, **kwargs)
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
diff --git a/src/earthkit/data/utils/xarray/profile.py b/src/earthkit/data/utils/xarray/profile.py
index dcf82dd..56ae896 100644
--- a/src/earthkit/data/utils/xarray/profile.py
+++ b/src/earthkit/data/utils/xarray/profile.py
@@ -145,6 +145,8 @@ class Profile:
         # generic
         self.decode_times = kwargs.pop("decode_times")
         self.decode_timedelta = kwargs.pop("decode_timedelta")
+        self.lazy_load = kwargs.pop("lazy_load")
+        self.release_source = kwargs.pop("release_source")
         self.strict = kwargs.pop("strict")
         self.errors = kwargs.pop("errors")
 
@@ -171,6 +173,7 @@ class Profile:
 
     @staticmethod
     def make(name_or_def, *args, **kwargs):
+        # print("name_or_def", name_or_def)
         if name_or_def is None:
             name_or_def = {}
 
@@ -281,14 +284,6 @@ class Profile:
         if not self.variables:
             raise ValueError(f"No metadata values found for variable key {self.variable_key}")
 
-    def update_dims(self, ds, attributes):
-        # variable keys cannot be dimensions
-        variable_keys = [self.variable_key]
-        # variable_keys = VARIABLE_KEYS + [self.variable_key]
-        # self.dims.remove(variable_keys, others=True)
-        self.dims.update(ds, attributes, variable_keys)
-        assert self.variable_key not in self.dim_keys
-
     def update(self, ds):
         """
         Parameters
@@ -314,9 +309,6 @@ class Profile:
         self.dims.update(ds)
         assert self.variable_key not in self.dim_keys
 
-        # self.update_variables(ds)
-        #         self.update_dims(ds, attributes)
-
         assert self.variable_key is not None
         assert self.variables
         assert self.variable_key not in self.dim_keys
diff --git a/src/earthkit/data/utils/xarray/splitter.py b/src/earthkit/data/utils/xarray/splitter.py
index 1085ebd..6528c1a 100644
--- a/src/earthkit/data/utils/xarray/splitter.py
+++ b/src/earthkit/data/utils/xarray/splitter.py
@@ -8,7 +8,6 @@
 #
 
 import logging
-import math
 from abc import ABCMeta
 from abc import abstractmethod
 
@@ -26,15 +25,13 @@ class Splitter(metaclass=ABCMeta):
         pass
 
     @staticmethod
-    def make(auto_split, split_dims):
+    def make(split_dims, auto_split=False):
+        """TODO: auto_split is not implemented"""
         if not auto_split and not split_dims:
             return NoSplitter()
         elif split_dims:
             return DimSplitter(split_dims)
-        elif auto_split:
-            return AutoSplitter()
-        else:
-            raise ValueError("Invalid split configuration")
+        raise ValueError("Invalid split configuration")
 
 
 class NoSplitter(Splitter):
@@ -68,54 +65,54 @@ class DimSplitter(Splitter):
             yield profile.dims.to_list(), ds_sel
 
 
-class AutoSplitter(Splitter):
-    def split(self, ds, profile):
-        grids = self.grids(ds)
-        dims = {k: ds.index(k) for k in profile.dim_keys}
-
-        for ds_gr, grid in self.ds_grids(grids, ds):
-            holes = {k: False for k in profile.dim_keys}
-            cube_vars = []
-            dims = {k: ds_gr.index(k) for k in profile.dim_keys}
-            t_dims = []
-            for d in dims:
-                if len(ds_gr.index(d)) > 1 or (not profile.dims.squeeze and len(ds_gr.index(d)) == 1):
-                    t_dims.append(d)
-
-            # try to see if dims form a cube
-            for v in profile.variables:
-                keys = {profile.variable_key: v}
-                ds_var = ds_gr.sel(**keys)
-
-                if len(ds_var) > 0:
-                    n = math.prod([len(ds_var.index(d)) for d in profile.dim_keys])
-                    # print(f" -> n={n} len={len(ds_var)}")
-                    if n == len(ds_var):
-                        cube_vars.append(v)
-                    else:
-                        for h in self.dims_hole(ds_var, t_dims):
-                            holes[h] = True
-
-            holes = [k for k in holes if holes[k]]
-            # print("holes=", holes)
-            if not holes:
-                yield t_dims, ds_gr
-            else:
-                splitter = DimSplitter(holes)
-                yield splitter.split(ds, profile)
-
-    def ds_grids(self, grids, ds):
-        if len(grids) == 1:
-            yield ds
-        elif len(grids) > 1:
-            for gr in grids:
-                ds_sel = ds.sel(md5GridSection=gr)
-                yield ds_sel, gr
-
-    def dims_hole(self, ds, full_dims):
-        holes = []
-        dims = {k: ds.index(k) for k in full_dims}
-        for d in dims:
-            if dims[d] != full_dims[d]:
-                holes.append(d)
-        return holes
+# class AutoSplitter(Splitter):
+#     def split(self, ds, profile):
+#         grids = self.grids(ds)
+#         dims = {k: ds.index(k) for k in profile.dim_keys}
+
+#         for ds_gr, grid in self.ds_grids(grids, ds):
+#             holes = {k: False for k in profile.dim_keys}
+#             cube_vars = []
+#             dims = {k: ds_gr.index(k) for k in profile.dim_keys}
+#             t_dims = []
+#             for d in dims:
+#                 if len(ds_gr.index(d)) > 1 or (not profile.dims.squeeze and len(ds_gr.index(d)) == 1):
+#                     t_dims.append(d)
+
+#             # try to see if dims form a cube
+#             for v in profile.variables:
+#                 keys = {profile.variable_key: v}
+#                 ds_var = ds_gr.sel(**keys)
+
+#                 if len(ds_var) > 0:
+#                     n = math.prod([len(ds_var.index(d)) for d in profile.dim_keys])
+#                     # print(f" -> n={n} len={len(ds_var)}")
+#                     if n == len(ds_var):
+#                         cube_vars.append(v)
+#                     else:
+#                         for h in self.dims_hole(ds_var, t_dims):
+#                             holes[h] = True
+
+#             holes = [k for k in holes if holes[k]]
+#             # print("holes=", holes)
+#             if not holes:
+#                 yield t_dims, ds_gr
+#             else:
+#                 splitter = DimSplitter(holes)
+#                 yield splitter.split(ds, profile)
+
+#     def ds_grids(self, grids, ds):
+#         if len(grids) == 1:
+#             yield ds
+#         elif len(grids) > 1:
+#             for gr in grids:
+#                 ds_sel = ds.sel(md5GridSection=gr)
+#                 yield ds_sel, gr
+
+#     def dims_hole(self, ds, full_dims):
+#         holes = []
+#         dims = {k: ds.index(k) for k in full_dims}
+#         for d in dims:
+#             if dims[d] != full_dims[d]:
+#                 holes.append(d)
+#         return holes


Test patch:
diff --git a/tests/grib/test_grib_url_stream.py b/tests/grib/test_grib_url_stream.py
index 4c88718..b388398 100644
--- a/tests/grib/test_grib_url_stream.py
+++ b/tests/grib/test_grib_url_stream.py
@@ -218,10 +218,6 @@ def test_grib_multi_url_stream_iter():
 
     assert isinstance(ds, StreamFieldList)
     assert len(ds._source.sources) == 2
-    assert ds._source._status() == [
-        {"reader": True, "stream": True},
-        {"reader": False, "stream": False},
-    ]
 
     # no fieldlist methods are available
     with pytest.raises((TypeError, NotImplementedError)):
@@ -245,11 +241,6 @@ def test_grib_multi_url_stream_iter():
     # stream consumed, no data is available
     assert sum([1 for _ in ds]) == 0
 
-    assert ds._source._status() == [
-        {"reader": True, "stream": True},
-        {"reader": True, "stream": True},
-    ]
-
 
 @pytest.mark.parametrize(
     "_kwargs,expected_meta",
diff --git a/tests/xr_engine/test_xr_engine.py b/tests/xr_engine/test_xr_engine.py
index edc71b9..6067176 100644
--- a/tests/xr_engine/test_xr_engine.py
+++ b/tests/xr_engine/test_xr_engine.py
@@ -9,6 +9,8 @@
 # nor does it submit to any jurisdiction.
 #
 
+import os
+import sys
 
 import numpy as np
 import pytest
@@ -16,6 +18,10 @@ import pytest
 from earthkit.data import from_source
 from earthkit.data.testing import earthkit_remote_test_data_file
 
+here = os.path.dirname(__file__)
+sys.path.insert(0, here)
+from xr_engine_fixtures import load_grib_data  # noqa: E402
+
 
 @pytest.mark.cache
 @pytest.mark.parametrize(
@@ -234,8 +240,12 @@ def test_xr_engine_detailed_check(api):
 
 
 @pytest.mark.cache
-def test_xr_engine_detailed_flatten_check():
-    ds_ek = from_source("url", earthkit_remote_test_data_file("test-data", "xr_engine", "level", "pl.grib"))
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("lazy_load", [False, True])
+@pytest.mark.parametrize("release_source", [False, True])
+def test_xr_engine_detailed_flatten_check(stream, lazy_load, release_source):
+    filename = "test-data/xr_engine/level/pl.grib"
+    ds_ek, ds_ek_ref = load_grib_data(filename, "url", stream=stream)
 
     kwargs = {
         "xarray_open_dataset_kwargs": {
@@ -245,6 +255,8 @@ def test_xr_engine_detailed_flatten_check():
                 "decode_timedelta": False,
                 "flatten_values": True,
                 "add_valid_time_coord": False,
+                "lazy_load": lazy_load,
+                "release_source": release_source,
             }
         }
     }
@@ -253,7 +265,7 @@ def test_xr_engine_detailed_flatten_check():
     assert ds is not None
 
     # dataset
-    ll = ds_ek[0].to_latlon(flatten=True)
+    ll = ds_ek_ref[0].to_latlon(flatten=True)
     lats = ll["lat"]
     lons = ll["lon"]
     data_vars = ["r", "t", "u", "v", "z"]
@@ -393,7 +405,7 @@ def test_xr_engine_detailed_flatten_check():
         ]
     )
 
-    v_ek = ds_ek.sel(param="t", time=0, levelist=500).to_numpy(flatten=True)
+    v_ek = ds_ek_ref.sel(param="t", time=0, levelist=500).to_numpy(flatten=True)
     assert np.allclose(r.values.flatten(), v_ek[:, [9 * 36, 10 * 36, 11 * 36]].flatten())
     assert np.allclose(r.values, vals_ref)
 
diff --git a/tests/xr_engine/test_xr_memory.py b/tests/xr_engine/test_xr_memory.py
new file mode 100644
index 0000000..ae9972e
--- /dev/null
+++ b/tests/xr_engine/test_xr_memory.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import os
+import sys
+
+import numpy as np
+import pytest
+
+here = os.path.dirname(__file__)
+sys.path.insert(0, here)
+from xr_engine_fixtures import load_grib_data  # noqa: E402
+
+
+class Counter:
+    n = 0
+
+    def increment(self):
+        self.n += 1
+
+    def get(self):
+        return self.n
+
+
+def count_calls(monkeypatch, cls, method):
+    counter = Counter()
+
+    def patched_method(self, *args, **kwargs):
+        if not self.released:
+            nonlocal counter
+            counter.increment()
+            return method(self, *args, **kwargs)
+
+    patched_method.__name__ = method.__name__
+    monkeypatch.setattr(cls, method.__name__, patched_method)
+    return counter
+
+
+@pytest.mark.cache
+@pytest.mark.parametrize(
+    "lazy_load,release_source,expected_result",
+    [
+        (False, False, {"call_count": 0}),
+        (False, True, {"call_count": 32}),
+        (True, False, {"call_count": 0}),
+        (True, True, {"call_count": 0}),
+    ],
+)
+def test_xr_engine_stream_release_source(lazy_load, release_source, expected_result, monkeypatch):
+    filename = "test-data/xr_engine/level/pl_small.grib"
+    ds_ek, _ = load_grib_data(filename, "url", stream=True)
+    ds_ek_ref, _ = load_grib_data(filename, "url", stream=False)
+
+    from earthkit.data.utils.xarray.fieldlist import ReleasableField
+
+    calls_cnt = count_calls(monkeypatch, ReleasableField, ReleasableField._release)
+
+    assert calls_cnt.get() == 0
+
+    kwargs = {
+        "lazy_load": lazy_load,
+        "release_source": release_source,
+    }
+
+    ds = ds_ek.to_xarray(**kwargs)
+    assert ds is not None
+
+    assert calls_cnt.get() == expected_result["call_count"]
+
+    # the stream is exhausted
+    with pytest.raises(StopIteration):
+        next(iter(ds_ek))
+
+    assert np.allclose(
+        ds_ek_ref.sel(param="t").order_by(["date", "time", "step", "levelist"]).to_numpy(),
+        ds["t"].values.reshape((16, 19, 36)),
+    )
+
+
+@pytest.mark.cache
+@pytest.mark.parametrize(
+    "lazy_load,release_source,expected_result",
+    [
+        (False, False, {"call_count": 0, "param": "t"}),
+        (False, True, {"call_count": 32, "param": None}),
+        (True, False, {"call_count": 0, "param": "t"}),
+        (True, True, {"call_count": 0, "param": "t"}),
+    ],
+)
+def test_xr_engine_array_field_release_source(lazy_load, release_source, expected_result, monkeypatch):
+    filename = "test-data/xr_engine/level/pl_small.grib"
+    ds_ek, _ = load_grib_data(filename, "url", stream=False)
+    ds_ek = ds_ek.to_fieldlist()
+
+    from earthkit.data.utils.xarray.fieldlist import ReleasableField
+
+    calls_cnt = count_calls(monkeypatch, ReleasableField, ReleasableField._release)
+
+    assert calls_cnt.get() == 0
+
+    kwargs = {
+        "lazy_load": lazy_load,
+        "release_source": release_source,
+    }
+
+    ds = ds_ek.to_xarray(**kwargs)
+    assert ds is not None
+
+    assert calls_cnt.get() == expected_result["call_count"]
+
+    if not lazy_load and release_source:
+        with pytest.raises(AttributeError):
+            ds_ek[0].metadata("param")
+    else:
+        assert np.allclose(
+            ds_ek.sel(param="t").order_by(["date", "time", "step", "levelist"]).to_numpy(),
+            ds["t"].values.reshape((16, 19, 36)),
+        )
+
+        assert ds_ek[0].metadata("param") == expected_result["param"]
diff --git a/tests/xr_engine/xr_engine_fixtures.py b/tests/xr_engine/xr_engine_fixtures.py
index eb08e79..713f95e 100644
--- a/tests/xr_engine/xr_engine_fixtures.py
+++ b/tests/xr_engine/xr_engine_fixtures.py
@@ -12,7 +12,10 @@
 import numpy as np
 
 from earthkit.data import from_source
-from earthkit.data.utils.xarray.fieldlist import WrappedFieldList
+from earthkit.data.testing import earthkit_examples_file
+from earthkit.data.testing import earthkit_remote_test_data_file
+from earthkit.data.testing import earthkit_test_data_file
+from earthkit.data.utils.xarray.fieldlist import XArrayInputFieldList
 
 
 def load_fieldlist(d):
@@ -47,7 +50,7 @@ def load_wrapped_fieldlist(d, profile, **kwargs):
     else:
         raise ValueError(f"Invalid data type={type(d)}")
 
-    return WrappedFieldList(ds, keys=profile.index_keys, **kwargs)
+    return XArrayInputFieldList(ds, keys=profile.index_keys, **kwargs)
 
 
 def compare_dims(ds, ref_coords, order_ref_var=None, sizes=False):
@@ -106,3 +109,30 @@ def compare_dim_size(ds, dims):
     for name, v in dims.items():
         assert name in ds.sizes, f"{name=} not in {ds.sizes}"
         assert ds.sizes[name] == v, f"{name=} {ds.sizes[name]} != {v}"
+
+
+def load_grib_data(filename, source, stream=False, folder="example"):
+    if source == "url":
+        path = earthkit_remote_test_data_file(filename)
+    else:
+        if folder == "example":
+            path = earthkit_examples_file(filename)
+        elif folder == "data":
+            path = earthkit_test_data_file(filename)
+        else:
+            raise ValueError("Invalid folder={folder}")
+
+    if source == "file":
+        ds = from_source("file", path)
+    elif source == "url":
+        ds = from_source("url", path)
+    else:
+        raise ValueError("Invalid source={source}")
+
+    ds_ref = ds
+
+    if stream:
+        f = open(ds.path, "rb")
+        ds = from_source("stream", f)
+
+    return ds, ds_ref