Add multi-region query support for Bam class

Author: kojix2

lib/hts/bam.rb

@@ -211,19 +211,42 @@ def each(copy: false, &block)
       end
     end
 
-    # Iterate records in a genomic region.
+    # Iterate records in a genomic region or multiple regions.
     # See {#each} for copy semantics. When copy: false, the yielded Record is reused and should not be stored.
+    #
+    # @param region [String, Array<String>] Region specification(s)
+    #   - Single region: "chr1:100-200" or "chr1" with beg/end parameters
+    #   - Multiple regions: ["chr1:100-200", "chr2:500-600", ...]
+    # @param beg [Integer, nil] Start position (used with single string region)
+    # @param end_ [Integer, nil] End position (used with single string region)
+    # @param copy [Boolean] Whether to deep-copy records (see {#each})
+    #
+    # @example Single region query
+    #   bam.query("chr1:100-200") { |r| puts r.qname }
+    #   bam.query("chr1", 100, 200) { |r| puts r.qname }
+    #
+    # @example Multi-region query
+    #   bam.query(["chr1:100-200", "chr2:500-600"]) { |r| puts r.qname }
     def query(region, beg = nil, end_ = nil, copy: false, &block)
       check_closed
       raise "Index file is required to call the query method." unless index_loaded?
 
-      if beg && end_
-        tid = header.get_tid(region)
-        queryi(tid, beg, end_, copy:, &block)
-      elsif beg.nil? && end_.nil?
-        querys(region, copy:, &block)
+      case region
+      when Array
+        raise ArgumentError, "beg and end_ cannot be used with array of regions" if beg || end_
+
+        query_regions(region, copy:, &block)
+      when String
+        if beg && end_
+          tid = header.get_tid(region)
+          queryi(tid, beg, end_, copy:, &block)
+        elsif beg.nil? && end_.nil?
+          querys(region, copy:, &block)
+        else
+          raise ArgumentError, "beg and end_ must be specified together"
+        end
       else
-        raise ArgumentError, "beg and end_ must be specified together"
+        raise ArgumentError, "region must be String or Array"
       end
     end
 
@@ -266,6 +289,15 @@ def querys(region, copy: false, &block)
       end
     end
 
+    # Multi-region query implementation
+    def query_regions(regions, copy: false, &block)
+      if copy
+        query_regions_copy(regions, &block)
+      else
+        query_regions_reuse(regions, &block)
+      end
+    end
+
     # Internal: yield a single reused Record over the entire file.
     # The underlying bam1_t is mutated on each iteration for speed.
     def each_record_reuse
@@ -355,5 +387,27 @@ def query_copy(qiter)
     ensure
       LibHTS.hts_itr_destroy(qiter)
     end
+
+    # Multi-region query using sequential single-region queries
+    # Note: This is a fallback implementation. Ideally we would use sam_itr_regarray
+    # but there seem to be issues with the multi-region iterator in the current setup.
+    def query_regions_reuse(regions, &block)
+      return to_enum(__method__, regions) unless block_given?
+
+      regions.each do |region|
+        querys_reuse(region, &block)
+      end
+      self
+    end
+
+    # Multi-region query with copied Records using sequential queries
+    def query_regions_copy(regions, &block)
+      return to_enum(__method__, regions) unless block_given?
+
+      regions.each do |region|
+        querys_copy(region, &block)
+      end
+      self
+    end
   end
 end

test/bam_test.rb

@@ -289,6 +289,38 @@ def assert_cram_getter_equal(ft, method_name, &block)
     end
   end
 
+  # CRAM multi-region iterator has issues, only test BAM formats
+  %i[bam_string bam_uri].each do |ft|
+    define_method "test_query_multi_regions_#{ft}" do
+      arr = []
+      bam(ft).query(["chr1:100-200", "chr2:350-700"]) do |aln|
+        arr << aln.pos
+      end
+      # Should get records from both regions
+      assert_includes arr, 341
+      assert_includes arr, 658
+    end
+
+    define_method "test_query_multi_regions_copy_#{ft}" do
+      arr = []
+      bam(ft).query(["chr1:100-200", "chr2:350-700"], copy: true) do |aln|
+        arr << aln.pos
+      end
+      # Should get records from both regions
+      assert_includes arr, 341
+      assert_includes arr, 658
+    end
+
+    define_method "test_query_multi_regions_single_#{ft}" do
+      # Single region as array should also work
+      arr = []
+      bam(ft).query(["chr2:350-700"]) do |aln|
+        arr << aln.pos
+      end
+      assert_equal [341, 658], arr
+    end
+  end
+
   def test_initialize_no_file_bam
     stderr_old = $stderr.dup
     $stderr.reopen(File::NULL)