bio2zarr 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

bio2zarr/vcf.py

@@ -40,6 +40,17 @@ FLOAT32_MISSING_AS_INT32, FLOAT32_FILL_AS_INT32 = np.array(
 )
 
 
+def display_number(x):
+    ret = "n/a"
+    if math.isfinite(x):
+        ret = f"{x: 0.2g}"
+    return ret
+
+
+def display_size(n):
+    return humanfriendly.format_size(n, binary=True)
+
+
 @dataclasses.dataclass
 class VcfFieldSummary:
     num_chunks: int = 0
@@ -53,11 +64,18 @@ class VcfFieldSummary:
     def update(self, other):
         self.num_chunks += other.num_chunks
         self.compressed_size += other.compressed_size
-        self.uncompressed_size = other.uncompressed_size
+        self.uncompressed_size += other.uncompressed_size
         self.max_number = max(self.max_number, other.max_number)
         self.min_value = min(self.min_value, other.min_value)
         self.max_value = max(self.max_value, other.max_value)
 
+    def asdict(self):
+        return dataclasses.asdict(self)
+
+    @staticmethod
+    def fromdict(d):
+        return VcfFieldSummary(**d)
+
 
 @dataclasses.dataclass
 class VcfField:
@@ -131,9 +149,14 @@ class VcfPartition:
     num_records: int = -1
 
 
+ICF_METADATA_FORMAT_VERSION = "0.2"
+ICF_DEFAULT_COMPRESSOR = numcodecs.Blosc(
+    cname="lz4", clevel=7, shuffle=numcodecs.Blosc.NOSHUFFLE
+).get_config()
+
+
 @dataclasses.dataclass
-class VcfMetadata:
-    format_version: str
+class IcfMetadata:
     samples: list
     contig_names: list
     contig_record_counts: dict
@@ -141,6 +164,10 @@ class VcfMetadata:
     fields: list
     partitions: list = None
     contig_lengths: list = None
+    format_version: str = None
+    compressor: dict = None
+    column_chunk_size: int = None
+    provenance: dict = None
 
     @property
     def info_fields(self):
@@ -164,12 +191,19 @@ class VcfMetadata:
 
     @staticmethod
     def fromdict(d):
+        if d["format_version"] != ICF_METADATA_FORMAT_VERSION:
+            raise ValueError(
+                "Intermediate columnar metadata format version mismatch: "
+                f"{d['format_version']} != {ICF_METADATA_FORMAT_VERSION}"
+            )
         fields = [VcfField.fromdict(fd) for fd in d["fields"]]
         partitions = [VcfPartition(**pd) for pd in d["partitions"]]
+        for p in partitions:
+            p.region = vcf_utils.Region(**p.region)
         d = d.copy()
         d["fields"] = fields
         d["partitions"] = partitions
-        return VcfMetadata(**d)
+        return IcfMetadata(**d)
 
     def asdict(self):
         return dataclasses.asdict(self)
@@ -220,16 +254,13 @@ def scan_vcf(path, target_num_partitions):
                     field.vcf_number = "."
                 fields.append(field)
 
-        metadata = VcfMetadata(
+        metadata = IcfMetadata(
             samples=vcf.samples,
             contig_names=vcf.seqnames,
             contig_record_counts=indexed_vcf.contig_record_counts(),
             filters=filters,
-            # TODO use the mapping dictionary
             fields=fields,
             partitions=[],
-            # FIXME do something systematic with this
-            format_version="0.1",
         )
         try:
             metadata.contig_lengths = vcf.seqlens
@@ -243,6 +274,8 @@ def scan_vcf(path, target_num_partitions):
         for region in regions:
             metadata.partitions.append(
                 VcfPartition(
+                    # TODO should this be fully resolving the path? Otherwise it's all
+                    # relative to the original WD
                     vcf_path=str(path),
                     region=region,
                 )
@@ -251,8 +284,19 @@ def scan_vcf(path, target_num_partitions):
         return metadata, vcf.raw_header
 
 
-def scan_vcfs(paths, show_progress, target_num_partitions, worker_processes=1):
-    logger.info(f"Scanning {len(paths)} VCFs")
+def scan_vcfs(
+    paths, show_progress, target_num_partitions, column_chunk_size, worker_processes=1
+):
+    logger.info(
+        f"Scanning {len(paths)} VCFs attempting to split into {target_num_partitions} partitions."
+    )
+    # An easy mistake to make is to pass the same file twice. Check this early on.
+    for path, count in collections.Counter(paths).items():
+        if not path.exists():  # NEEDS TEST
+            raise FileNotFoundError(path)
+        if count > 1:
+            raise ValueError(f"Duplicate path provided: {path}")
+
     progress_config = core.ProgressConfig(
         total=len(paths),
         units="files",
@@ -261,7 +305,7 @@ def scan_vcfs(paths, show_progress, target_num_partitions, worker_processes=1):
     )
     with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
         for path in paths:
-            pwm.submit(scan_vcf, path, target_num_partitions)
+            pwm.submit(scan_vcf, path, max(1, target_num_partitions // len(paths)))
         results = list(pwm.results_as_completed())
 
     # Sort to make the ordering deterministic
@@ -276,12 +320,12 @@ def scan_vcfs(paths, show_progress, target_num_partitions, worker_processes=1):
         contig_record_counts += metadata.contig_record_counts
         metadata.contig_record_counts.clear()
 
-    vcf_metadata, header = results[0]
+    icf_metadata, header = results[0]
     for metadata, _ in results[1:]:
-        if metadata != vcf_metadata:
+        if metadata != icf_metadata:
             raise ValueError("Incompatible VCF chunks")
 
-    vcf_metadata.contig_record_counts = dict(contig_record_counts)
+    icf_metadata.contig_record_counts = dict(contig_record_counts)
 
     # Sort by contig (in the order they appear in the header) first,
     # then by start coordinate
@@ -289,8 +333,15 @@ def scan_vcfs(paths, show_progress, target_num_partitions, worker_processes=1):
     all_partitions.sort(
         key=lambda x: (contig_index_map[x.region.contig], x.region.start)
     )
-    vcf_metadata.partitions = all_partitions
-    return vcf_metadata, header
+    icf_metadata.partitions = all_partitions
+    icf_metadata.format_version = ICF_METADATA_FORMAT_VERSION
+    icf_metadata.compressor = ICF_DEFAULT_COMPRESSOR
+    icf_metadata.column_chunk_size = column_chunk_size
+    # Bare minimum here for provenance - would be nice to include versions of key
+    # dependencies as well.
+    icf_metadata.provenance = {"source": f"bio2zarr-{provenance.__version__}"}
+    logger.info(f"Scan complete, resulting in {len(all_partitions)} partitions.")
+    return icf_metadata, header
 
 
 def sanitise_value_bool(buff, j, value):
@@ -385,7 +436,7 @@ def sanitise_value_float_2d(buff, j, value):
 
 def sanitise_int_array(value, ndmin, dtype):
     if isinstance(value, tuple):
-        value = [VCF_INT_MISSING if x is None else x for x in value]
+        value = [VCF_INT_MISSING if x is None else x for x in value]  #  NEEDS TEST
     value = np.array(value, ndmin=ndmin, copy=False)
     value[value == VCF_INT_MISSING] = -1
     value[value == VCF_INT_FILL] = -2
@@ -428,7 +479,7 @@ missing_value_map = {
 class VcfValueTransformer:
     """
     Transform VCF values into the stored intermediate format used
-    in the PickleChunkedVcf, and update field summaries.
+    in the IntermediateColumnarFormat, and update field summaries.
     """
 
     def __init__(self, field, num_samples):
@@ -516,29 +567,29 @@ class StringValueTransformer(VcfValueTransformer):
 class SplitStringValueTransformer(StringValueTransformer):
     def transform(self, vcf_value):
         if vcf_value is None:
-            return self.missing_value
+            return self.missing_value  # NEEDS TEST
         assert self.dimension == 1
         return np.array(vcf_value, ndmin=1, dtype="str")
 
 
-class PickleChunkedVcfField:
-    def __init__(self, pcvcf, vcf_field):
+def get_vcf_field_path(base_path, vcf_field):
+    if vcf_field.category == "fixed":
+        return base_path / vcf_field.name
+    return base_path / vcf_field.category / vcf_field.name
+
+
+class IntermediateColumnarFormatField:
+    def __init__(self, icf, vcf_field):
         self.vcf_field = vcf_field
-        self.path = self.get_path(pcvcf.path, vcf_field)
-        self.compressor = pcvcf.compressor
-        self.num_partitions = pcvcf.num_partitions
-        self.num_records = pcvcf.num_records
-        self.partition_record_index = pcvcf.partition_record_index
+        self.path = get_vcf_field_path(icf.path, vcf_field)
+        self.compressor = icf.compressor
+        self.num_partitions = icf.num_partitions
+        self.num_records = icf.num_records
+        self.partition_record_index = icf.partition_record_index
         # A map of partition id to the cumulative number of records
         # in chunks within that partition
         self._chunk_record_index = {}
 
-    @staticmethod
-    def get_path(base_path, vcf_field):
-        if vcf_field.category == "fixed":
-            return base_path / vcf_field.name
-        return base_path / vcf_field.category / vcf_field.name
-
     @property
     def name(self):
         return self.vcf_field.full_name
@@ -549,17 +600,17 @@ class PickleChunkedVcfField:
     def __repr__(self):
         partition_chunks = [self.num_chunks(j) for j in range(self.num_partitions)]
         return (
-            f"PickleChunkedVcfField(name={self.name}, "
+            f"IntermediateColumnarFormatField(name={self.name}, "
             f"partition_chunks={partition_chunks}, "
             f"path={self.path})"
         )
 
     def num_chunks(self, partition_id):
-        return len(self.chunk_cumulative_records(partition_id))
+        return len(self.chunk_record_index(partition_id)) - 1
 
     def chunk_record_index(self, partition_id):
         if partition_id not in self._chunk_record_index:
-            index_path = self.partition_path(partition_id) / "chunk_index.pkl"
+            index_path = self.partition_path(partition_id) / "chunk_index"
             with open(index_path, "rb") as f:
                 a = pickle.load(f)
             assert len(a) > 1
@@ -567,22 +618,27 @@ class PickleChunkedVcfField:
             self._chunk_record_index[partition_id] = a
         return self._chunk_record_index[partition_id]
 
-    def chunk_cumulative_records(self, partition_id):
-        return self.chunk_record_index(partition_id)[1:]
-
-    def chunk_num_records(self, partition_id):
-        return np.diff(self.chunk_cumulative_records(partition_id))
-
-    def chunk_files(self, partition_id, start=0):
-        partition_path = self.partition_path(partition_id)
-        for n in self.chunk_cumulative_records(partition_id)[start:]:
-            yield partition_path / f"{n}.pkl"
-
     def read_chunk(self, path):
         with open(path, "rb") as f:
             pkl = self.compressor.decode(f.read())
         return pickle.loads(pkl)
 
+    def chunk_num_records(self, partition_id):
+        return np.diff(self.chunk_record_index(partition_id))
+
+    def chunks(self, partition_id, start_chunk=0):
+        partition_path = self.partition_path(partition_id)
+        chunk_cumulative_records = self.chunk_record_index(partition_id)
+        chunk_num_records = np.diff(chunk_cumulative_records)
+        for count, cumulative in zip(
+            chunk_num_records[start_chunk:], chunk_cumulative_records[start_chunk + 1 :]
+        ):
+            path = partition_path / f"{cumulative}"
+            chunk = self.read_chunk(path)
+            if len(chunk) != count:
+                raise ValueError(f"Corruption detected in chunk: {path}")
+            yield chunk
+
     def iter_values(self, start=None, stop=None):
         start = 0 if start is None else start
         stop = self.num_records if stop is None else stop
@@ -603,9 +659,7 @@ class PickleChunkedVcfField:
             f"Read {self.vcf_field.full_name} slice [{start}:{stop}]:"
             f"p_start={start_partition}, c_start={start_chunk}, r_start={record_id}"
         )
-
-        for chunk_path in self.chunk_files(start_partition, start_chunk):
-            chunk = self.read_chunk(chunk_path)
+        for chunk in self.chunks(start_partition, start_chunk):
             for record in chunk:
                 if record_id == stop:
                     return
@@ -614,8 +668,7 @@ class PickleChunkedVcfField:
                 record_id += 1
         assert record_id > start
         for partition_id in range(start_partition + 1, self.num_partitions):
-            for chunk_path in self.chunk_files(partition_id):
-                chunk = self.read_chunk(chunk_path)
+            for chunk in self.chunks(partition_id):
                 for record in chunk:
                     if record_id == stop:
                         return
@@ -629,15 +682,11 @@ class PickleChunkedVcfField:
         ret = [None] * self.num_records
         j = 0
         for partition_id in range(self.num_partitions):
-            for chunk_path in self.chunk_files(partition_id):
-                chunk = self.read_chunk(chunk_path)
+            for chunk in self.chunks(partition_id):
                 for record in chunk:
                     ret[j] = record
                     j += 1
-        if j != self.num_records:
-            raise ValueError(
-                f"Corruption detected: incorrect number of records in {str(self.path)}."
-            )
+        assert j == self.num_records
         return ret
 
     def sanitiser_factory(self, shape):
@@ -674,7 +723,7 @@ class PickleChunkedVcfField:
 
 
 @dataclasses.dataclass
-class PcvcfFieldWriter:
+class IcfFieldWriter:
     vcf_field: VcfField
     path: pathlib.Path
     transformer: VcfValueTransformer
@@ -704,7 +753,7 @@ class PcvcfFieldWriter:
     def write_chunk(self):
         # Update index
         self.chunk_index.append(self.num_records)
-        path = self.path / f"{self.num_records}.pkl"
+        path = self.path / f"{self.num_records}"
         logger.debug(f"Start write: {path}")
         pkl = pickle.dumps(self.buff)
         compressed = self.compressor.encode(pkl)
@@ -723,37 +772,35 @@ class PcvcfFieldWriter:
         )
         if len(self.buff) > 0:
             self.write_chunk()
-        with open(self.path / "chunk_index.pkl", "wb") as f:
+        with open(self.path / "chunk_index", "wb") as f:
             a = np.array(self.chunk_index, dtype=int)
             pickle.dump(a, f)
 
 
-class PcvcfPartitionWriter(contextlib.AbstractContextManager):
+class IcfPartitionWriter(contextlib.AbstractContextManager):
     """
-    Writes the data for a PickleChunkedVcf partition.
+    Writes the data for a IntermediateColumnarFormat partition.
     """
 
     def __init__(
         self,
-        vcf_metadata,
+        icf_metadata,
         out_path,
         partition_index,
-        compressor,
-        *,
-        chunk_size=1,
     ):
         self.partition_index = partition_index
         # chunk_size is in megabytes
-        max_buffered_bytes = chunk_size * 2**20
+        max_buffered_bytes = icf_metadata.column_chunk_size * 2**20
         assert max_buffered_bytes > 0
+        compressor = numcodecs.get_codec(icf_metadata.compressor)
 
         self.field_writers = {}
-        num_samples = len(vcf_metadata.samples)
-        for vcf_field in vcf_metadata.fields:
-            field_path = PickleChunkedVcfField.get_path(out_path, vcf_field)
+        num_samples = len(icf_metadata.samples)
+        for vcf_field in icf_metadata.fields:
+            field_path = get_vcf_field_path(out_path, vcf_field)
             field_partition_path = field_path / f"p{partition_index}"
             transformer = VcfValueTransformer.factory(vcf_field, num_samples)
-            self.field_writers[vcf_field.full_name] = PcvcfFieldWriter(
+            self.field_writers[vcf_field.full_name] = IcfFieldWriter(
                 vcf_field,
                 field_partition_path,
                 transformer,
@@ -777,16 +824,23 @@ class PcvcfPartitionWriter(contextlib.AbstractContextManager):
         return False
 
 
-class PickleChunkedVcf(collections.abc.Mapping):
+# TODO rename to IntermediateColumnarFormat and move to icf.py
+
+
+class IntermediateColumnarFormat(collections.abc.Mapping):
     # TODO Check if other compressors would give reasonable compression
     # with significantly faster times
-    DEFAULT_COMPRESSOR = numcodecs.Blosc(cname="zstd", clevel=7)
 
-    def __init__(self, path, metadata, vcf_header):
-        self.path = path
-        self.metadata = metadata
-        self.vcf_header = vcf_header
-        self.compressor = self.DEFAULT_COMPRESSOR
+    def __init__(self, path):
+        self.path = pathlib.Path(path)
+        # TODO raise a more informative error here telling people this
+        # directory is either a WIP or the wrong format.
+        with open(self.path / "metadata.json") as f:
+            self.metadata = IcfMetadata.fromdict(json.load(f))
+        with open(self.path / "header.txt") as f:
+            self.vcf_header = f.read()
+
+        self.compressor = numcodecs.get_codec(self.metadata.compressor)
         self.columns = {}
         partition_num_records = [
             partition.num_records for partition in self.metadata.partitions
@@ -794,11 +848,15 @@ class PickleChunkedVcf(collections.abc.Mapping):
         # Allow us to find which partition a given record is in
         self.partition_record_index = np.cumsum([0] + partition_num_records)
         for field in self.metadata.fields:
-            self.columns[field.full_name] = PickleChunkedVcfField(self, field)
+            self.columns[field.full_name] = IntermediateColumnarFormatField(self, field)
+        logger.info(
+            f"Loaded IntermediateColumnarFormat(partitions={self.num_partitions}, "
+            f"records={self.num_records}, columns={self.num_columns})"
+        )
 
     def __repr__(self):
         return (
-            f"PickleChunkedVcf(fields={len(self)}, partitions={self.num_partitions}, "
+            f"IntermediateColumnarFormat(fields={len(self)}, partitions={self.num_partitions}, "
             f"records={self.num_records}, path={self.path})"
         )
 
@@ -812,15 +870,6 @@ class PickleChunkedVcf(collections.abc.Mapping):
         return len(self.columns)
 
     def summary_table(self):
-        def display_number(x):
-            ret = "n/a"
-            if math.isfinite(x):
-                ret = f"{x: 0.2g}"
-            return ret
-
-        def display_size(n):
-            return humanfriendly.format_size(n)
-
         data = []
         for name, col in self.columns.items():
             summary = col.vcf_field.summary
@@ -838,14 +887,6 @@ class PickleChunkedVcf(collections.abc.Mapping):
             data.append(d)
         return data
 
-    @functools.cached_property
-    def total_uncompressed_bytes(self):
-        total = 0
-        for col in self.columns.values():
-            summary = col.vcf_field.summary
-            total += summary.uncompressed_size
-        return total
-
     @functools.cached_property
     def num_records(self):
         return sum(self.metadata.contig_record_counts.values())
@@ -862,57 +903,121 @@ class PickleChunkedVcf(collections.abc.Mapping):
     def num_columns(self):
         return len(self.columns)
 
-    def mkdirs(self):
-        self.path.mkdir()
-        for col in self.columns.values():
-            col.path.mkdir(parents=True)
-            for j in range(self.num_partitions):
-                part_path = col.path / f"p{j}"
-                part_path.mkdir()
 
-    @staticmethod
-    def load(path):
-        path = pathlib.Path(path)
-        with open(path / "metadata.json") as f:
-            metadata = VcfMetadata.fromdict(json.load(f))
-        with open(path / "header.txt") as f:
-            header = f.read()
-        pcvcf = PickleChunkedVcf(path, metadata, header)
-        logger.info(
-            f"Loaded PickleChunkedVcf(partitions={pcvcf.num_partitions}, "
-            f"records={pcvcf.num_records}, columns={pcvcf.num_columns})"
-        )
-        return pcvcf
+class IntermediateColumnarFormatWriter:
+    def __init__(self, path):
+        self.path = pathlib.Path(path)
+        self.wip_path = self.path / "wip"
+        self.metadata = None
 
-    @staticmethod
-    def convert_partition(
-        vcf_metadata,
-        partition_index,
-        out_path,
+    @property
+    def num_partitions(self):
+        return len(self.metadata.partitions)
+
+    def init(
+        self,
+        vcfs,
         *,
         column_chunk_size=16,
+        worker_processes=1,
+        target_num_partitions=None,
+        show_progress=False,
     ):
-        partition = vcf_metadata.partitions[partition_index]
+        if self.path.exists():
+            shutil.rmtree(self.path)
+        vcfs = [pathlib.Path(vcf) for vcf in vcfs]
+        target_num_partitions = max(target_num_partitions, len(vcfs))
+
+        # TODO move scan_vcfs into this class
+        icf_metadata, header = scan_vcfs(
+            vcfs,
+            worker_processes=worker_processes,
+            show_progress=show_progress,
+            target_num_partitions=target_num_partitions,
+            column_chunk_size=column_chunk_size,
+        )
+        self.metadata = icf_metadata
+
+        self.mkdirs()
+
+        # Note: this is needed for the current version of the vcfzarr spec, but it's
+        # probably goint to be dropped.
+        # https://github.com/pystatgen/vcf-zarr-spec/issues/15
+        # May be useful to keep lying around still though?
+        logger.info(f"Writing VCF header")
+        with open(self.path / "header.txt", "w") as f:
+            f.write(header)
+
+        logger.info(f"Writing WIP metadata")
+        with open(self.wip_path / "metadata.json", "w") as f:
+            json.dump(self.metadata.asdict(), f, indent=4)
+        return self.num_partitions
+
+    def mkdirs(self):
+        # TODO add worker_processes here and do this with the ParallelWorkManager
+        logger.info(
+            f"Creating {len(self.metadata.fields) * self.num_partitions} directories"
+        )
+        self.path.mkdir()
+        self.wip_path.mkdir()
+        for field in self.metadata.fields:
+            col_path = get_vcf_field_path(self.path, field)
+            logger.debug(f"Make directories for {field.full_name} at {col_path}")
+            col_path.mkdir(parents=True)
+            for j in range(self.num_partitions):
+                part_path = col_path / f"p{j}"
+                part_path.mkdir()
+
+    def load_partition_summaries(self):
+        summaries = []
+        not_found = []
+        for j in range(self.num_partitions):
+            try:
+                with open(self.wip_path / f"p{j}_summary.json") as f:
+                    summary = json.load(f)
+                    for k, v in summary["field_summaries"].items():
+                        summary["field_summaries"][k] = VcfFieldSummary.fromdict(v)
+                    summaries.append(summary)
+            except FileNotFoundError:
+                not_found.append(j)
+        if len(not_found) > 0:
+            raise FileNotFoundError(
+                f"Partition metadata not found for {len(not_found)} partitions: {not_found}"
+            )
+        return summaries
+
+    def load_metadata(self):
+        if self.metadata is None:
+            with open(self.wip_path / f"metadata.json") as f:
+                self.metadata = IcfMetadata.fromdict(json.load(f))
+
+    def process_partition(self, partition_index):
+        self.load_metadata()
+        summary_path = self.wip_path / f"p{partition_index}_summary.json"
+        # If someone is rewriting a summary path (for whatever reason), make sure it
+        # doesn't look like it's already been completed.
+        # NOTE to do this properly we probably need to take a lock on this file - but
+        # this simple approach will catch the vast majority of problems.
+        if summary_path.exists():
+            summary_path.unlink()
+
+        partition = self.metadata.partitions[partition_index]
         logger.info(
             f"Start p{partition_index} {partition.vcf_path}__{partition.region}"
         )
-        info_fields = vcf_metadata.info_fields
+        info_fields = self.metadata.info_fields
         format_fields = []
         has_gt = False
-        for field in vcf_metadata.format_fields:
+        for field in self.metadata.format_fields:
             if field.name == "GT":
                 has_gt = True
             else:
                 format_fields.append(field)
 
-        compressor = PickleChunkedVcf.DEFAULT_COMPRESSOR
-
-        with PcvcfPartitionWriter(
-            vcf_metadata,
-            out_path,
+        with IcfPartitionWriter(
+            self.metadata,
+            self.path,
             partition_index,
-            compressor,
-            chunk_size=column_chunk_size,
         ) as tcw:
             with vcf_utils.IndexedVcf(partition.vcf_path) as ivcf:
                 num_records = 0
@@ -930,108 +1035,172 @@ class PickleChunkedVcf(collections.abc.Mapping):
                     if has_gt:
                         tcw.append("FORMAT/GT", variant.genotype.array())
                     for field in format_fields:
-                        val = None
-                        try:
-                            val = variant.format(field.name)
-                        except KeyError:
-                            pass
+                        val = variant.format(field.name)
                         tcw.append(field.full_name, val)
                     # Note: an issue with updating the progress per variant here like this
                     # is that we get a significant pause at the end of the counter while
                     # all the "small" fields get flushed. Possibly not much to be done about it.
                     core.update_progress(1)
+            logger.info(
+                f"Finished reading VCF for partition {partition_index}, flushing buffers"
+            )
 
+        partition_metadata = {
+            "num_records": num_records,
+            "field_summaries": {k: v.asdict() for k, v in tcw.field_summaries.items()},
+        }
+        with open(summary_path, "w") as f:
+            json.dump(partition_metadata, f, indent=4)
         logger.info(
             f"Finish p{partition_index} {partition.vcf_path}__{partition.region}="
             f"{num_records} records"
         )
-        return partition_index, tcw.field_summaries, num_records
 
-    @staticmethod
-    def convert(
-        vcfs, out_path, *, column_chunk_size=16, worker_processes=1, show_progress=False
+    def process_partition_slice(
+        self,
+        start,
+        stop,
+        *,
+        worker_processes=1,
+        show_progress=False,
     ):
-        out_path = pathlib.Path(out_path)
-        # TODO make scan work in parallel using general progress code too
-        target_num_partitions = max(1, worker_processes * 4)
-        vcf_metadata, header = scan_vcfs(
-            vcfs,
-            worker_processes=worker_processes,
-            show_progress=show_progress,
-            target_num_partitions=target_num_partitions,
-        )
-        pcvcf = PickleChunkedVcf(out_path, vcf_metadata, header)
-        pcvcf.mkdirs()
-
+        self.load_metadata()
+        if start == 0 and stop == self.num_partitions:
+            num_records = self.metadata.num_records
+        else:
+            # We only know the number of records if all partitions are done at once,
+            # and we signal this to tqdm by passing None as the total.
+            num_records = None
+        num_columns = len(self.metadata.fields)
+        num_samples = len(self.metadata.samples)
         logger.info(
-            f"Exploding {pcvcf.num_columns} columns {vcf_metadata.num_records} variants "
-            f"{pcvcf.num_samples} samples"
+            f"Exploding columns={num_columns} samples={num_samples}; "
+            f"partitions={stop - start} "
+            f"variants={'unknown' if num_records is None else num_records}"
         )
         progress_config = core.ProgressConfig(
-            total=vcf_metadata.num_records,
+            total=num_records,
             units="vars",
             title="Explode",
             show=show_progress,
         )
         with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-            for j, partition in enumerate(vcf_metadata.partitions):
-                pwm.submit(
-                    PickleChunkedVcf.convert_partition,
-                    vcf_metadata,
-                    j,
-                    out_path,
-                    column_chunk_size=column_chunk_size,
-                )
-            num_records = 0
-            partition_summaries = []
-            for index, summary, num_records in pwm.results_as_completed():
-                partition_summaries.append(summary)
-                vcf_metadata.partitions[index].num_records = num_records
-
-        total_records = sum(
-            partition.num_records for partition in vcf_metadata.partitions
+            for j in range(start, stop):
+                pwm.submit(self.process_partition, j)
+
+    def explode(self, *, worker_processes=1, show_progress=False):
+        self.load_metadata()
+        return self.process_partition_slice(
+            0,
+            self.num_partitions,
+            worker_processes=worker_processes,
+            show_progress=show_progress,
         )
-        assert total_records == pcvcf.num_records
 
-        for field in vcf_metadata.fields:
-            # Clear the summary to avoid problems when running in debug
-            # syncronous mode
-            field.summary = VcfFieldSummary()
+    def explode_partition(self, partition, *, show_progress=False, worker_processes=1):
+        self.load_metadata()
+        if partition < 0 or partition >= self.num_partitions:
+            raise ValueError(
+                "Partition index must be in the range 0 <= index < num_partitions"
+            )
+        return self.process_partition_slice(
+            partition,
+            partition + 1,
+            worker_processes=worker_processes,
+            show_progress=show_progress,
+        )
+
+    def finalise(self):
+        self.load_metadata()
+        partition_summaries = self.load_partition_summaries()
+        total_records = 0
+        for index, summary in enumerate(partition_summaries):
+            partition_records = summary["num_records"]
+            self.metadata.partitions[index].num_records = partition_records
+            total_records += partition_records
+        assert total_records == self.metadata.num_records
+
+        for field in self.metadata.fields:
             for summary in partition_summaries:
-                field.summary.update(summary[field.full_name])
+                field.summary.update(summary["field_summaries"][field.full_name])
 
-        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)
+        logger.info(f"Finalising metadata")
+        with open(self.path / "metadata.json", "w") as f:
+            json.dump(self.metadata.asdict(), f, indent=4)
+
+        logger.debug(f"Removing WIP directory")
+        shutil.rmtree(self.wip_path)
 
 
 def explode(
     vcfs,
-    out_path,
+    icf_path,
     *,
     column_chunk_size=16,
     worker_processes=1,
     show_progress=False,
 ):
-    out_path = pathlib.Path(out_path)
-    if out_path.exists():
-        shutil.rmtree(out_path)
-
-    PickleChunkedVcf.convert(
+    writer = IntermediateColumnarFormatWriter(icf_path)
+    num_partitions = writer.init(
         vcfs,
-        out_path,
+        # Heuristic to get reasonable worker utilisation with lumpy partition sizing
+        target_num_partitions=max(1, worker_processes * 4),
+        worker_processes=worker_processes,
+        show_progress=show_progress,
         column_chunk_size=column_chunk_size,
+    )
+    writer.explode(worker_processes=worker_processes, show_progress=show_progress)
+    writer.finalise()
+    return IntermediateColumnarFormat(icf_path)
+
+
+def explode_init(
+    icf_path,
+    vcfs,
+    *,
+    column_chunk_size=16,
+    target_num_partitions=1,
+    worker_processes=1,
+    show_progress=False,
+):
+    writer = IntermediateColumnarFormatWriter(icf_path)
+    return writer.init(
+        vcfs,
+        target_num_partitions=target_num_partitions,
         worker_processes=worker_processes,
         show_progress=show_progress,
+        column_chunk_size=column_chunk_size,
     )
-    return PickleChunkedVcf.load(out_path)
 
 
-def inspect(if_path):
+# NOTE only including worker_processes here so we can use the 0 option to get the
+# work done syncronously and so we can get test coverage on it. Should find a
+# better way to do this.
+def explode_partition(icf_path, partition, *, show_progress=False, worker_processes=1):
+    writer = IntermediateColumnarFormatWriter(icf_path)
+    writer.explode_partition(
+        partition, show_progress=show_progress, worker_processes=worker_processes
+    )
+
+
+def explode_finalise(icf_path):
+    writer = IntermediateColumnarFormatWriter(icf_path)
+    writer.finalise()
+
+
+def inspect(path):
+    path = pathlib.Path(path)
     # TODO add support for the Zarr format also
-    pcvcf = PickleChunkedVcf.load(if_path)
-    return pcvcf.summary_table()
+    if (path / "metadata.json").exists():
+        obj = IntermediateColumnarFormat(path)
+    elif (path / ".zmetadata").exists():
+        obj = VcfZarr(path)
+    else:
+        raise ValueError("Format not recognised")  # NEEDS TEST
+    return obj.summary_table()
+
+
+DEFAULT_ZARR_COMPRESSOR = numcodecs.Blosc(cname="zstd", clevel=7)
 
 
 @dataclasses.dataclass
@@ -1043,20 +1212,46 @@ class ZarrColumnSpec:
     dimensions: list
     description: str
     vcf_field: str
-    compressor: dict
+    compressor: dict = None
+    filters: list = None
     # TODO add filters
 
     def __post_init__(self):
         self.shape = tuple(self.shape)
         self.chunks = tuple(self.chunks)
         self.dimensions = tuple(self.dimensions)
+        self.compressor = DEFAULT_ZARR_COMPRESSOR.get_config()
+        self.filters = []
+        self._choose_compressor_settings()
+
+    def _choose_compressor_settings(self):
+        """
+        Choose compressor and filter settings based on the size and
+        type of the array, plus some hueristics from observed properties
+        of VCFs.
+
+        See https://github.com/pystatgen/bio2zarr/discussions/74
+        """
+        dt = np.dtype(self.dtype)
+        # Default is to not shuffle, because autoshuffle isn't recognised
+        # by many Zarr implementations, and shuffling can lead to worse
+        # performance in some cases anyway. Turning on shuffle should be a
+        # deliberate choice.
+        shuffle = numcodecs.Blosc.NOSHUFFLE
+        if dt.itemsize == 1:
+            # Any 1 byte field gets BITSHUFFLE by default
+            shuffle = numcodecs.Blosc.BITSHUFFLE
+        self.compressor["shuffle"] = shuffle
+
+
+ZARR_SCHEMA_FORMAT_VERSION = "0.2"
 
 
 @dataclasses.dataclass
-class ZarrConversionSpec:
+class VcfZarrSchema:
     format_version: str
-    chunk_width: int
-    chunk_length: int
+    samples_chunk_size: int
+    variants_chunk_size: int
     dimensions: list
     sample_id: list
     contig_id: list
@@ -1072,7 +1267,12 @@ class ZarrConversionSpec:
 
     @staticmethod
     def fromdict(d):
-        ret = ZarrConversionSpec(**d)
+        if d["format_version"] != ZARR_SCHEMA_FORMAT_VERSION:
+            raise ValueError(
+                "Zarr schema format version mismatch: "
+                f"{d['format_version']} != {ZARR_SCHEMA_FORMAT_VERSION}"
+            )
+        ret = VcfZarrSchema(**d)
         ret.columns = {
             key: ZarrColumnSpec(**value) for key, value in d["columns"].items()
         }
@@ -1080,19 +1280,20 @@ class ZarrConversionSpec:
 
     @staticmethod
     def fromjson(s):
-        return ZarrConversionSpec.fromdict(json.loads(s))
+        return VcfZarrSchema.fromdict(json.loads(s))
 
     @staticmethod
-    def generate(pcvcf, chunk_length=None, chunk_width=None):
-        m = pcvcf.num_records
-        n = pcvcf.num_samples
+    def generate(icf, variants_chunk_size=None, samples_chunk_size=None):
+        m = icf.num_records
+        n = icf.num_samples
         # FIXME
-        if chunk_width is None:
-            chunk_width = 1000
-        if chunk_length is None:
-            chunk_length = 10_000
-        logger.info(f"Generating schema with chunks={chunk_length, chunk_width}")
-        compressor = core.default_compressor.get_config()
+        if samples_chunk_size is None:
+            samples_chunk_size = 1000
+        if variants_chunk_size is None:
+            variants_chunk_size = 10_000
+        logger.info(
+            f"Generating schema with chunks={variants_chunk_size, samples_chunk_size}"
+        )
 
         def fixed_field_spec(
             name, dtype, vcf_field=None, shape=(m,), dimensions=("variants",)
@@ -1104,13 +1305,12 @@ class ZarrConversionSpec:
                 shape=shape,
                 description="",
                 dimensions=dimensions,
-                chunks=[chunk_length],
-                compressor=compressor,
+                chunks=[variants_chunk_size],
             )
 
-        alt_col = pcvcf.columns["ALT"]
+        alt_col = icf.columns["ALT"]
         max_alleles = alt_col.vcf_field.summary.max_number + 1
-        num_filters = len(pcvcf.metadata.filters)
+        num_filters = len(icf.metadata.filters)
 
         # # FIXME get dtype from lookup table
         colspecs = [
@@ -1153,7 +1353,7 @@ class ZarrConversionSpec:
         ]
 
         gt_field = None
-        for field in pcvcf.metadata.fields:
+        for field in icf.metadata.fields:
             if field.category == "fixed":
                 continue
             if field.name == "GT":
@@ -1162,11 +1362,11 @@ class ZarrConversionSpec:
             shape = [m]
             prefix = "variant_"
             dimensions = ["variants"]
-            chunks = [chunk_length]
+            chunks = [variants_chunk_size]
             if field.category == "FORMAT":
                 prefix = "call_"
                 shape.append(n)
-                chunks.append(chunk_width),
+                chunks.append(samples_chunk_size),
                 dimensions.append("samples")
             # TODO make an option to add in the empty extra dimension
             if field.summary.max_number > 1:
@@ -1181,14 +1381,13 @@ class ZarrConversionSpec:
                 chunks=chunks,
                 dimensions=dimensions,
                 description=field.description,
-                compressor=compressor,
             )
             colspecs.append(colspec)
 
         if gt_field is not None:
             ploidy = gt_field.summary.max_number - 1
             shape = [m, n]
-            chunks = [chunk_length, chunk_width]
+            chunks = [variants_chunk_size, samples_chunk_size]
             dimensions = ["variants", "samples"]
 
             colspecs.append(
@@ -1200,7 +1399,6 @@ class ZarrConversionSpec:
                     chunks=list(chunks),
                     dimensions=list(dimensions),
                     description="",
-                    compressor=compressor,
                 )
             )
             shape += [ploidy]
@@ -1214,7 +1412,6 @@ class ZarrConversionSpec:
                     chunks=list(chunks),
                     dimensions=list(dimensions),
                     description="",
-                    compressor=compressor,
                 )
             )
             colspecs.append(
@@ -1226,47 +1423,100 @@ class ZarrConversionSpec:
                     chunks=list(chunks),
                     dimensions=list(dimensions),
                     description="",
-                    compressor=compressor,
                 )
             )
 
-        return ZarrConversionSpec(
-            # TODO do something systematic
-            format_version="0.1",
-            chunk_width=chunk_width,
-            chunk_length=chunk_length,
+        return VcfZarrSchema(
+            format_version=ZARR_SCHEMA_FORMAT_VERSION,
+            samples_chunk_size=samples_chunk_size,
+            variants_chunk_size=variants_chunk_size,
             columns={col.name: col for col in colspecs},
             dimensions=["variants", "samples", "ploidy", "alleles", "filters"],
-            sample_id=pcvcf.metadata.samples,
-            contig_id=pcvcf.metadata.contig_names,
-            contig_length=pcvcf.metadata.contig_lengths,
-            filter_id=pcvcf.metadata.filters,
+            sample_id=icf.metadata.samples,
+            contig_id=icf.metadata.contig_names,
+            contig_length=icf.metadata.contig_lengths,
+            filter_id=icf.metadata.filters,
         )
 
 
-class SgvcfZarr:
+class VcfZarr:
     def __init__(self, path):
+        if not (path / ".zmetadata").exists():
+            raise ValueError("Not in VcfZarr format")  # NEEDS TEST
+        self.root = zarr.open(path, mode="r")
+
+    def __repr__(self):
+        return repr(self.root)  # NEEDS TEST
+
+    def summary_table(self):
+        data = []
+        arrays = [(a.nbytes_stored, a) for _, a in self.root.arrays()]
+        arrays.sort(key=lambda x: x[0])
+        for stored, array in reversed(arrays):
+            d = {
+                "name": array.name,
+                "dtype": str(array.dtype),
+                "stored": display_size(stored),
+                "size": display_size(array.nbytes),
+                "ratio": display_number(array.nbytes / stored),
+                "nchunks": str(array.nchunks),
+                "chunk_size": display_size(array.nbytes / array.nchunks),
+                "avg_chunk_stored": display_size(int(stored / array.nchunks)),
+                "shape": str(array.shape),
+                "chunk_shape": str(array.chunks),
+                "compressor": str(array.compressor),
+                "filters": str(array.filters),
+            }
+            data.append(d)
+        return data
+
+
+@dataclasses.dataclass
+class EncodingWork:
+    func: callable = dataclasses.field(repr=False)
+    start: int
+    stop: int
+    columns: list[str]
+    memory: int = 0
+
+
+class VcfZarrWriter:
+    def __init__(self, path, icf, schema):
         self.path = pathlib.Path(path)
-        self.root = None
+        self.icf = icf
+        self.schema = schema
+        store = zarr.DirectoryStore(self.path)
+        self.root = zarr.group(store=store)
 
-    def create_array(self, variable):
+    def init_array(self, variable):
         # print("CREATE", variable)
         object_codec = None
         if variable.dtype == "O":
             object_codec = numcodecs.VLenUTF8()
         a = self.root.empty(
-            variable.name,
+            "wip_" + variable.name,
             shape=variable.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,
         )
         a.attrs["_ARRAY_DIMENSIONS"] = variable.dimensions
 
-    def encode_column_slice(self, pcvcf, column, start, stop):
-        source_col = pcvcf.columns[column.vcf_field]
-        array = self.root[column.name]
+    def get_array(self, name):
+        return self.root["wip_" + name]
+
+    def finalise_array(self, variable_name):
+        source = self.path / ("wip_" + variable_name)
+        dest = self.path / variable_name
+        # Atomic swap
+        os.rename(source, dest)
+        logger.info(f"Finalised {variable_name}")
+
+    def encode_array_slice(self, column, start, stop):
+        source_col = self.icf.columns[column.vcf_field]
+        array = self.get_array(column.name)
         ba = core.BufferedArray(array, start)
         sanitiser = source_col.sanitiser_factory(ba.buff.shape)
 
@@ -1278,11 +1528,11 @@ class SgvcfZarr:
         ba.flush()
         logger.debug(f"Encoded {column.name} slice {start}:{stop}")
 
-    def encode_genotypes_slice(self, pcvcf, start, stop):
-        source_col = pcvcf.columns["FORMAT/GT"]
-        gt = core.BufferedArray(self.root["call_genotype"], start)
-        gt_mask = core.BufferedArray(self.root["call_genotype_mask"], start)
-        gt_phased = core.BufferedArray(self.root["call_genotype_phased"], start)
+    def encode_genotypes_slice(self, start, stop):
+        source_col = self.icf.columns["FORMAT/GT"]
+        gt = core.BufferedArray(self.get_array("call_genotype"), start)
+        gt_mask = core.BufferedArray(self.get_array("call_genotype_mask"), start)
+        gt_phased = core.BufferedArray(self.get_array("call_genotype_phased"), start)
 
         for value in source_col.iter_values(start, stop):
             j = gt.next_buffer_row()
@@ -1298,10 +1548,10 @@ class SgvcfZarr:
         gt_mask.flush()
         logger.debug(f"Encoded GT slice {start}:{stop}")
 
-    def encode_alleles_slice(self, pcvcf, start, stop):
-        ref_col = pcvcf.columns["REF"]
-        alt_col = pcvcf.columns["ALT"]
-        alleles = core.BufferedArray(self.root["variant_allele"], start)
+    def encode_alleles_slice(self, start, stop):
+        ref_col = self.icf.columns["REF"]
+        alt_col = self.icf.columns["ALT"]
+        alleles = core.BufferedArray(self.get_array("variant_allele"), start)
 
         for ref, alt in zip(
             ref_col.iter_values(start, stop), alt_col.iter_values(start, stop)
@@ -1313,10 +1563,10 @@ class SgvcfZarr:
         alleles.flush()
         logger.debug(f"Encoded alleles slice {start}:{stop}")
 
-    def encode_id_slice(self, pcvcf, start, stop):
-        col = pcvcf.columns["ID"]
-        vid = core.BufferedArray(self.root["variant_id"], start)
-        vid_mask = core.BufferedArray(self.root["variant_id_mask"], start)
+    def encode_id_slice(self, start, stop):
+        col = self.icf.columns["ID"]
+        vid = core.BufferedArray(self.get_array("variant_id"), start)
+        vid_mask = core.BufferedArray(self.get_array("variant_id_mask"), start)
 
         for value in col.iter_values(start, stop):
             j = vid.next_buffer_row()
@@ -1332,182 +1582,246 @@ class SgvcfZarr:
         vid_mask.flush()
         logger.debug(f"Encoded ID slice {start}:{stop}")
 
-    def encode_filters_slice(self, pcvcf, lookup, start, stop):
-        col = pcvcf.columns["FILTERS"]
-        var_filter = core.BufferedArray(self.root["variant_filter"], start)
+    def encode_filters_slice(self, lookup, start, stop):
+        col = self.icf.columns["FILTERS"]
+        var_filter = core.BufferedArray(self.get_array("variant_filter"), start)
 
         for value in col.iter_values(start, stop):
             j = var_filter.next_buffer_row()
             var_filter.buff[j] = False
-            try:
-                for f in value:
+            for f in value:
+                try:
                     var_filter.buff[j, lookup[f]] = True
-            except IndexError:
-                raise ValueError(f"Filter '{f}' was not defined in the header.")
+                except KeyError:
+                    raise ValueError(f"Filter '{f}' was not defined in the header.")
         var_filter.flush()
         logger.debug(f"Encoded FILTERS slice {start}:{stop}")
 
-    def encode_contig_slice(self, pcvcf, lookup, start, stop):
-        col = pcvcf.columns["CHROM"]
-        contig = core.BufferedArray(self.root["variant_contig"], start)
+    def encode_contig_slice(self, lookup, start, stop):
+        col = self.icf.columns["CHROM"]
+        contig = core.BufferedArray(self.get_array("variant_contig"), start)
 
         for value in col.iter_values(start, stop):
             j = contig.next_buffer_row()
-            try:
-                contig.buff[j] = lookup[value[0]]
-            except KeyError:
-                # TODO add advice about adding it to the spec
-                raise ValueError(f"Contig '{contig}' was not defined in the header.")
+            # 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()
         logger.debug(f"Encoded CHROM slice {start}:{stop}")
 
-    def encode_samples(self, pcvcf, sample_id, chunk_width):
-        if not np.array_equal(sample_id, pcvcf.metadata.samples):
-            raise ValueError("Subsetting or reordering samples not supported currently")
+    def encode_samples(self):
+        if not np.array_equal(self.schema.sample_id, self.icf.metadata.samples):
+            raise ValueError(
+                "Subsetting or reordering samples not supported currently"
+            )  # NEEDS TEST
         array = self.root.array(
             "sample_id",
-            sample_id,
+            self.schema.sample_id,
             dtype="str",
-            compressor=core.default_compressor,
-            chunks=(chunk_width,),
+            compressor=DEFAULT_ZARR_COMPRESSOR,
+            chunks=(self.schema.samples_chunk_size,),
         )
         array.attrs["_ARRAY_DIMENSIONS"] = ["samples"]
         logger.debug("Samples done")
 
-    def encode_contig_id(self, pcvcf, contig_names, contig_lengths):
+    def encode_contig_id(self):
         array = self.root.array(
             "contig_id",
-            contig_names,
+            self.schema.contig_id,
             dtype="str",
-            compressor=core.default_compressor,
+            compressor=DEFAULT_ZARR_COMPRESSOR,
         )
         array.attrs["_ARRAY_DIMENSIONS"] = ["contigs"]
-        if contig_lengths is not None:
+        if self.schema.contig_length is not None:
             array = self.root.array(
                 "contig_length",
-                contig_lengths,
+                self.schema.contig_length,
                 dtype=np.int64,
             )
             array.attrs["_ARRAY_DIMENSIONS"] = ["contigs"]
-        return {v: j for j, v in enumerate(contig_names)}
+        return {v: j for j, v in enumerate(self.schema.contig_id)}
 
-    def encode_filter_id(self, pcvcf, filter_names):
+    def encode_filter_id(self):
         array = self.root.array(
             "filter_id",
-            filter_names,
+            self.schema.filter_id,
             dtype="str",
-            compressor=core.default_compressor,
+            compressor=DEFAULT_ZARR_COMPRESSOR,
         )
         array.attrs["_ARRAY_DIMENSIONS"] = ["filters"]
-        return {v: j for j, v in enumerate(filter_names)}
+        return {v: j for j, v in enumerate(self.schema.filter_id)}
+
+    def init(self):
+        self.root.attrs["vcf_zarr_version"] = "0.2"
+        self.root.attrs["vcf_header"] = self.icf.vcf_header
+        self.root.attrs["source"] = f"bio2zarr-{provenance.__version__}"
+        for column in self.schema.columns.values():
+            self.init_array(column)
+
+    def finalise(self):
+        # for column in self.schema.columns.values():
+        #     self.finalise_array(column)
+        zarr.consolidate_metadata(self.path)
 
-    @staticmethod
     def encode(
-        pcvcf,
-        path,
-        conversion_spec,
-        *,
+        self,
         worker_processes=1,
         max_v_chunks=None,
         show_progress=False,
+        max_memory=None,
     ):
-        path = pathlib.Path(path)
-        # TODO: we should do this as a future to avoid blocking
-        if path.exists():
-            logger.warning(f"Deleting existing {path}")
-            shutil.rmtree(path)
-        write_path = path.with_suffix(path.suffix + f".{os.getpid()}.build")
-        store = zarr.DirectoryStore(write_path)
-        # FIXME, duplicating logic about the store
-        logger.info(f"Create zarr at {write_path}")
-        sgvcf = SgvcfZarr(write_path)
-        sgvcf.root = zarr.group(store=store, overwrite=True)
-        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
-        sgvcf.root.attrs["source"] = f"bio2zarr-{provenance.__version__}"
+        if max_memory is None:
+            # Unbounded
+            max_memory = 2**63
+        else:
+            # Value is specified in Mibibytes
+            max_memory *= 2**20  # NEEDS TEST
 
+        # TODO this will move into the setup logic later when we're making it possible
+        # to split the work by slice
         num_slices = max(1, worker_processes * 4)
         # Using POS arbitrarily to get the array slices
         slices = core.chunk_aligned_slices(
-            sgvcf.root["variant_position"], num_slices, max_chunks=max_v_chunks
+            self.get_array("variant_position"), num_slices, max_chunks=max_v_chunks
         )
         truncated = slices[-1][-1]
-        for array in sgvcf.root.values():
+        for array in self.root.values():
             if array.attrs["_ARRAY_DIMENSIONS"][0] == "variants":
                 shape = list(array.shape)
                 shape[0] = truncated
                 array.resize(shape)
 
-        chunked_1d = [
-            col for col in conversion_spec.columns.values() if len(col.chunks) <= 1
-        ]
+        total_bytes = 0
+        encoding_memory_requirements = {}
+        for col in self.schema.columns.values():
+            array = self.get_array(col.name)
+            # NOTE!! this is bad, we're potentially creating quite a large
+            # numpy array for basically nothing. We can compute this.
+            variant_chunk_size = array.blocks[0].nbytes
+            encoding_memory_requirements[col.name] = variant_chunk_size
+            logger.debug(
+                f"{col.name} requires at least {display_size(variant_chunk_size)} per worker"
+            )
+            total_bytes += array.nbytes
+
+        filter_id_map = self.encode_filter_id()
+        contig_id_map = self.encode_contig_id()
+
+        work = []
+        for start, stop in slices:
+            for col in self.schema.columns.values():
+                if col.vcf_field is not None:
+                    f = functools.partial(self.encode_array_slice, col)
+                    work.append(
+                        EncodingWork(
+                            f,
+                            start,
+                            stop,
+                            [col.name],
+                            encoding_memory_requirements[col.name],
+                        )
+                    )
+            work.append(
+                EncodingWork(self.encode_alleles_slice, start, stop, ["variant_allele"])
+            )
+            work.append(
+                EncodingWork(
+                    self.encode_id_slice, start, stop, ["variant_id", "variant_id_mask"]
+                )
+            )
+            work.append(
+                EncodingWork(
+                    functools.partial(self.encode_filters_slice, filter_id_map),
+                    start,
+                    stop,
+                    ["variant_filter"],
+                )
+            )
+            work.append(
+                EncodingWork(
+                    functools.partial(self.encode_contig_slice, contig_id_map),
+                    start,
+                    stop,
+                    ["variant_contig"],
+                )
+            )
+            if "call_genotype" in self.schema.columns:
+                variables = [
+                    "call_genotype",
+                    "call_genotype_phased",
+                    "call_genotype_mask",
+                ]
+                gt_memory = sum(
+                    encoding_memory_requirements[name] for name in variables
+                )
+                work.append(
+                    EncodingWork(
+                        self.encode_genotypes_slice, start, stop, variables, gt_memory
+                    )
+                )
+
+        # Fail early if we can't fit a particular column into memory
+        for wp in work:
+            if wp.memory >= max_memory:
+                raise ValueError(  # NEEDS TEST
+                    f"Insufficient memory for {wp.columns}: "
+                    f"{display_size(wp.memory)} > {display_size(max_memory)}"
+                )
+
         progress_config = core.ProgressConfig(
-            total=sum(sgvcf.root[col.name].nchunks for col in chunked_1d),
-            title="Encode 1D",
-            units="chunks",
+            total=total_bytes,
+            title="Encode",
+            units="B",
             show=show_progress,
         )
 
-        # Do these syncronously for simplicity so we have the mapping
-        filter_id_map = sgvcf.encode_filter_id(pcvcf, conversion_spec.filter_id)
-        contig_id_map = sgvcf.encode_contig_id(
-            pcvcf, conversion_spec.contig_id, conversion_spec.contig_length
-        )
+        used_memory = 0
+        max_queued = 4 * max(1, worker_processes)
+        encoded_slices = collections.Counter()
 
         with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-            pwm.submit(
-                sgvcf.encode_samples,
-                pcvcf,
-                conversion_spec.sample_id,
-                conversion_spec.chunk_width,
-            )
-            for start, stop in slices:
-                pwm.submit(sgvcf.encode_alleles_slice, pcvcf, start, stop)
-                pwm.submit(sgvcf.encode_id_slice, pcvcf, start, stop)
-                pwm.submit(
-                    sgvcf.encode_filters_slice, pcvcf, filter_id_map, start, stop
-                )
-                pwm.submit(sgvcf.encode_contig_slice, pcvcf, contig_id_map, start, stop)
-                for col in chunked_1d:
-                    if col.vcf_field is not None:
-                        pwm.submit(sgvcf.encode_column_slice, pcvcf, col, start, stop)
-
-        chunked_2d = [
-            col for col in conversion_spec.columns.values() if len(col.chunks) >= 2
-        ]
-        if len(chunked_2d) > 0:
-            progress_config = core.ProgressConfig(
-                total=sum(sgvcf.root[col.name].nchunks for col in chunked_2d),
-                title="Encode 2D",
-                units="chunks",
-                show=show_progress,
-            )
-            with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-                if "call_genotype" in conversion_spec.columns:
-                    logger.info(f"Submit encode call_genotypes in {len(slices)} slices")
-                    for start, stop in slices:
-                        pwm.submit(sgvcf.encode_genotypes_slice, pcvcf, start, stop)
-
-                for col in chunked_2d:
-                    if col.vcf_field is not None:
-                        logger.info(f"Submit encode {col.name} in {len(slices)} slices")
-                        for start, stop in slices:
-                            pwm.submit(
-                                sgvcf.encode_column_slice, pcvcf, col, start, stop
-                            )
-
-        zarr.consolidate_metadata(write_path)
-        # Atomic swap, now we've completely finished.
-        logger.info(f"Moving to final path {path}")
-        os.rename(write_path, path)
+            future = pwm.submit(self.encode_samples)
+            future_to_work = {future: EncodingWork(None, 0, 0, [])}
+
+            def service_completed_futures():
+                nonlocal used_memory
+
+                completed = pwm.wait_for_completed()
+                for future in completed:
+                    wp_done = future_to_work.pop(future)
+                    used_memory -= wp_done.memory
+                    logger.debug(
+                        f"Complete {wp_done}: used mem={display_size(used_memory)}"
+                    )
+                    for column in wp_done.columns:
+                        encoded_slices[column] += 1
+                        if encoded_slices[column] == len(slices):
+                            # Do this syncronously for simplicity. Should be
+                            # fine as the workers will probably be busy with
+                            # large encode tasks most of the time.
+                            self.finalise_array(column)
+
+            for wp in work:
+                if (
+                    used_memory + wp.memory > max_memory
+                    or len(future_to_work) > max_queued
+                ):
+                    service_completed_futures()
+                future = pwm.submit(wp.func, wp.start, wp.stop)
+                used_memory += wp.memory
+                logger.debug(f"Submit {wp}: used mem={display_size(used_memory)}")
+                future_to_work[future] = wp
+
+            logger.debug("All work submitted")
+            while len(future_to_work) > 0:
+                service_completed_futures()
 
 
 def mkschema(if_path, out):
-    pcvcf = PickleChunkedVcf.load(if_path)
-    spec = ZarrConversionSpec.generate(pcvcf)
+    icf = IntermediateColumnarFormat(if_path)
+    spec = VcfZarrSchema.generate(icf)
     out.write(spec.asjson())
 
 
@@ -1515,42 +1829,49 @@ def encode(
     if_path,
     zarr_path,
     schema_path=None,
-    chunk_length=None,
-    chunk_width=None,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
     max_v_chunks=None,
+    max_memory=None,
     worker_processes=1,
     show_progress=False,
 ):
-    pcvcf = PickleChunkedVcf.load(if_path)
+    icf = IntermediateColumnarFormat(if_path)
     if schema_path is None:
-        schema = ZarrConversionSpec.generate(
-            pcvcf,
-            chunk_length=chunk_length,
-            chunk_width=chunk_width,
+        schema = VcfZarrSchema.generate(
+            icf,
+            variants_chunk_size=variants_chunk_size,
+            samples_chunk_size=samples_chunk_size,
         )
     else:
         logger.info(f"Reading schema from {schema_path}")
-        if chunk_length is not None or chunk_width is not None:
-            raise ValueError("Cannot specify schema along with chunk sizes")
+        if variants_chunk_size is not None or samples_chunk_size is not None:
+            raise ValueError(
+                "Cannot specify schema along with chunk sizes"
+            )  # NEEDS TEST
         with open(schema_path, "r") as f:
-            schema = ZarrConversionSpec.fromjson(f.read())
-
-    SgvcfZarr.encode(
-        pcvcf,
-        zarr_path,
-        conversion_spec=schema,
+            schema = VcfZarrSchema.fromjson(f.read())
+    zarr_path = pathlib.Path(zarr_path)
+    if zarr_path.exists():
+        logger.warning(f"Deleting existing {zarr_path}")
+        shutil.rmtree(zarr_path)
+    vzw = VcfZarrWriter(zarr_path, icf, schema)
+    vzw.init()
+    vzw.encode(
         max_v_chunks=max_v_chunks,
         worker_processes=worker_processes,
+        max_memory=max_memory,
         show_progress=show_progress,
     )
+    vzw.finalise()
 
 
 def convert(
     vcfs,
     out_path,
     *,
-    chunk_length=None,
-    chunk_width=None,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
     worker_processes=1,
     show_progress=False,
     # TODO add arguments to control location of tmpdir
@@ -1565,8 +1886,8 @@ def convert(
         encode(
             if_dir,
             out_path,
-            chunk_length=chunk_length,
-            chunk_width=chunk_width,
+            variants_chunk_size=variants_chunk_size,
+            samples_chunk_size=samples_chunk_size,
             worker_processes=worker_processes,
             show_progress=show_progress,
         )
@@ -1744,16 +2065,14 @@ def validate(vcf_path, zarr_path, show_progress=False):
             name = colname.split("_", 1)[1]
             if name.isupper():
                 vcf_type = info_headers[name]["Type"]
-                # print(root[colname])
                 info_fields[name] = vcf_type, iter(root[colname])
-    # print(info_fields)
 
     first_pos = next(vcf).POS
     start_index = np.searchsorted(pos, first_pos)
     assert pos[start_index] == first_pos
     vcf = cyvcf2.VCF(vcf_path)
     if show_progress:
-        iterator = tqdm.tqdm(vcf, desc="   Verify", total=vcf.num_records)
+        iterator = tqdm.tqdm(vcf, desc=" Verify", total=vcf.num_records)  # NEEDS TEST
     else:
         iterator = vcf
     for j, row in enumerate(iterator, start_index):
@@ -1790,11 +2109,7 @@ def validate(vcf_path, zarr_path, show_progress=False):
                 assert_info_val_equal(vcf_val, zarr_val, vcf_type)
 
         for name, (vcf_type, zarr_iter) in format_fields.items():
-            vcf_val = None
-            try:
-                vcf_val = row.format(name)
-            except KeyError:
-                pass
+            vcf_val = row.format(name)
             zarr_val = next(zarr_iter)
             if vcf_val is None:
                 assert_format_val_missing(zarr_val, vcf_type)