Temporary fix to ensure coordinate realisation during pp load (#6767)

* put in a temporary fix to allow data realisation with derived coordinates during pp load

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* move and rename context manager

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add documentation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add whatsnew

* Update lib/iris/loading.py

Co-authored-by: Patrick Peglar <patrick.peglar@metoffice.gov.uk>

* address review comments

* refactor tests

* lint fix

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Patrick Peglar <patrick.peglar@metoffice.gov.uk>

Author: 3 people

docs/src/whatsnew/latest.rst

@@ -168,6 +168,12 @@ This document explains the changes made to Iris for this release
 #. `@ukmo-ccbunney`_ ``assert_CML`` now uses stricter array formatting to avoid
    changes in tests due to Numpy version changes. (:pull:`6743`)
 
+#. `@stephenworsley`_ added a private switch :obj:`~iris.loading._CONCRETE_DERIVED_LOADING`
+   for controlling laziness of coordinates from pp loading, avoiding a
+   slowdown due to merging. Note: this object is temporary and is likely
+   to be replaced by a permanent solution or else be renamed.
+   (:issue:`6755`, :pull:`6767`)
+
 
 .. comment
     Whatsnew author names (@github name) in alphabetical order. Note that,

lib/iris/fileformats/rules.py

@@ -14,6 +14,7 @@
 import iris.cube
 import iris.exceptions
 import iris.fileformats.um_cf_map
+from iris.loading import _CONCRETE_DERIVED_LOADING
 import iris.warnings
 
 Factory = collections.namedtuple("Factory", ["factory_class", "args"])
@@ -158,8 +159,12 @@ def _dereference_args(factory, reference_targets, regrid_cache, cube):
                 # match the grid of this cube.
                 src, cube = _ensure_aligned(regrid_cache, src, cube)
                 if src is not None:
+                    if _CONCRETE_DERIVED_LOADING:
+                        data = src.data
+                    else:
+                        data = src.core_data()
                     new_coord = iris.coords.AuxCoord(
-                        src.core_data(),
+                        data,
                         src.standard_name,
                         src.long_name,
                         src.var_name,

lib/iris/loading.py

@@ -4,6 +4,7 @@
 # See LICENSE in the root of the repository for full licensing details.
 """Iris general file loading mechanism."""
 
+from contextlib import contextmanager
 from dataclasses import dataclass
 import itertools
 import threading
@@ -15,6 +16,55 @@
 from iris.warnings import IrisLoadWarning
 
 
+class _ConcreteDerivedLoading(threading.local):
+    """An object to control whether factory references are loaded with lazy or real data.
+
+    Use via the run-time switch :const:`~iris.loading._CONCRETE_DERIVED_LOADING`.
+    Use :meth:`context` to temporarily activate or assign the property
+    `load_real_references` to True.
+
+    Notes
+    -----
+    This is intended as a temporary fix for a problem with PP loading causing
+    performance issues during loading (see https://github.com/SciTools/iris/issues/6755)
+    This is expected to be either removed or renamed in a future version.
+    """
+
+    def __init__(self):
+        self.load_real_references = False
+
+    def __bool__(self):
+        return self.load_real_references
+
+    @contextmanager
+    def context(self, load_real=True):
+        """Control whether references are loaded lazily.
+
+        Defines a context block within which the setting is applied.
+
+        Parameters
+        ----------
+        load_real : bool, default True
+            Sets whether references are loaded as concrete data, within the context.
+
+        Example
+        -------
+        >>> with _CONCRETE_DERIVED_LOADING.context(load_real=True):
+        ...     <code>
+        """
+        try:
+            old_value = self.load_real_references
+            self.load_real_references = load_real
+            yield
+        finally:
+            self.load_real_references = old_value
+
+
+# TODO: this is a temporary fix, either remove this when a permanent fix exists
+#  or else make this public if this is deemed necessary.
+_CONCRETE_DERIVED_LOADING = _ConcreteDerivedLoading()
+
+
 def _generate_cubes(uris, callback, constraints):
     import iris.io
 

lib/iris/tests/integration/test_pp.py

@@ -12,6 +12,7 @@
 from unittest import mock
 
 from cf_units import Unit
+import dask.array as da
 import numpy as np
 
 from iris.aux_factory import HybridHeightFactory, HybridPressureFactory
@@ -22,6 +23,7 @@
 from iris.fileformats.pp import load_pairs_from_fields
 import iris.fileformats.pp_load_rules
 from iris.fileformats.pp_save_rules import verify
+from iris.loading import _CONCRETE_DERIVED_LOADING
 import iris.util
 from iris.warnings import IrisUserWarning
 
@@ -144,9 +146,12 @@ def test_potential_temperature_level_round_trip(self):
         self.assertEqual(field.blev, potm_value)
 
     @staticmethod
-    def _field_with_data(scale=1, **kwargs):
+    def _field_with_data(scale=1, lazy=False, **kwargs):
         x, y = 40, 30
-        mock_data = np.arange(1200).reshape(y, x) * scale
+        if lazy:
+            mock_data = da.arange(1200).reshape(y, x) * scale
+        else:
+            mock_data = np.arange(1200).reshape(y, x) * scale
         mock_core_data = mock.MagicMock(return_value=mock_data)
         field = mock.MagicMock(
             core_data=mock_core_data,
@@ -253,6 +258,42 @@ def test_hybrid_pressure_round_trip(self):
         self.assertEqual(data_field.brlev, sigma_lower)
         self.assertEqual(data_field.brsvd, [sigma_upper, delta_upper])
 
+    def test_hybrid_pressure_lazy_load(self):
+        pressure_field = self._field_with_data(
+            10,
+            lazy=True,
+            stash=iris.fileformats.pp.STASH(1, 0, 409),
+            lbuser=[0, 0, 0, 409, 0, 0, 0],
+        )
+
+        # Make a fake data field which needs the reference surface.
+        model_level = 5678
+        sigma_lower, sigma, sigma_upper = 0.85, 0.9, 0.95
+        delta_lower, delta, delta_upper = 0.05, 0.1, 0.15
+        data_field = self._field_with_data(
+            lbvc=9,
+            lblev=model_level,
+            bhlev=delta,
+            bhrlev=delta_lower,
+            blev=sigma,
+            brlev=sigma_lower,
+            brsvd=[sigma_upper, delta_upper],
+        )
+
+        # Convert both fields to cubes.
+        load = mock.Mock(return_value=iter([pressure_field, data_field]))
+        with mock.patch("iris.fileformats.pp.load", new=load) as load:
+            pressure_cube, data_cube = iris.fileformats.pp.load_cubes("DUMMY")
+
+        assert data_cube.coord("surface_air_pressure").has_lazy_points()
+
+        # TODO: _CONCRETE_DERIVED_LOADING is a temporary fix, remove from test when a permanent fix exists
+        load = mock.Mock(return_value=iter([pressure_field, data_field]))
+        with mock.patch("iris.fileformats.pp.load", new=load) as load:
+            with _CONCRETE_DERIVED_LOADING.context():
+                _, realised_data_cube = iris.fileformats.pp.load_cubes("DUMMY")
+        assert not realised_data_cube.coord("surface_air_pressure").has_lazy_points()
+
     def test_hybrid_pressure_with_duplicate_references(self):
         # Make a fake reference surface field.
         pressure_field = self._field_with_data(