Repository: zarr-developers/VirtualiZarr

Issue:
Duplication of var names
During construction of `ChunkManifest`s, `find_var_names` is called to extract names of variables from `kerchunk` reference dicts. It appears that under the right conditions, this can be quite time consuming as later steps may end up doing the same work many times.

Discussion:

Patch:
diff --git a/virtualizarr/kerchunk.py b/virtualizarr/kerchunk.py
index 97f64b1..6e82067 100644
--- a/virtualizarr/kerchunk.py
+++ b/virtualizarr/kerchunk.py
@@ -165,8 +165,8 @@ def find_var_names(ds_reference_dict: KerchunkStoreRefs) -> list[str]:
     """Find the names of zarr variables in this store/group."""
 
     refs = ds_reference_dict["refs"]
-    found_var_names = [key.split("/")[0] for key in refs.keys() if "/" in key]
-    return found_var_names
+    found_var_names = {key.split("/")[0] for key in refs.keys() if "/" in key}
+    return list(found_var_names)
 
 
 def extract_array_refs(
diff --git a/virtualizarr/manifests/manifest.py b/virtualizarr/manifests/manifest.py
index 0297d17..cc196e6 100644
--- a/virtualizarr/manifests/manifest.py
+++ b/virtualizarr/manifests/manifest.py
@@ -71,8 +71,8 @@ class ChunkManifest:
     """
 
     _paths: np.ndarray[Any, np.dtypes.StringDType]  # type: ignore[name-defined]
-    _offsets: np.ndarray[Any, np.dtype[np.int32]]
-    _lengths: np.ndarray[Any, np.dtype[np.int32]]
+    _offsets: np.ndarray[Any, np.dtype[np.uint64]]
+    _lengths: np.ndarray[Any, np.dtype[np.uint64]]
 
     def __init__(self, entries: dict) -> None:
         """
@@ -100,8 +100,8 @@ class ChunkManifest:
 
         # Initializing to empty implies that entries with path='' are treated as missing chunks
         paths = np.empty(shape=shape, dtype=np.dtypes.StringDType())  # type: ignore[attr-defined]
-        offsets = np.empty(shape=shape, dtype=np.dtype("int32"))
-        lengths = np.empty(shape=shape, dtype=np.dtype("int32"))
+        offsets = np.empty(shape=shape, dtype=np.dtype("uint64"))
+        lengths = np.empty(shape=shape, dtype=np.dtype("uint64"))
 
         # populate the arrays
         for key, entry in entries.items():
@@ -128,8 +128,8 @@ class ChunkManifest:
     def from_arrays(
         cls,
         paths: np.ndarray[Any, np.dtype[np.dtypes.StringDType]],  # type: ignore[name-defined]
-        offsets: np.ndarray[Any, np.dtype[np.int32]],
-        lengths: np.ndarray[Any, np.dtype[np.int32]],
+        offsets: np.ndarray[Any, np.dtype[np.uint64]],
+        lengths: np.ndarray[Any, np.dtype[np.uint64]],
     ) -> "ChunkManifest":
         """
         Create manifest directly from numpy arrays containing the path and byte range information.
@@ -161,13 +161,13 @@ class ChunkManifest:
             raise ValueError(
                 f"paths array must have a numpy variable-length string dtype, but got dtype {paths.dtype}"
             )
-        if offsets.dtype != np.dtype("int32"):
+        if offsets.dtype != np.dtype("uint64"):
             raise ValueError(
-                f"offsets array must have 32-bit integer dtype, but got dtype {offsets.dtype}"
+                f"offsets array must have 64-bit unsigned integer dtype, but got dtype {offsets.dtype}"
             )
-        if lengths.dtype != np.dtype("int32"):
+        if lengths.dtype != np.dtype("uint64"):
             raise ValueError(
-                f"lengths array must have 32-bit integer dtype, but got dtype {lengths.dtype}"
+                f"lengths array must have 64-bit unsigned integer dtype, but got dtype {lengths.dtype}"
             )
 
         # check shapes


Test patch:
diff --git a/virtualizarr/tests/test_kerchunk.py b/virtualizarr/tests/test_kerchunk.py
index 22d6d7d..a6693e2 100644
--- a/virtualizarr/tests/test_kerchunk.py
+++ b/virtualizarr/tests/test_kerchunk.py
@@ -5,7 +5,7 @@ import ujson  # type: ignore
 import xarray as xr
 import xarray.testing as xrt
 
-from virtualizarr.kerchunk import FileType, _automatically_determine_filetype
+from virtualizarr.kerchunk import FileType, find_var_names, _automatically_determine_filetype, KerchunkStoreRefs
 from virtualizarr.manifests import ChunkManifest, ManifestArray
 from virtualizarr.xarray import dataset_from_kerchunk_refs
 
@@ -266,3 +266,14 @@ def test_FileType():
     assert "zarr" == FileType("zarr").name
     with pytest.raises(ValueError):
         FileType(None)
+
+
+def test_no_duplicates_find_var_names():
+    """Verify that we get a deduplicated list of var names"""
+    ref_dict = {
+        "refs": {
+            "x/something": {},
+            "x/otherthing": {}
+        }
+    }
+    assert len(find_var_names(ref_dict)) == 1
\ No newline at end of file