Variant Stats Direct Interface with IAF column (#597)

* Compute buffer sizes for variant stats when reading

* Populate buffers with expanded variant stats

* add direct variant stats reading to Reader class

* export direct variant stats via C API

* add retrieval of stats results in Python bindings

* add Python method to read variant stats

* add AF field to variant stats retrieval

* update existing Python interface to use C++ bindings

* buffer variant stats positions from map

* sort variant positions

* add AN field to variant stats retrieval

* Perform direct variant stats query without filter

* add test for direct variant stats interface

Author: awenocur

apis/python/src/tiledbvcf/allele_frequency.py

@@ -1,50 +1,14 @@
 import pandas
 
 
-def _calc_af(df) -> pandas.DataFrame:
-    """
-    Consolidate AC and compute AN, AF
-
-    :param pandas.Dataframe df
-    """
-    # Allele Count (AC) = sum of all AC at the same locus
-    # This step consolidates ACs from all ingested batches
-    df = df.groupby(["pos", "allele"], sort=True).sum(numeric_only=True)
-
-    # Allele Number (AN) = sum of AC at the same locus
-    an = df.groupby(["pos"], sort=True).ac.sum().rename("an")
-    df = df.join(an, how="inner").reset_index()
-
-    # Allele Frequency (AF) = AC / AN
-    df["af"] = df.ac / df.an
-    return df
-
-
 def read_allele_frequency(dataset_uri: str, region: str) -> pandas.DataFrame():
     """
     Read variant status
 
     :param dataset_uri: dataset URI
     :param region: genomics region to read
     """
-    import tiledb
-
-    # Get the variant stats uri
-    with tiledb.Group(dataset_uri) as g:
-        alleles_uri = g["variant_stats"].uri
-
-        try:
-            contig = region.split(":")[0]
-            start, end = map(int, region.split(":")[1].split("-"))
-            region_slice = slice(start, end)
-        except Exception as e:
-            raise ValueError(
-                f"Invalid region: {region}. Expected format: contig:start-end"
-            ) from e
-
-        with tiledb.open(alleles_uri) as A:
-            df = A.query(attrs=["ac", "allele"], dims=["pos", "contig"]).df[
-                contig, region_slice
-            ]
+    import tiledbvcf
 
-        return _calc_af(df)
+    ds = tiledbvcf.Dataset(uri=dataset_uri, mode="r")
+    return ds.read_variant_stats(region).to_pandas()

apis/python/src/tiledbvcf/binding/libtiledbvcf.cc

@@ -53,6 +53,7 @@ PYBIND11_MODULE(libtiledbvcf, m) {
       .def("get_materialized_attributes", &Reader::get_materialized_attributes)
       .def("get_sample_count", &Reader::get_sample_count)
       .def("get_sample_names", &Reader::get_sample_names)
+      .def("get_variant_stats_results", &Reader::get_variant_stats_results)
       .def("set_buffer_percentage", &Reader::set_buffer_percentage)
       .def(
           "set_tiledb_tile_cache_percentage",

apis/python/src/tiledbvcf/binding/reader.cc

@@ -27,9 +27,13 @@
  */
 
 #include <arrow/python/pyarrow.h>
+#include <stdint.h>
+#include <sys/types.h>
 #include <tiledbvcf/arrow.h>
+#include <cmath>
 #include <stdexcept>
 
+#include "arrow/type_fwd.h"
 #include "reader.h"
 
 namespace py = pybind11;
@@ -613,6 +617,88 @@ std::vector<std::string> Reader::get_sample_names() {
   return names;
 }
 
+py::object Reader::get_variant_stats_results() {
+  std::shared_ptr<arrow::Field> pos_field =
+      arrow::field("pos", arrow::uint32());
+  std::shared_ptr<arrow::Field> allele_field =
+      arrow::field("allele", arrow::large_utf8());
+  std::shared_ptr<arrow::Field> ac_field = arrow::field("ac", arrow::int32());
+  std::shared_ptr<arrow::Field> an_field = arrow::field("an", arrow::int32());
+  std::shared_ptr<arrow::Field> af_field = arrow::field("af", arrow::float32());
+  std::vector<std::shared_ptr<arrow::Field>> fields = {
+      pos_field, allele_field, ac_field, an_field, af_field};
+
+  size_t cardinality, alleles_size;
+  auto reader = ptr.get();
+
+  check_error(reader, tiledb_vcf_reader_prepare_variant_stats(reader));
+  check_error(
+      reader,
+      tiledb_vcf_reader_get_variant_stats_buffer_sizes(
+          reader, &cardinality, &alleles_size));
+
+  auto maybe_pos_buffer = arrow::AllocateBuffer(cardinality * sizeof(uint32_t));
+  auto maybe_allele_buffer = arrow::AllocateBuffer(alleles_size);
+  auto maybe_allele_offset_buffer =
+      arrow::AllocateBuffer((cardinality + 1) * sizeof(uint64_t));
+  auto maybe_ac_buffer = arrow::AllocateBuffer(cardinality * sizeof(int32_t));
+  auto maybe_an_buffer = arrow::AllocateBuffer(cardinality * sizeof(int32_t));
+  auto maybe_af_buffer = arrow::AllocateBuffer(cardinality * sizeof(float_t));
+  if (!(maybe_pos_buffer.ok() && maybe_allele_buffer.ok() &&
+        maybe_allele_offset_buffer.ok() && maybe_ac_buffer.ok() &&
+        maybe_an_buffer.ok() && maybe_af_buffer.ok())) {
+    throw std::runtime_error(
+        "TileDB-VCF-Py: buffer allocation failed for fetching stats");
+  }
+  std::shared_ptr<arrow::Buffer> pos_buffer = std::move(*maybe_pos_buffer);
+  std::shared_ptr<arrow::Buffer> allele_buffer =
+      std::move(*maybe_allele_buffer);
+  std::shared_ptr<arrow::Buffer> allele_offset_buffer =
+      std::move(*maybe_allele_offset_buffer);
+  std::shared_ptr<arrow::Buffer> ac_buffer = std::move(*maybe_ac_buffer);
+  std::shared_ptr<arrow::Buffer> an_buffer = std::move(*maybe_an_buffer);
+  std::shared_ptr<arrow::Buffer> af_buffer = std::move(*maybe_af_buffer);
+
+  check_error(
+      reader,
+      tiledb_vcf_reader_read_from_variant_stats(
+          reader,
+          reinterpret_cast<uint32_t*>(pos_buffer->mutable_data()),
+          reinterpret_cast<char*>(allele_buffer->mutable_data()),
+          reinterpret_cast<uint64_t*>(allele_offset_buffer->mutable_data()),
+          reinterpret_cast<int*>(ac_buffer->mutable_data()),
+          reinterpret_cast<int*>(an_buffer->mutable_data()),
+          reinterpret_cast<float_t*>(af_buffer->mutable_data())));
+
+  std::shared_ptr<arrow::Array> pos_array =
+      std::make_shared<arrow::UInt32Array>(cardinality, pos_buffer);
+  std::shared_ptr<arrow::Array> allele_array =
+      std::make_shared<arrow::LargeStringArray>(
+          cardinality, allele_offset_buffer, allele_buffer);
+  std::shared_ptr<arrow::Array> ac_array =
+      std::make_shared<arrow::Int32Array>(cardinality, ac_buffer);
+  std::shared_ptr<arrow::Array> an_array =
+      std::make_shared<arrow::Int32Array>(cardinality, an_buffer);
+  std::shared_ptr<arrow::Array> af_array =
+      std::make_shared<arrow::FloatArray>(cardinality, af_buffer);
+  std::vector<std::shared_ptr<arrow::Array>> arrays = {
+      pos_array, allele_array, ac_array, an_array, af_array};
+
+  std::shared_ptr<arrow::Schema> schema = arrow::schema(fields);
+  auto table = arrow::Table::Make(schema, arrays, cardinality);
+  check_arrow_error(table->Validate());
+
+  PyObject* obj = arrow::py::wrap_table(table);
+  if (obj == nullptr) {
+    PyErr_PrintEx(1);
+    throw std::runtime_error(
+        "TileDB-VCF-Py: Error converting stats export table to Arrow; null "
+        "Python object.");
+  }
+
+  return py::reinterpret_steal<py::object>(obj);
+}
+
 py::dtype Reader::to_numpy_dtype(tiledb_vcf_attr_datatype_t datatype) {
   switch (datatype) {
     case TILEDB_VCF_CHAR:

apis/python/src/tiledbvcf/binding/reader.h

@@ -133,6 +133,9 @@ class Reader {
   /** Retrieve list of registered samples names */
   std::vector<std::string> get_sample_names();
 
+  /** Get Stats Results **/
+  py::object get_variant_stats_results();
+
   /** Set reader verbose output mode */
   void set_verbose(bool verbose);
 

apis/python/src/tiledbvcf/dataset.py

@@ -275,6 +275,23 @@ def read_arrow(
 
         return self.continue_read_arrow()
 
+    def read_variant_stats(
+        self,
+        region: str = None,
+    ) -> pd.DataFrame:
+        """
+        Read variant stats from the dataset into a Pandas DataFrame
+
+        Parameters
+        ----------
+        region
+            Genomic region to be queried.
+        """
+        if self.mode != "r":
+            raise Exception("Dataset not open in read mode")
+        self.reader.set_regions(region)
+        return self.reader.get_variant_stats_results()
+
     def read(
         self,
         attrs: List[str] = DEFAULT_ATTRS,

apis/python/tests/test_tiledbvcf.py

@@ -1179,6 +1179,17 @@ def test_ingest_with_stats(tmp_path):
         ]["info_TILEDB_IAF"].iloc[0][0]
         == 0.9375
     )
+    ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r")
+    df = ds.read_variant_stats("chr1:1-10000")
+    assert df.shape == (11, 5)
+    df = tiledbvcf.allele_frequency.read_allele_frequency(
+        os.path.join(tmp_path, "stats_test"), "chr1:1-10000"
+    )
+    assert df.pos.is_monotonic_increasing
+    df["an_check"] = (df.ac / df.af).round(0).astype("int32")
+    assert df.an_check.equals(df.an)
+    df = ds.read_variant_stats("chr1:1-10000")
+    assert df.shape == (11, 5)
 
 
 # Ok to skip is missing bcftools in Windows CI job

libtiledbvcf/src/c_api/tiledbvcf.cc

@@ -954,6 +954,54 @@ int32_t tiledb_vcf_reader_set_scan_all_samples(
   return TILEDB_VCF_OK;
 }
 
+int32_t tiledb_vcf_reader_get_variant_stats_buffer_sizes(
+    tiledb_vcf_reader_t* reader, size_t* num_rows, size_t* alleles_size) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR || num_rows == nullptr ||
+      alleles_size == nullptr)
+    return TILEDB_VCF_ERR;
+
+  std::tuple<size_t, size_t> sizes;
+  if (SAVE_ERROR_CATCH(
+          reader, sizes = reader->reader_->variant_stats_buffer_sizes()))
+    return TILEDB_VCF_ERR;
+
+  *num_rows = std::get<0>(sizes);
+  *alleles_size = std::get<1>(sizes);
+
+  return TILEDB_VCF_OK;
+}
+
+int32_t tiledb_vcf_reader_prepare_variant_stats(tiledb_vcf_reader_t* reader) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+  if (SAVE_ERROR_CATCH(reader, reader->reader_->prepare_variant_stats()))
+    return TILEDB_VCF_ERR;
+  return TILEDB_VCF_OK;
+}
+
+int32_t tiledb_vcf_reader_read_from_variant_stats(
+    tiledb_vcf_reader_t* reader,
+    uint32_t* pos,
+    char* allele,
+    uint64_t* allele_offsets,
+    int* ac,
+    int* an,
+    float_t* af) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR || allele == nullptr ||
+      allele_offsets == nullptr || ac == nullptr) {
+    return TILEDB_VCF_ERR;
+  }
+
+  if (SAVE_ERROR_CATCH(
+          reader,
+          reader->reader_->read_from_variant_stats(
+              pos, allele, allele_offsets, ac, an, af))) {
+    return TILEDB_VCF_ERR;
+  }
+
+  return TILEDB_VCF_OK;
+}
+
 /* ********************************* */
 /*             BED FILE              */
 /* ********************************* */

libtiledbvcf/src/c_api/tiledbvcf.h

@@ -38,6 +38,7 @@
 
 #include <stdbool.h>
 #include <stdint.h>
+#include <sys/types.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -925,6 +926,39 @@ TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_set_af_filter(
 TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_set_scan_all_samples(
     tiledb_vcf_reader_t* reader, bool scan_all_samples);
 
+/**
+ * Returns the cardinality and aggregate allele length of expanded stats rows
+ * @param num_rows return number of rows by reference
+ * @param alleles_size return aggregate allele buffer size by reference
+ */
+TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_get_variant_stats_buffer_sizes(
+    tiledb_vcf_reader_t* reader, size_t* num_rows, size_t* alleles_size);
+
+/**
+ * Reads the contents of the stats array for the region, in preparation for
+ * conversion of the map to a data frame
+ */
+TILEDBVCF_EXPORT int32_t
+tiledb_vcf_reader_prepare_variant_stats(tiledb_vcf_reader_t* reader);
+
+/**
+ * Reads the expanded contents of the stats array into a set of buffers
+ * @param pos position buffer
+ * @param allele allele buffer
+ * @param allele_offsets allele offset buffer (n+1 cardinality)
+ * @param ac buffer of int representing allele count
+ * @param an buffer of int representing allele number
+ * @param af buffer of float representing internal allele frequency
+ */
+TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_read_from_variant_stats(
+    tiledb_vcf_reader_t* reader,
+    uint32_t* pos,
+    char* allele,
+    uint64_t* allele_offsets,
+    int* ac,
+    int* an,
+    float* af);
+
 /**
  * Sets export output directory
  * @param reader VCF reader object

libtiledbvcf/src/read/reader.cc

@@ -191,6 +191,11 @@ void Reader::init_af_filter() {
   }
 }
 
+void Reader::init_variant_stats_reader_for_export() {
+  Group group(*ctx_, dataset_->root_uri(), TILEDB_READ);
+  af_filter_ = std::make_unique<VariantStatsReader>(ctx_, group, false);
+}
+
 InMemoryExporter* Reader::set_in_memory_exporter() {
   // On the first call to set_buffer(), swap out any existing exporter with an
   // InMemoryExporter.
@@ -577,6 +582,12 @@ void Reader::init_for_reads_v4() {
   }
 }
 
+void Reader::init_for_variant_stats() {
+  assert(dataset_->metadata().version == TileDBVCFDataset::Version::V4);
+  init_variant_stats_reader_for_export();
+  LOG_DEBUG("Initializing reader for stats export");
+}
+
 bool Reader::next_read_batch() {
   if (dataset_->metadata().version == TileDBVCFDataset::V2 ||
       dataset_->metadata().version == TileDBVCFDataset::Version::V3)
@@ -930,6 +941,31 @@ bool Reader::next_read_batch_v4() {
   return true;
 }
 
+void Reader::prepare_variant_stats() {
+  init_for_variant_stats();
+  if (params_.regions.size() != 1) {
+    throw std::runtime_error(
+        "Error preparing variant stats: there should be exactly one region "
+        "specified");
+  }
+  af_filter_->add_region(params_.regions[0]);
+  af_filter_->compute_af();
+}
+
+void Reader::read_from_variant_stats(
+    uint32_t* pos,
+    char* allele,
+    uint64_t* allele_offsets,
+    int* ac,
+    int* an,
+    float_t* af) {
+  af_filter_->retrieve_variant_stats(pos, allele, allele_offsets, ac, an, af);
+}
+
+std::tuple<size_t, size_t> Reader::variant_stats_buffer_sizes() {
+  return af_filter_->variant_stats_buffer_sizes();
+}
+
 void Reader::init_exporter() {
   if (params_.export_to_disk) {
     if (params_.export_combined_vcf) {
@@ -1199,6 +1235,10 @@ bool Reader::process_query_results_v4() {
   if (read_state_.regions.empty())
     throw std::runtime_error(
         "Error processing query results; empty regions list.");
+  if (af_filter_ && params_.af_filter.empty()) {
+    throw std::runtime_error(
+        "Error processing query results; inconsistent AF filter state.");
+  }
 
   const auto& results = read_state_.query_results;
   const uint64_t num_cells = results.num_cells();