Improve API for VCF read_variant_stats() and read_allele_count()

apis/python/src/tiledbvcf/allele_frequency.py

@@ -11,4 +11,4 @@ def read_allele_frequency(dataset_uri: str, region: str) -> pandas.DataFrame:
     import tiledbvcf
 
     ds = tiledbvcf.Dataset(uri=dataset_uri, mode="r")
-    return ds.read_variant_stats(region).to_pandas()
+    return ds.read_variant_stats(region)

apis/python/src/tiledbvcf/dataset.py

@@ -2,10 +2,11 @@
 import shutil
 import warnings
 from collections import namedtuple
-from typing import List
+from typing import Generator, List
 
 import pandas as pd
 import pyarrow as pa
+import pyarrow.compute as pc
 
 from . import libtiledbvcf
 
@@ -100,6 +101,43 @@ class Dataset(object):
         TileDB configuration, alternative to `cfg.tiledb_config`.
     """
 
+    class Region(object):
+        """
+        Represents a 1-based inclusive region.
+
+        Parameters
+        ----------
+        region
+            A string in the form "<contig>:<start>-<end>".
+        """
+
+        def __init__(self, region: str):
+            try:
+                contig, interval = region.split(":")
+                start, end = map(int, interval.split("-"))
+            except Exception:
+                raise Exception(
+                    '"region" parameter must have format "<contig>:<start>-<end>"'
+                )
+            if contig == "":
+                raise Exception("Region contig cannot be empty")
+            if start <= 0:
+                raise Exception("Regions must be 1-based")
+            if end < start:
+                raise Exception(f'"{interval}" is not a valid region interval')
+            self.contig = contig
+            self.start = start
+            self.end = end
+
+        def __lt__(self, region):
+            return self.to_tuple() < region.to_tuple()
+
+        def __str__(self):
+            return f"{self.contig}:{self.start}-{self.end}"
+
+        def to_tuple(self):
+            return (self.contig, self.start, self.end)
+
     def __init__(
         self,
         uri: str,
@@ -219,6 +257,22 @@ def _set_write_cfg(self, cfg):
                     pass
             self.writer.set_tiledb_config(",".join(tiledb_config_list))
 
+    # parses and sorts regions, then generates consolidated regions one at a time
+    def _prepare_regions(self, regions: List[str]) -> Generator[Region, None, None]:
+        if not regions:
+            return
+        prev_region, *parsed_regions = sorted(map(self.Region, regions))
+        for r in parsed_regions:
+            if prev_region.contig != r.contig:
+                yield prev_region
+                prev_region = r
+            elif r.start <= prev_region.end + 1:
+                prev_region.end = max(r.end, prev_region.end)
+            else:  # the regions are non-overlapping
+                yield prev_region
+                prev_region = r
+        yield prev_region
+
     def read_arrow(
         self,
         attrs: List[str] = DEFAULT_ATTRS,
@@ -271,14 +325,18 @@ def read_arrow(
 
         if isinstance(regions, str):
             regions = [regions]
+        if isinstance(regions, list):
+            regions = map(str, self._prepare_regions(regions))
+        else:
+            regions = ""
+
         if isinstance(samples, str):
             samples = [samples]
 
         self.reader.reset()
         self.reader.set_export_to_disk(False)
         self._set_samples(samples, samples_file)
 
-        regions = "" if regions is None else regions
         self.reader.set_regions(",".join(regions))
         self.reader.set_attributes(attrs)
         self.reader.set_check_samples_exist(not skip_check_samples)
@@ -294,36 +352,173 @@ def read_arrow(
     def read_variant_stats(
         self,
         region: str = None,
+        drop_ref: bool = False,
+        regions: List[str] = None,
+        scan_all_samples: bool = False,
     ) -> pd.DataFrame:
         """
         Read variant stats from the dataset into a Pandas DataFrame
 
         Parameters
         ----------
+        drop_ref
+            Omit "ref" alleles from the results
+        regions
+            Genomic regions to be queried.
+        scan_all_samples
+            Scan all samples when computing internal allele frequency.
         region
-            Genomic region to be queried.
+            **DEPRECATED** - Genomic region to be queried.
         """
+        # TODO: deprecated region and parse regions like read()
+        if not (region or regions):
+            raise Exception('"region" or "regions" parameter is required')
+        if region and regions:
+            raise Exception('"region" and "regions" parameters are mutually exclusive')
+        if region:
+            warnings.warn(
+                '"region" parameter is deprecated, use "regions" instead',
+                DeprecationWarning,
+            )
+            regions = [region]
+
+        kwargs = {
+            "drop_ref": drop_ref,
+            "regions": regions,
+            "scan_all_samples": scan_all_samples,
+        }
+        return self.read_variant_stats_arrow(**kwargs).to_pandas()
+
+    def read_variant_stats_arrow(
+        self,
+        region: str = None,
+        drop_ref: bool = False,
+        regions: List[str] = None,
+        scan_all_samples: bool = False,
+    ) -> pa.Table:
+        """
+        Read variant stats from the dataset into a PyArrow Table
+
+        Parameters
+        ----------
+        drop_ref
+            Omit "ref" alleles from the results
+        regions
+            Genomic regions to be queried.
+        scan_all_samples
+            Scan all samples when computing internal allele frequency.
+        region
+            **DEPRECATED** - Genomic region to be queried.
+        """
+        # TODO: deprecated region and parse regions like read()
+        if not (region or regions):
+            raise Exception('"region" or "regions" parameter is required')
+        if region and regions:
+            raise Exception('"region" and "regions" parameters are mutually exclusive')
+        if region:
+            warnings.warn(
+                '"region" parameter is deprecated, use "regions" instead',
+                DeprecationWarning,
+            )
+            regions = [region]
         if self.mode != "r":
             raise Exception("Dataset not open in read mode")
-        self.reader.set_regions(region)
-        return self.reader.get_variant_stats_results()
+
+        self.reader.reset()
+        self.reader.set_scan_all_samples(scan_all_samples)
+
+        # generates stats, sorts the results, and adds contig column one region at a time
+        def variant_stats_generator(regions):
+            for r in regions:
+                self.reader.set_regions(str(r))
+                stats = self.reader.get_variant_stats_results()
+                stats = stats.sort_by([("pos", "ascending"), ("alleles", "ascending")])
+                n = stats.num_rows
+                contig_col = [r.contig] * n
+                yield stats.add_column(0, "contig", [contig_col])
+
+        # drop reference alleles from results
+        consolidated_regions = self._prepare_regions(regions)
+        stats_tbl = pa.concat_tables(variant_stats_generator(consolidated_regions))
+        if drop_ref:
+            expr = pc.field("alleles") != "ref"
+            return stats_tbl.filter(expr)
+        return stats_tbl
 
     def read_allele_count(
         self,
         region: str = None,
+        regions: List[str] = None,
     ) -> pd.DataFrame:
         """
         Read allele count from the dataset into a Pandas DataFrame
 
         Parameters
         ----------
+        regions
+            Genomic regions to be queried.
+        region
+            **DEPRECATED** - Genomic region to be queried.
+        """
+        # TODO: deprecated region and parse regions like read()
+        if not (region or regions):
+            raise Exception('"region" or "regions" parameter is required')
+        if region and regions:
+            raise Exception('"region" and "regions" parameters are mutually exclusive')
+        if region:
+            warnings.warn(
+                '"region" parameter is deprecated, use "regions" instead',
+                DeprecationWarning,
+            )
+            regions = [region]
+        if self.mode != "r":
+            raise Exception("Dataset not open in read mode")
+
+        return self.read_allele_count_arrow(regions=regions).to_pandas()
+
+    def read_allele_count_arrow(
+        self,
+        region: str = None,
+        regions: List[str] = None,
+    ) -> pa.Table:
+        """
+        Read allele count from the dataset into a Pandas DataFrame
+
+        Parameters
+        ----------
+        regions
+            Genomic regions to be queried.
         region
-            Genomic region to be queried.
+            **DEPRECATED** - Genomic region to be queried.
         """
+        # TODO: deprecated region and parse regions like read()
+        if not (region or regions):
+            raise Exception('"region" or "regions" parameter is required')
+        if region and regions:
+            raise Exception('"region" and "regions" parameters are mutually exclusive')
+        if region:
+            warnings.warn(
+                '"region" parameter is deprecated, use "regions" instead',
+                DeprecationWarning,
+            )
+            regions = [region]
         if self.mode != "r":
             raise Exception("Dataset not open in read mode")
-        self.reader.set_regions(region)
-        return self.reader.get_allele_count_results()
+
+        # generates counts and adds contig column one region at a time
+        def allele_count_generator(regions):
+            for r in regions:
+                self.reader.set_regions(str(r))
+                counts = self.reader.get_allele_count_results()
+                contigs = counts.sort_by(
+                    [("pos", "ascending"), ("ref", "ascending"), ("alt", "ascending")]
+                )
+                n = counts.num_rows
+                contig_col = [r.contig] * n
+                yield counts.add_column(0, "contig", [contig_col])
+
+        consolidated_regions = self._prepare_regions(regions)
+        return pa.concat_tables(allele_count_generator(consolidated_regions))
 
     def read(
         self,
@@ -377,14 +572,17 @@ def read(
 
         if isinstance(regions, str):
             regions = [regions]
+        if isinstance(regions, list):
+            regions = map(str, self._prepare_regions(regions))
+        else:
+            regions = ""
         if isinstance(samples, str):
             samples = [samples]
 
         self.reader.reset()
         self.reader.set_export_to_disk(False)
         self._set_samples(samples, samples_file)
 
-        regions = "" if regions is None else regions
         self.reader.set_regions(",".join(regions))
         self.reader.set_attributes(attrs)
         self.reader.set_check_samples_exist(not skip_check_samples)
@@ -446,14 +644,17 @@ def export(
 
         if isinstance(regions, str):
             regions = [regions]
+        if isinstance(regions, list):
+            regions = map(str, self._prepare_regions(regions))
+        else:
+            regions = ""
         if isinstance(samples, str):
             samples = [samples]
 
         self.reader.reset()
         self.reader.set_export_to_disk(True)
         self._set_samples(samples, samples_file)
 
-        regions = "" if regions is None else regions
         self.reader.set_regions(",".join(regions))
         self.reader.set_check_samples_exist(not skip_check_samples)
         self.reader.set_enable_progress_estimation(enable_progress_estimation)
@@ -501,13 +702,17 @@ def read_iter(
 
         if isinstance(regions, str):
             regions = [regions]
+        if isinstance(regions, list):
+            regions = map(str, self._prepare_regions(regions))
+        else:
+            regions = ""
         if isinstance(samples, str):
             samples = [samples]
 
         self.reader.reset()
 
         if not self.read_completed():
-            yield self.read(attrs, samples, regions, samples_file, bed_file)
+            yield self.read(attrs, samples, list(regions), samples_file, bed_file)
         while not self.read_completed():
             yield self.continue_read()
 
@@ -593,14 +798,18 @@ def count(
 
         if isinstance(regions, str):
             regions = [regions]
+        if isinstance(regions, list):
+            regions = map(str, self._prepare_regions(regions))
+        else:
+            regions = ""
         if isinstance(samples, str):
             samples = [samples]
+        elif samples is None:
+            samples = ""
 
         self.reader.reset()
         self.reader.set_export_to_disk(False)
 
-        samples = "" if samples is None else samples
-        regions = "" if regions is None else regions
         self.reader.set_samples(",".join(samples))
         self.reader.set_regions(",".join(regions))