bio2zarr 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

bio2zarr/vcf2zarr/vcz.py

@@ -1,3 +1,4 @@
+import contextlib
 import dataclasses
 import json
 import logging
@@ -12,6 +13,8 @@ import numcodecs
 import numpy as np
 import zarr
 
+from bio2zarr.zarr_utils import ZARR_FORMAT_KWARGS, zarr_v3
+
 from .. import constants, core, provenance
 from . import icf
 
@@ -20,18 +23,31 @@ logger = logging.getLogger(__name__)
 
 def inspect(path):
     path = pathlib.Path(path)
-    # TODO add support for the Zarr format also
+    if not path.exists():
+        raise ValueError(f"Path not found: {path}")
     if (path / "metadata.json").exists():
         obj = icf.IntermediateColumnarFormat(path)
+    # NOTE: this is too strict, we should support more general Zarrs, see #276
     elif (path / ".zmetadata").exists():
         obj = VcfZarr(path)
     else:
-        raise ValueError("Format not recognised")  # NEEDS TEST
+        raise ValueError(f"{path} not in ICF or VCF Zarr format")
     return obj.summary_table()
 
 
 DEFAULT_ZARR_COMPRESSOR = numcodecs.Blosc(cname="zstd", clevel=7)
 
+_fixed_field_descriptions = {
+    "variant_contig": "An identifier from the reference genome or an angle-bracketed ID"
+    " string pointing to a contig in the assembly file",
+    "variant_position": "The reference position",
+    "variant_length": "The length of the variant measured in bases",
+    "variant_id": "List of unique identifiers where applicable",
+    "variant_allele": "List of the reference and alternate alleles",
+    "variant_quality": "Phred-scaled quality score",
+    "variant_filter": "Filter status of the variant",
+}
+
 
 @dataclasses.dataclass
 class ZarrArraySpec:
@@ -46,6 +62,9 @@ class ZarrArraySpec:
     filters: list
 
     def __post_init__(self):
+        if self.name in _fixed_field_descriptions:
+            self.description = self.description or _fixed_field_descriptions[self.name]
+
         # Ensure these are tuples for ease of comparison and consistency
         self.shape = tuple(self.shape)
         self.chunks = tuple(self.chunks)
@@ -68,7 +87,7 @@ class ZarrArraySpec:
         num_samples,
         variants_chunk_size,
         samples_chunk_size,
-        variable_name=None,
+        array_name=None,
     ):
         shape = [num_variants]
         prefix = "variant_"
@@ -79,11 +98,12 @@ class ZarrArraySpec:
             shape.append(num_samples)
             chunks.append(samples_chunk_size)
             dimensions.append("samples")
-        if variable_name is None:
-            variable_name = prefix + vcf_field.name
+        if array_name is None:
+            array_name = prefix + vcf_field.name
         # TODO make an option to add in the empty extra dimension
-        if vcf_field.summary.max_number > 1:
+        if vcf_field.summary.max_number > 1 or vcf_field.full_name == "FORMAT/LAA":
             shape.append(vcf_field.summary.max_number)
+            chunks.append(vcf_field.summary.max_number)
             # TODO we should really be checking this to see if the named dimensions
             # are actually correct.
             if vcf_field.vcf_number == "R":
@@ -96,7 +116,7 @@ class ZarrArraySpec:
                 dimensions.append(f"{vcf_field.category}_{vcf_field.name}_dim")
         return ZarrArraySpec.new(
             vcf_field=vcf_field.full_name,
-            name=variable_name,
+            name=array_name,
             dtype=vcf_field.smallest_dtype(),
             shape=shape,
             chunks=chunks,
@@ -164,6 +184,62 @@ class ZarrArraySpec:
 ZARR_SCHEMA_FORMAT_VERSION = "0.4"
 
 
+def convert_local_allele_field_types(fields):
+    """
+    Update the specified list of fields to include the LAA field, and to convert
+    any supported localisable fields to the L* counterpart.
+
+    Note that we currently support only two ALT alleles per sample, and so the
+    dimensions of these fields are fixed by that requirement. Later versions may
+    use summry data storted in the ICF to make different choices, if information
+    about subsequent alleles (not in the actual genotype calls) should also be
+    stored.
+    """
+    fields_by_name = {field.name: field for field in fields}
+    gt = fields_by_name["call_genotype"]
+    if gt.shape[-1] != 2:
+        raise ValueError("Local alleles only supported on diploid data")
+
+    # TODO check if LA is already in here
+
+    shape = gt.shape[:-1]
+    chunks = gt.chunks[:-1]
+    dimensions = gt.dimensions[:-1]
+
+    la = ZarrArraySpec.new(
+        vcf_field=None,
+        name="call_LA",
+        dtype="i1",
+        shape=gt.shape,
+        chunks=gt.chunks,
+        dimensions=(*dimensions, "local_alleles"),
+        description=(
+            "0-based indices into REF+ALT, indicating which alleles"
+            " are relevant (local) for the current sample"
+        ),
+    )
+    ad = fields_by_name.get("call_AD", None)
+    if ad is not None:
+        # TODO check if call_LAD is in the list already
+        ad.name = "call_LAD"
+        ad.vcf_field = None
+        ad.shape = (*shape, 2)
+        ad.chunks = (*chunks, 2)
+        ad.dimensions = (*dimensions, "local_alleles")
+        ad.description += " (local-alleles)"
+
+    pl = fields_by_name.get("call_PL", None)
+    if pl is not None:
+        # TODO check if call_LPL is in the list already
+        pl.name = "call_LPL"
+        pl.vcf_field = None
+        pl.shape = (*shape, 3)
+        pl.chunks = (*chunks, 3)
+        pl.description += " (local-alleles)"
+        pl.dimensions = (*dimensions, "local_" + pl.dimensions[-1])
+    return [*fields, la]
+
+
 @dataclasses.dataclass
 class VcfZarrSchema(core.JsonDataclass):
     format_version: str
@@ -214,30 +290,38 @@ class VcfZarrSchema(core.JsonDataclass):
         return VcfZarrSchema.fromdict(json.loads(s))
 
     @staticmethod
-    def generate(icf, variants_chunk_size=None, samples_chunk_size=None):
+    def generate(
+        icf, variants_chunk_size=None, samples_chunk_size=None, local_alleles=None
+    ):
         m = icf.num_records
         n = icf.num_samples
-        # FIXME
         if samples_chunk_size is None:
-            samples_chunk_size = 1000
+            samples_chunk_size = 10_000
         if variants_chunk_size is None:
-            variants_chunk_size = 10_000
+            variants_chunk_size = 1000
+        if local_alleles is None:
+            local_alleles = False
         logger.info(
             f"Generating schema with chunks={variants_chunk_size, samples_chunk_size}"
         )
 
-        def spec_from_field(field, variable_name=None):
+        def spec_from_field(field, array_name=None):
             return ZarrArraySpec.from_field(
                 field,
                 num_samples=n,
                 num_variants=m,
                 samples_chunk_size=samples_chunk_size,
                 variants_chunk_size=variants_chunk_size,
-                variable_name=variable_name,
+                array_name=array_name,
             )
 
         def fixed_field_spec(
-            name, dtype, vcf_field=None, shape=(m,), dimensions=("variants",)
+            name,
+            dtype,
+            vcf_field=None,
+            shape=(m,),
+            dimensions=("variants",),
+            chunks=None,
         ):
             return ZarrArraySpec.new(
                 vcf_field=vcf_field,
@@ -246,13 +330,13 @@ class VcfZarrSchema(core.JsonDataclass):
                 shape=shape,
                 description="",
                 dimensions=dimensions,
-                chunks=[variants_chunk_size],
+                chunks=chunks or [variants_chunk_size],
             )
 
-        alt_col = icf.fields["ALT"]
-        max_alleles = alt_col.vcf_field.summary.max_number + 1
+        alt_field = icf.fields["ALT"]
+        max_alleles = alt_field.vcf_field.summary.max_number + 1
 
-        colspecs = [
+        array_specs = [
             fixed_field_spec(
                 name="variant_contig",
                 dtype=core.min_int_dtype(0, icf.metadata.num_contigs),
@@ -262,12 +346,14 @@ class VcfZarrSchema(core.JsonDataclass):
                 dtype="bool",
                 shape=(m, icf.metadata.num_filters),
                 dimensions=["variants", "filters"],
+                chunks=(variants_chunk_size, icf.metadata.num_filters),
             ),
             fixed_field_spec(
                 name="variant_allele",
                 dtype="O",
                 shape=(m, max_alleles),
                 dimensions=["variants", "alleles"],
+                chunks=(variants_chunk_size, max_alleles),
             ),
             fixed_field_spec(
                 name="variant_id",
@@ -280,28 +366,31 @@ class VcfZarrSchema(core.JsonDataclass):
         ]
         name_map = {field.full_name: field for field in icf.metadata.fields}
 
-        # Only two of the fixed fields have a direct one-to-one mapping.
-        colspecs.extend(
+        # Only three of the fixed fields have a direct one-to-one mapping.
+        array_specs.extend(
             [
-                spec_from_field(name_map["QUAL"], variable_name="variant_quality"),
-                spec_from_field(name_map["POS"], variable_name="variant_position"),
+                spec_from_field(name_map["QUAL"], array_name="variant_quality"),
+                spec_from_field(name_map["POS"], array_name="variant_position"),
+                spec_from_field(name_map["rlen"], array_name="variant_length"),
             ]
         )
-        colspecs.extend([spec_from_field(field) for field in icf.metadata.info_fields])
+        array_specs.extend(
+            [spec_from_field(field) for field in icf.metadata.info_fields]
+        )
 
         gt_field = None
         for field in icf.metadata.format_fields:
             if field.name == "GT":
                 gt_field = field
                 continue
-            colspecs.append(spec_from_field(field))
+            array_specs.append(spec_from_field(field))
 
-        if gt_field is not None:
-            ploidy = gt_field.summary.max_number - 1
+        if gt_field is not None and n > 0:
+            ploidy = max(gt_field.summary.max_number - 1, 1)
             shape = [m, n]
             chunks = [variants_chunk_size, samples_chunk_size]
             dimensions = ["variants", "samples"]
-            colspecs.append(
+            array_specs.append(
                 ZarrArraySpec.new(
                     vcf_field=None,
                     name="call_genotype_phased",
@@ -313,8 +402,9 @@ class VcfZarrSchema(core.JsonDataclass):
                 )
             )
             shape += [ploidy]
+            chunks += [ploidy]
             dimensions += ["ploidy"]
-            colspecs.append(
+            array_specs.append(
                 ZarrArraySpec.new(
                     vcf_field=None,
                     name="call_genotype",
@@ -325,7 +415,7 @@ class VcfZarrSchema(core.JsonDataclass):
                     description="",
                 )
             )
-            colspecs.append(
+            array_specs.append(
                 ZarrArraySpec.new(
                     vcf_field=None,
                     name="call_genotype_mask",
@@ -337,11 +427,14 @@ class VcfZarrSchema(core.JsonDataclass):
                 )
             )
 
+        if local_alleles:
+            array_specs = convert_local_allele_field_types(array_specs)
+
         return VcfZarrSchema(
             format_version=ZARR_SCHEMA_FORMAT_VERSION,
             samples_chunk_size=samples_chunk_size,
             variants_chunk_size=variants_chunk_size,
-            fields=colspecs,
+            fields=array_specs,
             samples=icf.metadata.samples,
             contigs=icf.metadata.contigs,
             filters=icf.metadata.filters,
@@ -434,6 +527,84 @@ class VcfZarrWriterMetadata(core.JsonDataclass):
         return ret
 
 
+def compute_la_field(genotypes):
+    """
+    Computes the value of the LA field for each sample given the genotypes
+    for a variant. The LA field lists the unique alleles observed for
+    each sample, including the REF.
+    """
+    v = 2**31 - 1
+    if np.any(genotypes >= v):
+        raise ValueError("Extreme allele value not supported")
+    G = genotypes.astype(np.int32)
+    if len(G) > 0:
+        # Anything < 0 gets mapped to -2 (pad) in the output, which comes last.
+        # So, to get this sorting correctly, we remap to the largest value for
+        # sorting, then map back. We promote the genotypes up to 32 bit for convenience
+        # here, assuming that we'll never have a allele of 2**31 - 1.
+        assert np.all(G != v)
+        G[G < 0] = v
+        G.sort(axis=1)
+        G[G[:, 0] == G[:, 1], 1] = -2
+        # Equal values result in padding also
+        G[G == v] = -2
+    return G.astype(genotypes.dtype)
+
+
+def compute_lad_field(ad, la):
+    assert ad.shape[0] == la.shape[0]
+    assert la.shape[1] == 2
+    lad = np.full((ad.shape[0], 2), -2, dtype=ad.dtype)
+    homs = np.where((la[:, 0] != -2) & (la[:, 1] == -2))
+    lad[homs, 0] = ad[homs, la[homs, 0]]
+    hets = np.where(la[:, 1] != -2)
+    lad[hets, 0] = ad[hets, la[hets, 0]]
+    lad[hets, 1] = ad[hets, la[hets, 1]]
+    return lad
+
+
+def pl_index(a, b):
+    """
+    Returns the PL index for alleles a and b.
+    """
+    return b * (b + 1) // 2 + a
+
+
+def compute_lpl_field(pl, la):
+    lpl = np.full((pl.shape[0], 3), -2, dtype=pl.dtype)
+
+    homs = np.where((la[:, 0] != -2) & (la[:, 1] == -2))
+    a = la[homs, 0]
+    lpl[homs, 0] = pl[homs, pl_index(a, a)]
+
+    hets = np.where(la[:, 1] != -2)[0]
+    a = la[hets, 0]
+    b = la[hets, 1]
+    lpl[hets, 0] = pl[hets, pl_index(a, a)]
+    lpl[hets, 1] = pl[hets, pl_index(a, b)]
+    lpl[hets, 2] = pl[hets, pl_index(b, b)]
+
+    return lpl
+
+
+@dataclasses.dataclass
+class LocalisableFieldDescriptor:
+    array_name: str
+    vcf_field: str
+    sanitise: callable
+    convert: callable
+
+
+localisable_fields = [
+    LocalisableFieldDescriptor(
+        "call_LAD", "FORMAT/AD", icf.sanitise_int_array, compute_lad_field
+    ),
+    LocalisableFieldDescriptor(
+        "call_LPL", "FORMAT/PL", icf.sanitise_int_array, compute_lpl_field
+    ),
+]
+
+
 @dataclasses.dataclass
 class VcfZarrWriteSummary(core.JsonDataclass):
     num_partitions: int
@@ -466,6 +637,12 @@ class VcfZarrWriter:
                 return True
         return False
 
+    def has_local_alleles(self):
+        for field in self.schema.fields:
+            if field.name == "call_LA" and field.vcf_field is None:
+                return True
+        return False
+
     #######################
     # init
     #######################
@@ -505,8 +682,7 @@ class VcfZarrWriter:
         )
 
         self.path.mkdir()
-        store = zarr.DirectoryStore(self.path)
-        root = zarr.group(store=store)
+        root = zarr.open(store=self.path, mode="a", **ZARR_FORMAT_KWARGS)
         root.attrs.update(
             {
                 "vcf_zarr_version": "0.2",
@@ -522,8 +698,7 @@ class VcfZarrWriter:
         self.wip_path.mkdir()
         self.arrays_path.mkdir()
         self.partitions_path.mkdir()
-        store = zarr.DirectoryStore(self.arrays_path)
-        root = zarr.group(store=store)
+        root = zarr.open(store=self.arrays_path, mode="a", **ZARR_FORMAT_KWARGS)
 
         total_chunks = 0
         for field in self.schema.fields:
@@ -547,7 +722,8 @@ class VcfZarrWriter:
             raise ValueError("Subsetting or reordering samples not supported currently")
         array = root.array(
             "sample_id",
-            [sample.id for sample in self.schema.samples],
+            data=[sample.id for sample in self.schema.samples],
+            shape=len(self.schema.samples),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
             chunks=(self.schema.samples_chunk_size,),
@@ -558,7 +734,8 @@ class VcfZarrWriter:
     def encode_contig_id(self, root):
         array = root.array(
             "contig_id",
-            [contig.id for contig in self.schema.contigs],
+            data=[contig.id for contig in self.schema.contigs],
+            shape=len(self.schema.contigs),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
         )
@@ -566,7 +743,8 @@ class VcfZarrWriter:
         if all(contig.length is not None for contig in self.schema.contigs):
             array = root.array(
                 "contig_length",
-                [contig.length for contig in self.schema.contigs],
+                data=[contig.length for contig in self.schema.contigs],
+                shape=len(self.schema.contigs),
                 dtype=np.int64,
                 compressor=DEFAULT_ZARR_COMPRESSOR,
             )
@@ -577,34 +755,42 @@ class VcfZarrWriter:
         # https://github.com/sgkit-dev/vcf-zarr-spec/issues/19
         array = root.array(
             "filter_id",
-            [filt.id for filt in self.schema.filters],
+            data=[filt.id for filt in self.schema.filters],
+            shape=len(self.schema.filters),
             dtype="str",
             compressor=DEFAULT_ZARR_COMPRESSOR,
         )
         array.attrs["_ARRAY_DIMENSIONS"] = ["filters"]
 
-    def init_array(self, root, variable, variants_dim_size):
-        object_codec = None
-        if variable.dtype == "O":
-            object_codec = numcodecs.VLenUTF8()
-        shape = list(variable.shape)
+    def init_array(self, root, array_spec, variants_dim_size):
+        kwargs = dict(ZARR_FORMAT_KWARGS)
+        filters = [numcodecs.get_codec(filt) for filt in array_spec.filters]
+        if array_spec.dtype == "O":
+            if zarr_v3():
+                filters = [*list(filters), numcodecs.VLenUTF8()]
+            else:
+                kwargs["object_codec"] = numcodecs.VLenUTF8()
+
+        if not zarr_v3():
+            kwargs["dimension_separator"] = self.metadata.dimension_separator
+
+        shape = list(array_spec.shape)
         # Truncate the variants dimension is max_variant_chunks was specified
         shape[0] = variants_dim_size
         a = root.empty(
-            variable.name,
+            name=array_spec.name,
             shape=shape,
-            chunks=variable.chunks,
-            dtype=variable.dtype,
-            compressor=numcodecs.get_codec(variable.compressor),
-            filters=[numcodecs.get_codec(filt) for filt in variable.filters],
-            object_codec=object_codec,
-            dimension_separator=self.metadata.dimension_separator,
+            chunks=array_spec.chunks,
+            dtype=array_spec.dtype,
+            compressor=numcodecs.get_codec(array_spec.compressor),
+            filters=filters,
+            **kwargs,
         )
         a.attrs.update(
             {
-                "description": variable.description,
+                "description": array_spec.description,
                 # Dimension names are part of the spec in Zarr v3
-                "_ARRAY_DIMENSIONS": variable.dimensions,
+                "_ARRAY_DIMENSIONS": array_spec.dimensions,
             }
         )
         logger.debug(f"Initialised {a}")
@@ -644,11 +830,15 @@ class VcfZarrWriter:
         self.encode_filters_partition(partition_index)
         self.encode_contig_partition(partition_index)
         self.encode_alleles_partition(partition_index)
-        for col in self.schema.fields:
-            if col.vcf_field is not None:
-                self.encode_array_partition(col, partition_index)
+        for array_spec in self.schema.fields:
+            if array_spec.vcf_field is not None:
+                self.encode_array_partition(array_spec, partition_index)
         if self.has_genotypes():
             self.encode_genotypes_partition(partition_index)
+            self.encode_genotype_mask_partition(partition_index)
+        if self.has_local_alleles():
+            self.encode_local_alleles_partition(partition_index)
+            self.encode_local_allele_fields_partition(partition_index)
 
         final_path = self.partition_path(partition_index)
         logger.info(f"Finalising {partition_index} at {final_path}")
@@ -658,95 +848,144 @@ class VcfZarrWriter:
         os.rename(partition_path, final_path)
 
     def init_partition_array(self, partition_index, name):
+        field_map = self.schema.field_map()
+        array_spec = field_map[name]
         # Create an empty array like the definition
-        src = self.arrays_path / name
+        src = self.arrays_path / array_spec.name
         # Overwrite any existing WIP files
-        wip_path = self.wip_partition_array_path(partition_index, name)
+        wip_path = self.wip_partition_array_path(partition_index, array_spec.name)
         shutil.copytree(src, wip_path, dirs_exist_ok=True)
-        store = zarr.DirectoryStore(self.wip_partition_path(partition_index))
-        wip_root = zarr.group(store=store)
-        array = wip_root[name]
-        logger.debug(f"Opened empty array {array.name} <{array.dtype}> @ {wip_path}")
-        return array
-
-    def finalise_partition_array(self, partition_index, name):
-        logger.debug(f"Encoded {name} partition {partition_index}")
+        array = zarr.open_array(store=wip_path, mode="a")
+        partition = self.metadata.partitions[partition_index]
+        ba = core.BufferedArray(array, partition.start, name)
+        logger.info(
+            f"Start partition {partition_index} array {name} <{array.dtype}> "
+            f"{array.shape} @ {wip_path}"
+        )
+        return ba
 
-    def encode_array_partition(self, column, partition_index):
-        array = self.init_partition_array(partition_index, column.name)
+    def finalise_partition_array(self, partition_index, buffered_array):
+        buffered_array.flush()
+        logger.info(
+            f"Completed partition {partition_index} array {buffered_array.name} "
+            f"max_memory={core.display_size(buffered_array.max_buff_size)}"
+        )
 
+    def encode_array_partition(self, array_spec, partition_index):
         partition = self.metadata.partitions[partition_index]
-        ba = core.BufferedArray(array, partition.start)
-        source_col = self.icf.fields[column.vcf_field]
-        sanitiser = source_col.sanitiser_factory(ba.buff.shape)
+        ba = self.init_partition_array(partition_index, array_spec.name)
+        source_field = self.icf.fields[array_spec.vcf_field]
+        sanitiser = source_field.sanitiser_factory(ba.buff.shape)
 
-        for value in source_col.iter_values(partition.start, partition.stop):
+        for value in source_field.iter_values(partition.start, partition.stop):
             # We write directly into the buffer in the sanitiser function
             # to make it easier to reason about dimension padding
             j = ba.next_buffer_row()
             sanitiser(ba.buff, j, value)
-        ba.flush()
-        self.finalise_partition_array(partition_index, column.name)
+        self.finalise_partition_array(partition_index, ba)
 
     def encode_genotypes_partition(self, partition_index):
-        gt_array = self.init_partition_array(partition_index, "call_genotype")
-        gt_mask_array = self.init_partition_array(partition_index, "call_genotype_mask")
-        gt_phased_array = self.init_partition_array(
-            partition_index, "call_genotype_phased"
-        )
-
         partition = self.metadata.partitions[partition_index]
-        gt = core.BufferedArray(gt_array, partition.start)
-        gt_mask = core.BufferedArray(gt_mask_array, partition.start)
-        gt_phased = core.BufferedArray(gt_phased_array, partition.start)
+        gt = self.init_partition_array(partition_index, "call_genotype")
+        gt_phased = self.init_partition_array(partition_index, "call_genotype_phased")
 
-        source_col = self.icf.fields["FORMAT/GT"]
-        for value in source_col.iter_values(partition.start, partition.stop):
+        source_field = self.icf.fields["FORMAT/GT"]
+        for value in source_field.iter_values(partition.start, partition.stop):
             j = gt.next_buffer_row()
-            icf.sanitise_value_int_2d(gt.buff, j, value[:, :-1])
+            icf.sanitise_value_int_2d(
+                gt.buff, j, value[:, :-1] if value is not None else None
+            )
             j = gt_phased.next_buffer_row()
-            icf.sanitise_value_int_1d(gt_phased.buff, j, value[:, -1])
+            icf.sanitise_value_int_1d(
+                gt_phased.buff, j, value[:, -1] if value is not None else None
+            )
+
+        self.finalise_partition_array(partition_index, gt)
+        self.finalise_partition_array(partition_index, gt_phased)
+
+    def encode_genotype_mask_partition(self, partition_index):
+        partition = self.metadata.partitions[partition_index]
+        gt_mask = self.init_partition_array(partition_index, "call_genotype_mask")
+        # Read back in the genotypes so we can compute the mask
+        gt_array = zarr.open_array(
+            store=self.wip_partition_array_path(partition_index, "call_genotype"),
+            mode="r",
+        )
+        for genotypes in core.first_dim_slice_iter(
+            gt_array, partition.start, partition.stop
+        ):
             # TODO check is this the correct semantics when we are padding
             # with mixed ploidies?
             j = gt_mask.next_buffer_row()
-            gt_mask.buff[j] = gt.buff[j] < 0
-        gt.flush()
-        gt_phased.flush()
-        gt_mask.flush()
+            gt_mask.buff[j] = genotypes < 0
+        self.finalise_partition_array(partition_index, gt_mask)
+
+    def encode_local_alleles_partition(self, partition_index):
+        partition = self.metadata.partitions[partition_index]
+        call_LA = self.init_partition_array(partition_index, "call_LA")
+
+        gt_array = zarr.open_array(
+            store=self.wip_partition_array_path(partition_index, "call_genotype"),
+            mode="r",
+        )
+        for genotypes in core.first_dim_slice_iter(
+            gt_array, partition.start, partition.stop
+        ):
+            la = compute_la_field(genotypes)
+            j = call_LA.next_buffer_row()
+            call_LA.buff[j] = la
+        self.finalise_partition_array(partition_index, call_LA)
 
-        self.finalise_partition_array(partition_index, "call_genotype")
-        self.finalise_partition_array(partition_index, "call_genotype_mask")
-        self.finalise_partition_array(partition_index, "call_genotype_phased")
+    def encode_local_allele_fields_partition(self, partition_index):
+        partition = self.metadata.partitions[partition_index]
+        la_array = zarr.open_array(
+            store=self.wip_partition_array_path(partition_index, "call_LA"),
+            mode="r",
+        )
+        # We got through the localisable fields one-by-one so that we don't need to
+        # keep several large arrays in memory at once for each partition.
+        field_map = self.schema.field_map()
+        for descriptor in localisable_fields:
+            if descriptor.array_name not in field_map:
+                continue
+            assert field_map[descriptor.array_name].vcf_field is None
+
+            buff = self.init_partition_array(partition_index, descriptor.array_name)
+            source = self.icf.fields[descriptor.vcf_field].iter_values(
+                partition.start, partition.stop
+            )
+            for la in core.first_dim_slice_iter(
+                la_array, partition.start, partition.stop
+            ):
+                raw_value = next(source)
+                value = descriptor.sanitise(raw_value, 2, raw_value.dtype)
+                j = buff.next_buffer_row()
+                buff.buff[j] = descriptor.convert(value, la)
+            self.finalise_partition_array(partition_index, buff)
 
     def encode_alleles_partition(self, partition_index):
-        array_name = "variant_allele"
-        alleles_array = self.init_partition_array(partition_index, array_name)
+        alleles = self.init_partition_array(partition_index, "variant_allele")
         partition = self.metadata.partitions[partition_index]
-        alleles = core.BufferedArray(alleles_array, partition.start)
-        ref_col = self.icf.fields["REF"]
-        alt_col = self.icf.fields["ALT"]
+        ref_field = self.icf.fields["REF"]
+        alt_field = self.icf.fields["ALT"]
 
         for ref, alt in zip(
-            ref_col.iter_values(partition.start, partition.stop),
-            alt_col.iter_values(partition.start, partition.stop),
+            ref_field.iter_values(partition.start, partition.stop),
+            alt_field.iter_values(partition.start, partition.stop),
         ):
             j = alleles.next_buffer_row()
             alleles.buff[j, :] = constants.STR_FILL
             alleles.buff[j, 0] = ref[0]
             alleles.buff[j, 1 : 1 + len(alt)] = alt
-        alleles.flush()
-
-        self.finalise_partition_array(partition_index, array_name)
+        self.finalise_partition_array(partition_index, alleles)
 
     def encode_id_partition(self, partition_index):
-        vid_array = self.init_partition_array(partition_index, "variant_id")
-        vid_mask_array = self.init_partition_array(partition_index, "variant_id_mask")
+        vid = self.init_partition_array(partition_index, "variant_id")
+        vid_mask = self.init_partition_array(partition_index, "variant_id_mask")
         partition = self.metadata.partitions[partition_index]
-        vid = core.BufferedArray(vid_array, partition.start)
-        vid_mask = core.BufferedArray(vid_mask_array, partition.start)
-        col = self.icf.fields["ID"]
+        field = self.icf.fields["ID"]
 
-        for value in col.iter_values(partition.start, partition.stop):
+        for value in field.iter_values(partition.start, partition.stop):
             j = vid.next_buffer_row()
             k = vid_mask.next_buffer_row()
             assert j == k
@@ -756,21 +995,17 @@ class VcfZarrWriter:
             else:
                 vid.buff[j] = constants.STR_MISSING
                 vid_mask.buff[j] = True
-        vid.flush()
-        vid_mask.flush()
 
-        self.finalise_partition_array(partition_index, "variant_id")
-        self.finalise_partition_array(partition_index, "variant_id_mask")
+        self.finalise_partition_array(partition_index, vid)
+        self.finalise_partition_array(partition_index, vid_mask)
 
     def encode_filters_partition(self, partition_index):
         lookup = {filt.id: index for index, filt in enumerate(self.schema.filters)}
-        array_name = "variant_filter"
-        array = self.init_partition_array(partition_index, array_name)
+        var_filter = self.init_partition_array(partition_index, "variant_filter")
         partition = self.metadata.partitions[partition_index]
-        var_filter = core.BufferedArray(array, partition.start)
 
-        col = self.icf.fields["FILTERS"]
-        for value in col.iter_values(partition.start, partition.stop):
+        field = self.icf.fields["FILTERS"]
+        for value in field.iter_values(partition.start, partition.stop):
             j = var_filter.next_buffer_row()
             var_filter.buff[j] = False
             for f in value:
@@ -780,28 +1015,24 @@ class VcfZarrWriter:
                     raise ValueError(
                         f"Filter '{f}' was not defined in the header."
                     ) from None
-        var_filter.flush()
 
-        self.finalise_partition_array(partition_index, array_name)
+        self.finalise_partition_array(partition_index, var_filter)
 
     def encode_contig_partition(self, partition_index):
         lookup = {contig.id: index for index, contig in enumerate(self.schema.contigs)}
-        array_name = "variant_contig"
-        array = self.init_partition_array(partition_index, array_name)
+        contig = self.init_partition_array(partition_index, "variant_contig")
         partition = self.metadata.partitions[partition_index]
-        contig = core.BufferedArray(array, partition.start)
-        col = self.icf.fields["CHROM"]
+        field = self.icf.fields["CHROM"]
 
-        for value in col.iter_values(partition.start, partition.stop):
+        for value in field.iter_values(partition.start, partition.stop):
             j = contig.next_buffer_row()
             # Note: because we are using the indexes to define the lookups
             # and we always have an index, it seems that we the contig lookup
             # will always succeed. However, if anyone ever does hit a KeyError
             # here, please do open an issue with a reproducible example!
             contig.buff[j] = lookup[value[0]]
-        contig.flush()
 
-        self.finalise_partition_array(partition_index, array_name)
+        self.finalise_partition_array(partition_index, contig)
 
     #######################
     # finalise
@@ -871,6 +1102,68 @@ class VcfZarrWriter:
         logger.info("Consolidating Zarr metadata")
         zarr.consolidate_metadata(self.path)
 
+    #######################
+    # index
+    #######################
+
+    def create_index(self):
+        """Create an index to support efficient region queries."""
+
+        root = zarr.open_group(store=self.path, mode="r+")
+
+        contig = root["variant_contig"]
+        pos = root["variant_position"]
+        length = root["variant_length"]
+
+        assert contig.cdata_shape == pos.cdata_shape
+
+        index = []
+
+        logger.info("Creating region index")
+        for v_chunk in range(pos.cdata_shape[0]):
+            c = contig.blocks[v_chunk]
+            p = pos.blocks[v_chunk]
+            e = p + length.blocks[v_chunk] - 1
+
+            # create a row for each contig in the chunk
+            d = np.diff(c, append=-1)
+            c_start_idx = 0
+            for c_end_idx in np.nonzero(d)[0]:
+                assert c[c_start_idx] == c[c_end_idx]
+                index.append(
+                    (
+                        v_chunk,  # chunk index
+                        c[c_start_idx],  # contig ID
+                        p[c_start_idx],  # start
+                        p[c_end_idx],  # end
+                        np.max(e[c_start_idx : c_end_idx + 1]),  # max end
+                        c_end_idx - c_start_idx + 1,  # num records
+                    )
+                )
+                c_start_idx = c_end_idx + 1
+
+        index = np.array(index, dtype=pos.dtype)
+        kwargs = {}
+        if not zarr_v3():
+            kwargs["dimension_separator"] = self.metadata.dimension_separator
+        array = root.array(
+            "region_index",
+            data=index,
+            shape=index.shape,
+            chunks=index.shape,
+            dtype=index.dtype,
+            compressor=numcodecs.Blosc("zstd", clevel=9, shuffle=0),
+            fill_value=None,
+            **kwargs,
+        )
+        array.attrs["_ARRAY_DIMENSIONS"] = [
+            "region_index_values",
+            "region_index_fields",
+        ]
+
+        logger.info("Consolidating Zarr metadata")
+        zarr.consolidate_metadata(self.path)
+
     ######################
     # encode_all_partitions
     ######################
@@ -880,8 +1173,8 @@ class VcfZarrWriter:
         Return the approximate maximum memory used to encode a variant chunk.
         """
         max_encoding_mem = 0
-        for col in self.schema.fields:
-            max_encoding_mem = max(max_encoding_mem, col.variant_chunk_nbytes)
+        for array_spec in self.schema.fields:
+            max_encoding_mem = max(max_encoding_mem, array_spec.variant_chunk_nbytes)
         gt_mem = 0
         if self.has_genotypes:
             gt_mem = sum(
@@ -921,9 +1214,9 @@ class VcfZarrWriter:
         num_workers = min(max_num_workers, worker_processes)
 
         total_bytes = 0
-        for col in self.schema.fields:
+        for array_spec in self.schema.fields:
             # Open the array definition to get the total size
-            total_bytes += zarr.open(self.arrays_path / col.name).nbytes
+            total_bytes += zarr.open(self.arrays_path / array_spec.name).nbytes
 
         progress_config = core.ProgressConfig(
             total=total_bytes,
@@ -936,9 +1229,21 @@ class VcfZarrWriter:
                 pwm.submit(self.encode_partition, partition_index)
 
 
-def mkschema(if_path, out):
+def mkschema(
+    if_path,
+    out,
+    *,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    local_alleles=None,
+):
     store = icf.IntermediateColumnarFormat(if_path)
-    spec = VcfZarrSchema.generate(store)
+    spec = VcfZarrSchema.generate(
+        store,
+        variants_chunk_size=variants_chunk_size,
+        samples_chunk_size=samples_chunk_size,
+        local_alleles=local_alleles,
+    )
     out.write(spec.asjson())
 
 
@@ -951,6 +1256,7 @@ def encode(
     max_variant_chunks=None,
     dimension_separator=None,
     max_memory=None,
+    local_alleles=None,
     worker_processes=1,
     show_progress=False,
 ):
@@ -963,6 +1269,7 @@ def encode(
         schema_path=schema_path,
         variants_chunk_size=variants_chunk_size,
         samples_chunk_size=samples_chunk_size,
+        local_alleles=local_alleles,
         max_variant_chunks=max_variant_chunks,
         dimension_separator=dimension_separator,
     )
@@ -973,6 +1280,7 @@ def encode(
         max_memory=max_memory,
     )
     vzw.finalise(show_progress)
+    vzw.create_index()
 
 
 def encode_init(
@@ -983,6 +1291,7 @@ def encode_init(
     schema_path=None,
     variants_chunk_size=None,
     samples_chunk_size=None,
+    local_alleles=None,
     max_variant_chunks=None,
     dimension_separator=None,
     max_memory=None,
@@ -995,6 +1304,7 @@ def encode_init(
             icf_store,
             variants_chunk_size=variants_chunk_size,
             samples_chunk_size=samples_chunk_size,
+            local_alleles=local_alleles,
         )
     else:
         logger.info(f"Reading schema from {schema_path}")
@@ -1032,22 +1342,34 @@ def convert(
     variants_chunk_size=None,
     samples_chunk_size=None,
     worker_processes=1,
+    local_alleles=None,
     show_progress=False,
-    # TODO add arguments to control location of tmpdir
+    icf_path=None,
 ):
-    with tempfile.TemporaryDirectory(prefix="vcf2zarr") as tmp:
-        if_dir = pathlib.Path(tmp) / "icf"
+    if icf_path is None:
+        cm = temp_icf_path(prefix="vcf2zarr")
+    else:
+        cm = contextlib.nullcontext(icf_path)
+
+    with cm as icf_path:
         icf.explode(
-            if_dir,
+            icf_path,
             vcfs,
             worker_processes=worker_processes,
             show_progress=show_progress,
         )
         encode(
-            if_dir,
+            icf_path,
             out_path,
             variants_chunk_size=variants_chunk_size,
             samples_chunk_size=samples_chunk_size,
             worker_processes=worker_processes,
             show_progress=show_progress,
+            local_alleles=local_alleles,
         )
+
+
+@contextlib.contextmanager
+def temp_icf_path(prefix=None):
+    with tempfile.TemporaryDirectory(prefix=prefix) as tmp:
+        yield pathlib.Path(tmp) / "icf"