Improve export run-time performance by caching sample headers (#881)

* Updated TileDBVCFDataset::SampleHeaders to support caching parsed, sample-specific headers

This functionality is implemented via TileDBVCFDataset::SampleHeaders::get_sample_header_shared(), which returns a shared_ptr with the option for SampleHeaders to cache a copy of the shared_ptr.

* Updated Reader::report_cell() to cache headers when exporting records

This significantly reduces run-time by reducing time spent parsing header strings and (de)allocating header objects.

Author: alancleary

libtiledbvcf/src/dataset/tiledbvcfdataset.cc

@@ -1374,10 +1374,11 @@ void TileDBVCFDataset::SampleHeaders::set_sample_header(
   }
 
   // Save the sample's unique header index and insertion order
-  sample_index_lookup.emplace(sample_name, Indexes{hdr_idx, sample_idx});
+  sample_index_lookup.emplace(
+      sample_name, Indexes{hdr_idx, sample_idx, nullptr});
 }
 
-SafeBCFHdr TileDBVCFDataset::SampleHeaders::get_sample_header(
+bcf_hdr_t* TileDBVCFDataset::SampleHeaders::parse_sample_header(
     const std::string& sample) const {
   bcf_hdr_t* hdr = bcf_hdr_init("r");
   if (!hdr) {
@@ -1411,9 +1412,30 @@ SafeBCFHdr TileDBVCFDataset::SampleHeaders::get_sample_header(
         "bcftools: failed to update VCF header.");
   }
 
+  return hdr;
+}
+
+SafeBCFHdr TileDBVCFDataset::SampleHeaders::get_sample_header(
+    const std::string& sample) const {
+  bcf_hdr_t* hdr = parse_sample_header(sample);
   return SafeBCFHdr(hdr, bcf_hdr_destroy);
 }
 
+SafeSharedBCFHdr TileDBVCFDataset::SampleHeaders::get_sample_header_shared(
+    std::string& sample, bool cache) {
+  Indexes& indexes = sample_index_lookup.at(sample);
+  if (cache) {
+    if (indexes.cached_header == nullptr) {
+      bcf_hdr_t* hdr = parse_sample_header(sample);
+      indexes.cached_header.reset(hdr);
+    }
+    return indexes.cached_header;
+  }
+
+  bcf_hdr_t* hdr = parse_sample_header(sample);
+  return SafeSharedBCFHdr(hdr);
+}
+
 std::vector<std::string>
 TileDBVCFDataset::SampleHeaders::header_ordered_samples() const {
   auto samples = samples_view();

libtiledbvcf/src/dataset/tiledbvcfdataset.h

@@ -220,6 +220,7 @@ class TileDBVCFDataset {
     struct Indexes {
       size_t header_idx;  // index in unique_headers
       size_t tiledb_idx;  // index in TileDB header array
+      SafeSharedBCFHdr cached_header;
     };
     // Save unique headers as strings for on-demand bcf_hdr_t parsing
     std::vector<std::string> unique_headers;
@@ -240,6 +241,14 @@ class TileDBVCFDataset {
     void set_sample_header(
         const char* hdr_str, const std::string& sample_name, size_t sample_idx);
 
+    /**
+     * Parses the header for the given sample name and returns a raw pointer.
+     *
+     * @param sample The name of the sample to get the header for
+     * @return The header for the sample
+     */
+    bcf_hdr_t* parse_sample_header(const std::string& sample) const;
+
    public:
     friend class TileDBVCFDataset;
 
@@ -275,13 +284,24 @@ class TileDBVCFDataset {
     SafeBCFHdr first() const;
 
     /**
-     * Gets the header for the given sample name.
+     * Gets the header for the given sample name and returns a unique pointer.
      *
      * @param sample The name of the sample to get the header for
-     * @return The first header
+     * @return The header for the sample
      */
     SafeBCFHdr get_sample_header(const std::string& sample) const;
 
+    /**
+     * Gets the header for the given sample name and returns a shared pointer,
+     * with optional caching.
+     *
+     * @param sample The name of the sample to get the header for
+     * @param cache If the shared pointer should be cached
+     * @return The header for the sample
+     */
+    SafeSharedBCFHdr get_sample_header_shared(
+        std::string& sample, bool cache = true);
+
     /**
      * Returns the sample names stored as a view of the internal lookup map.
      *

libtiledbvcf/src/read/reader.cc

@@ -1686,12 +1686,13 @@ bool Reader::report_cell(
     sample = SampleAndId{std::string(sample_name, size)};
   }
 
-  SafeBCFHdr hdr(nullptr, bcf_hdr_destroy);
+  bcf_hdr_t* hdr = nullptr;
   if (read_state_.need_headers) {
-    hdr = read_state_.current_hdrs.get_sample_header(sample.sample_name);
+    hdr = read_state_.current_hdrs.get_sample_header_shared(sample.sample_name)
+              .get();
   }
   if (!exporter_->export_record(
-          sample, hdr.get(), region, contig_offset, results, cell_idx))
+          sample, hdr, region, contig_offset, results, cell_idx))
     return false;
 
   // If no overflow, increment num records count.

libtiledbvcf/src/vcf/vcf_utils.h

@@ -46,6 +46,9 @@ namespace vcf {
 /** Alias for unique_ptr to bcf_hdr_t. */
 typedef std::unique_ptr<bcf_hdr_t, decltype(&bcf_hdr_destroy)> SafeBCFHdr;
 
+/** Alias for unique_ptr to bcf_hdr_t. */
+typedef std::shared_ptr<bcf_hdr_t> SafeSharedBCFHdr;
+
 /** Alias for unique_ptr to bcf1_t. */
 typedef std::unique_ptr<bcf1_t, decltype(&bcf_destroy)> SafeBCFRec;