Fix test_time_is_age test

Author: VeckoTheGecko

tests/test_particlefile.py

@@ -30,6 +30,7 @@
 from parcels.convert import copernicusmarine_to_sgrid
 from parcels.interpolators import XLinear, XLinear_Velocity
 from parcels.kernels import AdvectionRK4
+from tests import utils
 from tests.common_kernels import DoNothing
 
 
@@ -306,13 +307,13 @@ def IncreaseAge(particles, fieldset):  # pragma: no cover
 
     pset.execute(IncreaseAge, runtime=np.timedelta64(npart * 2, "s"), dt=np.timedelta64(1, "s"), output_file=ofile)
 
-    pytest.skip("# TODO: Need to figure out how times work with parquet output (#2386)")
-    ds = xr.open_zarr(tmp_parquet)
-    age = ds["age"][:].values.astype("timedelta64[s]")
-    ds_timediff = np.zeros_like(age)
-    for i in range(npart):
-        ds_timediff[i, :] = ds.time.values[i, :] - time[i]
-    np.testing.assert_equal(age, ds_timediff)
+    # df = pd.read_parquet(tmp_parquet)
+    df = utils.read_particlefile(tmp_parquet)
+
+    # Map sorted trajectory IDs to release times (0, 1, ..., npart-1 seconds)
+    for index, df_traj in df.groupby("trajectory"):
+        release_time = time[index]
+        np.testing.assert_equal(df_traj["age"].astype("timedelta64[s]").values, (df_traj["time"] - release_time).values)
 
 
 def test_reset_dt(fieldset, tmp_parquet):

tests/utils.py

@@ -8,6 +8,7 @@
 
 import cftime
 import numpy as np
+import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
 import xarray as xr
@@ -178,3 +179,35 @@ def assert_cftime_like_particlefile(parquet_path: Path) -> None:
         "CF-time values in Parquet did not get properly decoded. Are the attributes correct?"
     )
     return
+
+
+def read_particlefile(path: Path, decode_times: bool = True) -> pd.DataFrame:
+    assert path.suffix == ".parquet", "Only Parquet files are supported"
+
+    table = pq.read_table(path)
+
+    try:
+        time_field = table.field("time")
+    except KeyError as e:
+        raise ValueError(
+            f"Could not find 'time' column in parquet file. Are you sure {path=!r} is a particlefile?"
+        ) from e
+
+    try:
+        assert b"units" in time_field.metadata
+    except AssertionError as e:
+        raise ValueError(f"Could not find 'units' in the 'time' column metadata for parquet {path=!r}.") from e
+
+    attrs = {k.decode(): v.decode() for k, v in time_field.metadata.items()}
+
+    df = pd.read_parquet(path)
+    if not decode_times:
+        return df
+
+    values = table.column("time").to_numpy()
+    var = xr.Variable(("time",), values, attrs)
+    values = xr.coders.CFDatetimeCoder(time_unit="s").decode(var).values
+
+    df["time"] = values
+
+    return df