Merge pull request #75 from jeromekelleher/zarr-info

Zarr info

Author: jeromekelleher

bio2zarr/cli.py

@@ -31,7 +31,7 @@
     help="Chunk size in the samples dimension",
 )
 
-version = click.version_option(version=f"bio2zarr {provenance.__version__}")
+version = click.version_option(version=f"{provenance.__version__}")
 
 
 # Note: logging hasn't been implemented in the code at all, this is just

bio2zarr/vcf.py

@@ -40,6 +40,17 @@
 )
 
 
+def display_number(x):
+    ret = "n/a"
+    if math.isfinite(x):
+        ret = f"{x: 0.2g}"
+    return ret
+
+
+def display_size(n):
+    return humanfriendly.format_size(n)
+
+
 @dataclasses.dataclass
 class VcfFieldSummary:
     num_chunks: int = 0
@@ -812,15 +823,6 @@ def __len__(self):
         return len(self.columns)
 
     def summary_table(self):
-        def display_number(x):
-            ret = "n/a"
-            if math.isfinite(x):
-                ret = f"{x: 0.2g}"
-            return ret
-
-        def display_size(n):
-            return humanfriendly.format_size(n)
-
         data = []
         for name, col in self.columns.items():
             summary = col.vcf_field.summary
@@ -1028,10 +1030,16 @@ def explode(
     return PickleChunkedVcf.load(out_path)
 
 
-def inspect(if_path):
+def inspect(path):
+    path = pathlib.Path(path)
     # TODO add support for the Zarr format also
-    pcvcf = PickleChunkedVcf.load(if_path)
-    return pcvcf.summary_table()
+    if (path / "metadata.json").exists():
+        obj = PickleChunkedVcf.load(path)
+    elif (path / ".zmetadata").exists():
+        obj = VcfZarr(path)
+    else:
+        raise ValueError("Format not recognised")
+    return obj.summary_table()
 
 
 @dataclasses.dataclass
@@ -1244,6 +1252,40 @@ def fixed_field_spec(
         )
 
 
+class VcfZarr:
+    def __init__(self, path):
+        if not (path / ".zmetadata").exists():
+            raise ValueError("Not in VcfZarr format")
+        self.root = zarr.open(path, mode="r")
+
+    def __repr__(self):
+        return repr(self.root)
+
+    def summary_table(self):
+        data = []
+        arrays = [(a.nbytes_stored, a) for _, a in self.root.arrays()]
+        arrays.sort(key=lambda x: x[0])
+        for stored, array in reversed(arrays):
+            d = {
+                "name": array.name,
+                "dtype": str(array.dtype),
+                "stored": display_size(stored),
+                "size": display_size(array.nbytes),
+                "ratio": display_number(array.nbytes / stored),
+                "nchunks": str(array.nchunks),
+                "chunk_size": display_size(array.nbytes / array.nchunks),
+                "avg_chunk_stored": display_size(int(stored / array.nchunks)),
+                "shape": str(array.shape),
+                "chunk_shape": str(array.chunks),
+                "compressor": str(array.compressor),
+                "filters": str(array.filters),
+            }
+            data.append(d)
+        return data
+
+
+# TODO refactor this into a VcfZarrWriter class, and get rid of the
+# statis methods.
 class SgvcfZarr:
     def __init__(self, path):
         self.path = pathlib.Path(path)

tests/test_cli.py

@@ -127,6 +127,7 @@ def test_convert_plink(self):
             )
 
 
+
 class TestVcfPartition:
     def test_num_parts(self):
         path = "tests/data/vcf/NA12878.prod.chr20snippet.g.vcf.gz"

tests/test_vcf_examples.py

@@ -334,6 +334,19 @@ def test_worker_processes(self, ds, tmp_path, worker_processes):
         ds2 = sg.load_dataset(out)
         xt.assert_equal(ds, ds2)
 
+    def test_inspect(self, tmp_path):
+        # TODO pretty weak test, we should be doing this better somewhere else
+        out = tmp_path / "example.vcf.zarr"
+        vcf.convert(
+            [self.data_path],
+            out,
+            chunk_length=3,
+        )
+        data = vcf.inspect(out)
+        assert len(data) > 0
+        for row in data:
+            assert "name" in row
+
 
 class Test1000G2020Example:
     data_path = "tests/data/vcf/1kg_2020_chrM.vcf.gz"