Repository: openradar/xradar

Issue:
NEXRAD issue - IndexError index out of bounds
* xradar version: 0.6.4
* xarray version: 2024.7.0
* Python version: 3.12.4
* Operating System: Linux

### Description

I have found a puzzling bug that only comes up in certain situations with Dask

### What I Did

```python
import xradar
import xarray as xr
import pooch

# download and open a NEXRAD2 file from S3
url = "https://noaa-nexrad-level2.s3.amazonaws.com/2024/09/01/FOP1/FOP120240901_000347_V06"
local_file = pooch.retrieve(url, known_hash=None)
ds = xr.open_dataset(local_file, group="sweep_0", engine="nexradlevel2")

# create a chunked version 
dsc = ds.chunk()
# load one variable - IMPORTANT - skipping this step makes the next line work
dsc.DBZH.load()
# load the entire dataset
dsc.load()
# - > IndexError: index 140 is out of bounds for axis 0 with size 38
# try all the variables
for v in dsc:
    print(v)
    try:
        dsc[v].load()  # also fails with dsc!
        print("ok")
    except Exception as e:
        print(e)
# DBZH
# ok
# ZDR
# index 212 is out of bounds for axis 0 with size 212
# PHIDP
# index 140 is out of bounds for axis 0 with size 130
# RHOHV
# index 140 is out of bounds for axis 0 with size 22
# CCORH
# index 140 is out of bounds for axis 0 with size 73
# sweep_mode
# ok
# sweep_number
# ok
# prt_mode
# ok
# follow_mode
# ok
# sweep_fixed_angle
# ok
```

Possibly related to #180.

Experience tells me this has something to do with Dask task tokenization. 


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/openradar/xradar/pull/269?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) Report
Attention: Patch coverage is `96.29630%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 93.16%. Comparing base [(`fc43e60`)](https://app.codecov.io/gh/openradar/xradar/commit/fc43e6051a3d483521482c667b30a435054bc3ea?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) to head [(`4076e29`)](https://app.codecov.io/gh/openradar/xradar/commit/4076e29459d1b6d915a8726971dc9cf993fe13a0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar).

| [Files with missing lines](https://app.codecov.io/gh/openradar/xradar/pull/269?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | Patch % | Lines |
|---|---|---|
| [xradar/io/backends/nexrad\_level2.py](https://app.codecov.io/gh/openradar/xradar/pull/269?src=pr&el=tree&filepath=xradar%2Fio%2Fbackends%2Fnexrad_level2.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar#diff-eHJhZGFyL2lvL2JhY2tlbmRzL25leHJhZF9sZXZlbDIucHk=) | 96.29% | [1 Missing :warning: ](https://app.codecov.io/gh/openradar/xradar/pull/269?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) |

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


```diff
@@            Coverage Diff             @@
##             main     #269      +/-   ##
==========================================
- Coverage   93.57%   93.16%   -0.41%     
==========================================
  Files          26       26              
  Lines        5073     5076       +3     
==========================================
- Hits         4747     4729      -18     
- Misses        326      347      +21     
```

| [Flag](https://app.codecov.io/gh/openradar/xradar/pull/269/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | Coverage Δ | |
|---|---|---|
| [notebooktests](https://app.codecov.io/gh/openradar/xradar/pull/269/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | `?` | |
| [unittests](https://app.codecov.io/gh/openradar/xradar/pull/269/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar) | `93.16% <96.29%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/openradar/xradar/pull/269?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openradar).   
: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=openradar).


Patch:
diff --git a/docs/history.md b/docs/history.md
index 1946f0a..2999ec2 100644
--- a/docs/history.md
+++ b/docs/history.md
@@ -12,6 +12,7 @@
 * FIX: Correct retrieval of intermediate records in nexrad level2 reader ({issue}`259`) ({pull}`261`) by [@kmuehlbauer](https://github.com/kmuehlbauer).
 * FIX: Test for magic number BZhX1AY&SY (where X is any number between 0..9) when retrieving BZ2 record indices in nexrad level2 reader ({issue}`264`) ({pull}`266`) by [@kmuehlbauer](https://github.com/kmuehlbauer).
 * ENH: Add message type 1 decoding to nexrad level 2 reader ({issue}`256`) ({pull}`267`) by [@kmuehlbauer](https://github.com/kmuehlbauer).
+* ENH: Introduce file locks for nexrad level2 and iris backend ({issue}`207`) ({pull}`268`) by [@kmuehlbauer](https://github.com/kmuehlbauer).
 
 ## 0.8.0 (2024-11-04)
 
diff --git a/xradar/io/backends/iris.py b/xradar/io/backends/iris.py
index 59f0d12..ee45f9f 100644
--- a/xradar/io/backends/iris.py
+++ b/xradar/io/backends/iris.py
@@ -46,6 +46,7 @@ import xarray as xr
 from xarray import DataTree
 from xarray.backends.common import AbstractDataStore, BackendArray, BackendEntrypoint
 from xarray.backends.file_manager import CachingFileManager
+from xarray.backends.locks import SerializableLock, ensure_lock
 from xarray.backends.store import StoreBackendEntrypoint
 from xarray.core import indexing
 from xarray.core.utils import FrozenDict
@@ -72,6 +73,9 @@ from .common import (
     _get_subgroup,
 )
 
+IRIS_LOCK = SerializableLock()
+
+
 #: mapping from IRIS names to CfRadial2/ODIM
 iris_mapping = {
     "DB_DBT": "DBTH",
@@ -3786,9 +3790,10 @@ class IrisArrayWrapper(BackendArray):
             self.shape = (nrays, nbins)
 
     def _getitem(self, key):
-        # read the data and put it into dict
-        self.datastore.root.get_moment(self.group, self.name)
-        return self.datastore.ds["sweep_data"][self.name][key]
+        with self.datastore.lock:
+            # read the data and put it into dict
+            self.datastore.root.get_moment(self.group, self.name)
+            return self.datastore.ds["sweep_data"][self.name][key]
 
     def __getitem__(self, key):
         return indexing.explicit_indexing_adapter(
@@ -3805,16 +3810,19 @@ class IrisStore(AbstractDataStore):
     Ported from wradlib.
     """
 
-    def __init__(self, manager, group=None):
+    def __init__(self, manager, group=None, lock=IRIS_LOCK):
         self._manager = manager
         self._group = int(group[6:]) + 1
         self._filename = self.filename
         self._need_time_recalc = False
+        self.lock = ensure_lock(lock)
 
     @classmethod
-    def open(cls, filename, mode="r", group=None, **kwargs):
+    def open(cls, filename, mode="r", group=None, lock=None, **kwargs):
+        if lock is None:
+            lock = IRIS_LOCK
         manager = CachingFileManager(IrisRawFile, filename, mode=mode, kwargs=kwargs)
-        return cls(manager, group=group)
+        return cls(manager, group=group, lock=lock)
 
     @property
     def filename(self):
@@ -3991,6 +3999,7 @@ class IrisBackendEntrypoint(BackendEntrypoint):
         use_cftime=None,
         decode_timedelta=None,
         group=None,
+        lock=None,
         first_dim="auto",
         reindex_angle=False,
         fix_second_angle=False,
@@ -4000,6 +4009,7 @@ class IrisBackendEntrypoint(BackendEntrypoint):
         store = IrisStore.open(
             filename_or_obj,
             group=group,
+            lock=lock,
             loaddata=False,
         )
 
diff --git a/xradar/io/backends/nexrad_level2.py b/xradar/io/backends/nexrad_level2.py
index 27c1485..1cf9509 100644
--- a/xradar/io/backends/nexrad_level2.py
+++ b/xradar/io/backends/nexrad_level2.py
@@ -44,6 +44,7 @@ import xarray as xr
 from xarray import DataTree
 from xarray.backends.common import AbstractDataStore, BackendArray, BackendEntrypoint
 from xarray.backends.file_manager import CachingFileManager
+from xarray.backends.locks import SerializableLock, ensure_lock
 from xarray.backends.store import StoreBackendEntrypoint
 from xarray.core import indexing
 from xarray.core.utils import FrozenDict, close_on_error
@@ -78,6 +79,9 @@ from .iris import (
     string_dict,
 )
 
+NEXRADL2_LOCK = SerializableLock()
+
+
 #: mapping from NEXRAD names to CfRadial2/ODIM
 nexrad_mapping = {
     "REF": "DBZH",
@@ -1324,36 +1328,32 @@ class NexradLevel2ArrayWrapper(BackendArray):
         self.shape = (nrays, nbins)
 
     def _getitem(self, key):
-        # read the data if not available
-        try:
-            data = self.datastore.ds["sweep_data"][self.name]["data"]
-            print("ZZZZZAA:", self.name, data, key)
-        except KeyError:
-            print("XXXXX:", self.group, self.name)
-            self.datastore.root.get_data(self.group, self.name)
-            data = self.datastore.ds["sweep_data"][self.name]["data"]
-            print("ZZZZZBB:", self.name, data, key)
-            print("YY0:", self.name, len(data), len(data[0]))
-        # see 3.2.4.17.6 Table XVII-I Data Moment Characteristics and Conversion for Data Names
-        word_size = self.datastore.ds["sweep_data"][self.name]["word_size"]
-        if self.name == "PHI" and word_size == 16:
-            # 10 bit mask, but only for 2 byte data
-            x = np.uint16(0x3FF)
-        elif self.name == "ZDR" and word_size == 16:
-            # 11 bit mask, but only for 2 byte data
-            x = np.uint16(0x7FF)
-        else:
-            x = np.uint8(0xFF)
-        print("YY1:", self.name, len(data[0]), self.shape)
-        if len(data[0]) < self.shape[1]:
-            return np.pad(
-                np.vstack(data) & x,
-                ((0, 0), (0, self.shape[1] - len(data[0]))),
-                mode="constant",
-                constant_values=0,
-            )[key]
-        else:
-            return (np.vstack(data) & x)[key]
+        with self.datastore.lock:
+            # read the data if not available
+            try:
+                data = self.datastore.ds["sweep_data"][self.name]["data"]
+            except KeyError:
+                self.datastore.root.get_data(self.group, self.name)
+                data = self.datastore.ds["sweep_data"][self.name]["data"]
+            # see 3.2.4.17.6 Table XVII-I Data Moment Characteristics and Conversion for Data Names
+            word_size = self.datastore.ds["sweep_data"][self.name]["word_size"]
+            if self.name == "PHI" and word_size == 16:
+                # 10 bit mask, but only for 2 byte data
+                x = np.uint16(0x3FF)
+            elif self.name == "ZDR" and word_size == 16:
+                # 11 bit mask, but only for 2 byte data
+                x = np.uint16(0x7FF)
+            else:
+                x = np.uint8(0xFF)
+            if len(data[0]) < self.shape[1]:
+                return np.pad(
+                    np.vstack(data) & x,
+                    ((0, 0), (0, self.shape[1] - len(data[0]))),
+                    mode="constant",
+                    constant_values=0,
+                )[key]
+            else:
+                return (np.vstack(data) & x)[key]
 
     def __getitem__(self, key):
         return indexing.explicit_indexing_adapter(
@@ -1365,24 +1365,29 @@ class NexradLevel2ArrayWrapper(BackendArray):
 
 
 class NexradLevel2Store(AbstractDataStore):
-    def __init__(self, manager, group=None):
+    def __init__(self, manager, group=None, lock=NEXRADL2_LOCK):
         self._manager = manager
         self._group = int(group[6:])
         self._filename = self.filename
+        self.lock = ensure_lock(lock)
 
     @classmethod
-    def open(cls, filename, mode="r", group=None, **kwargs):
+    def open(cls, filename, mode="r", group=None, lock=None, **kwargs):
+        if lock is None:
+            lock = NEXRADL2_LOCK
         manager = CachingFileManager(
             NEXRADLevel2File, filename, mode=mode, kwargs=kwargs
         )
-        return cls(manager, group=group)
+        return cls(manager, group=group, lock=lock)
 
     @classmethod
-    def open_groups(cls, filename, groups, mode="r", **kwargs):
+    def open_groups(cls, filename, groups, mode="r", lock=None, **kwargs):
+        if lock is None:
+            lock = NEXRADL2_LOCK
         manager = CachingFileManager(
             NEXRADLevel2File, filename, mode=mode, kwargs=kwargs
         )
-        return {group: cls(manager, group=group) for group in groups}
+        return {group: cls(manager, group=group, lock=lock) for group in groups}
 
     @property
     def filename(self):
@@ -1534,6 +1539,7 @@ class NexradLevel2BackendEntrypoint(BackendEntrypoint):
         use_cftime=None,
         decode_timedelta=None,
         group=None,
+        lock=None,
         first_dim="auto",
         reindex_angle=False,
         fix_second_angle=False,
@@ -1543,6 +1549,7 @@ class NexradLevel2BackendEntrypoint(BackendEntrypoint):
         store = NexradLevel2Store.open(
             filename_or_obj,
             group=group,
+            lock=lock,
             loaddata=False,
         )
 
@@ -1610,6 +1617,7 @@ def open_nexradlevel2_datatree(
     fix_second_angle=False,
     site_coords=True,
     optional=True,
+    lock=None,
     **kwargs,
 ):
     """Open a NEXRAD Level2 dataset as an `xarray.DataTree`.
@@ -1732,6 +1740,7 @@ def open_nexradlevel2_datatree(
         fix_second_angle=fix_second_angle,
         site_coords=site_coords,
         optional=optional,
+        lock=lock,
         **kwargs,
     )
     ls_ds: list[xr.Dataset] = [sweep_dict[sweep] for sweep in sweep_dict.keys()]
@@ -1771,10 +1780,12 @@ def open_sweeps_as_dict(
     fix_second_angle=False,
     site_coords=True,
     optional=True,
+    lock=None,
     **kwargs,
 ):
     stores = NexradLevel2Store.open_groups(
         filename=filename_or_obj,
+        lock=lock,
         groups=sweeps,
     )
     groups_dict = {}


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index eaee33b..0638785 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -3,3 +3,17 @@
 # Distributed under the MIT License. See LICENSE for more info.
 
 """Unit test package for xradar."""
+
+import importlib
+
+import pytest
+
+
+def skip_import(name):
+    try:
+        importlib.import_module(name)
+        found = True
+    except ImportError:
+        found = False
+
+    return pytest.mark.skipif(not found, reason=f"requires {name}")
diff --git a/tests/io/test_io.py b/tests/io/test_io.py
index 1b7c871..67a9e5a 100644
--- a/tests/io/test_io.py
+++ b/tests/io/test_io.py
@@ -14,6 +14,7 @@ import pytest
 import xarray as xr
 
 import xradar.io
+from tests import skip_import
 from xradar.io import (
     open_cfradial1_datatree,
     open_datamet_datatree,
@@ -1101,3 +1102,20 @@ def test_open_nexradlevel2_datatree(nexradlevel2_files):
         }
         assert np.round(ds.elevation.mean().values.item(), 1) == elevations[i]
         assert ds.sweep_number.values == int(grp[6:])
+
+
+@skip_import("dask")
+@pytest.mark.parametrize(
+    "nexradlevel2_files", ["nexradlevel2_gzfile", "nexradlevel2_bzfile"], indirect=True
+)
+def test_nexradlevel2_dask_load(nexradlevel2_files):
+    ds = xr.open_dataset(nexradlevel2_files, group="sweep_0", engine="nexradlevel2")
+    dsc = ds.chunk()
+    dsc.load()
+
+
+@skip_import("dask")
+def test_iris_dask_load(iris0_file):
+    ds = xr.open_dataset(iris0_file, group="sweep_0", engine="iris")
+    dsc = ds.chunk()
+    dsc.load()