Fix progress bars and missing partition errors

Author: benjeffery

bio2zarr/cli.py

@@ -81,16 +81,18 @@ def explode(vcfs, out_path, verbose, worker_processes, column_chunk_size):
 @click.command
 @click.argument("vcfs", nargs=-1, required=True)
 @click.argument("out_path", type=click.Path())
+@click.argument("num_partitions", type=int, required=True)
 @verbose
 @worker_processes
-def explode_init(vcfs, out_path, verbose, worker_processes):
+def explode_init(vcfs, out_path, num_partitions, verbose, worker_processes):
     """
     Initial step for parallel conversion of VCF(s) to columnar intermediate format
     """
     setup_logging(verbose)
     vcf.explode_init(
         vcfs,
         out_path,
+        num_partitions=num_partitions,
         worker_processes=worker_processes,
         show_progress=True,
     )

bio2zarr/vcf.py

@@ -896,12 +896,20 @@ def write_header(self):
 
     def load_partition_summaries(self):
         summaries = []
+        not_found = []
         for j in range(self.num_partitions):
-            with open(self.path / f"partition_{j}_metadata.json") as f:
-                summary = json.load(f)
-                for k, v in summary['field_summaries'].items():
-                    summary['field_summaries'][k] = VcfFieldSummary(**v)
-                summaries.append(summary)
+            try:
+                with open(self.path / f"partition_{j}_metadata.json") as f:
+                    summary = json.load(f)
+                    for k, v in summary['field_summaries'].items():
+                        summary['field_summaries'][k] = VcfFieldSummary(**v)
+                    summaries.append(summary)
+            except FileNotFoundError:
+                not_found.append(j)
+        if not_found:
+            raise FileNotFoundError(
+                f"Partition metadata not found for {len(not_found)} partitions: {not_found}"
+            )
         return summaries
 
 
@@ -1008,8 +1016,13 @@ def convert_slice(self, start, stop, *, worker_processes=1, show_progress=False,
             f"Exploding {self.num_columns} columns {self.metadata.num_records} variants "
             f"{self.num_samples} samples"
         )
+        if start < 0:
+            raise ValueError(f"start={start} must be non-negative")
+        if stop > self.num_partitions:
+            raise ValueError(f"stop={stop} must be less than the number of partitions")
+        num_records_to_progress = sum([partition.num_records for partition in self.metadata.partitions[start:stop]])
         progress_config = core.ProgressConfig(
-            total=self.metadata.num_records,
+            total=num_records_to_progress,
             units="vars",
             title="Explode",
             show=show_progress,