Accept object dtype for variant_allele (don't change to string dtype)

Author: tomwhite

sgkit/api.py

@@ -33,7 +33,7 @@ def create_genotype_call_dataset(
         The (index of the) contig for each variant.
     variant_position : array_like, int
         The reference position of the variant.
-    variant_alleles : array_like, zero-terminated bytes, e.g. "S1"
+    variant_alleles : array_like, zero-terminated bytes, e.g. "S1", or object
         The possible alleles for the variant.
     sample_id : array_like, str
         The unique identifier of the sample.
@@ -55,7 +55,7 @@ def create_genotype_call_dataset(
     """
     check_array_like(variant_contig, kind="i", ndim=1)
     check_array_like(variant_position, kind="i", ndim=1)
-    check_array_like(variant_alleles, kind="S", ndim=2)
+    check_array_like(variant_alleles, kind={"S", "O"}, ndim=2)
     check_array_like(sample_id, kind="U", ndim=1)
     check_array_like(call_genotype, kind="i", ndim=3)
     data_vars: Dict[Hashable, Any] = {
@@ -102,7 +102,7 @@ def create_genotype_dosage_dataset(
         The (index of the) contig for each variant.
     variant_position : array_like, int
         The reference position of the variant.
-    variant_alleles : array_like, S1
+    variant_alleles : array_like, zero-terminated bytes, e.g. "S1", or object
         The possible alleles for the variant.
     sample_id : array_like, str
         The unique identifier of the sample.
@@ -120,7 +120,7 @@ def create_genotype_dosage_dataset(
     """
     check_array_like(variant_contig, kind="i", ndim=1)
     check_array_like(variant_position, kind="i", ndim=1)
-    check_array_like(variant_alleles, kind="S", ndim=2)
+    check_array_like(variant_alleles, kind={"S", "O"}, ndim=2)
     check_array_like(sample_id, kind="U", ndim=1)
     check_array_like(call_dosage, kind="f", ndim=2)
     data_vars: Dict[Hashable, Any] = {

sgkit/io/vcfzarr_reader.py

@@ -1,7 +1,4 @@
-from typing import Any
-
 import dask.array as da
-import numpy as np
 import xarray as xr
 import zarr
 
@@ -45,22 +42,11 @@ def read_vcfzarr(path: PathType) -> xr.Dataset:
     variant_contig_names = list(variant_contig_names)
 
     # For variant alleles, combine REF and ALT into a single array
-    # and calculate the number of alleles so we can set the dtype correctly
     variants_ref = da.from_zarr(vcfzarr["variants/REF"])
     variants_alt = da.from_zarr(vcfzarr["variants/ALT"])
-
-    def max_str_len(arr: ArrayLike) -> Any:
-        return arr.map_blocks(
-            lambda s: np.char.str_len(s.astype(str)), dtype=np.int8
-        ).max()
-
-    max_allele_length = max(
-        da.compute(max_str_len(variants_ref), max_str_len(variants_alt))
-    )
-    variants_ref_alt = da.concatenate(
+    variant_alleles = da.concatenate(
         [_ensure_2d(variants_ref), _ensure_2d(variants_alt)], axis=1
     )
-    variant_alleles = variants_ref_alt.astype(f"S{max_allele_length}")
 
     variants_id = da.from_zarr(vcfzarr["variants/ID"]).astype(str)
 

sgkit/tests/test_vcfzarr_reader.py

@@ -24,15 +24,15 @@ def test_read_vcfzarr(shared_datadir):
     assert_array_equal(
         ds["variant_allele"],
         [
-            [b"A", b"C", b"", b""],
-            [b"A", b"G", b"", b""],
-            [b"G", b"A", b"", b""],
-            [b"T", b"A", b"", b""],
-            [b"A", b"G", b"T", b""],
-            [b"T", b"", b"", b""],
-            [b"G", b"GA", b"GAC", b""],
-            [b"T", b"", b"", b""],
-            [b"AC", b"A", b"ATG", b"C"],
+            ["A", "C", "", ""],
+            ["A", "G", "", ""],
+            ["G", "A", "", ""],
+            ["T", "A", "", ""],
+            ["A", "G", "T", ""],
+            ["T", "", "", ""],
+            ["G", "GA", "GAC", ""],
+            ["T", "", "", ""],
+            ["AC", "A", "ATG", "C"],
         ],
     )
     assert_array_equal(