Improvements to stats run-time performance (#877)

* Added total_size() methods to allele count and variant stats classes

These methods return the total size of the classes' buffers in bytes.

* Improved AlleleCount process and flush run-time

This was done by using an unordered set for tracking sample names.

* Improved SampleStats process and flush run-time

This was done by using an unordered map for saving sample stats, minimizing map lookups via memoization, and replacing a nested map with structs.

* Improved VariantStats process and flush run-time

This was done by reusing vectors for (missing)GT values, using an unordered set for tracking sample names, minimizing map lookups via memoization, and creating separate codepaths for updating v2 and v3 arrays, the prior of which can be done much more efficiently via appending instead of using the v3 insertion strategy.

Author: alancleary

libtiledbvcf/src/stats/allele_count.cc

@@ -371,6 +371,16 @@ void AlleleCount::flush(bool finalize) {
   }
 }
 
+size_t AlleleCount::total_size() const {
+  return contig_buffer_.size() + sizeof(uint64_t) * contig_offsets_.size() +
+         sizeof(uint32_t) * pos_buffer_.size() + ref_buffer_.size() +
+         sizeof(uint64_t) * ref_offsets_.size() + alt_buffer_.size() +
+         sizeof(uint64_t) * alt_offsets_.size() + filter_buffer_.size() +
+         sizeof(uint64_t) * filter_offsets_.size() + gt_buffer_.size() +
+         sizeof(uint64_t) * gt_offsets_.size() +
+         sizeof(int32_t) * count_buffer_.size();
+}
+
 void AlleleCount::process(
     const bcf_hdr_t* hdr,
     const std::string& sample_name,
@@ -525,13 +535,13 @@ void AlleleCount::update_results() {
 std::tuple<size_t, size_t, size_t, size_t, size_t>
 AlleleCountReader::allele_count_buffer_sizes() {
   size_t ref = 0, alt = 0, filter = 0, gt = 0;
-  for (std::pair<AlleleCountKey, int32_t> ac : AlleleCountGroupBy) {
+  for (std::pair<AlleleCountKey, int32_t> ac : alleleCountGroupBy) {
     ref += ac.first.ref.size();
     alt += ac.first.alt.size();
     filter += ac.first.filter.size();
     gt += ac.first.gt.size();
   }
-  return {AlleleCountGroupBy.size(), ref, alt, filter, gt};
+  return {alleleCountGroupBy.size(), ref, alt, filter, gt};
 }
 
 void AlleleCountReader::prepare_allele_count(Region region) {
@@ -554,7 +564,7 @@ void AlleleCountReader::prepare_allele_count(Region region) {
       // TODO: add function to generate key by concatenating pos, ref, alt,
       // filter, gt
       AlleleCountKey key(pos, ref, alt, filter, gt);
-      AlleleCountGroupBy[key] += mq.data<uint32_t>("count")[i];
+      alleleCountGroupBy[key] += mq.data<uint32_t>("count")[i];
     }
   }
 }
@@ -581,7 +591,7 @@ void AlleleCountReader::read_from_allele_count(
   filter_offsets[0] = 0;
   gt_offsets[0] = 0;
   size_t i = 0;
-  for (std::pair<AlleleCountKey, int32_t> ac : AlleleCountGroupBy) {
+  for (std::pair<AlleleCountKey, int32_t> ac : alleleCountGroupBy) {
     pos[i] = ac.first.pos;
     std::strncpy(
         ref + ref_offsets[i], ac.first.ref.data(), ac.first.ref.size());

libtiledbvcf/src/stats/allele_count.h

@@ -30,7 +30,9 @@
 #include <atomic>
 #include <map>
 #include <mutex>
+#include <set>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include <htslib/vcf.h>
@@ -187,6 +189,11 @@ class AlleleCount {
    */
   void flush(bool finalize = false);
 
+  /**
+   * Returns the sum of sizes of all buffers (in bytes).
+   */
+  size_t total_size() const;
+
  private:
   // give the AlleleCountReader access to AlleleCount's private members
   friend class AlleleCountReader;
@@ -252,7 +259,7 @@ class AlleleCount {
   int count_delta_ = 1;
 
   // Set of sample names in this query (per thread)
-  std::set<std::string> sample_names_;
+  std::unordered_set<std::string> sample_names_;
 
   // Counts grouped by "key" at the current locus.
   // Use map to keep dimension keys sorted and maintain global order.
@@ -395,7 +402,7 @@ class AlleleCountReader {
   AlleleCountReader(std::shared_ptr<Context> ctx, const Group& group);
 
  private:
-  std::map<AlleleCountKey, int32_t> AlleleCountGroupBy;
+  std::map<AlleleCountKey, int32_t> alleleCountGroupBy;
   std::shared_ptr<Array> array_;
 };
 

libtiledbvcf/src/stats/sample_stats.cc

@@ -24,6 +24,8 @@
  * THE SOFTWARE.
  */
 
+#include <ranges>
+
 #include <htslib/vcf.h>
 #include <htslib/vcfutils.h>
 
@@ -206,31 +208,35 @@ void SampleStats::process(
     const std::string& contig,
     uint32_t pos,
     bcf1_t* rec) {
+  Stats& stats = stats_[sample];
+
   if (bcf_get_format_int32(hdr, rec, "DP", &dst_, &ndst_) > 0) {
     if (dst_[0] != bcf_int32_missing) {
-      auto dp = static_cast<uint64_t>(dst_[0]);
-      stats_[sample]["dp_sum"] += dp;
-      stats_[sample]["dp_sum2"] += dp * dp;
-      stats_[sample]["dp_count"] += 1;
-      if (!stats_[sample].contains("dp_min")) {
-        stats_[sample]["dp_min"] = std::numeric_limits<uint64_t>::max();
+      std::unique_ptr<DPStats>& dp_stats = stats.dp_stats;
+      if (dp_stats == nullptr) {
+        dp_stats.reset(new DPStats());
       }
-      stats_[sample]["dp_min"] = std::min(stats_[sample]["dp_min"], dp);
-      stats_[sample]["dp_max"] = std::max(stats_[sample]["dp_max"], dp);
+      auto dp = static_cast<uint64_t>(dst_[0]);
+      dp_stats->dp_sum += dp;
+      dp_stats->dp_sum2 += dp * dp;
+      dp_stats->dp_count += 1;
+      dp_stats->dp_min = std::min(dp_stats->dp_min, dp);
+      dp_stats->dp_max = std::max(dp_stats->dp_max, dp);
     }
   }
 
   if (bcf_get_format_int32(hdr, rec, "GQ", &dst_, &ndst_) > 0) {
     if (dst_[0] != bcf_int32_missing) {
-      auto gq = static_cast<uint64_t>(dst_[0]);
-      stats_[sample]["gq_sum"] += gq;
-      stats_[sample]["gq_sum2"] += gq * gq;
-      stats_[sample]["gq_count"] += 1;
-      if (!stats_[sample].contains("gq_min")) {
-        stats_[sample]["gq_min"] = std::numeric_limits<uint64_t>::max();
+      std::unique_ptr<GQStats>& gq_stats = stats.gq_stats;
+      if (gq_stats == nullptr) {
+        gq_stats.reset(new GQStats());
       }
-      stats_[sample]["gq_min"] = std::min(stats_[sample]["gq_min"], gq);
-      stats_[sample]["gq_max"] = std::max(stats_[sample]["gq_max"], gq);
+      auto gq = static_cast<uint64_t>(dst_[0]);
+      gq_stats->gq_sum += gq;
+      gq_stats->gq_sum2 += gq * gq;
+      gq_stats->gq_count += 1;
+      gq_stats->gq_min = std::min(gq_stats->gq_min, gq);
+      gq_stats->gq_max = std::max(gq_stats->gq_max, gq);
     }
   }
 
@@ -307,19 +313,19 @@ void SampleStats::process(
     n_singleton += count == 1;
   }
 
-  stats_[sample]["n_records"] += 1;
-  stats_[sample]["n_called"] += !is_missing;
-  stats_[sample]["n_not_called"] += is_missing;
-  stats_[sample]["n_hom_ref"] += is_hom_ref;
-  stats_[sample]["n_het"] += is_het;
-  stats_[sample]["n_singleton"] += n_singleton;
-  stats_[sample]["n_snp"] += n_ti + n_tv;
-  stats_[sample]["n_transition"] += n_ti;
-  stats_[sample]["n_transversion"] += n_tv;
-  stats_[sample]["n_insertion"] += n_ins;
-  stats_[sample]["n_deletion"] += n_del;
-  stats_[sample]["n_star"] += n_star;
-  stats_[sample]["n_multiallelic"] += is_multi;
+  stats.n_records += 1;
+  stats.n_called += !is_missing;
+  stats.n_not_called += is_missing;
+  stats.n_hom_ref += is_hom_ref;
+  stats.n_het += is_het;
+  stats.n_singleton += n_singleton;
+  stats.n_snp += n_ti + n_tv;
+  stats.n_transition += n_ti;
+  stats.n_transversion += n_tv;
+  stats.n_insertion += n_ins;
+  stats.n_deletion += n_del;
+  stats.n_star += n_star;
+  stats.n_multiallelic += is_multi;
 }
 
 void SampleStats::flush(bool finalize) {
@@ -351,30 +357,54 @@ void SampleStats::flush(bool finalize) {
     add_buffer_(name);
   }
 
-  bool found_dp = false;
-  bool found_gq = false;
-
-  for (const auto& [sample, stat] : stats_) {
+  // Sort the stats keys and iterate them in order
+  auto keys_view = stats_ | std::views::keys;
+  std::vector<std::string> samples(keys_view.begin(), keys_view.end());
+  std::sort(samples.begin(), samples.end());
+  for (const std::string& sample : samples) {
+    Stats& stats = stats_[sample];
     buffers["sample"]->push_back(sample);
 
     // Add stats to buffers
-    for (const auto& [field, value] : stat) {
-      buffers[field]->push_back(value);
-      found_dp |= field == "dp_sum";
-      found_gq |= field == "gq_sum";
-    }
-
-    // Add nulls for missing DP fields
-    if (!found_dp) {
+    buffers["n_records"]->push_back(stats.n_records);
+    buffers["n_called"]->push_back(stats.n_called);
+    buffers["n_not_called"]->push_back(stats.n_not_called);
+    buffers["n_hom_ref"]->push_back(stats.n_hom_ref);
+    buffers["n_het"]->push_back(stats.n_het);
+    buffers["n_singleton"]->push_back(stats.n_singleton);
+    buffers["n_snp"]->push_back(stats.n_snp);
+    buffers["n_transition"]->push_back(stats.n_transition);
+    buffers["n_transversion"]->push_back(stats.n_transversion);
+    buffers["n_insertion"]->push_back(stats.n_insertion);
+    buffers["n_deletion"]->push_back(stats.n_deletion);
+    buffers["n_star"]->push_back(stats.n_star);
+    buffers["n_multiallelic"]->push_back(stats.n_multiallelic);
+
+    // Add DP fields
+    const std::unique_ptr<DPStats>& dp_stats = stats.dp_stats;
+    if (dp_stats != nullptr) {
+      buffers["dp_sum"]->push_back(dp_stats->dp_sum);
+      buffers["dp_sum2"]->push_back(dp_stats->dp_sum2);
+      buffers["dp_count"]->push_back(dp_stats->dp_count);
+      buffers["dp_min"]->push_back(dp_stats->dp_min);
+      buffers["dp_max"]->push_back(dp_stats->dp_max);
+    } else {
       buffers["dp_sum"]->push_null();
       buffers["dp_sum2"]->push_null();
       buffers["dp_count"]->push_null();
       buffers["dp_min"]->push_null();
       buffers["dp_max"]->push_null();
     }
 
-    // Add nulls for missing GQ fields
-    if (!found_gq) {
+    // Add GQ fields
+    const std::unique_ptr<GQStats>& gq_stats = stats.gq_stats;
+    if (gq_stats != nullptr) {
+      buffers["gq_sum"]->push_back(gq_stats->gq_sum);
+      buffers["gq_sum2"]->push_back(gq_stats->gq_sum2);
+      buffers["gq_count"]->push_back(gq_stats->gq_count);
+      buffers["gq_min"]->push_back(gq_stats->gq_min);
+      buffers["gq_max"]->push_back(gq_stats->gq_max);
+    } else {
       buffers["gq_sum"]->push_null();
       buffers["gq_sum2"]->push_null();
       buffers["gq_count"]->push_null();

libtiledbvcf/src/stats/sample_stats.h

@@ -187,6 +187,40 @@ class SampleStats {
   void flush(bool finalize = false);
 
  private:
+  struct DPStats {
+    uint64_t dp_sum = 0;
+    uint64_t dp_sum2 = 0;
+    uint64_t dp_count = 0;
+    uint64_t dp_min = std::numeric_limits<uint64_t>::max();
+    uint64_t dp_max = 0;
+  };
+
+  struct GQStats {
+    uint64_t gq_sum = 0;
+    uint64_t gq_sum2 = 0;
+    uint64_t gq_count = 0;
+    uint64_t gq_min = std::numeric_limits<uint64_t>::max();
+    uint64_t gq_max = 0;
+  };
+
+  struct Stats {
+    uint64_t n_records = 0;
+    uint64_t n_called = 0;
+    uint64_t n_not_called = 0;
+    uint64_t n_hom_ref = 0;
+    uint64_t n_het = 0;
+    uint64_t n_singleton = 0;
+    uint64_t n_snp = 0;
+    uint64_t n_transition = 0;
+    uint64_t n_transversion = 0;
+    uint64_t n_insertion = 0;
+    uint64_t n_deletion = 0;
+    uint64_t n_star = 0;
+    uint64_t n_multiallelic = 0;
+    std::unique_ptr<DPStats> dp_stats = nullptr;
+    std::unique_ptr<GQStats> gq_stats = nullptr;
+  };
+
   // Array URI basename
   inline static const std::string SAMPLE_STATS_ARRAY = "sample_stats";
 
@@ -208,8 +242,8 @@ class SampleStats {
   // Current contig
   std::string contig_;
 
-  // Aggregate stats for each sample. map: sample -> (map: field -> count)
-  std::map<std::string, std::unordered_map<std::string, uint64_t>> stats_;
+  // Aggregate stats for each sample.
+  std::unordered_map<std::string, Stats> stats_;
 
   // Get the URI for the array from the root group
   static std::string group_uri(const Group& group) {