Add min_int_dtype function

Author: jeromekelleher

bio2zarr/core.py

@@ -16,6 +16,16 @@
 numcodecs.blosc.use_threads = False
 
 
+def min_int_dtype(min_value, max_value):
+    if min_value > max_value:
+        raise ValueError("min_value must be <= max_value")
+    for a_dtype in ["i1", "i2", "i4", "i8"]:
+        info = np.iinfo(a_dtype)
+        if info.min <= min_value and max_value <= info.max:
+            return a_dtype
+    raise OverflowError("Integer cannot be represented")
+
+
 def chunk_aligned_slices(z, n, max_chunks=None):
     """
     Returns at n slices in the specified zarr array, aligned

bio2zarr/vcf.py

@@ -123,13 +123,7 @@ def smallest_dtype(self):
         if self.vcf_type == "Float":
             ret = "f4"
         elif self.vcf_type == "Integer":
-            dtype = "i4"
-            for a_dtype in ["i1", "i2"]:
-                info = np.iinfo(a_dtype)
-                if info.min <= s.min_value and s.max_value <= info.max:
-                    dtype = a_dtype
-                    break
-            ret = dtype
+            ret = core.min_int_dtype(s.min_value, s.max_value)
         elif self.vcf_type == "Flag":
             ret = "bool"
         elif self.vcf_type == "Character":
@@ -187,6 +181,14 @@ def format_fields(self):
                 fields.append(field)
         return fields
 
+    @property
+    def num_contigs(self):
+        return len(self.contig_names)
+
+    @property
+    def num_filters(self):
+        return len(self.filters)
+
     @property
     def num_records(self):
         return sum(self.contig_record_counts.values())
@@ -1388,24 +1390,22 @@ def fixed_field_spec(
 
         alt_col = icf.columns["ALT"]
         max_alleles = alt_col.vcf_field.summary.max_number + 1
-        num_filters = len(icf.metadata.filters)
 
-        # # FIXME get dtype from lookup table
         colspecs = [
             fixed_field_spec(
                 name="variant_contig",
-                dtype="i2",  # FIXME
+                dtype=core.min_int_dtype(0, icf.metadata.num_contigs),
             ),
             fixed_field_spec(
                 name="variant_filter",
                 dtype="bool",
-                shape=(m, num_filters),
+                shape=(m, icf.metadata.num_filters),
                 dimensions=["variants", "filters"],
             ),
             fixed_field_spec(
                 name="variant_allele",
                 dtype="str",
-                shape=[m, max_alleles],
+                shape=(m, max_alleles),
                 dimensions=["variants", "alleles"],
             ),
             fixed_field_spec(

tests/test_core.py

@@ -5,6 +5,55 @@
 from bio2zarr import core
 
 
+class TestMinIntDtype:
+    @pytest.mark.parametrize(
+        ("min_value", "max_value", "dtype"),
+        [
+            (0, 1, "i1"),
+            (0, 0, "i1"),
+            (0, 127, "i1"),
+            (127, 128, "i2"),
+            (-127, 0, "i1"),
+            (-127, -126, "i1"),
+            (0, 2**15 - 1, "i2"),
+            (-(2**15), 2**15 - 1, "i2"),
+            (0, 2**15, "i4"),
+            (-(2**15), 2**15, "i4"),
+            (0, 2**31 - 1, "i4"),
+            (-(2**31), 2**31 - 1, "i4"),
+            (2**31 - 1, 2**31 - 1, "i4"),
+            (0, 2**31, "i8"),
+            (0, 2**32, "i8"),
+        ],
+    )
+    def test_values(self, min_value, max_value, dtype):
+        assert core.min_int_dtype(min_value, max_value) == dtype
+
+    @pytest.mark.parametrize(
+        ("min_value", "max_value"),
+        [
+            (0, 2**63),
+            (-(2**63) - 1, 0),
+            (0, 2**65),
+        ],
+    )
+    def test_overflow(self, min_value, max_value):
+        with pytest.raises(OverflowError, match="Integer cannot"):
+            core.min_int_dtype(min_value, max_value)
+
+    @pytest.mark.parametrize(
+        ("min_value", "max_value"),
+        [
+            (1, 0),
+            (-1, -2),
+            (2**31, 2**31 - 1),
+        ],
+    )
+    def test_bad_min_max(self, min_value, max_value):
+        with pytest.raises(ValueError, match="must be <="):
+            core.min_int_dtype(min_value, max_value)
+
+
 class TestParallelWorkManager:
     @pytest.mark.parametrize("total", [1, 10, 2**63])
     @pytest.mark.parametrize("workers", [0, 1])

tests/test_vcf.py

@@ -225,7 +225,7 @@ def test_filter_id(self, schema):
     def test_variant_contig(self, schema):
         assert schema["columns"]["variant_contig"] == {
             "name": "variant_contig",
-            "dtype": "i2",
+            "dtype": "i1",
             "shape": [9],
             "chunks": [10000],
             "dimensions": ["variants"],