perf(gribjump): compute coordinates once instead of per-field

Extract grid coordinates before the field loop instead of inside it,
reducing O(N) calls to geography methods to O(1).

Also make shared grid indices explicit in ExtractionRequestCollection
and replace assert with proper ValueError check.

Author: andreas-grafberger

src/earthkit/data/sources/gribjump.py

@@ -37,10 +37,6 @@ def split_mars_requests(request: dict[str, Any]) -> list[dict[str, Any]]:
     returns result arrays without metadata, so each field must be requested individually
     to map outputs correctly.
 
-    NOTE: Parsing of MARS requests should ideally not be handled here but in a dedicated
-    component like pymetkit. Consider updating this function once something appropriate
-    is available.
-
     Parameters
     ----------
     request : dict[str, Any]
@@ -129,17 +125,13 @@ def mask_to_ranges(mask: np.ndarray) -> list[tuple[int, int]]:
 
 @dataclasses.dataclass
 class ExtractionRequest:
-    """
-    Simple wrapper of pygribjump.ExtractionRequest and the original FDB request dict.
-
-    Can be removed once pygribjump.ExtractionRequest provides a reference to the request dictionary
-    with original MARS keyword types.
+    """Wrapper around pygribjump.ExtractionRequest that also stores the original request dict.
 
     Parameters
     ----------
     extraction_request : pygj.ExtractionRequest
         The GribJump extraction request object.
-    request : dict[str, str]
+    request : dict[str, Any]
         The original request dictionary used to create the extraction request.
     """
 
@@ -162,8 +154,7 @@ def build_extraction_request(
     mask: Optional[np.ndarray] = None,
     indices: Optional[np.ndarray] = None,
 ) -> ExtractionRequest:
-    """
-    Builds an ExtractionRequest from the given request dictionary and optional parameters.
+    """Builds an ExtractionRequest from the given request dictionary and optional parameters.
 
     Parameters
     ----------
@@ -201,6 +192,17 @@ def build_extraction_request(
 
 
 class ExtractionRequestCollection(UserList):
+    """Collection of extraction requests sharing the same grid indices."""
+
+    def __init__(self, extraction_requests: list[ExtractionRequest], indices: np.ndarray):
+        """Internal constructor. Use from_mars_requests() to create instances."""
+        super().__init__(extraction_requests)
+        self._indices = indices
+
+    @property
+    def indices(self) -> np.ndarray:
+        """Grid indices shared by all extraction requests."""
+        return self._indices
 
     @classmethod
     def from_mars_requests(
@@ -210,25 +212,7 @@ def from_mars_requests(
         mask: Optional[np.ndarray] = None,
         indices: Optional[np.ndarray] = None,
     ) -> "ExtractionRequestCollection":
-        """Creates an ExtractionRequestCollection from MARS requests.
-
-        One of the parameters `ranges`, `mask`, or `indices` must be provided.
-
-        Parameters
-        ----------
-        mars_requests : list[dict[str, str]]
-            List of MARS requests, each represented as a dictionary of keywords.
-        ranges : Optional[list[tuple[int, int]]], optional
-            The ranges for the extraction requests, by default None.
-        mask : Optional[np.ndarray], optional
-            The mask for the extraction requests, by default None.
-        indices : Optional[np.ndarray], optional
-            The indices for the extraction requests, by default None.
-        Returns
-        -------
-        ExtractionRequestCollection
-            A collection of ExtractionRequest objects created from the MARS requests.
-        """
+        """Create collection from MARS requests. Exactly one of ranges/mask/indices must be provided."""
 
         if sum(opt is not None for opt in (ranges, mask, indices)) != 1:
             raise ValueError(
@@ -243,34 +227,39 @@ def from_mars_requests(
             mask = None
 
         extraction_requests = [build_extraction_request(req, ranges, mask, indices) for req in mars_requests]
-        return cls(extraction_requests)
+
+        # All requests share the same indices; get from the first one
+        if not extraction_requests:
+            raise ValueError("Cannot create ExtractionRequestCollection from empty mars_requests list")
+        canonical_indices = extraction_requests[0].indices()
+
+        return cls(extraction_requests, canonical_indices)
 
 
 class FieldExtractList(SimpleFieldList):
     """Lazily loaded representation of points extracted from multiple fields using GribJump.
 
-    .. warning::
-        This implementation is **not thread-safe**. Concurrent access from multiple threads
-        may result in race conditions during lazy loading. Use appropriate synchronization
-        if accessing from multiple threads.
-
     .. note::
         This class should not be instantiated directly. Use the ``gribjump`` source instead:
         ``earthkit.data.from_source("gribjump", request, ranges=ranges)``
 
-    This class inherits from SimpleFieldList and provides lazy loading of grid point
-    extractions from GRIB fields via GribJump. FieldList operations like ``sel()``,
-    ``group_by()``, etc. might work but are not guaranteed to be fully functional.
+    .. warning::
+        This implementation is **not thread-safe**. Concurrent access from multiple threads
+        may result in race conditions during lazy loading.
+
+    This class provides lazy loading of grid point extractions from GRIB fields via GribJump.
+    FieldList operations like ``sel()``, ``group_by()``, etc. might work but are not guaranteed
+    to be fully functional.
 
     Known Limitations
     -----------------
-    * **No validation**: Grid indices are not validated against actual field grids.
-        Incorrect indices may return unexpected grid points.
-    * **Not thread-safe**: Concurrent access may cause race conditions during lazy loading.
-    * **Limited metadata**: Only metadata from the request dictionary is available,
-        except for latitude/longitude coordinates when ``fetch_coords_from_fdb=True`` is used.
-    * **No efficient slicing**: Lazy loading of selections/slices is not supported.
-    * **Serialization issues**: Pickling/unpickling might not work reliably.
+    * Grid indices are not validated against actual field grids. Incorrect indices may return
+      unexpected grid points.
+    * Not thread-safe - concurrent access may cause race conditions during lazy loading.
+    * Limited metadata - only metadata from the request dictionary is available, except for
+      latitude/longitude coordinates when ``fetch_coords_from_fdb=True`` is used.
+    * No efficient slicing - lazy loading of selections/slices is not supported.
+    * Serialization issues - pickling/unpickling might not work reliably.
     """
 
     def __init__(
@@ -283,7 +272,6 @@ def __init__(
         self._requests = requests
         self._fdb_retriever = fdb_retriever
 
-        # These attributes are set lazily after loading the data.
         self._loaded = False
         self._grid_indices = None
         self._reference_metadata: Optional[GribMetadata] = None
@@ -307,23 +295,26 @@ def _load(self):
         context = {"origin": "earthkit-data"}
         extraction_results = self._gj.extract(extraction_requests, ctx=context)
 
+        # Get the shared indices from the collection
+        indices = self._requests.indices
+
+        # Pre-compute grid coordinates once for all fields
+        coords = self._get_grid_coordinates(indices)
+
         fields = []
-        indices = None
-        ranges = None
         for request, result in zip(self._requests, extraction_results):
-            if ranges is None:
-                ranges = request.ranges
-                indices = request.indices()
-            else:
-                if request.ranges != ranges:
-                    raise ValueError(
-                        f"Extraction request has different ranges than the first request: {request.ranges} != {ranges}"
-                    )
             arr = result.values_flat
             shape = arr.shape
 
             metadata = UserMetadata(request.request, shape=shape)
-            metadata = self._enrich_metadata_with_coordinates(indices, metadata)
+            if coords is not None:
+                grid_latitudes, grid_longitudes = coords
+                metadata = metadata.override(
+                    {
+                        "latitudes": grid_latitudes,
+                        "longitudes": grid_longitudes,
+                    }
+                )
 
             field = ArrayField(arr, metadata)
             fields.append(field)
@@ -350,21 +341,26 @@ def _load_reference_metadata(self):
         self._reference_metadata = metadata
         return metadata
 
-    def _enrich_metadata_with_coordinates(self, indices: np.ndarray, metadata: UserMetadata) -> UserMetadata:
-        """Enriches the metadata with coordinates if reference metadata is available."""
+    def _get_grid_coordinates(self, indices: np.ndarray) -> Optional[tuple[np.ndarray, np.ndarray]]:
+        """Get latitude/longitude coordinates at the specified grid indices.
+
+        Parameters
+        ----------
+        indices : np.ndarray
+            Grid indices to extract coordinates for.
+
+        Returns
+        -------
+        Optional[tuple[np.ndarray, np.ndarray]]
+            A tuple of (latitudes, longitudes) arrays, or None if no reference metadata is available.
+        """
         if (reference_metadata := self._load_reference_metadata()) is None:
-            return metadata
+            return None
 
         reference_geography = reference_metadata.geography
         grid_latitudes = reference_geography.latitudes()[indices]
         grid_longitudes = reference_geography.longitudes()[indices]
-        metadata = metadata.override(
-            {
-                "latitudes": grid_latitudes,
-                "longitudes": grid_longitudes,
-            }
-        )
-        return metadata
+        return (grid_latitudes, grid_longitudes)
 
     def to_xarray(self, *args, **kwargs):
         kwargs = kwargs.copy()
@@ -424,11 +420,12 @@ def __init__(
             A 1D array of grid indices to retrieve, by default None.
         fetch_coords_from_fdb : bool, optional
             If set to True, loads the first field's metadata from the FDB to extract the coordinates
-            at the specified indices.
+            at the specified indices, by default False.
         fdb_kwargs : Optional[dict[str, Any]], optional
             Only used when `fetch_coords_from_fdb=True`. A dict of
             keyword arguments passed to the `pyfdb.FDB` constructor. These arguments are only passed
-            to the FDB when fetching coordinates and is not used by GribJump for the extraction itself.
+            to the FDB when fetching coordinates and is not used by GribJump for the extraction itself,
+            by default None.
         """
 
         super().__init__(**kwargs)