Backport PR #5581 on branch yt-4.4.x ([BUG]: Fix reloaded sphere crashing due to chunking)

yt/frontends/ytdata/io.py

@@ -212,10 +212,11 @@ def _read_particle_data_file(self, data_file, ptf, selector):
         data_return = {}
 
         with h5py.File(data_file.filename, mode="r") as f:
+            # NOTE: f represents the full-length file
             index_mask = slice(data_file.start, data_file.end)
             for ptype, field_list in sorted(ptf.items()):
                 if selector is None or getattr(selector, "is_all_data", False):
-                    mask = index_mask
+                    mask = None
                 else:
                     units = _get_position_array_units(ptype, f, "x")
                     x, y, z = (
@@ -231,8 +232,10 @@ def _read_particle_data_file(self, data_file, ptf, selector):
                         continue
 
                 for field in field_list:
-                    data = f[ptype][field][mask].astype("float64", copy=False)
-                    data_return[ptype, field] = data
+                    data = f[ptype][field][index_mask]  # aligns field to chunk size
+                    if mask is not None:
+                        data = data[mask]  # boolean mask is chunk-length
+                    data_return[ptype, field] = data.astype("float64", copy=False)
 
         return data_return
 

yt/frontends/ytdata/tests/test_data_reload.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from yt.loaders import load_uniform_grid
-from yt.testing import requires_module_pytest
+from yt.loaders import load, load_uniform_grid
+from yt.testing import fake_amr_ds, requires_module_pytest
 from yt.utilities.on_demand_imports import _h5py as h5py
 
 
@@ -24,3 +24,34 @@ def test_save_as_data_unit_system(tmp_path):
 
     with h5py.File(fi, mode="r") as f:
         assert f.attrs["unit_system_name"] == "code"
+
+
+@requires_module_pytest("h5py")
+def test_save_as_data_chunk(tmp_path):
+    """
+    Tests whether a saved (sphere) dataset matches what was saved, and whether
+    accessing the saved data via different paths produces the same result.
+    """
+    sphere_path = tmp_path / "test_sphere.h5"
+    ds = fake_amr_ds(length_unit="kpc")
+    sp = ds.sphere(ds.domain_center, (0.25, "kpc"))
+    original_data = sp["stream", "Density"]
+    sp.save_as_dataset(sphere_path, fields=[("stream", "Density")])
+
+    sp_ds = load(sphere_path)  # should produce a 7-chunk dataset
+    assert len(sp_ds.index.data_files) > 1, "Test data not chunked."
+
+    reloaded_data = sp_ds.data["grid", "Density"]  # previously crashed with chunking
+    all_reloaded_data = sp_ds.all_data()["grid", "Density"]
+
+    np.testing.assert_array_equal(
+        original_data,
+        reloaded_data,
+        err_msg="Reloaded sphere produces different data to what was saved.",
+    )
+
+    np.testing.assert_array_equal(
+        all_reloaded_data,
+        reloaded_data,
+        err_msg="Reloaded sphere produces different results from .data vs .all_data().",
+    )