Repository: chrovis/cljam
Pull request #322 (2024-08-29 06:46:36)

Fine-grained partition of CRAM containers and slices
## Overview

This PR introduces changes to fine-tune how containers and slices are partitioned when writing CRAM files.

The CRAM specification has flexibility in how containers and slices are partitioned, and it’s up to each CRAM encoder to decide how to do it. In CRAM index files, index entries are created for each slice, so how containers and slices are partitioned affects the performance of index access in CRAM files.

In particular, having _multiple-reference_ (or _multi-ref_) slices can make index access less efficient. So, CRAM encoders should aim to avoid creating multi-ref slices whenever possible, and even when they are necessary, the size of such slices should be kept to a minimum.

The change introduced in this PR is a bit complex, and this complexity mainly comes from these considerations.

## Specification

The size of containers and slices can be controlled by the following three options provided to the CRAM writer:

- `records-per-slice`
- `slices-per-container`
- `min-single-ref-slice-size`

Basically, the CRAM writer tries to pack the number of records specified by `records-per-slice` into one slice and the number of slices specified by `slices-per-container` into one container.

Additionally, the CRAM writer adjusts the sizes of containers and slices to minimize the occurrence of multi-ref containers and slices as much as possible:

- Records mapped to different references are placed in different slices, ensuring that each slice remains single-ref
- However, if the size of a slice falls below the value specified by `min-single-ref-slice-size`, the records mapped to different references will be put together into one multi-ref slice
- A multi-ref slice will be placed in its own container, and no other slices will be included in the same container
- If two single-ref slices consist of records mapped to different references, they will not be placed in the same container
- There are also some other minor adjustments included


Patch:
diff --git a/src/cljam/io/cram.clj b/src/cljam/io/cram.clj
index e557b4d..f618eb6 100644
--- a/src/cljam/io/cram.clj
+++ b/src/cljam/io/cram.clj
@@ -58,6 +58,12 @@
         a sequence reader that reads sequences from the reference file.
         This may be omitted only when the CRAM file to be read does not require
         a reference file.
+    - records-per-slice: The maximum number of records a slice may contain.
+        Defaults to 10000.
+    - slices-per-container: The maximum number of slices a container may contain.
+        Defaults to 1.
+    - min-single-ref-slice-size: The minimum number of records required to emit
+        a single-reference slice. Defaults to 1000.
     - ds-compressor-overrides: A function to override data series compressors.
         Given a data series keyword, returns a keyword or a set of keywords
         representing compression method. It may return another function to add
diff --git a/src/cljam/io/cram/encode/partitioning.clj b/src/cljam/io/cram/encode/partitioning.clj
new file mode 100644
index 0000000..3cffa4b
--- /dev/null
+++ b/src/cljam/io/cram/encode/partitioning.clj
@@ -0,0 +1,89 @@
+(ns cljam.io.cram.encode.partitioning
+  (:require [cljam.io.sam.util.header :as sam.header])
+  (:import [java.util ArrayList List]))
+
+(defn- slice-record-collector
+  [header {:keys [^long records-per-slice ^long min-single-ref-slice-size]
+           :or {records-per-slice 10000, min-single-ref-slice-size 1000}}]
+  (let [sort-by-coord? (= (sam.header/sort-order header) sam.header/order-coordinate)
+        ;; CRAM files sorted by coord can easily get back from a multi-ref state
+        ;; to a single-ref state. To support this, multi-ref slices should be kept
+        ;; as small as possible
+        multi-upper-bound (if sort-by-coord?
+                            min-single-ref-slice-size
+                            records-per-slice)]
+    (fn [empty-container? alns]
+      (let [slice-records (ArrayList.)]
+        (loop [ref-state nil, alns alns]
+          (if (seq alns)
+            (let [[{:keys [rname] :as aln} & more] alns
+                  n-records (.size slice-records)]
+              (case ref-state
+                nil
+                (let [ref-state' (if (= rname "*") :unmapped rname)]
+                  (.add slice-records aln)
+                  (recur ref-state' more))
+
+                :unmapped
+                (if (< n-records records-per-slice)
+                  (let [ref-state' (cond (= rname "*") ref-state
+                                         sort-by-coord? (throw
+                                                         (ex-info
+                                                          (str "Unmapped records "
+                                                               "must be last")
+                                                          {}))
+                                         :else :multi-ref)]
+                    (.add slice-records aln)
+                    (recur ref-state' more))
+                  [ref-state slice-records alns])
+
+                :multi-ref
+                (if (< n-records multi-upper-bound)
+                  (do (.add slice-records aln)
+                      (recur ref-state more))
+                  [ref-state slice-records alns])
+
+                (if (= ref-state rname)
+                  (if (< n-records records-per-slice)
+                    (do (.add slice-records aln)
+                        (recur ref-state more))
+                    [ref-state slice-records alns])
+                  ;; If the container already contains one or more single-ref
+                  ;; slices, instead of creating and adding a new multi-ref slice
+                  ;; to that container, add the accumulated single-ref slice
+                  ;; to the container and add the current record to the next slice
+                  (if (and empty-container? (< n-records min-single-ref-slice-size))
+                    (do (.add slice-records aln)
+                        (recur :multi-ref more))
+                    [ref-state slice-records alns]))))
+            [ref-state slice-records alns]))))))
+
+(defn with-each-container
+  "Partitions the given alignment records into containers, which are represented
+  as a List of Lists of record maps, and calls the function f with them."
+  [header options alns f]
+  (let [{:keys [^long slices-per-container] :or {slices-per-container 1}} options
+        collect-fn (slice-record-collector header options)
+        container-records (ArrayList.)]
+    (loop [ref-state nil, written 0, ready 0, alns alns]
+      (if (seq alns)
+        (let [n-slices (.size container-records)
+              [ref-state' ^List slice-records alns'] (collect-fn (zero? n-slices) alns)]
+          (if (or (= n-slices slices-per-container)
+                  (= ref-state' :multi-ref)
+                  (and (some? ref-state) (not= ref-state ref-state')))
+            (let [written' (+ written ready)]
+              (when-not (.isEmpty container-records)
+                (f written container-records))
+              (.clear container-records)
+              (if (= ref-state' :multi-ref)
+                (do (f written' (doto (ArrayList.) (.add slice-records)))
+                    (recur nil (+ written' (.size slice-records)) 0 alns'))
+                (let [ready' (.size slice-records)]
+                  (.add container-records slice-records)
+                  (recur ref-state' written' ready' alns'))))
+            (let [ready' (+ ready (.size slice-records))]
+              (.add container-records slice-records)
+              (recur ref-state' written ready' alns'))))
+        (when-not (.isEmpty container-records)
+          (f written container-records))))))
diff --git a/src/cljam/io/cram/encode/record.clj b/src/cljam/io/cram/encode/record.clj
index 64685d4..bcd10e4 100644
--- a/src/cljam/io/cram/encode/record.clj
+++ b/src/cljam/io/cram/encode/record.clj
@@ -1,11 +1,11 @@
 (ns cljam.io.cram.encode.record
   (:require [cljam.io.cram.encode.alignment-stats :as stats]
+            [cljam.io.cram.encode.tag-dict :as tag-dict]
+            [cljam.io.cram.seq-resolver.protocol :as resolver]
             [cljam.io.sam.util.cigar :as sam.cigar]
             [cljam.io.sam.util.flag :as sam.flag]
-            [cljam.io.sam.util.option :as sam.option]
-            [cljam.io.cram.seq-resolver.protocol :as resolver]
-            [cljam.io.cram.encode.tag-dict :as tag-dict])
-  (:import [java.util Arrays]))
+            [cljam.io.sam.util.option :as sam.option])
+  (:import [java.util Arrays List]))
 
 (defn- ref-index [rname->idx rname]
   (if  (= rname "*")
@@ -202,9 +202,9 @@
   "Preprocesses slice records to calculate some record fields prior to record
   encoding that are necessary for the CRAM writer to generate some header
   components."
-  [{:keys [rname->idx subst-mat seq-resolver tag-dict-builder]} ^objects records]
-  (dotimes [i (alength records)]
-    (let [record (aget records i)
+  [{:keys [rname->idx subst-mat seq-resolver tag-dict-builder]} ^List records]
+  (dotimes [i (.size records)]
+    (let [record (.get records i)
           ;; these flag bits of CF are hard-coded at the moment:
           ;; - 0x01: quality scores stored as array (true)
           ;; - 0x02: detached (true)
@@ -217,4 +217,4 @@
           record' (assoc record
                          ::flag cf ::ref-index ri ::end end
                          ::features fs ::tags-index tags-id)]
-      (aset records i record'))))
+      (.set records i record'))))
diff --git a/src/cljam/io/cram/writer.clj b/src/cljam/io/cram/writer.clj
index 02aafc4..66b2dd9 100644
--- a/src/cljam/io/cram/writer.clj
+++ b/src/cljam/io/cram/writer.clj
@@ -2,6 +2,7 @@
   (:require [cljam.io.crai :as crai]
             [cljam.io.cram.encode.alignment-stats :as stats]
             [cljam.io.cram.encode.context :as context]
+            [cljam.io.cram.encode.partitioning :as partition]
             [cljam.io.cram.encode.record :as record]
             [cljam.io.cram.encode.structure :as struct]
             [cljam.io.cram.seq-resolver.protocol :as resolver]
@@ -54,9 +55,7 @@
                                                       preservation-map
                                                       seq-resolver)
         {:keys [ds-compressor-overrides tag-compressor-overrides]} options]
-    (dotimes [i (alength container-records)]
-      (let [slice-records (aget container-records i)]
-        (record/preprocess-slice-records container-ctx slice-records)))
+    (run! (partial record/preprocess-slice-records container-ctx) container-records)
     (context/finalize-container-context container-ctx
                                         ds-compressor-overrides
                                         tag-compressor-overrides)))
@@ -180,8 +179,7 @@
         compression-header-block (generate-compression-header-block container-ctx)
         container-header (generate-container-header compression-header-block slices)
         ^DataOutputStream out (.-stream wtr)
-        container-offset (.size out)
-        counter' (:counter (peek slices))]
+        container-offset (.size out)]
     (struct/encode-container-header out (assoc container-header :counter counter))
     (.write out compression-header-block)
     (run! (fn [{:keys [^bytes header-block data-blocks]}]
@@ -190,17 +188,11 @@
           slices)
     (when-let [index-writer (.-index-writer wtr)]
       (write-index-entries index-writer container-offset container-header
-                           slices container-records))
-    counter'))
-
-(defn- partition-alignments [slices-per-container records-per-slice alns]
-  (->> alns
-       (partition-all records-per-slice)
-       (map object-array)
-       (partition-all slices-per-container)
-       (map object-array)))
+                           slices container-records))))
 
 (defn write-alignments
   "Writes all the given alignments, which is a sequence of alignment maps."
-  [wtr alns header]
-  (reduce (partial write-container wtr header) 0 (partition-alignments 1 10000 alns)))
+  [^CRAMWriter wtr alns header]
+  (partition/with-each-container header (.-options wtr) alns
+    (fn [counter container-records]
+      (write-container wtr header counter container-records))))


Test patch:
diff --git a/test/cljam/io/cram/encode/partitioning_test.clj b/test/cljam/io/cram/encode/partitioning_test.clj
new file mode 100644
index 0000000..a6c9e2f
--- /dev/null
+++ b/test/cljam/io/cram/encode/partitioning_test.clj
@@ -0,0 +1,201 @@
+(ns cljam.io.cram.encode.partitioning-test
+  (:require [cljam.io.cram.encode.partitioning :as partition]
+            [clojure.test :refer [are deftest testing]]))
+
+(defn- partition-alignments [header options alns]
+  (let [acc (volatile! [])]
+    (partition/with-each-container header options alns
+      (fn [counter container-records]
+        (vswap! acc conj [counter (mapv vec container-records)])))
+    @acc))
+
+(deftest with-each-container-test
+  (let [options {:slices-per-container 2
+                 :records-per-slice 3
+                 :min-single-ref-slice-size 2}]
+    (testing "sorted by coord"
+      (let [header {:HD {:SO "coordinate"}}]
+        (are [input expected]
+             (= expected
+                (partition-alignments header options input))
+          []
+          []
+
+          [{:qname "q1", :rname "chr1"}]
+          [[0 [[{:qname "q1", :rname "chr1"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "chr1"}
+           {:qname "q4", :rname "chr1"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}
+                {:qname "q3", :rname "chr1"}]
+               [{:qname "q4", :rname "chr1"}]]]]
+
+          (map (fn [i] {:qname (str \q (inc (long i))), :rname "chr1"}) (range 6))
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}
+                {:qname "q3", :rname "chr1"}]
+               [{:qname "q4", :rname "chr1"}
+                {:qname "q5", :rname "chr1"}
+                {:qname "q6", :rname "chr1"}]]]]
+
+          (map (fn [i] {:qname (str \q (inc (long i))), :rname "chr1"}) (range 7))
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}
+                {:qname "q3", :rname "chr1"}]
+               [{:qname "q4", :rname "chr1"}
+                {:qname "q5", :rname "chr1"}
+                {:qname "q6", :rname "chr1"}]]]
+           [6 [[{:qname "q7", :rname "chr1"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}]]]
+           [2 [[{:qname "q3", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "chr1"}
+           {:qname "q4", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}
+                {:qname "q3", :rname "chr1"}]]]
+           [3 [[{:qname "q4", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "chr2"}
+           {:qname "q4", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}]]]
+           [2 [[{:qname "q3", :rname "chr2"}
+                {:qname "q4", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr2"}
+           {:qname "q3", :rname "chr2"}
+           {:qname "q4", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr2"}]]]
+           [2 [[{:qname "q3", :rname "chr2"}
+                {:qname "q4", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "chr1"}
+           {:qname "q4", :rname "chr1"}
+           {:qname "q5", :rname "chr2"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}
+                {:qname "q3", :rname "chr1"}]
+               [{:qname "q4", :rname "chr1"}]]]
+           [4 [[{:qname "q5", :rname "chr2"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr2"}
+           {:qname "q3", :rname "chr2"}
+           {:qname "q4", :rname "chr3"}
+           {:qname "q5", :rname "chr4"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr2"}]]]
+           [2 [[{:qname "q3", :rname "chr2"}
+                {:qname "q4", :rname "chr3"}]]]
+           [4 [[{:qname "q5", :rname "chr4"}]]]]
+
+          [{:qname "q1", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "*"}
+           {:qname "q2", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "*"}
+                {:qname "q2", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "*"}
+           {:qname "q2", :rname "*"}
+           {:qname "q3", :rname "*"}
+           {:qname "q4", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "*"}
+                {:qname "q2", :rname "*"}
+                {:qname "q3", :rname "*"}]
+               [{:qname "q4", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}]]]
+           [2 [[{:qname "q3", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "chr1"}
+           {:qname "q3", :rname "*"}
+           {:qname "q4", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "chr1"}]]]
+           [2 [[{:qname "q3", :rname "*"}
+                {:qname "q4", :rname "*"}]]]]
+
+          [{:qname "q1", :rname "chr1"}
+           {:qname "q2", :rname "*"}
+           {:qname "q3", :rname "*"}
+           {:qname "q4", :rname "*"}]
+          [[0 [[{:qname "q1", :rname "chr1"}
+                {:qname "q2", :rname "*"}]]]
+           [2 [[{:qname "q3", :rname "*"}
+                {:qname "q4", :rname "*"}]]]])
+        (are [input] (thrown? Exception (partition-alignments header options input))
+          [{:qname "q1", :rname "*"}
+           {:qname "q2", :rname "chr1"}]
+
+          [{:qname "q1", :rname "*"}
+           {:qname "q2", :rname "*"}
+           {:qname "q3", :rname "chr1"}])))
+    (testing "unsorted"
+      (are [input expected]
+           (= expected (partition-alignments {:HD {:SO "unsorted"}} options input))
+        [{:qname "q1", :rname "chr1"}
+         {:qname "q2", :rname "chr2"}
+         {:qname "q3", :rname "chr1"}]
+        [[0 [[{:qname "q1", :rname "chr1"}
+              {:qname "q2", :rname "chr2"}
+              {:qname "q3", :rname "chr1"}]]]]
+
+        [{:qname "q1", :rname "chr1"}
+         {:qname "q2", :rname "chr2"}
+         {:qname "q3", :rname "chr1"}
+         {:qname "q4", :rname "chr2"}]
+        [[0 [[{:qname "q1", :rname "chr1"}
+              {:qname "q2", :rname "chr2"}
+              {:qname "q3", :rname "chr1"}]]]
+         [3 [[{:qname "q4", :rname "chr2"}]]]]
+
+        [{:qname "q1", :rname "*"}
+         {:qname "q2", :rname "chr1"}]
+        [[0 [[{:qname "q1", :rname "*"}
+              {:qname "q2", :rname "chr1"}]]]]
+
+        [{:qname "q1", :rname "*"}
+         {:qname "q2", :rname "*"}
+         {:qname "q3", :rname "chr1"}]
+        [[0 [[{:qname "q1", :rname "*"}
+              {:qname "q2", :rname "*"}
+              {:qname "q3", :rname "chr1"}]]]]))))
diff --git a/test/cljam/io/cram/encode/record_test.clj b/test/cljam/io/cram/encode/record_test.clj
index 3fda2a6..b00f050 100644
--- a/test/cljam/io/cram/encode/record_test.clj
+++ b/test/cljam/io/cram/encode/record_test.clj
@@ -6,7 +6,8 @@
             [cljam.io.sequence :as cseq]
             [cljam.test-common :as common]
             [clojure.test :refer [are deftest is testing]]
-            [clojure.walk :as walk]))
+            [clojure.walk :as walk])
+  (:import [java.util ArrayList]))
 
 (def ^:private test-seq-resolver
   (let [seqs (with-open [r (cseq/reader common/test-fa-file)]
@@ -78,7 +79,7 @@
 
 (deftest preprocess-slice-records-test
   (let [cram-header {:SQ [{:SN "ref"}]}
-        records (object-array
+        records (ArrayList.
                  [{:rname "ref", :pos 1, :cigar "5M", :seq "AGAAT", :qual "HFHHH"
                    :options [{:RG {:type "Z", :value "rg001"}}
                              {:MD {:type "Z", :value "2C2"}}
@@ -132,7 +133,7 @@
              ::record/flag 0x0b, ::record/ref-index -1, ::record/end 10, ::record/tags-index 1
              ::record/features []}]
            (walk/prewalk #(if (.isArray (class %)) (vec %) %)
-                         records)))
+                         (vec records))))
     (is (= [[{:tag :MD, :type \Z} {:tag :NM, :type \c}]
             []]
            (:tag-dict container-ctx)))
@@ -158,7 +159,7 @@
                        :RG
                        [{:ID "rg001"}
                         {:ID "rg002"}]}
-          records (object-array
+          records (ArrayList.
                    [{:qname "q001", :flag 99, :rname "ref", :pos 1, :end 5, :mapq 0,
                      :cigar "5M", :rnext "=", :pnext 151, :tlen 150, :seq "AGAAT", :qual "HFHHH"
                      :options [{:RG {:type "Z", :value "rg001"}}
@@ -337,7 +338,7 @@
     (let [cram-header {:SQ
                        [{:SN "ref"}
                         {:SN "ref2"}]}
-          records (object-array
+          records (ArrayList.
                    [{:qname "q001", :flag 77, :rname "*", :pos 0, :end 0, :mapq 0,
                      :cigar "*", :rnext "*", :pnext 0, :tlen 0, :seq "AATCC", :qual "CCFFF"
                      :options []}