feat(sort): add CB (cellular barcode) to template-coordinate sort key

Port fgbio PR #1142 to include the CB tag in the template-coordinate
sort order for single-cell data. Without CB in the sort, templates from
different cells at the same locus are interleaved, breaking grouper
adjacency assumptions.

Comparison order is now:
  tid1 → tid2 → pos1 → pos2 → neg1 → neg2 → CB → MI → name → library → is_upper

Changes:
- Add cb_hash (u64) to TemplateKey between secondary and tertiary,
  expanding packed key from 32 to 40 bytes and inline header from 40
  to 48 bytes
- Add cid (String) to TemplateCoordinateKey with length-first then
  lexicographic comparison matching fgbio
- Add cell_tag parameter to compare_template_coordinate_raw() for raw
  byte comparison path
- Add cell_tag field and builder method to RawExternalSorter
- Add --cell-tag / -c CLI argument to sort command (default "CB")
- Update group and dedup help text to mention --cell-tag in sort
  suggestions

Author: nh13

crates/fgumi-raw-bam/src/sort.rs

@@ -3,7 +3,7 @@ use std::cmp::Ordering;
 use crate::cigar::mate_unclipped_5prime;
 use crate::cigar::unclipped_5prime_sort;
 use crate::fields::{self, flags, mate_pos, mate_ref_id, pos, read_name, ref_id};
-use crate::tags::{find_mc_tag_in_record, find_mi_tag_in_record};
+use crate::tags::{find_mc_tag_in_record, find_mi_tag_in_record, find_string_tag_in_record};
 
 #[must_use]
 pub fn compare_coordinate_raw(a: &[u8], b: &[u8]) -> Ordering {
@@ -57,9 +57,15 @@ pub fn compare_queryname_raw(a: &[u8], b: &[u8]) -> Ordering {
 /// Compare for template-coordinate ordering using raw bytes.
 ///
 /// This matches samtools' template-coordinate sorting which uses unclipped 5' positions.
+/// When `cell_tag` is `Some`, the CB (cellular barcode) tag is included in the comparison
+/// between neg2 and MI, matching fgbio's sort order.
 #[inline]
 #[must_use]
-pub fn compare_template_coordinate_raw(a: &[u8], b: &[u8]) -> Ordering {
+pub fn compare_template_coordinate_raw(
+    a: &[u8],
+    b: &[u8],
+    cell_tag: Option<&[u8; 2]>,
+) -> Ordering {
     // Extract all needed fields from both records
     let a_tid = ref_id(a);
     let a_pos = pos(a);
@@ -133,6 +139,16 @@ pub fn compare_template_coordinate_raw(a: &[u8], b: &[u8]) -> Ordering {
             (false, true) => Ordering::Greater,
             _ => Ordering::Equal,
         })
+        .then_with(|| {
+            // CB (cellular barcode): length-first, then lexicographic (matching fgbio)
+            if let Some(tag) = cell_tag {
+                let a_cb = find_string_tag_in_record(a, tag).unwrap_or(b"");
+                let b_cb = find_string_tag_in_record(b, tag).unwrap_or(b"");
+                a_cb.len().cmp(&b_cb.len()).then_with(|| a_cb.cmp(b_cb))
+            } else {
+                Ordering::Equal
+            }
+        })
         .then_with(|| compare_mi_tags_raw(a, b))
         .then_with(|| compare_names_raw(a, b))
         .then_with(|| a_upper.cmp(&b_upper))
@@ -201,6 +217,7 @@ fn compare_mi_tags_raw(a: &[u8], b: &[u8]) -> Ordering {
 mod tests {
     use super::*;
     use crate::testutil::*;
+    use rstest::rstest;
     use std::cmp::Ordering;
 
     // ========================================================================
@@ -319,7 +336,7 @@ mod tests {
             -1,
             &[],
         );
-        assert_eq!(compare_template_coordinate_raw(&rec, &rec), Ordering::Equal);
+        assert_eq!(compare_template_coordinate_raw(&rec, &rec, None), Ordering::Equal);
     }
 
     #[test]
@@ -329,8 +346,8 @@ mod tests {
         let cigar = &[(10 << 4) | 0]; // 10M
         let a = make_bam_bytes(0, 100, 0, b"rea", cigar, 10, -1, -1, &[]);
         let b = make_bam_bytes(2, 100, 0, b"rea", cigar, 10, -1, -1, &[]);
-        assert_eq!(compare_template_coordinate_raw(&a, &b), Ordering::Less);
-        assert_eq!(compare_template_coordinate_raw(&b, &a), Ordering::Greater);
+        assert_eq!(compare_template_coordinate_raw(&a, &b, None), Ordering::Less);
+        assert_eq!(compare_template_coordinate_raw(&b, &a, None), Ordering::Greater);
     }
 
     #[test]
@@ -339,7 +356,7 @@ mod tests {
         let a = make_bam_bytes(-1, -1, flags::UNMAPPED, b"aaa", &[], 0, -1, -1, &[]);
         let b = make_bam_bytes(-1, -1, flags::UNMAPPED, b"zzz", &[], 0, -1, -1, &[]);
         // Both fully unmapped, compare by name
-        assert_eq!(compare_template_coordinate_raw(&a, &b), Ordering::Less);
+        assert_eq!(compare_template_coordinate_raw(&a, &b, None), Ordering::Less);
     }
 
     #[test]
@@ -350,7 +367,7 @@ mod tests {
         let b = make_bam_bytes(0, 200, flags::PAIRED, b"rea", cigar, 10, 0, 100, &[]);
         // a is unmapped with mate at tid=0,pos=100; b is mapped at tid=0,pos=200
         // a sorts by mate position (100) which is < b's position (200)
-        let cmp = compare_template_coordinate_raw(&a, &b);
+        let cmp = compare_template_coordinate_raw(&a, &b, None);
         assert_ne!(cmp, Ordering::Equal);
     }
 
@@ -381,7 +398,7 @@ mod tests {
             &[],
         );
         // Both mapped, mates unmapped. a at pos 100 < b at pos 200
-        assert_eq!(compare_template_coordinate_raw(&a, &b), Ordering::Less);
+        assert_eq!(compare_template_coordinate_raw(&a, &b, None), Ordering::Less);
     }
 
     #[test]
@@ -394,7 +411,7 @@ mod tests {
         let b = make_bam_bytes(0, 100, flags::PAIRED, b"rea", cigar, 10, 0, 200, &[]);
         // After canonical ordering, both should produce same (tid1=0,tid2=0,pos1=~100,pos2=~200)
         // so they compare equal on positions, differentiated by is_upper
-        let cmp = compare_template_coordinate_raw(&a, &b);
+        let cmp = compare_template_coordinate_raw(&a, &b, None);
         // a is_upper=true, b is_upper=false -> a > b (true > false)
         assert_eq!(cmp, Ordering::Greater);
     }
@@ -428,7 +445,7 @@ mod tests {
             &[],
         );
         // Reverse strand sorts before forward in samtools convention (neg1=true < neg1=false)
-        let cmp = compare_template_coordinate_raw(&a, &b);
+        let cmp = compare_template_coordinate_raw(&a, &b, None);
         assert_ne!(cmp, Ordering::Equal);
     }
 
@@ -467,7 +484,38 @@ mod tests {
             &aux_b,
         );
         // MI 10 < MI 20
-        assert_eq!(compare_template_coordinate_raw(&a, &b), Ordering::Less);
+        assert_eq!(compare_template_coordinate_raw(&a, &b, None), Ordering::Less);
+    }
+
+    // ========================================================================
+    // compare_template_coordinate_raw: CB (cell barcode) tag tests
+    // ========================================================================
+
+    /// CB ordering: parameterized cases for compare_template_coordinate_raw.
+    #[rstest]
+    // Same-length CBs: AAAA < BBBB lexicographically
+    #[case(b"CBZ\x41\x41\x41\x41\x00".as_slice(), b"CBZ\x42\x42\x42\x42\x00".as_slice(), Some(b"CB"), Ordering::Less, "AAAA < BBBB")]
+    // CB ignored when cell_tag is None
+    #[case(b"CBZ\x41\x41\x41\x41\x00".as_slice(), b"CBZ\x42\x42\x42\x42\x00".as_slice(), None, Ordering::Equal, "CB ignored without cell_tag")]
+    // Missing CB (empty) < present CB by length
+    #[case(b"".as_slice(), b"CBZ\x41\x41\x41\x41\x00".as_slice(), Some(b"CB"), Ordering::Less, "no CB < has CB")]
+    // Shorter CB < longer CB by length
+    #[case(b"CBZA\x00".as_slice(), b"CBZAA\x00".as_slice(), Some(b"CB"), Ordering::Less, "A < AA by length")]
+    fn test_compare_template_coordinate_raw_cb_ordering(
+        #[case] aux_a: &[u8],
+        #[case] aux_b: &[u8],
+        #[case] cell_tag: Option<&[u8; 2]>,
+        #[case] expected: Ordering,
+        #[case] msg: &str,
+    ) {
+        let cigar = &[(10 << 4) | 0]; // 10M
+        let a = make_bam_bytes(
+            0, 100, flags::PAIRED | flags::MATE_UNMAPPED, b"rea", cigar, 10, -1, -1, aux_a,
+        );
+        let b = make_bam_bytes(
+            0, 100, flags::PAIRED | flags::MATE_UNMAPPED, b"rea", cigar, 10, -1, -1, aux_b,
+        );
+        assert_eq!(compare_template_coordinate_raw(&a, &b, cell_tag), expected, "{msg}");
     }
 
     // ========================================================================

src/commands/dedup.rs

@@ -1039,7 +1039,7 @@ is selected as the representative; all others are marked as duplicates.
 # Input Requirements
 
 - Must be processed with `fgumi zipper` (adds `pa` tag for secondary/supplementary reads)
-- Must be sorted with `fgumi sort --order template-coordinate`
+- Must be sorted with `fgumi sort --order template-coordinate --cell-tag CB`
 - UMI tags on reads (default: RX tag)
 
 Note: Using `samtools sort` will NOT work correctly because it doesn't use the
@@ -1188,7 +1188,9 @@ impl Command for MarkDuplicates {
                 "Input BAM must be template-coordinate sorted.\n\n\
                 To prepare your BAM file, run:\n  \
                 fgumi zipper -u unmapped.bam -r reference.fa -o merged.bam\n  \
-                fgumi sort -i merged.bam -o sorted.bam --order template-coordinate"
+                fgumi sort -i merged.bam -o sorted.bam --order template-coordinate \
+                --cell-tag {}",
+                self.cell_tag
             );
         }
         info!("Template-coordinate sorted");
@@ -1398,9 +1400,11 @@ impl Command for MarkDuplicates {
                 `fgumi zipper` during the merge of unmapped and mapped BAMs.\n\n\
                 To fix this, re-run your pipeline starting from `fgumi zipper`:\n  \
                 fgumi zipper --unmapped unmapped.bam --mapped aligned.bam -o merged.bam\n  \
-                fgumi sort -i merged.bam -o sorted.bam --order template-coordinate\n  \
+                fgumi sort -i merged.bam -o sorted.bam --order template-coordinate \
+                --cell-tag {}\n  \
                 fgumi dedup -i sorted.bam -o deduped.bam",
-                final_metrics.missing_pa_tag
+                final_metrics.missing_pa_tag,
+                self.cell_tag
             );
         }
 

src/commands/group.rs

@@ -675,8 +675,9 @@ Accepts reads in any order (including unsorted) and outputs reads sorted by:
    4. Read Name
 
 It is recommended to sort the reads into template-coordinate order prior to running
-this tool to avoid re-sorting the input. Use `fgumi sort --order template-coordinate` for
-the pre-sorting. The output will always be written in template-coordinate order.
+this tool to avoid re-sorting the input. Use `fgumi sort --order template-coordinate`
+(with matching `--cell-tag`) for the pre-sorting. The output will always be written
+in template-coordinate order.
 
 During grouping, reads and templates are filtered out as follows:
 
@@ -924,7 +925,9 @@ impl Command for GroupReadsByUmi {
             bail!(
                 "Input BAM must be template-coordinate sorted.\n\n\
                 To sort your BAM file, run:\n  \
-                fgumi sort -i input.bam -o sorted.bam --order template-coordinate"
+                fgumi sort -i input.bam -o sorted.bam --order template-coordinate \
+                --cell-tag {}",
+                self.cell_tag
             );
         }
         info!("template coordinate sorted");

src/commands/sort.rs

@@ -25,7 +25,7 @@ use clap::{Parser, ValueEnum};
 use fgumi_lib::bam_io::create_bam_reader;
 use fgumi_lib::logging::OperationTimer;
 use fgumi_lib::sort::{RawExternalSorter, SortOrder};
-use fgumi_lib::validation::validate_file_exists;
+use fgumi_lib::validation::{string_to_tag, validate_file_exists};
 use log::info;
 use std::path::PathBuf;
 
@@ -176,6 +176,15 @@ pub struct Sort {
     /// position tracking.
     #[arg(long = "write-index", default_value = "false")]
     pub write_index: bool,
+
+    /// Cell barcode tag for template-coordinate sort.
+    ///
+    /// When sorting in template-coordinate order, this tag is included in the
+    /// sort key so that templates from different cells at the same locus are
+    /// not interleaved. Use the same value as `--cell-tag` in `fgumi group`
+    /// and `fgumi dedup`. Only used for template-coordinate sort.
+    #[arg(short = 'c', long = "cell-tag", default_value = "CB")]
+    pub cell_tag: String,
 }
 
 /// Parse memory size string (e.g., "512M", "1G", "2G").
@@ -264,6 +273,17 @@ impl Command for Sort {
 }
 
 impl Sort {
+    /// Parse the cell tag for template-coordinate sort/verify, returning `None`
+    /// for other sort orders.
+    fn parse_cell_tag(&self) -> Result<Option<[u8; 2]>> {
+        if matches!(self.order, SortOrderArg::TemplateCoordinate) {
+            let tag = string_to_tag(&self.cell_tag, "cell-tag")?;
+            Ok(Some([tag.as_ref()[0], tag.as_ref()[1]]))
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Execute sort mode: read, sort, and write output.
     fn execute_sort(&self, command_line: &str) -> Result<()> {
         let output = self.output.as_ref().expect("output required for sort mode");
@@ -286,10 +306,15 @@ impl Sort {
             self.max_memory
         };
 
+        let cell_tag = self.parse_cell_tag()?;
+
         info!("Starting Sort");
         info!("Input: {}", self.input.display());
         info!("Output: {}", output.display());
         info!("Sort order: {:?}", self.order);
+        if let Some(ct) = cell_tag {
+            info!("Cell tag: {}{}", ct[0] as char, ct[1] as char);
+        }
         if self.memory_per_thread {
             info!(
                 "Max memory: {} MB ({} MB/thread × {} threads)",
@@ -318,6 +343,10 @@ impl Sort {
             .write_index(self.write_index)
             .pg_info(crate::version::VERSION.to_string(), command_line.to_string());
 
+        if let Some(ct) = cell_tag {
+            sorter = sorter.cell_tag(ct);
+        }
+
         if let Some(ref tmp) = self.tmp_dir {
             sorter = sorter.temp_dir(tmp.clone());
         }
@@ -349,11 +378,16 @@ impl Sort {
         use std::cmp::Ordering;
         use std::fs::File;
 
+        let cell_tag = self.parse_cell_tag()?;
+
         let timer = OperationTimer::new("Verifying BAM sort order");
 
         info!("Starting Sort Verification");
         info!("Input: {}", self.input.display());
         info!("Expected order: {:?}", self.order);
+        if let Some(ct) = cell_tag {
+            info!("Cell tag: {}{}", ct[0] as char, ct[1] as char);
+        }
 
         // Get header using noodles reader, then use raw reader for records
         let (_, header) = create_bam_reader(&self.input, 1)?;
@@ -383,7 +417,7 @@ impl Sort {
                 let lib_lookup = LibraryLookup::from_header(&header);
                 verify_sort_order(
                     raw_reader,
-                    |bam| extract_template_key_inline(bam, &lib_lookup),
+                    |bam| extract_template_key_inline(bam, &lib_lookup, cell_tag.as_ref()),
                     // Use core_cmp to ignore name_hash tie-breaker differences
                     // This allows both fgumi and samtools sorted files to pass
                     |key, prev| key.core_cmp(prev) == Ordering::Less,