Add standard_dimensions for VCZ and fix chunk size

Sets default chunk size to min(size, default) for large dimensions.

Closes #368

Author: jeromekelleher

bio2zarr/vcf.py

@@ -1057,40 +1057,41 @@ def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
     def generate_schema(
         self, variants_chunk_size=None, samples_chunk_size=None, local_alleles=None
     ):
-        m = self.num_records
-        n = self.num_samples
         if local_alleles is None:
             local_alleles = False
 
         max_alleles = max(self.fields["ALT"].vcf_field.summary.max_number + 1, 2)
-        dimensions = {
-            "variants": vcz.VcfZarrDimension(
-                size=m, chunk_size=variants_chunk_size or vcz.DEFAULT_VARIANT_CHUNK_SIZE
-            ),
-            "samples": vcz.VcfZarrDimension(
-                size=n, chunk_size=samples_chunk_size or vcz.DEFAULT_SAMPLE_CHUNK_SIZE
-            ),
-            # ploidy and genotypes added conditionally below
-            "alleles": vcz.VcfZarrDimension(size=max_alleles),
-            "alt_alleles": vcz.VcfZarrDimension(size=max_alleles - 1),
-            "filters": vcz.VcfZarrDimension(size=self.metadata.num_filters),
-        }
 
         # Add ploidy and genotypes dimensions only when needed
         max_genotypes = 0
         for field in self.metadata.format_fields:
             if field.vcf_number == "G":
                 max_genotypes = max(max_genotypes, field.summary.max_number)
+
+        ploidy = None
+        genotypes_size = None
         if self.gt_field is not None:
             ploidy = max(self.gt_field.summary.max_number - 1, 1)
-            dimensions["ploidy"] = vcz.VcfZarrDimension(size=ploidy)
-            max_genotypes = math.comb(max_alleles + ploidy - 1, ploidy)
-            dimensions["genotypes"] = vcz.VcfZarrDimension(size=max_genotypes)
+            # NOTE: it's not clear why we're computing this, when we must have had
+            # at least one number=G field to require it anyway?
+            genotypes_size = math.comb(max_alleles + ploidy - 1, ploidy)
+            # assert max_genotypes == genotypes_size
         else:
             if max_genotypes > 0:
                 # there is no GT field, but there is at least one Number=G field,
                 # so need to define genotypes dimension
-                dimensions["genotypes"] = vcz.VcfZarrDimension(size=max_genotypes)
+                genotypes_size = max_genotypes
+
+        dimensions = vcz.standard_dimensions(
+            variants_size=self.num_records,
+            variants_chunk_size=variants_chunk_size,
+            samples_size=self.num_samples,
+            samples_chunk_size=samples_chunk_size,
+            alleles_size=max_alleles,
+            filters_size=self.metadata.num_filters,
+            ploidy_size=ploidy,
+            genotypes_size=genotypes_size,
+        )
 
         schema_instance = vcz.VcfZarrSchema(
             format_version=vcz.ZARR_SCHEMA_FORMAT_VERSION,
@@ -1173,7 +1174,7 @@ def fixed_field_spec(name, dtype, source=None, dimensions=("variants",)):
                 continue
             array_specs.append(spec_from_field(field))
 
-        if self.gt_field is not None and n > 0:
+        if self.gt_field is not None and self.num_samples > 0:
             array_specs.append(
                 vcz.ZarrArraySpec(
                     name="call_genotype_phased",

bio2zarr/vcz.py

@@ -121,6 +121,51 @@ def fromdict(cls, d):
             chunk_size=d.get("chunk_size", d["size"]),
         )
 
+    @classmethod
+    def unchunked(cls, size):
+        return cls(size, size)
+
+
+def standard_dimensions(
+    *,
+    variants_size,
+    samples_size,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    alleles_size=None,
+    filters_size=None,
+    ploidy_size=None,
+    genotypes_size=None,
+):
+    """
+    Returns a dictionary mapping dimension names to definition for the standard
+    fields in a VCF.
+    """
+    if variants_chunk_size is None:
+        variants_chunk_size = max(1, min(variants_size, DEFAULT_VARIANT_CHUNK_SIZE))
+    if samples_chunk_size is None:
+        samples_chunk_size = max(1, min(samples_size, DEFAULT_SAMPLE_CHUNK_SIZE))
+
+    dimensions = {
+        "variants": VcfZarrDimension(variants_size, variants_chunk_size),
+        "samples": VcfZarrDimension(samples_size, samples_chunk_size),
+    }
+
+    if alleles_size is not None:
+        dimensions["alleles"] = VcfZarrDimension.unchunked(alleles_size)
+        dimensions["alt_alleles"] = VcfZarrDimension.unchunked(alleles_size - 1)
+
+    if filters_size is not None:
+        dimensions["filters"] = VcfZarrDimension.unchunked(filters_size)
+
+    if ploidy_size is not None:
+        dimensions["ploidy"] = VcfZarrDimension.unchunked(ploidy_size)
+
+    if genotypes_size is not None:
+        dimensions["genotypes"] = VcfZarrDimension.unchunked(genotypes_size)
+
+    return dimensions
+
 
 @dataclasses.dataclass
 class ZarrArraySpec:
@@ -1117,7 +1162,6 @@ def __init__(self, path):
     def create_index(self):
         """Create an index to support efficient region queries."""
         root = zarr.open_group(store=self.path, mode="r+")
-        print(list(root.keys()))
         if (
             "variant_contig" not in root
             or "variant_position" not in root

tests/test_vcz.py

@@ -187,8 +187,8 @@ def test_chunk_sizes(self, icf_path, samples_chunk_size, variants_chunk_size):
     def test_default_chunk_size(self, icf_path):
         icf = vcf_mod.IntermediateColumnarFormat(icf_path)
         schema = icf.generate_schema()
-        assert schema.dimensions["samples"].chunk_size == 10_000
-        assert schema.dimensions["variants"].chunk_size == 1000
+        assert schema.dimensions["samples"].chunk_size == 3
+        assert schema.dimensions["variants"].chunk_size == 9
 
 
 class TestSchemaJsonRoundTrip:
@@ -329,7 +329,6 @@ def test_chunk_not_too_large(self, schema, size):
         field = schema.field_map()["variant_H2"]
         schema.dimensions[field.dimensions[-1]].size = size
         schema.dimensions[field.dimensions[-1]].chunk_size = size
-        print(schema.dimensions)
         schema.validate()
 
 
@@ -338,8 +337,8 @@ def test_format_version(self, schema):
         assert schema.format_version == vcz.ZARR_SCHEMA_FORMAT_VERSION
 
     def test_chunk_size(self, schema):
-        assert schema.dimensions["samples"].chunk_size == 10000
-        assert schema.dimensions["variants"].chunk_size == 1000
+        assert schema.dimensions["samples"].chunk_size == 3
+        assert schema.dimensions["variants"].chunk_size == 9
 
     def test_variant_contig(self, schema):
         assert get_field_dict(schema, "variant_contig") == {