BCF and header enhancements: introduce error classes for better error handling; implement sample subsetting functionality in Bcf::Header; add tests for subset behavior and error conditions

Author: kojix2

lib/hts/bcf.rb

@@ -3,6 +3,7 @@
 require_relative "../htslib"
 
 require_relative "hts"
+require_relative "bcf/errors"
 require_relative "bcf/header"
 require_relative "bcf/info"
 require_relative "bcf/format"
@@ -38,16 +39,16 @@ def self.build_index(file_name, index_name = nil, min_shift = 14, threads = 0, v
 
       case LibHTS.bcf_index_build3(file_name, index_name, min_shift, threads)
       when 0 # successful
-      when -1 then raise "indexing failed"
-      when -2 then raise "opening #{file_name} failed"
-      when -3 then raise "format not indexable"
-      when -4 then raise "failed to create and/or save the index"
-      else raise "unknown error"
+      when -1 then raise IndexError, "Indexing failed for #{file_name}"
+      when -2 then raise IndexError, "Opening #{file_name} failed while building the index"
+      when -3 then raise IndexError, "#{file_name} is not in an indexable format"
+      when -4 then raise IndexError, "Failed to create or save the index for #{file_name}"
+      else raise IndexError, "Unknown index build error for #{file_name}"
       end
     end
 
     def initialize(file_name, mode = "r", index: nil, threads: nil,
-                   build_index: false)
+                   build_index: false, subset: nil)
       if block_given?
         message = "HTS::Bcf.new() does not take block; Please use HTS::Bcf.open() instead"
         raise message
@@ -61,13 +62,16 @@ def initialize(file_name, mode = "r", index: nil, threads: nil,
       @nthreads   = threads
       @hts_file   = LibHTS.hts_open(@file_name, mode)
 
-      raise Errno::ENOENT, "Failed to open #{@file_name}" if @hts_file.null?
+      raise OpenError, "Failed to open #{@file_name}" if @hts_file.null?
 
       set_threads(threads) if threads
 
+      raise SubsetError, "Sample subsetting is only available when reading BCF/VCF files" if subset && @mode[0] == "w"
+
       return if @mode[0] == "w"
 
-      @header = Bcf::Header.new(@hts_file)
+      @read_header = Bcf::Header.new(@hts_file)
+      @header = subset ? @read_header.subset(subset) : @read_header
       build_index(index) if build_index
       @idx = load_index(index)
       @start_position = tell
@@ -219,8 +223,8 @@ def each(copy: false, &block)
     def query(region, beg = nil, end_ = nil, copy: false, &block)
       check_closed
 
-      raise "query is only available for BCF files" unless file_format == "bcf"
-      raise "Index file is required to call the query method." unless index_loaded?
+      raise QueryError, "Query is only available for BCF files" unless file_format == "bcf"
+      raise MissingIndexError, "Index file is required to call the query method for #{@file_name}" unless index_loaded?
 
       case region
       when Array
@@ -269,7 +273,7 @@ def queryi_reuse(tid, beg, end_, &block)
       return to_enum(__method__, tid, beg, end_) unless block_given?
 
       qiter = LibHTS.bcf_itr_queryi(@idx, tid, beg, end_)
-      raise "Failed to query region #{tid} #{beg} #{end_}" if qiter.null?
+      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
 
       query_reuse_yield(qiter, &block)
       self
@@ -278,8 +282,8 @@ def queryi_reuse(tid, beg, end_, &block)
     def querys_reuse(region, &block)
       return to_enum(__method__, region) unless block_given?
 
-      qiter = LibHTS.bcf_itr_querys(@idx, header, region)
-      raise "Failed to query region #{region}" if qiter.null?
+      qiter = LibHTS.bcf_itr_querys(@idx, read_header, region)
+      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
 
       query_reuse_yield(qiter, &block)
       self
@@ -303,6 +307,7 @@ def query_reuse_yield(qiter)
           break if slen == -1
           raise if slen < -1
 
+          apply_subset!(record)
           yield record
         end
       ensure
@@ -314,7 +319,7 @@ def queryi_copy(tid, beg, end_, &block)
       return to_enum(__method__, tid, beg, end_) unless block_given?
 
       qiter = LibHTS.bcf_itr_queryi(@idx, tid, beg, end_)
-      raise "Failed to query region #{tid} #{beg} #{end_}" if qiter.null?
+      raise QueryError, "Failed to query region #{tid}:#{beg}-#{end_} in #{@file_name}" if qiter.null?
 
       query_copy_yield(qiter, &block)
       self
@@ -323,8 +328,8 @@ def queryi_copy(tid, beg, end_, &block)
     def querys_copy(region, &block)
       return to_enum(__method__, region) unless block_given?
 
-      qiter = LibHTS.bcf_itr_querys(@idx, header, region)
-      raise "Failed to query region #{region}" if qiter.null?
+      qiter = LibHTS.bcf_itr_querys(@idx, read_header, region)
+      raise QueryError, "Failed to query region #{region.inspect} in #{@file_name}" if qiter.null?
 
       query_copy_yield(qiter, &block)
       self
@@ -346,7 +351,9 @@ def query_copy_yield(qiter)
         break if slen == -1
         raise if slen < -1
 
-        yield Record.new(header, bcf1)
+        record = Record.new(header, bcf1)
+        apply_subset!(record)
+        yield record
       end
     ensure
       LibHTS.bcf_itr_destroy(qiter)
@@ -359,7 +366,10 @@ def each_record_reuse
 
       bcf1 = LibHTS.bcf_init
       record = Record.new(header, bcf1)
-      yield record while LibHTS.bcf_read(@hts_file, header, bcf1) != -1
+      while LibHTS.bcf_read(@hts_file, read_header, bcf1) != -1
+        apply_subset!(record)
+        yield record
+      end
       self
     end
 
@@ -368,11 +378,25 @@ def each_record_copy
 
       return to_enum(__method__) unless block_given?
 
-      while LibHTS.bcf_read(@hts_file, header, bcf1 = LibHTS.bcf_init) != -1
+      while LibHTS.bcf_read(@hts_file, read_header, bcf1 = LibHTS.bcf_init) != -1
         record = Record.new(header, bcf1)
+        apply_subset!(record)
         yield record
       end
       self
     end
+
+    def read_header
+      @read_header || header
+    end
+
+    def apply_subset!(record)
+      return unless header.subset?
+
+      rc = LibHTS.bcf_subset(header.struct, record.struct, header.subset_sample_count, header.subset_imap_pointer || ::FFI::Pointer::NULL)
+      return if rc >= 0
+
+      raise SubsetError, "Failed to subset samples #{header.subset_samples.inspect} while reading #{@file_name}"
+    end
   end
 end

lib/hts/bcf/errors.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module HTS
+  class Bcf < Hts
+    class Error < HTS::Error; end
+
+    class OpenError < Error; end
+    class IndexError < Error; end
+    class MissingIndexError < IndexError; end
+    class QueryError < Error; end
+    class HeaderError < Error; end
+    class SubsetError < HeaderError; end
+    class UnknownSampleError < SubsetError; end
+  end
+end

lib/hts/bcf/header.rb

@@ -35,6 +35,9 @@ def initialize(arg = nil)
 
         @sync_depth = 0
         @sync_needed = false
+        @subset_samples = nil
+        @subset_imap = nil
+        @subset_imap_pointer = nil
 
         yield self if block_given?
       end
@@ -87,6 +90,44 @@ def samples
           .map(&:read_string)
       end
 
+      attr_reader :subset_samples
+
+      def subset?( )
+        !@subset_imap.nil?
+      end
+
+      def subset_sample_count
+        subset? ? @subset_samples.length : 0
+      end
+
+      def subset_imap_pointer
+        @subset_imap_pointer
+      end
+
+      def subset(samples)
+        subset_samples = normalize_subset_samples(samples)
+        validate_subset_samples!(subset_samples)
+
+        sample_pointers = nil
+        imap_pointer = nil
+        if subset_samples.empty?
+          subset_hdr = LibHTS.bcf_hdr_subset(@bcf_hdr, 0, ::FFI::Pointer::NULL, ::FFI::Pointer::NULL)
+        else
+          encoded_samples = subset_samples.map { |name| FFI::MemoryPointer.from_string(name) }
+          sample_pointers = FFI::MemoryPointer.new(:pointer, subset_samples.length)
+          sample_pointers.write_array_of_pointer(encoded_samples)
+          imap_pointer = FFI::MemoryPointer.new(:int, subset_samples.length)
+          subset_hdr = LibHTS.bcf_hdr_subset(@bcf_hdr, subset_samples.length, sample_pointers, imap_pointer)
+        end
+
+        raise SubsetError, "Failed to subset BCF header samples #{subset_samples.inspect}" if subset_hdr.to_ptr.null?
+
+        composed_imap = compose_subset_imap(read_subset_imap(imap_pointer, subset_samples.length))
+        self.class.new(subset_hdr).tap do |header|
+          header.send(:set_subset_state, subset_samples, composed_imap)
+        end
+      end
+
       def add_sample(sample, sync: true)
         rc = LibHTS.bcf_hdr_add_sample(@bcf_hdr, sample)
         raise "Failed to add sample #{sample}" if rc.negative?
@@ -341,6 +382,58 @@ def initialize_copy(orig)
         @bcf_hdr = LibHTS.bcf_hdr_dup(orig.struct)
         @sync_depth = 0
         @sync_needed = false
+        set_subset_state(orig.subset_samples, orig.send(:subset_imap))
+      end
+
+      protected
+
+      attr_reader :subset_imap
+
+      def set_subset_state(samples, imap)
+        @subset_samples = samples&.dup
+        @subset_imap = imap&.dup
+        @subset_imap_pointer = build_subset_imap_pointer(@subset_imap)
+      end
+
+      private
+
+      def normalize_subset_samples(samples)
+        case samples
+        when String
+          [samples]
+        else
+          Array(samples).map(&:to_s)
+        end
+      rescue TypeError
+        raise SubsetError, "Sample subset must be a String or an Array of sample names"
+      end
+
+      def validate_subset_samples!(subset_samples)
+        duplicates = subset_samples.group_by(&:itself).select { |_name, group| group.length > 1 }.keys
+        raise SubsetError, "Duplicate sample names in subset: #{duplicates.join(', ')}" unless duplicates.empty?
+
+        missing = subset_samples.reject { |name| samples.include?(name) }
+        raise UnknownSampleError, "Unknown sample names: #{missing.join(', ')}" unless missing.empty?
+      end
+
+      def read_subset_imap(pointer, length)
+        return [] if length.zero?
+
+        pointer.read_array_of_int(length)
+      end
+
+      def compose_subset_imap(imap)
+        base_imap = @subset_imap || Array.new(samples.length, &:itself)
+        imap.map { |index| base_imap.fetch(index) }
+      end
+
+      def build_subset_imap_pointer(imap)
+        return nil unless imap
+        return nil if imap.empty?
+
+        FFI::MemoryPointer.new(:int, imap.length).tap do |pointer|
+          pointer.write_array_of_int(imap)
+        end
       end
     end
   end

test/bcf/header_test.rb

@@ -68,6 +68,30 @@ def test_add_sample
     assert_equal %w[A B kojix2 kojix3], hdr2.samples
   end
 
+  def test_subset_returns_new_header
+    subset = @hdr.subset(["B"])
+
+    assert_equal %w[A B], @hdr.samples
+    assert_equal ["B"], subset.samples
+    assert_equal 1, subset.nsamples
+  end
+
+  def test_subset_rejects_unknown_samples
+    error = assert_raises(HTS::Bcf::UnknownSampleError) do
+      @hdr.subset(["missing"])
+    end
+
+    assert_match(/missing/, error.message)
+  end
+
+  def test_subset_rejects_duplicates
+    error = assert_raises(HTS::Bcf::SubsetError) do
+      @hdr.subset(["A", "A"])
+    end
+
+    assert_match(/Duplicate sample names/, error.message)
+  end
+
   def test_sync
     hdr2 = @hdr.clone
     hdr2.add_sample("kojix1", sync: false)

test/bcf_test.rb

@@ -7,6 +7,10 @@ def test_bcf_path
     File.expand_path("../htslib/test/index.vcf", __dir__)
   end
 
+  def test_multi_sample_bcf_path
+    File.expand_path("../htslib/test/tabix/vcf_file.bcf", __dir__)
+  end
+
   def setup
     @bcf = HTS::Bcf.new(test_bcf_path)
   end
@@ -128,11 +132,45 @@ def test_format
   def test_initialize_no_file_bcf
     stderr_old = $stderr.dup
     $stderr.reopen(File::NULL)
-    assert_raises(Errno::ENOENT) { HTS::Bcf.new("/tmp/no_such_file") }
+    assert_raises(HTS::Bcf::OpenError) { HTS::Bcf.new("/tmp/no_such_file") }
     $stderr.flush
     $stderr.reopen(stderr_old)
   end
 
+  def test_initialize_with_subset
+    bcf = HTS::Bcf.new(test_multi_sample_bcf_path, subset: ["B"])
+
+    assert_equal ["B"], bcf.samples
+    assert_equal 1, bcf.nsamples
+    assert_equal ["0/1"], bcf.first.format("GT")
+  ensure
+    bcf&.close
+  end
+
+  def test_query_requires_index
+    bcf = HTS::Bcf.new(Fixtures["test.bcf"], index: "/tmp/no_such_test_bcf_index.csi")
+
+    error = assert_raises(HTS::Bcf::MissingIndexError) do
+      bcf.query("poo:4000-4100").first
+    end
+
+    assert_match(/Index file is required/, error.message)
+  ensure
+    bcf&.close
+  end
+
+  def test_query_invalid_region_raises_query_error
+    bcf = HTS::Bcf.open(Fixtures["test.bcf"])
+
+    error = assert_raises(HTS::Bcf::QueryError) do
+      bcf.query("unknown:1-10").first
+    end
+
+    assert_match(/unknown:1-10/, error.message)
+  ensure
+    bcf&.close
+  end
+
   def test_query
     bcf = HTS::Bcf.open(Fixtures["test.bcf"])
     assert_equal 4021, bcf.query("poo", 4000, 4100).first.pos + 1