Merge pull request #488 from martindurant/append_deep

More thorough store handling during combine/append
fsspec · Aug 2, 2024 · 3ae8939 · 3ae8939
2 parents bc1275c + 429d1df
commit 3ae8939
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/kerchunk/combine.py b/kerchunk/combine.py
@@ -143,7 +143,7 @@ def __init__(
             raise ValueError("Values being mapped cannot also be identical")
         self.preprocess = preprocess
         self.postprocess = postprocess
-        self.out = out or {}
+        self.out = out if out is not None else {}
         self.coos = None
         self.done = set()
 
@@ -383,7 +383,9 @@ def store_coords(self):
         """
         Write coordinate arrays into the output
         """
-        group = zarr.open(self.out)
+        kv = {}
+        store = zarr.storage.KVStore(kv)
+        group = zarr.open(store)
         m = self.fss[0].get_mapper("")
         z = zarr.open(m)
         for k, v in self.coos.items():
@@ -435,6 +437,7 @@ def store_coords(self):
                 else:
                     arr.attrs.update(self.cf_units[k])
             # TODO: rewrite .zarray/.zattrs with ujson to save space. Maybe make them by hand anyway.
+        self.out.update(kv)
         logger.debug("Written coordinates")
         for fn in [".zgroup", ".zattrs"]:
             # top-level group attributes from first input

diff --git a/kerchunk/tests/test_combine.py b/kerchunk/tests/test_combine.py
@@ -475,29 +475,29 @@ def test_lazy_filler(tmpdir, refs):
     mzz.second_pass()
 
     # actual references don't show
-    assert set(out) == {
+    assert set(out.zmetadata) == {
         ".zattrs",
         ".zgroup",
-        ".zmetadata",
         "data/.zarray",
         "data/.zattrs",
         "static/.zarray",
         "static/.zattrs",
         "time/.zarray",
         "time/.zattrs",
     }
+    assert out._items
     out.flush()
-    assert set(out) == {
+    assert set(out.zmetadata) == {
         ".zattrs",
         ".zgroup",
-        ".zmetadata",
         "data/.zarray",
         "data/.zattrs",
         "static/.zarray",
         "static/.zattrs",
         "time/.zarray",
         "time/.zattrs",
     }
+    assert set(out._items) == {".zmetadata"}
     allfiles = fs.find(tmpdir)
     assert [
         f"{tmpdir}/{a}" in allfiles for a in ["static/refs.0.parq", "data/refs.0.parq"]