Add VCF header and format version

Closes #14

Author: jeromekelleher

bio2zarr/vcf.py

@@ -178,6 +178,7 @@ def make_field_def(name, vcf_type, vcf_number):
 def scan_vcfs(paths, show_progress):
     partitions = []
     vcf_metadata = None
+    header = None
     logger.info(f"Scanning {len(paths)} VCFs")
     for path in tqdm.tqdm(paths, desc="Scan ", disable=not show_progress):
         vcf = cyvcf2.VCF(path)
@@ -215,6 +216,9 @@ def scan_vcfs(paths, show_progress):
 
         if vcf_metadata is None:
             vcf_metadata = metadata
+            # We just take the first header, assuming the others
+            # are compatible.
+            header = vcf.raw_header
         else:
             if metadata != vcf_metadata:
                 raise ValueError("Incompatible VCF chunks")
@@ -230,7 +234,7 @@ def scan_vcfs(paths, show_progress):
         )
     partitions.sort(key=lambda x: x.first_position)
     vcf_metadata.partitions = partitions
-    return vcf_metadata
+    return vcf_metadata, header
 
 
 def sanitise_value_bool(buff, j, value):
@@ -668,9 +672,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 class PickleChunkedVcf(collections.abc.Mapping):
-    def __init__(self, path, metadata):
+    def __init__(self, path, metadata, vcf_header):
         self.path = path
         self.metadata = metadata
+        self.vcf_header = vcf_header
 
         self.columns = {}
         for field in self.metadata.fields:
@@ -753,7 +758,9 @@ def load(path):
         path = pathlib.Path(path)
         with open(path / "metadata.json") as f:
             metadata = VcfMetadata.fromdict(json.load(f))
-        return PickleChunkedVcf(path, metadata)
+        with open(path / "header.txt") as f:
+            header = f.read()
+        return PickleChunkedVcf(path, metadata, header)
 
     @staticmethod
     def convert_partition(
@@ -820,8 +827,8 @@ def convert(
     ):
         out_path = pathlib.Path(out_path)
         # TODO make scan work in parallel using general progress code too
-        vcf_metadata = scan_vcfs(vcfs, show_progress=show_progress)
-        pcvcf = PickleChunkedVcf(out_path, vcf_metadata)
+        vcf_metadata, header = scan_vcfs(vcfs, show_progress=show_progress)
+        pcvcf = PickleChunkedVcf(out_path, vcf_metadata, header)
         pcvcf.mkdirs()
 
         total_variants = sum(
@@ -855,6 +862,8 @@ def convert(
 
         with open(out_path / "metadata.json", "w") as f:
             json.dump(vcf_metadata.asdict(), f, indent=4)
+        with open(out_path / "header.txt", "w") as f:
+            f.write(header)
         return pcvcf
 
 
@@ -1214,7 +1223,6 @@ def encode_contig(self, pcvcf, contig_names, contig_lengths):
         logger.debug("Contig done")
 
     def encode_filters(self, pcvcf, filter_names):
-        self.root.attrs["filters"] = filter_names
         array = self.root.array(
             "filter_id",
             filter_names,
@@ -1277,6 +1285,9 @@ def convert(
         for column in conversion_spec.columns.values():
             sgvcf.create_array(column)
 
+        sgvcf.root.attrs["vcf_zarr_version"] = "0.2"
+        sgvcf.root.attrs["vcf_header"] = pcvcf.vcf_header
+
         progress_config = core.ProgressConfig(
             total=pcvcf.total_uncompressed_bytes,
             title="Encode",

tests/test_vcf_examples.py

@@ -3,6 +3,7 @@
 import xarray.testing as xt
 import pytest
 import sgkit as sg
+import cyvcf2
 
 from bio2zarr import vcf
 
@@ -33,6 +34,10 @@ def test_filters(self, ds):
             ],
         )
 
+    def test_header(self, ds):
+        vcf = cyvcf2.VCF(self.data_path)
+        assert ds.attrs["vcf_header"] == vcf.raw_header
+
     def test_contigs(self, ds):
         nt.assert_array_equal(ds["contig_id"], ["19", "20", "X"])
         assert "contig_length" not in ds